← Back to Mesh-based solvers Module · scimba_jax

FEM solver

Continuous, node-based finite elements: a degree of freedom lives at a node and is shared by every cell that touches it, so the discrete solution is continuous across cell boundaries by construction.

1. A 2D Laplacian, node by node

\(-\Delta u = f\)

A Mesh (see Domains & meshes), a basis of local shape functions, and VariablesFE to glue them into one global vector of nodal values. EllipticFEscheme assembles the weak form and solves — matrix-free here, by conjugate gradient, since a pure diffusion operator is symmetric positive definite. 

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.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
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_fe import VariablesFE
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 value per NODE, shared between cells -- sparse-and-symmetric where DG is block-diagonal.
variables = VariablesFE(basis=basis, nb_variables=1)


def source(x):  # so that u(x, y) = sin(pi x) sin(pi y)
    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))
scheme = EllipticFEscheme(model, variables)
scheme = EllipticFEscheme.solve(scheme, matrix_free=True, tol=1e-10)
# Q2 on 24x24 cells: 1.8e-05 relative L2 error against the manufactured solution.

2. A nonlinear coefficient needs Newton

\(-\nabla\cdot(A(u)\nabla u) = f\)

As soon as the diffusion coefficient A depends on u itself, the problem is no longer linear — a bare CG solve has nothing to converge to. NLEllipticWeakForm differentiates A(u) automatically, and NewtonSolver takes it from there. See Nonlinear solvers for what else is on offer once a problem needs it. 

from scimba_jax.physical_models.classical_weakform.nl_diffusion_advection_reaction_weak_form import (  # noqa: E501
    NLEllipticWeakForm,
)
from scimba_jax.linear_approximation.solvers.newton import NewtonSolver

# A(u) = 1 + 10 u^2, not a constant -- see Nonlinear solvers for what changes.
pde = NLEllipticWeakForm(
    dim=2, A=lambda x, u: (1.0 + 10.0 * u**2) * jnp.eye(2), f=source
)
model = AbstractPhysicalWeakModel.from_weak_form(pde, dirichlet=lambda x: jnp.zeros(1))
scheme = EllipticFEscheme(model, variables)
scheme, report = EllipticFEscheme.solve(
    scheme, solver=NewtonSolver(max_iter=20, tol=1e-10), return_report=True
)
# 6 Newton iterations, residual 6.4e-11.

3. Time-dependent problems

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

TimeDiscreteFEscheme turns the elliptic scheme above into one stage of a Butcher tableau, and re-solves it at every stage of every time step — the basis, the mesh and the Dirichlet lift are untouched, only the right-hand side changes. Explicit Euler is the cheapest choice per step but conditionally stable; implicit Euler and the second-order Pareschi–Russo scheme trade a (still linear, still solved by the same machinery) implicit stage for a time step that does not shrink with the mesh. 

from scimba_jax.linear_approximation.galerkin.fem.time_discrete_fe_scheme import (
    TimeDiscreteFEscheme,
)
from scimba_jax.time_discrete.butcher_tableau import build_implicit_euler_tableau

# du/dt - u_xx = 0, u(., 0) = sin(pi x). TimeDiscreteFEscheme re-solves the
# SAME elliptic weak form at every stage of a Butcher tableau.
mesh_1d = Mesh(dim=1, n_cells=(16,), ref_quad=UnitSquareTensorized(dim=1, order=4),
               mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]))
basis_1d = AnalyticBasis(
    nb_basis=2, out_dim=1, mesh=mesh_1d,
    local_basis=lambda y, i, m: local_lagrange_basis(y, i, m, order=1, out_dim=1),
    basis_type="scalar",
)
spatial_pde = LaplacianWeakForm(dim=1, f=lambda x: jnp.zeros(()))
scheme = TimeDiscreteFEscheme(
    spatial_weak_form_factory=spatial_pde,
    variables=VariablesFE(basis=basis_1d, nb_variables=1),
    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)
# Implicit Euler: unconditionally stable, first order. 8 steps to t=0.1,
# 16 P1 cells: 5.5e-02 relative L2 error (dominated by the coarse mesh).

A fine mesh makes each Krylov iteration slower to converge, not just slower to run — see Multigrid for the fix that keeps the iteration count flat as the mesh refines.