laplace

#!/usr/bin/env python

from __future__ import division
from numpy import *
from matplotlib import pyplot
from mpl_toolkits.mplot3d import Axes3D
import scipy.linalg

# We're going to solve the 2D Laplace equation on an n by n grid, with sinusoidal Dirichlet boundary conditions
n = 52
dx = 1/(n+1)

# Pick random sinusoids.  Our boundary conditions will be z(x,y) = sum_i coeffs[i]*sin(freq[i,0]*x+freq[i,1]*y)
m = 10
random.seed(138131)
coeffs = random.randn(10)
freq = 2*pi*random.randn(10,2)

# Set up dense matrix, treating it initially as a rank 4 tensor.  We're only going to fill in the upper triangle.
A = zeros((n,n,n,n))
i = arange(n).reshape(-1,1)
j = i.reshape(1,-1)
A[i,j,i,j] = -4
A[i[:-1],j,i[1:],j] = 1
A[i,j[:,:-1],i,j[:,1:]] = 1
A = A.reshape((n**2,n**2)) # Smash A down from rank 4 to rank 2

# Compute the right hand side
b = zeros((n,n))
x = dx*arange(n).reshape(-1,1)
b[:,0] += (coeffs*sin(freq[:,0]*x)).sum(axis=-1)
b[:,-1] += (coeffs*sin(freq[:,0]*x+freq[:,1])).sum(axis=-1)
b[0,:] += (coeffs*sin(freq[:,1]*x)).sum(axis=-1)
b[-1,:] += (coeffs*sin(freq[:,1]*x+freq[:,0])).sum(axis=-1)
b = b.reshape(n**2) # Smash b down from rank 2 to rank 1

# Solve the linear system, using only the upper triangle of A
z = scipy.linalg.solve(-A,b,sym_pos=1)
z = z.reshape(n,n)
print 'z =\n%s'%z

# Plot
x = y = dx*arange(n)
x,y = meshgrid(x,y)
fig = pyplot.figure()
axes = fig.add_subplot(111,projection='3d')
axes.plot_surface(x,y,z,rstride=4,cstride=4)
pyplot.show()