Finite element solver for the 1D wave equation

A finite-element solver for the 1D wave equation.

We use finite elements to obtain the mass \(M\) and stiffness \(K\) matrices. We use bilinear discretization for the time stepping.

\[ \begin{bmatrix} \dot{\mathbf{x}} \\ \ddot{\mathbf{x}} \end{bmatrix} = \begin{bmatrix} 0 & \mathbf{I} \\ -\mathbf{M}^{-1}\mathbf{K} & 0 \end{bmatrix} \begin{bmatrix} \mathbf{x} \\ \dot{\mathbf{x}} \end{bmatrix} + \begin{bmatrix} 0 \\ \mathbf{M}^{-1} \end{bmatrix} \mathbf{f} \]

---

source

discretize

 discretize (A, step)

Jax compatible bilinear discretization from https://github.com/srush/annotated-s4

---

source

Wave1dSolverFE

 Wave1dSolverFE (sampling_rate:float, final_time:float, length:float,
                 n_gridpoints:int, wave_speed:float=1)

This class solves the 1D wave equation using finite elements and state space discretization.

	Type	Default	Details
sampling_rate	float		sampling rate in Hz
final_time	float		final time in seconds
length	float		length of the string in meters
n_gridpoints	int		number of points in the string
wave_speed	float	1	wave speed in m/s

Test

from physmodjax.solver.generator import Gaussian

n_gridpoints = 200
solver = Wave1dSolverFE(
    sampling_rate=2000,
    final_time=1,
    length=1,
    n_gridpoints=n_gridpoints,
    wave_speed=10,
)

u0 = Gaussian(num_points=n_gridpoints)()
v0 = np.zeros_like(u0)

t, u, v = solver.solve(u0, v0)

dx: 0.005025125628140704 in meters
dt: 0.0005 in seconds
number of points (n_gridpoints): (200,)
time in samples (nt): (2000,)

plt.plot(solver.grid, u[50], label="initial")

# show the solution viewed from above
plt.figure(figsize=(5, 10))
plt.pcolormesh(solver.grid, t, u)
plt.xlabel("x")
plt.ylabel("t")
plt.colorbar()
plt.show()