Code demonstrates solving an elliptic equation using an iterative solver

poisson.py

"""
Author: Jonathan Zrake, NYU CCPP

Code demonstrates solving an elliptic equation using an iterative solver.
"""

import itertools
import numpy as np
import matplotlib.pyplot as plt


# Describe the physical domain
# ------------------------------------------------------------------------------
Nx = 50
Ny = 50
x0 = 0.0
x1 = 1.5
y0 = 0.0
y1 = 2.0
dx = (x1 - x0) / (Nx - 1)
dy = (y1 - y0) / (Ny - 1)
X, Y = np.ogrid[x0:x1:Nx*1j,y0:y1:Ny*1j]


def set_boundary(U):
    """
    For this problem, set the boundary conditions by fixing the potential to be
    U=0 outside the quadrilateral (Y >= 1.5 - 2*X) and (Y <= 2.75 - 1.5*X) and
    on the domain boundary.
    """
    U[:, 0] = 0.0
    U[:,-1] = 0.0
    U[ 0,:] = 0.0
    U[-1,:] = 0.0
    U[np.where(1 - (Y >= 1.5 - 2*X) * (Y <= 2.75 - 1.5*X))] = 0.0


def gauss_seidel_iteration(U, f, dt):
    """
    Updates the potential `U` toward a solution of the Poisson equation with
    source `f`, with step size `dt`.

    where U is the potential and f is the source, i.e. solves the diffusion
    equation

    dU/dt = del^2 U - f

    The grid spacings are allowed to be different in x and y.
    """
    s = 2*dt * (1.0/dx**2 + 1.0/dy**2)
    for i,j in itertools.product(range(1,Nx-1), range(1,Ny-1)):
        U[i,j] = (U[i,j] * (1.0 - s) +
                  dt/dx**2 * (U[i+1,j] + U[i-1,j]) +
                  dt/dy**2 * (U[i,j+1] + U[i,j-1]) - dt * f[i,j])


def solve_poisson():
    """
    Solves the Poisson equation del^2 U = f using an iterative relaxation
    method.
    """
    fig1 = plt.figure(1)
    fig2 = plt.figure(2)
    ax1 = fig1.add_subplot('111')
    ax2 = fig2.add_subplot('111')

    Z = np.zeros([Nx,Ny])
    f = np.ones([Nx,Ny]) # source term

    dt = dx*dy / 4.0
    errors = [ ]
    for n in range(1000):
        U = Z.copy()
        gauss_seidel_iteration(U, f, dt)
        set_boundary(U)
        e = abs(U - Z).mean()
        Z[:,:] = U
        errors.append(e)
        print "n =", n, "error =", e

    ax1.semilogy(errors)
    ax2.imshow(Z, interpolation='nearest')
    plt.show()


if __name__ == "__main__":
    solve_poisson()