2d linear Partial Differential Equation Solver using finite differences

Question

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}
    figure['layout']['updatemenus'] = [{
                'buttons': [{'args': [None,
                                     {'frame': {'duration': 500},
                                      'fromcurrent': True,
                                      'transition': {'easing': 'quadratic-in-out'}}],
                                    '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))
    ''

Gareth Rees · Accepted Answer · 2018-06-05 18:40:02Z

There's no docstring for fd_solve. The post contains a description that would make a good start.
The docstring has zs[2:, 1:-1, 1:-1] = 0 but this doesn't match the code zs[1:,1:-1,1:-1] = 0.
eqn_parse does not seem to parse anything. I think you could pick a better name, for example pde_kernel.
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.

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))

Instead of:
```
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')
```
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?
The slice 1:-1 appears in many places; I suggest making a global constant, for example INTERIOR = slice(1, -1).

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.)

Stack Exchange Network

2d linear Partial Differential Equation Solver using finite differences

1 Answer 1

You must log in to answer this question.

Hot Network Questions

2d linear Partial Differential Equation Solver using finite differences

1 Answer 1

You must log in to answer this question.

Related

Hot Network Questions