DG advection with a vortical velocity

\[\begin{split} \mathbb{S}_u \begin{cases} \Omega = [0, 1] \times [0, 1] \\ u_0(x,y) = U_0x^2(1-x)^4y^2(1-y)^4 \\ \textbf{a}(x,y) = \sin(\pi x)\cos(\pi y)\textbf{e}_x -\cos(\pi x)\sin(\pi y)\textbf{e}_y \\ \end{cases} \end{split}\]

import numpy as np
from ufl import SpatialCoordinate, as_vector, sin, cos
from lucifex.mesh import rectangle_mesh
from lucifex.fem import Constant
from lucifex.fdm import ( 
    FiniteDifference, FunctionSeries, ConstantSeries,
    BE, FE, advective_timestep)
from lucifex.solver import ibvp , BoundaryConditions
from lucifex.sim import run, Simulation
from lucifex.plt import (
    plot_colormap, save_figure, 
    plot_streamlines, create_multifigure,
)       
from lucifex.utils.py_utils import nested_dict
from lucifex.utils.fenicsx_utils import is_continuous_lagrange
from lucifex.utils.npy_utils import as_index
from lucifex.pde.advection import advection, dg_advection


def create_simulation(
    element: tuple[str, int],
    Lx: float,
    Ly: float,
    Nx: int,  
    Ny: float,  
    courant: float,
    D_adv: FiniteDifference,
    mu: float,
):
    mesh = rectangle_mesh(Lx, Ly, Nx, Ny)
    t = ConstantSeries(mesh, name='t', ics=0.0)
    x = SpatialCoordinate(mesh)
    a = as_vector(
        (sin(np.pi * x[0]) * cos(np.pi * x[1]), -cos(np.pi * x[0]) * sin(np.pi * x[1]))
    )
    dt = advective_timestep(a, 'hmin', courant, mesh=mesh)
    dt = Constant(mesh, dt, name='dt')
    u = FunctionSeries((mesh, *element), name='u', store=1)
    ics = lambda x: mu * (x[0] ** 2) * (x[1] ** 2) * ((1 - x[0]) ** 4) * ((1 - x[1]) ** 4)
    bcs = BoundaryConditions(
        ('dirichlet', lambda x: x[0], 0.0),
    )
    if is_continuous_lagrange(u.function_space):
        u_solver = ibvp(advection, ics, bcs)(u, dt, a, D_adv)
    else:
        u_solver = ibvp(dg_advection, ics)(u, dt, a, D_adv, bcs=bcs)
    return Simulation(u_solver, t, dt, auxiliary=[('a', a)])


Lx = 1.0
Ly = 1.0
Nx = 64
Ny = 64
h = Lx / Nx
mu = 1e3
courant = 0.8

elem_opts = [
    ('DP', 1),
    ('P', 1),
]
D_adv_opts = (BE, FE)

simulations = nested_dict((FiniteDifference, tuple, Simulation))

for elem in elem_opts:
    for D_adv in D_adv_opts:
        simulations[elem][D_adv] = create_simulation(elem, Lx, Ly, Nx, Ny, courant, D_adv, mu)

n_stop = 100
for elem in elem_opts:
    for D_adv in D_adv_opts:
        run(simulations[elem][D_adv], n_stop)

for elem in elem_opts:
    fam, deg = elem
    for D_adv in D_adv_opts:
        u = simulations[elem][D_adv]['u']
        a = simulations[elem][D_adv]['a']
        time_indices = as_index(u.time_series, (0, 0.5, -1),)
        suptitle = f'{fam}$_{deg}$\n $\mathcal{{D}}_{{\mathbf{{a}}, u}}=\mathrm{{{D_adv}}}$'
        mfig, axs_main, axs_cbar = create_multifigure(
            n_cols=len(time_indices), 
            cbars=True, 
            suptitle=suptitle,
        )
        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, mesh=u.function_space.mesh, density=0.75, color='cyan')
        save_figure(f'{u.name}(x,y,t)_{fam}{deg}_{D_adv}')(mfig)

        idx = time_indices[1]
        ui = u.series[idx]
        ti = u.time_series[idx]
        title = f'{suptitle}\n$u(t={ti:.3f})$'
        fig, ax = plot_colormap(ui, title=title)
        plot_streamlines(fig, ax, a, mesh=u.function_space.mesh, density=0.75, color='cyan')
        thumbnail = (elem == ('DP', 1) and D_adv is BE)
        save_figure(f'{u.name}(x,y,t={ti:.3f})_{fam}{deg}_{D_adv}', thumbnail=thumbnail)(fig, close=True)