# 2d linear Partial Differential Equation Solver using finite differences

This is code that solves partial differential equations on a rectangular domain using partial differences. fd_solve takes an equation, a partially filled in output, and a tuple of the x, y, and t steps to use. eqn_parse turns a representation of an equation to a lambda equation that can be easily used. in __main__, I have created two examples that use this code, one for the wave equation, and the other for the heat equation. Running this code requires numpy, scipy, and plotly.

from plotly.offline import init_notebook_mode, plot
import plotly.graph_objs as go

from scipy.signal import convolve2d
import numpy as np

def fd_solve(eqn, zs, boundary):
# Coordinates for evaluating.
xs, ys, ts = boundary

offset = np.any(zs[1, -1:1, -1:1]) # 1 if u_tt != 0
# Loop through time
for i, t in enumerate(ts[:-1], start=offset):
zs[i+1, 1:-1,1:-1] = eqn(zs, i)

# Magic to make it plot
nframe = 100
frames = [{'name': str(t), 'data': [{'z': zs[i], 'type':'surface'}]} for i,t in enumerate(ts) if i % nframe == 0]
steps = [{'method': 'animate', 'label': str(t)[0:5],
'args': [[str(t)], {'frame': {'duration': 100}, 'mode': 'immediate'}]} for t in ts if t % nframe == 0]
print(len(frames))
figure = {'data': [go.Surface(x=xs, y=ys, z=zs[0])],
'layout': {},
'frames': frames}
'buttons': [{'args': [None,
{'frame': {'duration': 500},
'fromcurrent': True,
'label': 'Play',
'method': 'animate'},
{'args': [[None],
{'frame': {'duration': 0,},
'mode': 'immediate',
'transition': {'duration': 0}}],
'label': 'Pause',
'method': 'animate'}],
'type': 'buttons', 'x': 0.1, 'y': .05}]
sliders_dict = {'active': 0,
'currentvalue': {'font': {'size': 20}, 'prefix': 't: ',
'xanchor': 'right'},
'transition': {'duration': 100},
'steps': steps}
figure['layout']['sliders'] = [sliders_dict]
plot(figure)

def eqn_parse(dt=.01, dx=.05, ts=(0,1,0), xs=(0,0,1), ys=(0,0,1)):
''' Numerical kernel for PDEs of the form
ts[1]*u_t + ts[2]*u_tt = xs[1]*u_x + xs[2]*u_xx + ys[1]*u_y + ys[2]'''

assert xs[0]==ys[0]==0

# 2nd deriv
kernel = np.array([[0,        ys[2],         0],
[xs[2], -2*ys[2]-2*xs[2], xs[2]],
[0,        ys[2],         0]])
kernel /= dx

# 1st deriv
kernel += .5 * np.array([[0,    -ys[1], 0],
[-xs[1], 0,     xs[1]],
[ 0,     ys[1], 0]])
kernel /= dx

if ts[0]==ts[2]==0:
kernel *= dt * ts[1]
return lambda zs, i, : zs[i, 1:-1,1:-1] + convolve2d(zs[i], kernel, mode='valid')

if ts[0]==ts[1]==0:
kernel *= dt*dt*ts[2]
return lambda zs, i, : (-zs[i-1, 1:-1,1:-1]
+ 2 * zs[i, 1:-1,1:-1]
+ convolve2d(zs[i], kernel, mode='valid'))

if __name__ == '__main__':
'''
# Wave eqn
dx,c =.01,.5
dt = dx/(10*c) # dt < dx/(c) for stability
eqn = eqn_parse(dt=dt, dx=dx, ts=(0,0,1), xs=(0,0,c), ys=(0,0,c))

xs, ys, ts = np.arange(-1,1,dx), np.arange(-1,1,dx), np.arange(0,10,dt)

# Boundary conditions
bc = lambda t, x, y : (x-x+y-y)*(t-t) #(x**2-y**3)*np.cos(t)
zs = bc(ts[:,None,None],xs[None,:,None], ys[None,None,:])

# Initial Conditions
zs[0, 1:-1, 1:-1] = 0
zs[0, 1:-1, 1:-1] = dt*np.exp(-(xs[1:-1, None]**2+ys[None, 1:-1]**2)**.5)
zs[2:, 1:-1, 1:-1] = 0

fd_solve(eqn, zs, (xs,ys,ts))
'''
# Heat eqn
dx,k =.05,1.0
dt = dx*dx/(16*k) # dt < dx*dx/(4*k) for stability

eqn = eqn_parse(dt=dt, dx=dx, ts=(0,k,0), xs=(0,0,k), ys=(0,0,k))
xs, ys, ts = np.arange(-1,1,dx), np.arange(-1,1,dx), np.arange(0,2,dt)

# Boundary conditions
bc = lambda t, x, y : (x**2+y**2)*np.cos(3*t)
zs = bc(ts[:,None,None], xs[None,:,None], ys[None,None,:])

# Initial Conditions
ic = lambda t, x, y : (t-t)+1
zs[0,1:-1,1:-1] = ic(ts[:1,None,None],xs[None,1:-1,None], ys[None,None,1:-1])
zs[1:,1:-1,1:-1] = 0

fd_solve(eqn, zs, (xs,ys,ts))
''


1. There's no docstring for fd_solve. The post contains a description that would make a good start.

2. The docstring has zs[2:, 1:-1, 1:-1] = 0 but this doesn't match the code zs[1:,1:-1,1:-1] = 0.

3. eqn_parse does not seem to parse anything. I think you could pick a better name, for example pde_kernel.

4. eqn_parse might return None in some circumstances. Is this the expected behaviour? It would be better to raise an exception if the input is incorrect.

5. When building a dictionary, it is often clearer to call dict instead of using a dictionary literal. When calling dict you don't have to quote the keys, which reduces the amount of syntactic clutter. For example:

sliders_dict = dict(active=0,
currentvalue=dict(font=dict(size=20),
prefix='t: ',
xanchor='right'),
steps=steps,
transition=dict(duration=100))


ts[:,None,None], xs[None,:,None], ys[None,None,:]


I suggest using numpy.meshgrid to make it clear that you are constructing a grid of coordinates:

np.meshgrid(ts, xs, ys, sparse=True, indexing='ij')

7. Instead of (t-t)+1, I suggest using np.ones_like(t) to make it clear that you are constructing an array of ones with the same shape and dtype as t. (See numpy.ones_like.) But since this is constant, why not just write zs[0,1:-1,1:-1] = 1 and avoid the need for ic?

8. The slice 1:-1 appears in many places; I suggest making a global constant, for example INTERIOR = slice(1, -1).

9. The boundary condition computation is wasteful as most of the values computed will be overwritten in the subsequent steps. It might be better to start with

zs = np.zeros(ts.shape + xs.shape + ys.shape)


and then fill in the values only in the boundary planes:

BOUNDARY = [0, -1]
INTERIOR = slice(1, -1)
zs[0, INTERIOR, INTERIOR] = 1
zs[:, BOUNDARY, :] = bc(*np.meshgrid(ts, xs[BOUNDARY], ys, sparse=True, indexing='ij'))
zs[:, INTERIOR, BOUNDARY] = bc(*np.meshgrid(ts, xs[INTERIOR], ys[BOUNDARY], sparse=True, indexing='ij'))


(The repetition in the last two lines can easily be eliminated using a helper function.)