
# customize a color map
import numpy as np
from matplotlib import cm
from matplotlib.colors import LinearSegmentedColormap

rainbow = cm.get_cmap('rainbow')
nc = 180
colors = rainbow( np.linspace(0,1,nc) )  # get nc colors from the colormap
colors[nc//2-0:nc//2+1] = [1,1,1,1]  # mark the half-way point white
mycmap = LinearSegmentedColormap.from_list('colormap', colors) # create custom colormap with modified colors


# 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 0*np.ones_like(x)#sin(5*x)
def bottom(x): return 100*np.ones_like(x)#x+1

a = 0.; b = 6.; M = 120#60#16 # number of horizontal grid intervals
c = 0.; d = 5.; N = 100#50 #15 # 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')

# form the matrix

np.set_printoptions(linewidth=300)

a = np.eye(mn)*(-2/h**2 - 2/k**2)  # start with the main diagonal in place

ia = np.arange(mn)
# make the sub- and super-diagonal (which are the same)
# by Jeremy's method
od = np.ones(m); od[-1] = 0
tmp = np.tile(od,n)[:-1]
#print(tmp)
a[ia[1: ],ia[:-1]] = tmp/h**2
a[ia[:-1],ia[1: ]] = tmp/h**2

# still need to put in the other super and sub diagonals  
a[ia[:-m],ia[m:]] = 1/k**2
a[ia[m:],ia[:-m]] = 1/k**2
#display(tmp)
#display(np.array(a,dtype=int))
#imshow(a,interpolation='nearest'); colorbar()

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,:] = 100 #bottom(x)
#display(u)
figure(figsize=(10,10))
imshow(u,interpolation='nearest',cmap=mycmap); colorbar();

h,k 0.05 0.05
Memory requirement for matrix: 1110.335688 MB
Solving the 11781x11781 linear system
took 9.87 seconds


# 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


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,extent=(0,1,0,1));
plt.colorbar();

Solving the banded linear system
took 0.08 seconds

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

Taking advantage of the banded structure¶

Exploiting banded structure¶