FastWave_Maybe.py

# -*- coding: utf-8 -*-
"""
Created on Mon May 29 02:45:53 2017

@author: Chronum
"""

import numpy as np
import matplotlib.pyplot as plt
from numba import jit

#@jit
def fd(matrix_new, matrix, matrix_old, rowindices, colindices, args):
    """Simulates one integration step of the wave equation"""
    wave_speed = args["wave_speed"]
    time_step = args["time_step"]
    space_step = args["space_step"]
    
    c = wave_speed**2 * time_step**2 / space_step**2
    
    # matrix_new = np.zeros(matrix.shape)
    i = rowindices# np.array(range(1, matrix.shape[0]-1))
    j  =colindices# np.array(range(1, matrix.shape[1]-1))
    
    matrix_new[i,j] = 2*(1 - 2*c)*matrix[i,j] - matrix_old[i,j] + \
    c*(matrix[i+1,j] + matrix[i-1,j] + matrix[i,j+1] + matrix[i,j-1])
    # return matrix_new

#@jit
def sim2(xspan, yspan, tspan, args):
    x = np.arange(xspan[0], xspan[1], args["space_step"])
    y = np.arange(yspan[0], yspan[1], args["space_step"])
    t = np.arange(tspan[0], tspan[1], args["time_step"])
    
    xx, yy = np.meshgrid(x, y)
    result = np.zeros((x.size, y.size, t.size))
    
    # Set initial conditions
    result[:, :, 0] = np.exp(-(xx**2 + yy**2))
    result[:, :, 1] = np.exp(-(xx**2 + yy**2))
    
    rowIndices = np.array(range(1, result.shape[0]-1))
    colIndices = np.array(range(1, result.shape[1]-1))
    # Integrate
    for i in range(1, t.size-1):
         fd(result[:, :, i+1], result[:, :, i], result[:, :, i-1],
            rowIndices, colIndices, args)
    return result

args = {}
args["wave_speed"] = 1
args["time_step"] = 0.01
args["space_step"] = 0.05
xspan = [-5, 5]
yspan = [-5, 5]
tspan = [0, 10]

kek = sim2(xspan, yspan, tspan, args)