Advection-diffusion of a Gaussian in an interval

\[\begin{split} \mathbb{S} \begin{cases} \Omega = [0, L_x] \\ u_0(x)=\exp\left(-\frac{(x - x_0)^2}{\sigma^2}\right) \\ u_{\text{N}}(x=0)=0 \\ u_{\text{D}}(x=L_x)=0 \\ \textbf{a}(\textbf{x})=x(L_x - x)\textbf{e}_x \\ \end{cases} \end{split}\]

import numpy as np
from lucifex.mesh import interval_mesh
from lucifex.fdm import (
    CN, AM1, AB1, AB2, FiniteDifference, advective_timestep, diffusive_timestep,
    FunctionSeries, ConstantSeries, finite_difference_order,
)
from lucifex.fem import Function, Constant
from lucifex.solver import ibvp, evaluation, extrema, BoundaryConditions
from lucifex.sim import run, Simulation
from lucifex.viz import plot_line, plot_twin_lines, plot_stacked_lines, save_figure
from lucifex.utils import nested_dict
from lucifex.pde.advection_diffusion import advection_diffusion


def create_simulation(
    Lx: float,
    Nx: int,
    dt: float,
    a_max: float,
    d: float,
    D_adv: FiniteDifference,
    D_diff: FiniteDifference,
) -> Simulation:
    order = finite_difference_order(D_adv, D_diff)
    mesh = interval_mesh(Lx, Nx)
    t = ConstantSeries(mesh, name='t', ics=0.0)
    dt = Constant(mesh, dt, name='dt')
    a = Function((mesh, 'P', 1, 1), name='a')
    ax = lambda x: 4 * a_max * x[0] * (Lx - x[0]) / Lx**2
    a.interpolate(lambda x: (ax(x), ))
    d = Constant(mesh, 0.1, name='d')
    u = FunctionSeries((mesh, 'P', 1), name='u', order=order, store=1)
    u_bcs = BoundaryConditions(
        ('neumann', lambda x: x[0], 0.0),
        ('dirichlet', lambda x: x[0] - Lx, 0.0),
    )
    u_ics = lambda x: np.exp(-(x[0] - 0.2 * Lx)**2 / (0.001 * Lx))
    u_solver = ibvp(advection_diffusion, u_ics, u_bcs)(u, dt, a, d, D_adv, D_diff)
    uMinMax = ConstantSeries(mesh, name='uMinMax', shape=(2,), store=1)
    uMinMax_solver = evaluation(uMinMax, extrema)(u[0])
    solvers = [u_solver, uMinMax_solver]
    return Simulation(solvers, t, dt)


Lx = 2.0
Nx = 200
dt = 0.01
a_max = 1.0
d = 0.1

h = Lx/ Nx
dt_diffusive = diffusive_timestep(d, h)
dt_advective = advective_timestep(a_max, h)
courant_diffusive = dt / dt_diffusive
courant_advective = dt / dt_advective

D_adv = AB2
D_diff_opts = (AB1, AM1, CN)
simulations = nested_dict((FiniteDifference, Simulation))

for D_diff in D_diff_opts:
    simulations[D_diff] = create_simulation(Lx, Nx, dt, a_max, d, D_adv, D_diff)
    
n_init = 5
n_stop = 150
for D_diff in D_diff_opts:
    run(simulations[D_diff], n_stop, n_init=n_init)

slc = slice(0, None, 2)
texstr = lambda D: str(D).replace('◦', '\circ ').replace('₁', '_1').replace('₂', '_2')
minmax_labels = ('$\min(u)$', '$\max(u)$')
D_adv_tex = texstr(D_adv)

for D_diff in D_diff_opts:
    u = simulations[D_diff]['u']
    s = texstr(D_diff)
    title_texts = [
        f"$\mathcal{{D}}_{{\mathbf{{a}}, u}}=\mathrm{{{D_adv_tex}}}$",
        f"$\mathcal{{D}}_{{\mathsf{{D}}, u}}=\mathrm{{{texstr(D_diff)}}}$",
        f"$C_{{{{\mathbf{{a}}}}}}={courant_advective}$",
        f"$C_{{\mathsf{{D}}}}={courant_diffusive}$",
    ]
    title = "\n".join(title_texts)
    legend_labels=(min(u.time_series[slc]), max(u.time_series[slc]))
    fig, ax = plot_line(u.series[slc], legend_labels, '$t$', cyc='jet', x_label='$x$', y_label='$u$', title=title)
    save_figure(f'u(x,t)_{D_diff.name}', thumbnail=(D_diff is CN))(fig)
    uMinMax = simulations[D_diff]['uMinMax']
    uMin, uMax = uMinMax.sub(0), uMinMax.sub(1)
    fig, ax = plot_twin_lines(
        uMinMax.time_series,
        (uMin.value_series, uMax.value_series),
        minmax_labels,
        x_label='$t$',
        title=title,
    )
    save_figure(f'uMinMax(t)_{D_diff.name}_twinned')(fig)
    plot_stacked_lines(
        [(uMin.time_series, uMin.value_series), (uMax.time_series, uMax.value_series)],
        '$t$',
        minmax_labels,
        title=title,
    )
    save_figure(f'uMinMax(t)_{D_diff.name}_stacked')(fig)