Poisson equation solved directly (i.e. by Gaussian elimination)

image.png

Hugely inefficient to ignore the banded form.

In [9]:
# customize a color map
import numpy as np
from matplotlib import cm
from matplotlib.colors import LinearSegmentedColormap

rainbow = cm.get_cmap('rainbow')
nc = 80
colors = rainbow( np.linspace(0,1,nc) )  # get nc colors from the colormap
colors[nc//2-1:nc//2+2] = [1,1,1,1]  # mark the half-way point white
mycmap = LinearSegmentedColormap.from_list('colormap', colors) # create custom colormap with modified colors
In [30]:
# Poisson equation solved directly
# Note: this full direct solution is prohibitively inefficient

import numpy as np
from pylab import imshow,show,figure,colorbar
import matplotlib.pyplot as plt
import time

def f(x,y): return 0*x

def left  (y): return 0*np.ones_like(y)#y
def right (y): return 0*np.ones_like(y)#y*y
def top   (x): return 100*np.ones_like(x)#sin(5*x)
def bottom(x): return 0*np.ones_like(x)#x+1

a = 0.; b = 6.; M = 6 # number of horizontal grid intervals
c = 0.; d = 5.; N = 5 # number of vertical   grid intervals

h = (b-a)/float(M); m = M-1; x = np.linspace(a,b,m+2); xinterior = x[1:-1]
k = (d-c)/float(N); n = N-1; y = np.linspace(c,d,n+2); yinterior = y[1:-1]
print('h,k',h,k)

mn = m*n
print('Memory requirement for matrix:',8*mn**2*1.e-6,'MB')

# matrix
a = np.eye(mn)*(-2/h**2 -2/k**2)  # diagonal
for i in range(mn-1): 
    a[i  ,i+1] = 1/h**2  # right neighbor
    a[i+1,i  ] = 1/h**2  # left  neighbor
for i in range(1,n): # erase terms where no right neighbor or left neighbor exists
    a[m*i-1,m*i  ] = 0
    a[m*i  ,m*i-1] = 0
for i in range(mn-m): # write in the vertical neighbor bands
    a[i  ,i+m] = 1/k**2  # neighbor below
    a[i+m,i  ] = 1/k**2  # neighbor above

figure(figsize=(10,10))
imshow(a,interpolation='nearest'); colorbar()

# rhs
X,Y = np.meshgrid(xinterior,yinterior)
Xflat = X.reshape(mn)
Yflat = Y.reshape(mn)
b = f(Xflat,Yflat)

for j in range(n):  # left and right edges
    b[m*j    ] -= left ( yinterior[j] )/h**2 
    b[m*j+m-1] -= right( yinterior[j] )/h**2

for i in range(m):  # top and bottom edges
    b[i        ] -= top   ( xinterior[i] )/k**2 
    b[i+(n-1)*m] -= bottom( xinterior[i] )/k**2

print(f'Solving the {mn}x{mn} linear system')
tic = time.time()
uinterior = np.linalg.solve(a,b).reshape((n,m))
toc = time.time()
print('took',round(toc-tic,2),'seconds')

u = np.empty((n+2,m+2))  # make picture
u[1:-1,1:-1] = uinterior
u[:, 0] = left  (y)
u[:,-1] = right (y)
u[ 0,:] = top   (x)
u[-1,:] = bottom(x)
figure(figsize=(10,10))
imshow(u,interpolation='nearest',cmap=mycmap); colorbar();
h,k 1.0 1.0
Memory requirement for matrix: 0.0031999999999999997 MB
Solving the 20x20 linear system
took 0.0 seconds

Taking advantage of the banded structure

image.png

In [28]:
# code I found on the internet to prepare matrix in scipy banded format
import numpy as np
def diagonal_form(a, upper = 1, lower= 1):
    """
    a is a numpy square matrix
    this function converts a square matrix to diagonal ordered form
    returned matrix in ab shape which can be used directly for scipy.linalg.solve_banded
    """
    n = a.shape[1]
    assert(np.all(a.shape ==(n,n)))
    
    ab = np.zeros((2*n-1, n))
    
    for i in range(n):
        ab[i,(n-1)-i:] = np.diagonal(a,(n-1)-i)
        
    for i in range(n-1): 
        ab[(2*n-2)-i,:i+1] = np.diagonal(a,i-(n-1))

    mid_row_inx = int(ab.shape[0]/2)
    upper_rows = [mid_row_inx - i for i in range(1, upper+1)]
    upper_rows.reverse()
    upper_rows.append(mid_row_inx)
    lower_rows = [mid_row_inx + i for i in range(1, lower+1)]
    keep_rows = upper_rows+lower_rows
    ab = ab[keep_rows,:]

    return ab
In [31]:
from scipy.linalg import solve_banded
abanded = diagonal_form(a,m,m)
figure(figsize=(10,10))

imshow(abanded,interpolation='nearest'); colorbar()

print('Solving the banded linear system')
tic = time.time()
uinterior = solve_banded((m,m),abanded,b).reshape((n,m))
toc = time.time()
print('took',round(toc-tic,2),'seconds')
u = np.empty((n+2,m+2))
u[1:-1,1:-1] = uinterior
u[:, 0] = left  (y)
u[:,-1] = right (y)
u[ 0,:] = top   (x)
u[-1,:] = bottom(x)
figure(figsize=(10,10))
imshow(u,interpolation='nearest',cmap=mycmap);
plt.colorbar();
Solving the banded linear system
took 0.0 seconds
In [ ]: