Wave equation in an interval

from typing import Callable
import numpy as np
from lucifex.mesh import interval_mesh
from lucifex.fem import Constant, Function
from lucifex.fdm import (
    FiniteDifference, FiniteDifferenceArgwise, FunctionSeries, ConstantSeries, 
    BE, finite_difference_order, advective_timestep,
)
from lucifex.solver import ibvp, BoundaryConditions
from lucifex.sim import Simulation, run
from lucifex.plt import plot_line, save_figure, create_animation, display_animation
from lucifex.pde.wave import wave

def create_simulation(
    Lx: float,
    Nx: int,
    dt: float,
    d: float,
    D_diff: FiniteDifference | FiniteDifferenceArgwise,
    u_zero: Callable,
    u_dot_zero: float | Callable,
) -> Simulation:
    order = finite_difference_order(D_diff, minimum=2)
    mesh = interval_mesh(Lx, Nx)
    t = ConstantSeries(mesh, name='t', ics=0.0)
    dt = Constant(mesh, dt, name='dt')
    d = Constant(mesh, d, name='d')
    u = FunctionSeries((mesh, 'P', 1), name='u', order=order, store=1)
    u0 = Function(
        u.function_space, 
        u_zero,
    )
    u0_dot = Function(
        u.function_space, 
        u_dot_zero,
    )
    u.update(u0 - dt * u0_dot)
    u.forward(-float(dt))    
    u_bcs = BoundaryConditions(
        ('dirichlet', lambda x: x[0], 0.0),
        ('dirichlet', lambda x: x[0] - Lx, 0.0),
    )
    u_solver = ibvp(wave, ics=u_zero, bcs=u_bcs)(u, dt, d, D_diff)
    return Simulation(u_solver, t, dt)


Nx = 300
Lx = 1.0
sigma = 0.01 * Lx
u_zero = lambda x: np.exp(-(x[0] - Lx/2)**2 / sigma)
u_dot_zero = 0.0
dt = 0.0005
d = 1.0
D_diff = BE
dt_advective = advective_timestep(np.sqrt(d), Lx/Nx)
courant = dt / dt_advective

simulation = create_simulation(Lx, Nx, dt, d, D_diff, u_zero, u_dot_zero)

t_stop = 1.0
n_stop = 1000
run(simulation, n_stop, t_stop) 

u = simulation['u']

slc = slice(1, None, 20)
legend_labels=(min(u.time_series[slc]), max(u.time_series[slc]))
title = "\n".join(
    (
        f"$\mathcal{{D}}_{{\mathsf{{D}}, u}}=\mathrm{{{D_diff}}}$",
        "$C_{\sqrt{\mathsf{D}}}=X$".replace('X', str(courant)),
    )
)
fig, ax = plot_line(
    u.series[slc], 
    legend_labels, 
    '$t$', 
    cyc='jet',
    x_label='$x$', 
    y_label='$u$',
    title=title,
)
save_figure('u(x,t)', thumbnail=True)(fig)

y_label_series = [f'$u(t={t:.3f})$' for t in u.time_series[slc]]

anim = create_animation(
    plot_line,
    x_label='$x$',
)(u.series[slc], y_label=y_label_series)
anim_path = save_figure('u(x,t)', return_path=True)(anim)

display_animation(anim_path)