1. 2D Laplacian problem
\(-\Delta u = f\)
The solver requires a Mesh (see
Domains & meshes), a
basis of local shape functions,
and VariablesFE, which combines them into a single
global vector of nodal values. EllipticFEscheme then
assembles the weak form and solves the linear system. Here, the
system is solved matrix-free with the conjugate gradient method,
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: the matrix is sparse and symmetric.
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. Nonlinear diffusion coefficient
\(-\nabla\cdot(A(u)\nabla u) = f\)
When the diffusion coefficient A depends on
u, the problem becomes nonlinear, and a single CG solve
is no longer sufficient. NLEllipticWeakForm computes
the derivative of A(u) automatically, and
NewtonSolver solves the nonlinear problem. Other
nonlinear solvers are presented on the
Nonlinear solvers page.
from scimba_jax.physical_models.classical_weakform.nl_diffusion_advection_reaction_weak_form import (
NLEllipticWeakForm,
)
from scimba_jax.linear_approximation.solvers.newton import NewtonSolver
# Nonlinear diffusion coefficient A(u) = 1 + 10 u^2 (see Nonlinear solvers).
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 uses the elliptic scheme above as
one stage of a Butcher tableau, and solves it at every stage of
every time step. The basis, the mesh and the Dirichlet lift stay
the same, and only the right-hand side changes. Explicit Euler has
the lowest cost per time step, but it is only conditionally stable.
Implicit Euler and the second-order Pareschi-Russo scheme require
solving an implicit stage, which is still linear and is solved with
the same tools. In exchange, their time step does not decrease when
the mesh is refined.
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 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).
On a finer mesh, each Krylov iteration is more expensive, and the number of iterations also increases. The Multigrid page presents a method that keeps the number of iterations constant as the mesh is refined.