Advection-diffusion of a Gaussian in a rectangle

import numpy as np
from ufl import SpatialCoordinate, as_vector
from lucifex.mesh import rectangle_mesh, mesh_boundary
from lucifex.fdm import (
    CN, AM1, AB1, AB3, FiniteDifference, FiniteDifferenceArgwise,
    FunctionSeries, ConstantSeries, finite_difference_order, as_grid_function_series,
)
from lucifex.fem import Constant
from lucifex.solver import ibvp, BoundaryConditions, evaluation
from lucifex.sim import Simulation, run
from lucifex.utils.npy_utils import as_index
from lucifex.plt import plot_colormap, plot_line, save_figure, create_multifigure, plot_streamlines
from lucifex.pde.advection_diffusion import advection_diffusion

def velocity_modulation(
    t: Constant | float,
    eps: float,
    omega: float,
) -> float:
    return 1 + eps * np.cos(omega * float(t))

def create_simulation(
    Lx: float,
    Ly: float,
    Nx: int,
    Ny: int,
    dt: float,
    d: float,
    eps: float,
    omega: float,
    D_adv: FiniteDifference | FiniteDifferenceArgwise,
    D_diff: FiniteDifference,
) -> Simulation:
    order = finite_difference_order(D_adv, D_diff)
    mesh = rectangle_mesh(Lx, Ly, Nx, Ny)
    boundary = mesh_boundary(
        mesh, 
        {
            "left": lambda x: x[0],
            "right": lambda x: x[0] - Lx,
            "lower": lambda x: x[1],
            "upper": lambda x: x[1] - Ly,
        },
    )
    x = SpatialCoordinate(mesh)
    t = ConstantSeries(mesh, name='t', ics=0.0)
    dt = Constant(mesh, dt, name='dt')
    d = Constant(mesh, d, name='d')
    at = ConstantSeries(
        mesh, 
        name='at', 
        order=order,
        store=1,
        ics=velocity_modulation(0.0, eps, omega),
    )
    axy = as_vector((0.5 * Ly -x[1], x[0] - 0.5 * Lx))
    a = at * axy
    at_solver = evaluation(at, velocity_modulation, future=True)(
        t[1], eps, omega,
    )
    u = FunctionSeries((mesh, 'P', 1), name='u', store=1)
    ics = lambda x: np.exp(-((x[0] - 0.75 * Lx)**2 + (x[1] - 0.5* Ly)**2)/ (0.01 * Lx))
    bcs = BoundaryConditions(
        ("dirichlet", boundary.union, 0.0),  
    )
    u_solver = ibvp(advection_diffusion, ics, bcs)(u, dt, axy, d, D_adv, D_diff)
    solvers = (at_solver, u_solver)
    auxiliary = (('axy', axy), ('a', a), at, d)
    return Simulation(solvers, t, dt, auxiliary)


Nx = 64
Ny = 64
Lx = 1.0
Ly = 1.0

d = 0.005
eps = 0.5
dt_period = 0.4
omega = 2 * np.pi / dt_period

dt = 0.05 * dt_period

D_adv_opts = (AB1, AB1 @ CN, AB3 @ CN, AM1)
D_diff = CN

simulations: dict[FiniteDifference | FiniteDifferenceArgwise, Simulation] = {}

for D_adv in D_adv_opts:
    simulations[D_adv] = create_simulation(Lx, Ly, Nx, Ny, dt, d, eps, omega, D_adv, D_diff)

t_stop = 20 * dt_period
for D_adv in D_adv_opts:
    run(simulations[D_adv], t_stop=t_stop)

as_tex_str = lambda fd: (
    repr(fd).replace('@', '\circ ').replace('1', '_1').replace('2', '_2').replace('3', '_3')
)
grid_index = (int(0.25 * Nx), int(0.5 * Ny))
uij_lines = []
uij_labels = []

for D_adv, sim in simulations.items():
    u, at, a = sim['u', 'at', 'a']
    u_grid = as_grid_function_series(u)
    uij_series = [v[grid_index] for v in u_grid.value_series]
    uij_lines.append((u_grid.time_series, uij_series))
    uij_labels.append(f'$\mathrm{{{as_tex_str(D_adv)}}}$')
    time_indices = as_index(u.time_series, (0, 1/3, 2/3, -1), fraction=True)
    mfig, axs_main, axs_cbar = create_multifigure(
        n_cols=len(time_indices), 
        cbars=True, 
        suptitle=f'$\mathcal{{D}}_{{\mathbf{{a}}, u}}=\mathrm{{{as_tex_str(D_adv)}}}$',
    )
    for i, axm, axc in zip(time_indices, axs_main, axs_cbar):
        title = f'$u(t={u.time_series[i]:.3f})$'
        plot_colormap(mfig, axm, u.series[i], title=title, cbar_ax=axc)
        plot_streamlines(mfig, axm, a.series[i], mesh=u.function_space.mesh, density=0.75, color='cyan')
    if D_adv is D_adv_opts[-1]:
        xi, yj = (u_grid.mesh.x_axis[i] for i in grid_index)
        plot_line(
            (at.time_series, at.value_series),
            x_label='$t$',
            y_label='$a_{\\omega}$',
        )

y_label=f'$u(x={xi:.2f}, y={yj:.2f})$'
fig, ax = plot_line(
    uij_lines, 
    y_label=y_label,
    x_label='$t$',
    legend_title=f"$\mathcal{{D}}_{{\mathbf{{a}}, u}}$",
    legend_labels=uij_labels,
)
save_figure(y_label.replace(')', ',t)').replace(' ', ''), thumbnail=True)(fig)