In a region with no free charge, the electrostatic potential obeys the Laplace equation:
`grad^2 phi = 0`
Jacobi iteration is a simple numerical technique to find solutions to problems like this. For example, if you have a 2D equal-spaced cartesian mesh of points, the numerical gradient
`(del^2 phi)/(del x^2) + (del^2 phi)/(del y^2)` can be approximated as:
(phi[i+1,j] + phi[i-1,j] - 2*phi[i,j]) + (phi[i,j+1] + phi[i,j-1] - 2*phi[i,j])
If that is set to zero, then one can iterate the expression:
next_phi[i,j] = 0.25 * (phi[i+1,j] + phi[i-1,j] + phi[i,j+1] + phi[i,j-1])
It is slow, but here's a basic example. Scroll down to see cool ways to make it much faster!
Shown as a contour using mostly Matplotlib defaults:
1
2
| CS = ax.contour(phi, origin = 'lower')
ax.clabel(CS, inline=1, fontsize=11)
|
The script to set up the problem (
phi,
do_me, and
sphere) and then iterate is quite short. Instead of looping over all the indices in the array as you would in a compiled language like C or FORTRAN, in Python we try to just use existing NumPy methods like
np.roll in line 29.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
| import numpy as np
import matplotlib.pyplot as plt
nx, ny = 101, 101
phi = np.zeros((ny, nx), dtype = 'float')
do_me = np.ones_like(phi, dtype='bool')
x0, y0, r0 = 40, 65, 12
x = np.arange(nx, dtype = 'float')[None,:]
y = np.arange(ny, dtype = 'float')[:,None]
rsq = (x-x0)**2 + (y-y0)**2
circle = rsq <= r0**2
phi[circle] = 1.0
do_me[circle] = False
do_me[0,:], do_me[-1,:], do_me[:,0], do_me[:,-1] = False, False, False, False
n, nper = 100, 100
phi_hold = np.zeros((n+1, ny, nx))
phi_hold[0] = phi
for i in range(n):
for j in range(nper):
phi2 = 0.25*(np.roll(phi, 1, axis=0) +
np.roll(phi, -1, axis=0) +
np.roll(phi, 1, axis=1) +
np.roll(phi, -1, axis=1) )
phi[do_me] = phi2[do_me]
phi_hold[i+1] = phi
change = phi_hold[1:] - phi_hold[:-1]
cmin, cmax = [], []
for thing in change:
cmin.append(thing[do_me].min())
cmax.append(thing[do_me].max())
|
The actual calculation (lines 26 to 36) takes about 2.2 seconds on my laptop. That's 100*100*101*101 multiplications plus three times that many sums, among other things, or about 6ns per flop. Since a lap-top-CPU-flop is more like about 1 or 1.5ns, there's definitely some room for improvement. (That GPU is just sitting there hungry for work, and could probably do it much faster!)
There's an upcoming post on lap-top-CPU-flops with NumPy; stay tuned!
So, let's make the problem a little harder by making it 3D, Instead of ~10,000 mesh points we'll have ~1,000,000. At 8 bytes per point, that's beyond the cache, so it's going to be slower still.
and two slices below the sphere:
This took 525 seconds. That's 100*100*101*101*101 multiplications plus five summations, or about 9ns per flop. It's creeping up there! ,And we're still changing at about 3E-04 per 100 iterations. The next
Relax Laplace! post will cover some ways to bring that back down.
The 3D version of the script for reference:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
| nx, ny, nz = 101, 101, 101
phi = np.zeros((nz, ny, nx), dtype = 'float')
do_me = np.ones_like(phi, dtype='bool')
x0, y0, z0, r0 = 40, 65, 50, 12
x = np.arange(nx, dtype = 'float')[None,None,:]
y = np.arange(ny, dtype = 'float')[None,:,None]
z = np.arange(ny, dtype = 'float')[:,None,None]
rsq = (x-x0)**2 + (y-y0)**2 + (z-z0)**2
sphere = rsq <= r0**2
phi[sphere] = 1.0
do_me[sphere] = False
do_me[0,:,:], do_me[-1,:,:] = False, False
do_me[:,0,:], do_me[:,-1,:] = False, False
do_me[:,:,0], do_me[:,:,-1] = False, False
n, nper = 100, 100
phi_hold = np.zeros((n+1, nz, ny, nx))
phi_hold[0] = phi
one_sixth = (1./6.)
for i in range(n):
for j in range(nper):
phi2 = one_sixth*(np.roll(phi, 1, axis=0) +
np.roll(phi, -1, axis=0) +
np.roll(phi, 1, axis=1) +
np.roll(phi, -1, axis=1) +
np.roll(phi, 1, axis=2) +
np.roll(phi, -1, axis=2) )
phi[do_me] = phi2[do_me]
print i
phi_hold[i+1] = phi
|
