The physics is written with ParamFunctions
Residuals and weak forms are written with a single object, the
ParamFunction: a function of the model and of the
physical variables, combined and differentiated like the
expression it stands for. A physical model is a set of such
expressions, one per label (the interior, each boundary, the
initial condition).
A ParamFunction wraps fn(model, *physical_args):
model is the nonlinear part (containing the trainable
parameters θ) and *physical_args are the
physical variables (space, time, parameters). In the linear setting
(FEM, DG, etc.), there is also a linear part: the local basis DOFs,
called "dofsl" in the code. Algebra and differentiation
never evaluate eagerly: every operator returns a new
ParamFunction, vmapped over collocation points at
evaluation time.
1. ParamFunction
A PDE unknown is represented as \(u(x;\theta)\) for a PINN, or
\(u(x;\theta,\alpha)\) in the DG/FEM setting, where α are the
local linear degrees of freedom of a basis expansion. Both live behind
the same interface: a ParamFunction(dims, fn, f_type)
wrapping fn(model, *args). The first argument is always the model, and the
physical variables (named by f_type, e.g.
"x_mu") follow.
The method .freeze stops the gradient. More specifically, .freeze("model")
stops it through θ
(the nonlinear part); .freeze("dofsl") stops it through α instead (linear
variables only).
This is the nonlinear/linear split
that the math suggests, which is expressed as two gradient targets on a single object.
from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import (
ParamScalarFunction,
)
# u(x; theta) = model(x), where theta are the (nonlinear)
# network parameters
u = ParamScalarFunction(
{"x": 2}, fn=lambda model, x: model(x), f_type="x"
)
# stop gradient w.r.t. theta (the model)
u_frozen = u.freeze()
# linear variables only: stop gradient w.r.t. the DOFs
u_frozen = u.freeze("dofsl")
2. Function algebra
Binary operators promote f_type to the union of both
operands' variables (e.g. f(x,μ) * g(t,x,μ) → h(t,x,μ))
and return a new ParamFunction. Overloaded operators include
+ - * / (elementwise), << (composition,
f << g meaning f∘g), and @
(matrix/vector product, on ParamFieldFunction /
ParamMatrixFunction).
# elementwise algebra builds a new ParamFunction, without
# evaluating it
residual = u - u_exact # __sub__
scaled = 2.0 * u # __rmul__
ratio = u / v # __truediv__
# composition: h(*args) = f(g(*args))
h = f << g # __lshift__
# matrix-vector product (ParamFieldFunction /
# ParamMatrixFunction only)
diffusion_flux = A @ grad_u # __matmul__
3. Automatic differentiation
ParamFunction implements several usual differential operators as methods:
gradient_x(), etc.
Every differential operator applies jax.grad /
jacrev / jacfwd / hessian to the
ParamFunction's __call__, with
argnums resolved from f_type. Hence,
.gradient_x() differentiates w.r.t. whichever argument is
named "x". The result is wrapped back into a
ParamFunction,
which can be differentiated or otherwise manipulated further.
# jax.grad(u, argnums=argnums["x"]) -> ParamFieldFunction
grad_u = u.gradient_x()
# trace(hessian(u, x)) -> ParamScalarFunction
lap_u = u.laplacian_x()
# trace(jacobian(q, x)) -> ParamScalarFunction
div_q = q.divergence_x()
# jax.jacrev(u, argnums=argnums["t"]) -> ParamScalarFunction
dt_u = u.d_t()
# full Jacobian -> ParamFieldFunction
J = u.jacobian_x()
4. Kinds of ParamFunctions
ParamFunction is the base class, containing most overloaded and differential operators.
ParamVecFunction extends this base class for fixed-size vector output;
ParamFieldFunction and ParamScalarFunction further extend
those (a field's size is the spatial dimension, a scalar's is
always 1). ParamMatrixFunction is a separate sibling, directly
under ParamFunction, for matrix-valued output (e.g. a
Jacobian).
u = ParamScalarFunction(
{"x": 2}, fn=lambda model, x: model(x), f_type="x"
)
n = ParamFieldFunction(
{"x": 2},
fn=lambda *a: normal(a[1]),
f_type="x",
main_var="x",
)
w = ParamVecFunction(
3, {"x": 2}, fn=lambda model, x: model(x), f_type="x"
)
# Jacobian of the model w.r.t. theta, built in evaluators.py
J = ParamMatrixFunction(func.dims, jacrev_theta, func.f_type)
ParamScalarFunction
gradient_x() / laplacian_x()
\(\nabla u\) and \(\Delta u\), which are a field and a scalar, respectively.
anisotropic_laplacian_x(A)
\(-\nabla\cdot(A\nabla u)\) for a matrix-valued A.
advection_x(v) / bracket_x(g)
\(v\cdot\nabla u\), and the Poisson bracket \((\nabla u\times\nabla g)\cdot e\).
hessian_x() / det_hessian_x()
Full Hessian, and its (log-)determinant (Monge-Ampère).
ParamFieldFunction
dot(other)
Dot product with another field → a ParamScalarFunction.
divergence_x() / curl_x()
\(\nabla\cdot F\), and \(\nabla\times F\) (scalar in 2D, vector in 3D).
jacobian_x()
\(\partial F/\partial x\) → a ParamMatrixFunction.
laplacian_x()
Vector Laplacian, component-wise trace of the Hessian.
ParamVecFunction
component(i) / components()
Extract one (or all) scalar component(s).
cat(funcs)
Classmethod: concatenate several Param*Functions into one.
@ (__matmul__)
Matrix/vector product with a field or matrix.
ParamMatrixFunction
det() / abs_det()
(Absolute) determinant → a ParamScalarFunction.
not vector-like
No .component()/.cat(), since it is a
sibling of ParamVecFunction.
5. Strong-form residual
As usual, a Residual outputs an unevaluated function. Its
construct_residual(*rho) takes the model's
ParamFunction output(s) and returns another
ParamFunction. This is the pointwise residual, evaluated later at
every sampled point.
class AbstractResidual(ScimbaPytree):
@abstractmethod
def construct_residual(
self, *rho: PARAM_FUNC_TYPE, precomputed: dict | None = None
) -> PARAM_FUNC_TYPE:
"""Build the residual from the model's outputs.
Evaluated later, vmapped over sampled collocation points.
"""
class DirichletResidual(BoundaryResidual):
def construct_residual(self, *vars):
rho = vars[0]
return rho # compared with f_rhs = g(x) in the loss
class NeumannResidual(BoundaryResidual):
def construct_residual(self, *vars):
rho = vars[0]
grad_x_rho = rho.gradient_x()
n_model = ParamFieldFunction(rho.dims, lambda *args: args[rho.argnums["n"]], rho.f_type)
return grad_x_rho.dot(n_model) # d(u)/dn
6. Linking residuals and labels
A model's physical_residuals dict maps exactly the labels
the sampler produces ("interior", "bc N",
"bc D"..., see Samplers, §4)
to a Residual. Projector iterates over this dict,
looks up each label in the sample dict, and squares/weights every
residual into its own loss term.
class Projection(AbstractPhysicalModel):
def __init__(self, main_domain, size, model_type, f_rhs):
super().__init__(main_domain=main_domain)
self.physical_residuals = {
"interior": IdResidual(domain=main_domain, size=1, model_type="x_mu", f_rhs=f_rhs),
}
# boundary residuals use the sampler labels ("bc N", "bc D", ...,
# see Samplers, §4); the Projector looks up each residual by label
7. AbstractLinearWeakForm: bilinear_form + linear_form
For the classical DG/FEM solver, the model is written in its weak form. It comprises a
bilinear form \(a(u,v)\) and a linear form \(l(v)\), functions
of both a trial and a test ParamFunction,
and assembled into a linear system. This is a different case than a pointwise residual squared into
a PINN loss.
class AbstractLinearWeakForm(ScimbaPytree):
def __init__(self, dim: int): ...
@abstractmethod
def bilinear_form(
self, u: list[PARAM_FUNC_TYPE], v: list[PARAM_FUNC_TYPE]
) -> PARAM_FUNC_TYPE:
"""Compute the bilinear form a(u, v)."""
@abstractmethod
def linear_form(self, v: list[PARAM_FUNC_TYPE]) -> PARAM_FUNC_TYPE:
"""Compute the linear form l(v)."""
8. Examples of strong-form models
We present below several InteriorResidual subclasses from
physical_models/.
The same pattern is used every time: first get the
ParamFunction(s) from *vars, then combine them with
algebra and differentiation, and finally return the residual.
LaplacianResidual
\(-\Delta u = f\)
class LaplacianResidual(InteriorResidual):
def __init__(self, domain, f_rhs=None, model_type="x"):
super().__init__(
domain=domain,
size=1,
model_type=model_type,
f_rhs=f_rhs,
)
def construct_residual(self, *vars):
rho = vars[0]
return -rho.laplacian_x()
GeneralEllipticResidual
\(-\nabla\cdot(A\nabla u) + \mathbf{b}\cdot\nabla u + cu = f\)
class GeneralEllipticResidual(InteriorResidual):
def __init__(
self,
domain,
model_type="x",
f_rhs=None,
A=None,
b=None,
c=None,
):
super().__init__(
domain=domain,
size=1,
model_type=model_type,
f_rhs=f_rhs,
)
self.A, self.b, self.c = A, b, c
def construct_residual(self, *vars):
rho = vars[0]
result = rho.anisotropic_laplacian_x(self.A)
if self.b is not None:
result = result + rho.advection_x(self.b)
if self.c is not None:
result = result + rho * self.c
return result
This is also the time-independent anisotropic diffusion
residual, where A is any matrix field. For pure
diffusion, leave b and c unset.
NavierStokesResidual
\((\mathbf{u}\cdot\nabla)\mathbf{u} - \nu\Delta\mathbf{u} + \nabla p = \mathbf{f},\quad \nabla\cdot\mathbf{u} = 0\)
class NavierStokesResidual(InteriorResidual):
def __init__(
self,
domain,
dim=2,
nu=1.0,
mode="vec",
f_rhs=None,
model_type="x",
):
super().__init__(
domain=domain,
size=dim + 1,
model_type=model_type,
f_rhs=f_rhs,
)
self.nu, self.dim, self.mode = nu, dim, mode
def construct_residual(self, *vars):
u, p = _extract_u_p(vars, self.mode, self.dim)
lap_u = u.laplacian_x()
momentum = [
u.component(i).advection_x(u)
- self.nu * lap_u.component(i)
+ p.partial_derivative_x(i)
for i in range(self.dim)
]
return ParamVecFunction.cat(
momentum + [u.divergence_x()]
)
WaveResidual
\(\partial_{tt} u - \Delta u = f\)
class WaveResidual(InteriorResidual):
def __init__(
self, domain, time_domain, model_type="t_x", f_rhs=None
):
super().__init__(
domain=domain,
size=1,
model_type=model_type,
f_rhs=f_rhs,
time_domain=time_domain,
)
def construct_residual(self, *vars):
rho = vars[0]
return rho.d_tt() - rho.laplacian_x()
SteadyEulerResidual1D
\(\partial_t W + \partial_x F(W) = 0,\quad W = (\rho, q, E)\)
class SteadyEulerResidual1D(InteriorResidual):
def __init__(
self,
domain,
time_domain,
gamma=1.4,
f_rhs=None,
model_type="x",
):
super().__init__(
domain=domain,
size=domain.dim + 2, # (rho, q, E)
model_type=model_type,
f_rhs=f_rhs,
time_domain=time_domain,
)
self.gamma = gamma
def construct_residual(self, *vars):
rho, q, e = vars[0].components()
u = q / rho
p = (self.gamma - 1) * (e - 0.5 * rho * u**2)
flux_rho = q.gradient_x()
flux_q = (rho * u**2 + p).gradient_x()
flux_e = (u * (e + p)).gradient_x()
return ParamVecFunction.cat((flux_rho, flux_q, flux_e))
9. Examples of weak-form models
This is the same idea as in the previous section, but in weak form. They are used by the classical DG/FEM solver.
LaplacianWeakForm
\(a(u,v) = \int \nabla u\cdot\nabla v\,dx,\quad l(v) = \int f v\,dx\)
class LaplacianWeakForm(EllipticWeakForm):
def __init__(self, dim, f=None):
super().__init__(
dim=dim,
A=lambda x: jnp.eye(dim),
b=lambda x: jnp.zeros(dim),
c=lambda x: jnp.zeros(()),
f=f,
)
def bilinear_form(self, u, v):
return u.gradient_x().dot(v.gradient_x())
# linear_form(v) = f * v is inherited from EllipticWeakForm
EllipticWeakForm
\(a(u,v) = \int (A\nabla u)\cdot\nabla v + (\mathbf{b}\cdot\nabla u)v + cuv\,dx,\quad l(v) = \int f v\,dx\)
class EllipticWeakForm(AbstractLinearWeakForm):
def __init__(self, dim, A=None, b=None, c=None, f=None):
super().__init__(dim=dim)
self.A, self.b, self.c, self.f = A, b, c, f
def bilinear_form(self, u, v):
fields = self.get_fields()
grad_u, grad_v = u.gradient_x(), v.gradient_x()
diffusion = grad_v.dot(fields["A"] @ grad_u)
advection = grad_u.dot(fields["b"]) * v
reaction = fields["c"] * u * v
return diffusion + advection + reaction
def linear_form(self, v):
return self.get_fields()["f"] * v
# A(x) can be any dim x dim matrix, e.g. anisotropic and time-
# independent:
# A = lambda x: jnp.eye(2) * (2.0 + jnp.sin(x[0]))
More PDE models (Grad-Shafranov, Monge-Ampère, parametric
Dirichlet/Neumann...) are in physical_models/, on
GitLab ↗.