# +----------------------------------------------------+
# |                  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]
#    Testing convergence with N->inf. for optimal SOR method

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=8
n_pl=5
pl_n=np.zeros(n_pl)
pl_u=np.zeros(n_pl)
for ii in range(n_pl):
    print(N)
    pl_n[ii]=N
    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)
    pl_u[ii]=u[N/4,N/4]
    N=N*2
    
plt.plot(np.log10(pl_n),np.log10(abs(pl_u-upoint(0.25,0.25))),'ro')
plt.plot(np.log10(pl_n),np.log10(pl_n**(-2)),'r')
plt.show()

plt.contourf(u)
plt.show()