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.57,0.57,0.57)]

max=5
l=0.2
lmax = int(l*max+0.5)
v=zeros((2*max+1,2*max+1,2*max+1),Float)
for i in arange(max-lmax,max+lmax+1,1):
    for j in arange(max-lmax,max+lmax+1,1):
        for k in arange(max-lmax,max+lmax+1,1):
            v[i][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 or i>max+lmax or j<max-lmax or j>max+lmax
                    or k<max-lmax or k>max+lmax):
                    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.05
for i in arange(1,2*max,1):
    for j in arange(1,2*max,1):
        for k in arange(1,2*max,1):
            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
            if ((i==max-lmax or i==max+lmax) and j>=max-lmax and j<=max+lmax
                and k>=max-lmax and k<=max+lmax):
                x=2.*x
            if ((j==max-lmax or j==max+lmax) and i>=max-lmax and i<=max+lmax
                and k>=max-lmax and k<=max+lmax):
                y=2.*y
            if ((k==max-lmax or k==max+lmax) and j>=max-lmax and j<=max+lmax
                and i>=max-lmax and i<=max+lmax):
                z=2.*z
            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)
side=2.*float(lmax)/float(max)
box(pos=(0,0,0),size=(side,side,side),color=(0,0,1))
box(pos=(1,0,0),size=(0.01,2,2))
