# +----------------------------------------------------+
# |                  Lekce $ - demo 2                  |
# +----------------------------------------------------+
#  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]
#    Comparison of Jacobi, Gauss-Sidel, SOR

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]=upoint(x[0],x[jj])
        u[N,jj]=upoint(x[N],x[jj])
        u[jj,0]=upoint(x[jj],x[0])
        u[jj,N]=upoint(x[jj],x[N])
    for j1 in range(1,N):
        for j2 in range(1,N):
            if (((x[j1]-0.5)**2+(x[j2]-0.5)**2)<0.03):
                u[j1,j2]=np.log(np.sqrt(0.03))
            else:
                if (not Omega(x[j1],x[j2])):
                    u[j1,j2]=upoint(x[j1],x[j2])
    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=40
nstep=500

pl_i=np.array(range(nstep))
pl_Jacobi=np.zeros(nstep)
u,x = Ini_Ux(N)
for i in range(nstep):
    pl_Jacobi[i]=u[N/4,N/4]
    Sweep_Jacobi(u,N)

pl_GS=np.zeros(nstep)
u,x = Ini_Ux(N)
for i in range(nstep):
    pl_GS[i]=u[N/4,N/4]
    Sweep_GS(u,N)

pl_SOR=np.zeros(nstep)
u,x = Ini_Ux(N)
for i in range(nstep):
    pl_SOR[i]=u[N/4,N/4]
    Sweep_SOR(u,N,1.5)

#lmbd=1-0.5*(np.pi/(N+1))**2
lmbd=0.95 #(=1-0.05)

lmbd=1.0-0.05/4.0
lmbd_SOR=(lmbd/(1.0+np.sqrt(1.0-lmbd**2)))**2
w=2.0/(1+np.sqrt(1.0-lmbd**2))

pl_SOR2=np.zeros(nstep)
u,x = Ini_Ux(N)
for i in range(nstep):
    pl_SOR2[i]=u[N/4,N/4]
    Sweep_SOR(u,N,w)


plt.plot(pl_i,np.log(abs(pl_Jacobi-pl_SOR2[-1]))/np.log(10.0),'ro')
plt.plot(pl_i,np.log(abs(pl_GS-pl_SOR2[-1]))/np.log(10.0),'bo')
plt.plot(pl_i,np.log(abs(pl_SOR-pl_SOR2[-1]))/np.log(10.0),'go')
plt.plot(pl_i,np.log(abs(pl_SOR2-pl_SOR2[-1]))/np.log(10.0),'go')
plt.plot(pl_i,np.log((lmbd)**pl_i)/np.log(10.0),'r')
plt.plot(pl_i[:70],np.log((lmbd_SOR)**pl_i[:70])/np.log(10.0),'g')

plt.show()