Splitting method for the advection-reaction equation

\[\begin{split} \mathbb{S}_u \begin{cases} \Omega = [0, L_x] \\ u_0(x)=\exp\left(-\frac{(x - x_0)^2}{\sigma^2}\right) \\ u_{\text{I}}(x=0) = 0 \\ \textbf{a}(t)=a\,\textbf{e}_x & \text{constant velocity} \\ \Sigma(u) = k u(1-u) & \text{nonlinear reaction}\\ \end{cases} \end{split}\]

import numpy as np
from dolfinx.fem import FunctionSpace
from lucifex.mesh import interval_mesh
from lucifex.fem import Constant
from lucifex.fdm import (
    FiniteDifference, FunctionSeries, ConstantSeries, 
    FE, CN, finite_difference_order,
)
from lucifex.solver import ibvp, BoundaryConditions, interpolation, evaluation
from lucifex.sim import Simulation, run
from lucifex.plt import plot_line, save_figure
from lucifex.pde.advection import advection_reaction, advection
from lucifex.pde.evolution import evolution_rhs

def create_simulation(
    Lx: float,
    Nx: int,
    dt: float,
    a: float,
    k: float,
    D_adv: FiniteDifference,
    D_src: FiniteDifference,
    x0: float,
    sigma: tuple[float, float],
    splitting: str | None
) -> Simulation:
    order = finite_difference_order(D_adv, D_src)
    mesh = interval_mesh(Lx, Nx)
    t = ConstantSeries(mesh, name='t', ics=0.0)
    dt = Constant(mesh, dt, name='dt')

    a = Constant(mesh, (a, ), 'a')
    k = Constant(mesh, k, 'k')
    source = lambda u: k * u * (1 - u)

    sigma_minus, sigma_plus = sigma
    gaussian = lambda x, s: np.exp(-(x - x0)**2 / s**2)
    u_ics = lambda x: (
        (gaussian(x[0], sigma_minus)) * (x[0] <= x0)
        + (gaussian(x[0], sigma_plus)) * (x[0] > x0)
    )
    u_bcs = BoundaryConditions(('dirichlet', lambda x: x[0], 0.0))

    fs = FunctionSpace(mesh, ('P', 1))
    u = FunctionSeries(fs, 'u', order=order, store=1, ics=u_ics)

    if splitting is None:
        u_solver = ibvp(advection_reaction, bcs=u_bcs)(
            u, dt, a, j=source(u), D_adv=D_adv, D_src=D_src,
        )
    elif splitting == 'lie':
        uTilde1 = FunctionSeries(fs, 'uTilde1', order=order)
        uTilde2 = FunctionSeries(fs, 'uTilde2', order=order)
        evaluator = lambda u: u
        # FIXME for higher order finite differences also setting from u[-1], u[-2], ...
        uTilde1_setter = evaluation(uTilde1, evaluator, future=False, overwrite=True)(
            u[0]
        )
        uTilde1_solver = ibvp(advection, bcs=u_bcs, future=True, overwrite=True)(
            uTilde1, dt, a, D_adv
        )
        uTilde2_setter = evaluation(uTilde2, evaluator, future=False, overwrite=True)(
            uTilde1[1]
        )
        uTilde2_solver = interpolation(uTilde2, evolution_rhs, future=True)(
            uTilde2, dt, r=source(uTilde2), D_rhs=D_src,
        )
        u_setter = evaluation(u, evaluator, future=True, overwrite=True)(
            uTilde2[1]
        )
        u_solver = (
            uTilde1_setter, uTilde1_solver, 
            uTilde2_setter, uTilde2_solver, 
            u_setter,
        ) 
    elif splitting == 'strang':
        dt_half = Constant(mesh, 0.5 * float(dt), name='dt')
        uTilde1 = FunctionSeries(fs, 'uTilde1', order=order)
        uTilde2 = FunctionSeries(fs, 'uTilde2', order=order) 
        uTilde3 = FunctionSeries(fs, 'uTilde3', order=order) 
        evaluator = lambda u: u
        uTilde1_setter = evaluation(uTilde1, evaluator, future=False, overwrite=True)(
            u[0]
        )
        uTilde1_solver = ibvp(advection, bcs=u_bcs, future=True, overwrite=True)(
            uTilde1, dt_half, a, D_adv
        )
        uTilde2_setter = evaluation(uTilde2, evaluator, future=False, overwrite=True)(
            uTilde1[1]
        )
        uTilde2_solver = interpolation(uTilde2, evolution_rhs, future=True)(
            uTilde2, dt, r=source(uTilde2), D_rhs=D_src,
        )
        uTilde3_setter = evaluation(uTilde3, evaluator, future=False, overwrite=True)(
            uTilde2[1]
        )
        uTilde3_solver = ibvp(advection, bcs=u_bcs, future=True, overwrite=True)(
            uTilde3, dt_half, a, D_adv
        )
        u_setter = evaluation(u, evaluator, future=True, overwrite=True)(
            uTilde3[1]
        )
        u_solver = (
            uTilde1_setter, uTilde1_solver, 
            uTilde2_setter, uTilde2_solver, 
            uTilde3_setter, uTilde3_solver, 
            u_setter,
        ) 
    else:
        raise ValueError(splitting)

    return Simulation(u_solver, t, dt)


splitting_opts = ('lie', 'strang', None)
simulations: dict[None | str, Simulation] = {}

Lx = 1.0
Nx = 100
dt = 1e-3
a = 1.0
k = 10.0
D_adv = CN
D_src = FE

x0 = 0.3
sigma_minus = 0.1
sigma_plus = 5 * sigma_minus
sigma = (sigma_minus, sigma_plus)

for splitting in splitting_opts:
    simulations[splitting] = create_simulation(
        Lx, Nx, dt, a, k, D_adv, D_src, x0, sigma, splitting,
    )

t_stop = 1.0
n_stop = 100
for splitting in splitting_opts:
    run(simulations[splitting], n_stop, t_stop) 

slc = slice(0, None, 10)

for splitting, sim in simulations.items():
    u = sim['u']
    legend_title = f"{str(splitting).capitalize()}\n\n$t$"
    legend_labels = [f"{time:.3f}" for time in u.time_series[slc]]
    fig, ax = plot_line(
        u.series[slc], 
        legend_labels,
        legend_title,
        x_label='$x$',
        y_label='$u$',
        cyc='jet',
    )
    save_figure(f'u(x,t)_{splitting}', thumbnail=(splitting is splitting_opts[0]))(fig)