#!/usr/bin/env python3
'''
Program for solving Poisson equation (Laplace(x) = b) on a rectangular
grid (Nx x Ny).
'''
import numpy as np
import matplotlib.pyplot as plt
def getLinearizedIdx(x,y):
'''
Returns index in vector x, which corresponds with point (x,y) on grid.
'''
return x+y*Nx
def initMatrix():
'''
Creates matrix A and vector b for eqn Ax=b. Three possibilities
are involved:
1) Approximation of Laplace operator by finite differences
x[(i+1,j)]+x[(i-1),j]+x[(i,j+1)]+x[(i,j-1)]-4*x[(i,j)]=h*h*Laplace(x),
where (i,j) denotes point in grid (=getLinearizedIdx(i,j)).
2) Dirichlet boundary condition: x[(i,j)] = 0 (more generally x[k]=C,
where C is constant).
3) von Neumann boundary condition: d/dx(x[(i,j)]) = 0 or d/dy(x[(i,j)])=0
approximated by (d/dy(x[(i,j)]) = 0) <=> (x[(i,j)]=x[(i,j+1)]).
'''
tmp = 1./(h*h)
N = Nx*Ny
A = np.zeros((N,N))
b = np.zeros(N)
for i in range(Nx):
for j in range(Ny):
k = getLinearizedIdx(i,j)
# upper plate
if j==0: #Dirichlet
A[k,k] = 1.0
b[k] = 1.0 # nonzero potential
# lower plate
elif j==Ny-1: #Dirichlet
A[k,k] = 1.0
# circle
elif (j-Ny/4)**2+(i-Nx/2)**2<=Ny/5: #Dirichlet
A[k,k] = 1.0
# sides
elif i==0: #von Neumann
A[k,k] = 1.0
A[k,getLinearizedIdx(i+1,j)] = -1.0
elif i==Nx-1: #von Neumann
A[k,k] = 1.0
A[k,getLinearizedIdx(i-1,j)] = -1.0
# Laplace operator approximation
else:
A[k,k] = -4.0*tmp
A[k,getLinearizedIdx(i-1,j)] = 1.0*tmp
A[k,getLinearizedIdx(i+1,j)] = 1.0*tmp
A[k,getLinearizedIdx(i,j-1)] = 1.0*tmp
A[k,getLinearizedIdx(i,j+1)] = 1.0*tmp
return A,b
# parameters
Nx,Ny = 100,100 #grid size
h = 0.1 #size of grid cell in meters (in both dimensions)
# boundaries of area (used only for graph legend)
X0,Y0 = -5.0, 0.0
X1,Y1 = X0+h*Nx, Y0+h*Ny
# actual solving
A,b = initMatrix()
x = np.linalg.solve(A,b)
# displaying of results
plt.imshow(x.reshape((Ny,Nx)), extent=(X0,X1,Y0,Y1))
plt.colorbar().set_label("$\Phi$ [V]")
plt.xlabel("$x$ [m]")
plt.ylabel("$y$ [m]")
plt.show()