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

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,
                                          prefix='t: ',
  6. 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')
  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.)


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.