# /usr/bin/env python
import sys
import copy
from math import *
import numpy as np
from numpy.random import default_rng
import matplotlib.pyplot as plt
import matplotlib.animation as ani
[docs]def print_progress(step, total):
"""
Prints a progress bar.
Args:
step (int): progress counter
total (int): counter at completion
"""
message = "simulation progress ["
total_bar_length = 60
percentage = int(step / total * 100)
bar_fill = int(step / total * total_bar_length)
for i in range(total_bar_length):
if i < bar_fill:
message += "|"
else:
message += " "
message += "] "+str(percentage)+" %"
if step < total:
print(message, end="\r")
else:
print(message)
[docs]def draw(frame, wave):
"""
Draws the wave as a colormap.
Used for animation.
Args:
frame (int): index of the frame to be drawn
wave (list): list of waves as a time series
"""
plt.clf()
ax = plt.axes()
ax.set_aspect('equal')
vector_to_matrix( wave[frame] )
plt.pcolormesh(wave[frame], vmin=-0.5, vmax=0.5, cmap='RdBu')
[docs]def draw_3d(frame, wave, xx, yy):
"""
Draws the wave as a surface.
Used for animation.
Args:
frame (int): index of the frame to be drawn
wave (list): list of waves as a time series
xx (array): meshgrid with x coordinates
yy (array): meshgrid with y coordinates
"""
plt.clf()
ax = plt.axes(projection ='3d')
ax.set_zlim(-5, 5)
vector_to_matrix( wave[frame] )
ax.plot_surface(xx, yy, wave[frame], vmin=-0.5, vmax=0.5)
[docs]def animate(wave):
"""
Animates the simulation as a colormap.
Args:
wave (list): list of waves as a time series
"""
nframes = len(wave)
fig = plt.figure()
motion = ani.FuncAnimation(fig, draw, nframes, interval=10, fargs=( wave , ) )
plt.show()
[docs]def animate_3d(wave):
"""
Animates the simulation as a surface.
Args:
wave (list): list of waves as a time series
"""
nframes = len(wave)
fig = plt.figure()
n = int( sqrt( wave[0].size ) )
x = np.linspace(0,1,n)
y = np.linspace(0,1,n)
xx, yy = np.meshgrid(x,y)
motion = ani.FuncAnimation(fig, draw_3d, nframes, interval=10, fargs=( wave , xx, yy ) )
plt.show()
[docs]def propagator(n, v, dt, dx,
top_boundary_free = False,
bottom_boundary_free = False,
left_boundary_free = False,
right_boundary_free = False):
"""
Calculates the matrix describing the time evolution of the system.
The discretized wave equation can be written in matrix form as
.. math ::
u(t + \\Delta t) = M u(t) - u(t - \\Delta t).
Here :math:`M` is a sparse matrix which is built out of :math:`N \\times N` blocks
on its diagonal.
Each block controls wave dynamics on a horizontal line of grid points in the 2D system.
They contain the diagonal elements
.. math ::
M_{i,i} = 2 - 4 v^2 \\left( \\frac{\\Delta t}{\\Delta x} \\right)^2
and off-diagonals
.. math ::
M_{i,i\\pm 1} = v^2 \\left( \\frac{\\Delta t}{\\Delta x} \\right)^2
except for the first and last row of the block where the element outisde
the block remains zero.
In addition, adjacent blocks are also coupled by the elements
.. math ::
M_{i,i\\pm N} = v^2 \\left( \\frac{\\Delta t}{\\Delta x} \\right)^2.
These elements represent the coupling of grid points on a vertical line in the 2D system.
If the medium is free to move at the boundaries,
the first or last row off-diagonals of each block must be adjusted to
:math:`2 v^2 \\left( \\frac{\\Delta t}{\\Delta x} \\right)^2`.
This is also true for top and bottom boundaries and the top and bottom off-blocks.
Args:
n_nodes (int): number of grid points :math:`N` (i.e., the number of elements in the vector :math:`u`)
v (float): wave speed
dt (float): time step, :math:`\\Delta t`
dx (float): distance between grid points, :math:`\\Delta x`
top_boundary_free (bool): If True, free medium boundary condition is applied at the top bound
bottom_boundary_free (bool): If True, free medium boundary condition is applied at the bottom bound
left_boundary_free (bool): If True, free medium boundary condition is applied at the left bound
right_boundary_free (bool): If True, free medium boundary condition is applied at the right bound
Returns:
array: the matrix :math:`M`
"""
n_nodes = n*n
# initialize the array to zeros
m = np.zeros([n_nodes, n_nodes])
# calculate and store v^2 dt^2 / cx^2
c = v*v*dt*dt/(dx*dx)
#
# example: the matrix should look like this for
# N = 3, c = 0.1 and fixed bounds
#
# 1.6 0.2 0.0 0.2 0.0 0.0 0.0 0.0 0.0
# 0.2 1.6 0.2 0.0 0.2 0.0 0.0 0.0 0.0
# 0.0 0.2 1.6 0.0 0.0 0.2 0.0 0.0 0.0
# 0.2 0.0 0.0 1.6 0.2 0.0 0.2 0.0 0.0
# 0.0 0.2 0.0 0.2 1.6 0.2 0.0 0.2 0.0
# 0.0 0.0 0.2 0.0 0.2 1.6 0.0 0.0 0.2
# 0.0 0.0 0.0 0.2 0.0 0.0 1.6 0.2 0.0
# 0.0 0.0 0.0 0.0 0.2 0.0 0.2 1.6 0.2
# 0.0 0.0 0.0 0.0 0.0 0.2 0.0 0.2 1.6
#
for i in range(n_nodes):
# diagonals
m[i,i] = 2 - 4*c
# off diagonals
if i%n != n-1 and i < n_nodes-1:
m[i,i+1] = c
if i%n == 0 and left_boundary_free:
m[i,i+1] = 2*c
if i%n != 0 and i > 0:
m[i,i-1] = c
if i%n == n-1 and right_boundary_free:
m[i,i-1] = 2*c
# off blocks
if i>n-1:
m[i,i-n] = c
if i >= n_nodes-n-1 and top_boundary_free:
m[i,i-n] = 2*c
if i<n_nodes-n-1:
m[i,i+n] = c
if i <= n-1 and bottom_boundary_free:
m[i,i+n] = 2*c
return m
[docs]def add_gaussian(film, a, center_x, center_y):
"""
Initialize the wave to a gaussian initial shape.
The function adjusts the shape of the wave by adding
a narrow gaussian function to the wave function.
Since this operation is additive,
it can be called several times to add several pulses.
Args:
film (array): discretized wave function to be modified
a (float): pulse height
center_x (float): fractional position of the peak in x direction, should be between 0 and 1
center_y (float): fractional position of the peak in y direction, should be between 0 and 1
"""
n_nodes = film.shape[0]
for i in range(n_nodes):
for j in range(n_nodes):
film[i,j] += a*exp( - ( 5.0/float(n_nodes) ) * ((i-n_nodes*center_y)**2 + (j-n_nodes*center_x)**2 ))
[docs]def first_step(initial_wave, matrix, initial_velocity = 0, dt = 0):
"""
Performs the first time step.
Since the wave equation is a second degree partial differential equation
with respect to time, one needs to know the wave function at two
time steps in order to simulate the wave dynamics.
Often one knows instead the initial shape of the wave, :math:`u(0)`, and the
initial velocity of the medium, :math:`\\partial_t u(0)`.
Especially if the medium is initially at rest, :math:`\\partial_t u(0) = 0`.
This function calculates the shape of the wave after the first as
.. math ::
u(\\Delta t) = \\frac{1}{2} M u(0) + \\partial_t u(0) \\Delta t.
Args:
initial_wave (array): discretized wave function in the beginning, :math:`u(0)`
matrix (array): matrix :math:`M` - should be pre-calculated using :meth:`propagator`
initial_velocity (array): velocity of the wave medium in the beginning, :math:`\\partial_t u(0)`
dt (float): time step :math:`\\Delta t`
Returns:
array: discretized wave function after the first time step, :math:`u(\\Delta t)`
"""
return 0.5 * matrix @ initial_wave + dt*initial_velocity
[docs]def run_simulation(initial_wave, next_wave, matrix, dt, time, recording_dt = 0):
"""
Run a dynamic wave simulation.
The discretized wave equation can be written in matrix form as
.. math ::
u(t + \\Delta t) = M u(t) - u(t - \\Delta t).
This function iterates this step for the required time.
The function returns a list containing the wave function at different times.
.. note ::
This function is incomplete!
Args:
initial_wave (array): discretized wave function at the start, :math:`u(0)`
next_wave (array): discretized wave function after first step, :math:`u(\\Delta t)`
matrix (array): matrix :math:`M` - should be pre-calculated using :meth:`propagator`
dt (float): time step :math:`\\Delta t`
time (float): total simulation time
recording_dt (float): time interval between saved wavefunctions
Returns:
list: time evolution of the wave function
"""
n_steps = int(time/dt)
if recording_dt < dt:
recording_interval = n_steps
else:
recording_interval = int(recording_dt/dt)
history = []
current_wave = next_wave
past_wave = initial_wave
for i in range(n_steps):
# apply dynamics
# move u(j) to u(j-1) and u(j+1) to u(j)
# todo
print_progress(i+1,n_steps)
if i%recording_interval == 0:
history.append(current_wave)
return history
[docs]def matrix_to_vector(matrix):
"""
Reshapes a square matrix to a vector.
The wave functions are treated as :math:`N \\times N` matrices
e.g. during plotting, but as :math:`1 \\times N^2` vectors
during calculation.
This function reshapes the matrix into a vector.
Args:
matrix (array): the matrix to reshape
"""
n = matrix.shape[0]
nn = n*n
matrix.shape = (nn,)
[docs]def vector_to_matrix(vector):
"""
Reshapes a vector to a square matrix.
The wave functions are treated as :math:`N \\times N` matrices
e.g. during plotting, but as :math:`1 \\times N^2` vectors
during calculation.
This function reshapes the vector into a matrix.
Note: the length of the vector must be a square number.
Args:
vector (array): the vector to reshape
"""
nn = vector.size
n = int ( sqrt( nn ) )
vector.shape = (n,n)
[docs]def main():
"""
The main program. Simulates a 2D wave.
"""
# Simulation parameters
# number of grid points in both axis directions
# (n x n points in total)
n_nodes = 51
# length of the simulated area in x and y
length = 1.0
# distance between grid points in x (and y) and t
dx = length/(n_nodes-1)
dt = 0.005
# timing
simulation_time = 2
sample_dt = 0.01
# wave speed
v = 1.0
# Calculate the matrix M for time dynamics
matrix = propagator(n_nodes, v, dt, dx,
top_boundary_free = False,
bottom_boundary_free = False,
left_boundary_free = False,
right_boundary_free = False)
# initial conditions
wave_initial = np.zeros( [n_nodes, n_nodes] )
v_initial = np.zeros( [n_nodes, n_nodes] )
add_gaussian(wave_initial, a = 1, center_x = 0.2, center_y = 0.3)
add_gaussian(wave_initial, a = 2, center_x = 0.4, center_y = 0.7)
matrix_to_vector( wave_initial )
matrix_to_vector( v_initial )
wave_dt = first_step(wave_initial, matrix, dt, v_initial)
# record evolution of the wave
wavefunction = [wave_initial]
#
# Run simulation
#
wavefunction += run_simulation(wave_initial, wave_dt, matrix, dt, simulation_time, sample_dt)
# Show the result
animate(wavefunction)
animate_3d(wavefunction)
if __name__ == "__main__":
main()