#  Author:
#
#    Christian Hill
#
import numpy as np

class WaveEqn2D:
    def __init__( self, nx = 500, ny = 500, c = 0.2, h = 1, dt = 1,
                 use_mur_abc = True ):
#
#  init() initializes the simulation.
#
#  nx and ny are the dimension of the domain;
#  c is the wave speed;
#  h and dt are the space and time grid spacings;
#  If use_mur_abc is True, the Mur absorbing boundary
#  conditions will be used; if False, the Dirichlet
#  (reflecting) boundary conditions are used.
#  
        self.nx = nx
        self.ny = ny
        self.c = c
        self.h = 1
        self.dt = 1
        self.use_mur_abc = use_mur_abc
        self.alpha = self.c * self.dt / self.h
        self.alpha2 = self.alpha**2

        self.u = np.zeros ( ( 3, ny, nx ) )

    def update ( self ):
#
## update() updates the simulation by one time step.
#
#  The simulation stores 3 copies of the solution:
#  u[0]: solution at the current time.
#  u[1]: solution at previous time
#  u[2]: solution two time steps ago.
#
#  On calling update, we overwrite u[2] with u[1], and u[1] with u[0].
#  Then we calculate a new u[0] using a discretized version of the PDE.
#
        u = self.u
        nx = self.nx
        ny = self.ny

        u[2] = u[1]
        u[1] = u[0]
#
#  Calculate U at the new time.
#
#  ( u0 - 2 u1 + u2 ) / dt^2 = c^2 * ( u1(left) - u1(center) + u1(right) ) / dx^2
# 
#  Solve for u0.
#  Leave boundary values at 0.
#  On initial step, u0 is set to initial condition.
#  On next step, the solution two steps back is assumed to have been zero.
#  Thereafter, we have enough data.
#
        u[0, 1:ny-1, 1:nx-1]  = self.alpha2 * (
                      u[1, 0:ny-2, 1:nx-1]
                  +   u[1, 2:ny,   1:nx-1]
                  +   u[1, 1:ny-1, 0:nx-2]
                  +   u[1, 1:ny-1, 2:nx]
                  - 4*u[1, 1:ny-1, 1:nx-1] ) \
                  + ( 2 * u[1, 1:ny-1, 1:nx-1]
                        - u[2, 1:ny-1, 1:nx-1] )

        if self.use_mur_abc:
            # Mur absorbing boundary conditions.
            kappa = (1 - self.alpha) / (1 + self.alpha) 
            u[0, 0, 1:nx-1] = (u[1, 1, 1:nx-1]
                               - kappa * (
                                     u[0, 1, 1:nx-1]
                                   - u[1, 0, 1:nx-1])
                              )
            u[0, ny-1, 1:nx-1] = (u[1, ny-2, 1:nx-1]
                               + kappa * (
                                   u[1, ny-1, 1:nx-1]
                                 - u[0, ny-2, 1:nx-1])
                              )
            u[0, 1:ny-1, 0] = (u[1, 1:ny-1, 1]
                               - kappa * (
                                   u[0, 1:ny-1, 1]
                               - u[1, 1:ny-1, 0])
                              )
            u[0, 1:ny-1, nx-1] = (u[1, 1:ny-1, nx-2]
                               + kappa * (
                                   u[1, 1:ny-1, nx-1]
                                 - u[0, 1:ny-1, nx-2])
                              )