from visual import *

print """
Right button drag to rotate view.
Left button drag up or down to move in or out.
"""
scene.title = "Electric field by Gauss-Seidel"
scene.center = vector(0,0,0)
scene.lights = [vector(0.5,0.5,0.5)]
scene.ambient = 0.3

max=5
d=0.2
lmax = int(d*max+0.5)
w=2*lmax
v=zeros((2*max+1,2*max+1,2*max+1),Float)
for j in arange(max-w,max+w+1,1):
    for k in arange(max-w,max+w+1,1):
        v[max-lmax][j][k]=1.
        v[max+lmax][j][k]=-1.

iter=0
diff=1
while (iter<21 or diff>1e-5):
    diff=0.
    iter=iter+1
    for i in arange(1,2*max,1):
        for j in arange(1,2*max,1):
            for k in arange(1,2*max,1):
                if ((i!=max-lmax and i!=max+lmax) or j<max-w or j>max+w
                    or k<max-w or k>max+w):
                    tmp=v[i][j][k]
                    v[i][j][k]=(v[i-1][j][k]+v[i+1][j][k]
                                +v[i][j+1][k]+v[i][j-1][k]
                                +v[i][j][k-1]+v[i][j][k+1])/6.
                    diff=diff+abs(tmp-v[i][j][k])
    diff=diff/(2.*max+1.)**3

print iter,diff
scale=0.04
for i in arange(1,2*max,1):
    for j in arange(1,2*max,1):
        for k in arange(1,2*max,1):
            if ((i==max-lmax or i==max+lmax) and j>=max-w
                and j<=max+w and k>=max-w and k<=max+w): break
            x=-(v[i+1][j][k]-v[i-1][j][k])*max*0.5
            y=-(v[i][j+1][k]-v[i][j-1][k])*max*0.5
            z=-(v[i][j][k+1]-v[i][j][k-1])*max*0.5
            px=float(i)/float(max)-1
            py=float(j)/float(max)-1
            pz=float(k)/float(max)-1
            arrow(pos=(px,py,pz),axis=(scale*x,scale*y,scale*z),shaftwidth=0.02)
depth=float(lmax)/float(max)
width=2.*float(w)/float(max)
box(pos=(-depth,0,0),size=(0.01,width,width),color=(1,0,0))
box(pos=(depth,0,0),size=(0.01,width,width),color=(0,0,1))
box(pos=(1,0,0),size=(0.01,2,2))