DG advection-diffusion in a rectangle

import numpy as np
from ufl import SpatialCoordinate, as_vector
from lucifex.mesh import rectangle_mesh
from lucifex.fem import Constant, Function
from lucifex.fdm import (
    BE, CN, advective_timestep, 
    FiniteDifference, FunctionSeries, ConstantSeries, finite_difference_order,
)
from lucifex.solver import ibvp , BoundaryConditions, evaluation, extrema
from lucifex.sim import Simulation, run
from lucifex.plt import (
    plot_streamlines, plot_line, save_figure, 
    plot_colormap, plot_colormap_multifigure,
)
from lucifex.utils.fenicsx_utils import is_continuous_lagrange
from lucifex.pde.advection_diffusion import dg_advection_diffusion, advection_diffusion

def create_simulation(
    element: tuple[str, int],
    Lx: float, 
    Ly: float,
    Nx: int,
    Ny: int,
    dt: float,
    d: float,
    D_adv: FiniteDifference,
    D_diff: FiniteDifference,
    alpha: float | tuple[float, float] = 10.0,
    dg_kws: tuple = (None, None),
) -> Simulation:
    order = finite_difference_order(D_adv, D_diff)
    mesh = rectangle_mesh(Lx, Ly, Nx, Ny)
    t = ConstantSeries(mesh, name='t', ics=0.0)
    dt = Constant(mesh, dt, name='dt')
    x = SpatialCoordinate(mesh)
    a = as_vector((0.5 * Ly -x[1], x[0] - 0.5 * Lx))
    d = Constant(mesh, d, name='d')
    u = FunctionSeries(
        (mesh, *element), 
        name='u', 
        order=order, 
        store=1, 
        ics=0.0,
    )
    uMinMax = ConstantSeries(
        mesh, 
        name='uMinMax', 
        order=order,
        shape=(2,),
        store=1,
    )
    u_lower = lambda x: (
        np.exp(-((x[0] - 0.25 * Lx) / (0.05 * Lx))**2)
        + np.exp(-((x[0] - 0.75 * Lx) / (0.05 * Lx))**2)
    )
    bcs = BoundaryConditions(
        ('dirichlet', lambda x: x[0], 0.0),
        ('dirichlet', lambda x: x[0] - Lx, 0.0),
        ('dirichlet', lambda x: x[1], u_lower),
        ('neumann', lambda x: x[1] - Ly, 0.0),
    )
    if is_continuous_lagrange(u.function_space):
        u_solver = ibvp(advection_diffusion, bcs=bcs)(
            u, dt, a, d, D_adv, D_diff,
        )
    else:
        u_solver = ibvp(dg_advection_diffusion)(
            u, dt, alpha, a, d, D_adv, D_diff, 
            bcs=bcs, dg_kws=dg_kws,
        )
    uMinMax_solver = evaluation(uMinMax, extrema)(u[0])
    solvers = (u_solver, uMinMax_solver)
    auxiliary = (('a', a), d)
    return Simulation(solvers, t, dt, auxiliary)


Lx = 1.0
Ly = 1.0
Nx = 40
Ny = 40
hx = Lx / Nx
hy = Ly / Ny
hMin = min(hx, hx)

d = 1e-3
aMax = np.sqrt(Lx**2 + Ly**2)
courant = 1.0
dt = advective_timestep(aMax, hMin, courant)

D_adv = BE
D_diff = BE

elem_opts = (
    ('P', 1),
    ('DP', 1),
    ('DP', 2),
)
simulations: dict[None | float, Simulation] = {}

for elem in elem_opts:
    simulations[elem] = create_simulation(
        elem, Lx, Ly, Nx, Ny, dt, d, D_adv, D_diff,
    )

n_stop = 50
for elem in elem_opts:
    run(simulations[elem], n_stop)

time_indices = (1, 10, 20, 30, 40, -1)

for elem in elem_opts:
    sim = simulations[elem]
    u, uMinMax, a = sim['u', 'uMinMax', 'a']
    fam, deg = elem
    elem_title = f'{fam}$_{deg}$'
    mfig, axs, _ = plot_colormap_multifigure(
        n_rows=2, 
        cbars=True, 
        suptitle=elem_title,
    )(
        [u.series[i] for i in time_indices],
        title=[f'$u(t={u.time_series[i]:.3f})$' for i in time_indices],
    )
    plot_streamlines(mfig, axs[0], a, mesh=u.function_space.mesh, density=0.75, color='cyan')
    save_figure(f'u(x,y,t)_{fam}{deg}')(mfig)
    fig, ax = plot_line(
        (uMinMax.time_series, uMinMax.value_series), 
        cyc='black',
        x_label='$t$',
        legend_labels=['$\min_{\\mathbf{x}}u$', '$\max_{\\mathbf{x}}u$'],
        title=elem_title,
    )
    save_figure(f'uMinMax(t)_{fam}{deg}', thumbnail=(elem == elem_opts[-1]))(fig)