# +----------------------------------------------------+
# |                  Lekce $ - demo 4                  |
# +----------------------------------------------------+
#  FINITE DIFFERENCE METHODS for boundary value problem
#   
#  PROBLEM: to solve Laplace equation with Dirichlet condition
#    solution on <0,1> X <0,1> with some part missing
#    solution using various relax. methods
#  THIS VERSION: test for analytically solvable simple problem
#    solution in global array u[0:N,0:N]
#    Final solution for condensator

import numpy as np
#import scipy.linalg as linalg
import matplotlib.pyplot as plt

#Logical function, returns True, if (x,y) in Omega = domain of problem
def Omega(x,y):
    epsilon=1e-8  #this value should be smaller than any used grid stepsize
    return (abs(x-0.5)+abs(y-0.5)>0.25+epsilon)

#Boundary condition for test with poit charge
def upoint(x,y):
    return np.log(np.sqrt((x-0.5)**2+(y-0.5)**2))

#Definition of initial estimate, including correct boundary condition 
def Ini_Ux(N):
    u=np.zeros((N+1,N+1))
    x=np.linspace(0.0,1.0,N+1)
    for jj in range(N+1):          #Boundary condition
        u[0,jj]=0.0
        u[N,jj]=0.0
        u[jj,0]=0.0
        u[jj,N]=0.0
    for j1 in range(1,N):
        for j2 in range(1,N):
            if (not Omega(x[j1],x[j2])):
                u[j1,j2]=1.0
    return u,x

#One sweep of Jacobi method
def Sweep_Jacobi(u,N):
    v=u.copy()
    for j1 in range(1,N):
        for j2 in range(1,N):
            if (Omega(x[j1],x[j2])):
                u[j1,j2]=0.25*(v[j1+1,j2]+v[j1-1,j2]+v[j1,j2+1]+v[j1,j2-1])

#One sweep of Gauss-Sidel
def Sweep_GS(u,N):
    for j1 in range(1,N):
        for j2 in range(1,N):
            if (Omega(x[j1],x[j2])):
                u[j1,j2]=0.25*(u[j1+1,j2]+u[j1-1,j2]+u[j1,j2+1]+u[j1,j2-1])

#One sweep SOR
def Sweep_SOR(u,N,w):
    for j1 in range(1,N):
        for j2 in range(1,N):
            if (Omega(x[j1],x[j2])):
                u[j1,j2]=w*0.25*(u[j1+1,j2]+u[j1-1,j2]+u[j1,j2+1]+u[j1,j2-1])+(1-w)*u[j1,j2]


N=64
nstep=4*N
rho_J=1-0.05*(20/N)**2
w=2.0/(1+np.sqrt(1.0-rho_J**2))
u,x = Ini_Ux(N)
for i in range(nstep):
    Sweep_SOR(u,N,w)
    
plt.contourf(u)
plt.show()