DG solver — Scimba Jax

1. The same 2D Laplacian, one block per cell

\(-\Delta u = f\)

Same mesh, same basis family, a different Variables class: VariablesDG keeps one full set of degrees of freedom per cell instead of sharing nodes. What replaces the sharing is a flux — here SIPGFlux, symmetric interior penalty — added on every interior and boundary face.

import jax.numpy as jnp
from scimba_jax.linear_approximation.basis.analytic_bases import local_lagrange_basis
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.galerkin.dg.elliptic_dg_scheme import (
    EllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import SIPGFlux
from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.variables.variables_dg import VariablesDG
from scimba_jax.mapping.mapping import InvertibleFunction, Mapping
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.laplacian_weak_form import (
    LaplacianWeakForm,
)

order = 2
mesh = Mesh(
    dim=2,
    n_cells=(24, 24),
    ref_quad=UnitSquareTensorized(dim=2, order=2 * order + 1),
    mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]),
)
basis = AnalyticBasis(
    nb_basis=(order + 1) ** 2,
    out_dim=1,
    mesh=mesh,
    local_basis=lambda y, i, m: local_lagrange_basis(y, i, m, order=order, out_dim=1),
    basis_type="scalar",
)
# One block of DOFs per CELL, no sharing -- a jump across a face is the
# unknown the flux term controls, not an error.
variables = VariablesDG(basis=basis, nb_variables=1)


def source(x):  # same manufactured solution as the FEM page
    return 2 * jnp.pi**2 * jnp.sin(jnp.pi * x[0]) * jnp.sin(jnp.pi * x[1])


pde = LaplacianWeakForm(dim=2, f=source)
model = AbstractPhysicalWeakModel.from_weak_form(pde, dirichlet=lambda x: jnp.zeros(1))
# Continuity is PENALISED, not assumed -- strength must scale like 1/h.
flux = SIPGFlux(sigma=10.0 * order * (order + 1))
scheme = EllipticDGscheme(model, variables, flux, use_scan_quad=False)
scheme = EllipticDGscheme.solve(scheme, matrix_free=True, tol=1e-10)
# Same problem, order and mesh as the FEM page: 1.7e-05 relative L2 error.

2. What a flux actually is

\(\widehat{F} = \{\nabla u\}\cdot n - \dfrac{\sigma}{h}\,[\![u]\!]\)

Every DG scheme reads a flux the same way: given the two sides of a face (their trace, gradient and outward normal), return what each side owes the other. SIPGFlux is the default for a diffusion term; an upwind flux is the same contract applied to advection, and a Riemann solver for a hyperbolic system.

class AbstractFlux(ScimbaPytree):
    """What a DG scheme asks of a flux."""

    def __call__(self, varL, varR):
        """Two sides' contributions on one INTERIOR face."""
        raise NotImplementedError

    def boundary_call(self, var_interior, n_interior, dirichlet_val):
        """The interior side's contribution on one BOUNDARY face."""
        raise NotImplementedError

# SIPGFlux is consistent, symmetric (keeps the operator SPD), and penalised
# by a LOCAL sigma/h -- a global h is wrong as soon as the mesh isn't uniform.
# An upwind flux is the same contract, for advection instead of diffusion.

3. Time-dependent problems

\(\partial_t u - \Delta u = 0\)

TimeDiscreteDGscheme re-solves the same elliptic DG problem at every stage of a Butcher tableau, exactly like TimeDiscreteFEscheme on the FEM page. The one thing worth knowing before choosing a tableau: DG's interior-penalty flux adds stiffness beyond the bare operator, so an explicit tableau needs a noticeably smaller time step than FEM would at the same mesh — an implicit tableau sidesteps that restriction entirely.

from scimba_jax.linear_approximation.basis.analytic_bases import local_taylor_basis
from scimba_jax.linear_approximation.galerkin.dg.time_discrete_dg_scheme import (
    TimeDiscreteDGscheme,
)
from scimba_jax.time_discrete.butcher_tableau import build_implicit_euler_tableau

# Same du/dt - u_xx = 0 as the FEM time page. DG's explicit-Euler threshold
# is ~5x stricter than FEM's here (the penalty adds stiffness), so implicit
# is the natural default for DG diffusion.
n_cells, order = 16, 1
mesh_1d = Mesh(dim=1, n_cells=(n_cells,), ref_quad=UnitSquareTensorized(dim=1, order=4),
               mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]))
basis_1d = AnalyticBasis(
    nb_basis=order + 1, out_dim=1, mesh=mesh_1d,
    local_basis=lambda y, i, m: local_taylor_basis(y, i, m, order=order, out_dim=1),
    basis_type="scalar",
)
scheme = TimeDiscreteDGscheme(
    spatial_weak_form_factory=LaplacianWeakForm(dim=1, f=lambda x: jnp.zeros(())),
    variables=VariablesDG(basis=basis_1d, nb_variables=1),
    flux=SIPGFlux(sigma=(order + 1) * (order + 2), h=1.0 / n_cells),
    butcher_tableau=build_implicit_euler_tableau(),
    dt=0.1 / 8,
    dirichlet=lambda x: jnp.zeros(1),
)
dofsl_init = scheme.initialize(lambda x: jnp.sin(jnp.pi * x))
dofsl_final, history = scheme.solve(dofsl_init, t0=0.0, nt=8)
# Same 8 steps, order-1 Taylor basis: 5.5e-02 relative L2 error -- matches
# FEM's at this resolution, as expected for the same order and mesh.

DG's block structure changes what an efficient solver looks like too — see Multigrid for the block-Jacobi smoother that reads one cell's dense block per sweep instead of a single diagonal entry, and Nonlinear solvers for a nonlinear flux.