1. Training on data
Given pairs (f_i, F_i) of a source term and its known
solution, a PhysicNOProjector is optimized in the same
way as a PINN, with a data residual
instead of a PDE residual. Each training example is a physical
model with its own f, so a batch is a batch
of models.
from scimba_jax.neural_operator.physic_no.abstract_deeponet import AbstractDeepONet
from scimba_jax.nonlinear_approximation.approximation_spaces.physic_no_approximation_spaces import (
PhysicNOApproximationSpace,
)
from scimba_jax.nonlinear_approximation.numerical_solvers.physic_no_projectors import (
PhysicNOProjector,
)
from scimba_jax.physical_models.ode.anti_derivative import AntiDerivative1DWithData
# A DeepONet learns the map f -> F, and each training example is a physical model.
# pre_encoder gives the f of a model to the branch network.
class DeepONetAntiDerivative(AbstractDeepONet):
def pre_encoder_size(self):
return self.grid_size
def pre_encoder(self, physical_model):
return physical_model.data_residuals["interior"].data_vals[0][..., 0]
no = DeepONetAntiDerivative(grid_size=20)
space = PhysicNOApproximationSpace(dims={"x": 1}, list_models=[no], model_type="x")
# Training on data: each pde has its own (f, F) pair, and only_data=True
# turns off the physics terms.
projector = PhysicNOProjector(pdes, space, sampler, only_data=True, weights=custom_weights)
key, projector = projector.project(key, space, n_epochs=20, batch_size=4, n_dl=20)
# 4 antiderivative instances, grid of 20, 20 epochs: loss 0.18 -> 0.048.
2. Training on physics (PhysicNO)
When the data residual is replaced by a PDE residual, the
architecture stays the same. Only the physical model changes, and
its pre_encoder reads f as a callable
instead of a precomputed array. No labelled solution is needed: the
network learns to satisfy the equation at sampled points, as a PINN
does, but for a whole family of right-hand sides at once.
from scimba_jax.physical_models.ode.anti_derivative import AntiDerivative1D
# Same architecture: pre_encoder now reads f as a function, and the physical
# model contains the ODE residual u' - f (no data on F).
pdes = [AntiDerivative1D(main_domain=domain_x, f_rhs=f_i) for f_i in sampled_fs]
projector = PhysicNOProjector(pdes, space, sampler) # only_data defaults to False
key, projector = projector.project(
key, space, n_epochs=20, batch_size=4, n_colloc=50, n_bc_colloc=1
)
# Same 4 instances, 20 epochs, with collocation: loss 0.70 -> 0.50. The loss
# decreases more slowly than with data, but no solution F is needed.
Batching over simulations
A PINN batches the points of one problem, while a PhysicNO
batches problems. Every training instance is a physical
model of the same class, with its own f and its own
data. The N instances are stacked once
with create_batch: every array leaf of the model gains
a leading axis of length N, and the domain, the
residual definitions and the other shared data are stored a single
time. A mini-batch is then a slice of this stack by index.
Building a new stack at each step would trigger a recompilation.
At each step, the sampler draws one set of
collocation points, shared by the whole mini-batch
(in_axes=None), while each model has its own
f and its own observations (in_axes=0).
The residual is vmap-ed over the models, and the loss
is its mean over the simulations. The shapes
(batch_size, n_colloc) are fixed, so the step is
compiled once. The data common to all instances (for instance, an
observation grid) is stored in the model as a numpy
array, and is not duplicated batch_size times. Only
the values that differ between simulations have the batch axis.
from scimba_jax.physical_models.model_sampler import (
ModelSampler,
)
# N simulations = N physical models of one class, each with
# its own f.
pdes = [
AntiDerivative1D(main_domain=domain_x, f_rhs=f_i)
for f_i in sampled_fs
]
# Stacked once (create_batch): each array leaf gets a leading
# axis of length N,
# and the domain and residual definitions are stored once.
sampler = ModelSampler(pdes)
# A mini-batch is a slice of the stack (one pytree with 4
# models). A new stack
# would give a new treedef, and thus a compilation at each step.
key, mini_batch, idx = sampler.sample(key, batch_size=4)
# At each step, the collocation points are shared and the
# residual is vmapped over f:
# vmap(residual, in_axes=(None, None, 0))(space, points,
# mini_batch)
# PhysicNOProjector does all of this from the list of models.
projector = PhysicNOProjector(pdes, space, sampler_x)
key, projector = projector.project(
key, space, n_epochs=20, batch_size=4, n_colloc=50
)
3. Architectures
The architectures differ by the geometries they can handle and by the way they transfer information across the domain. They all use the training loop described above.
FNO
Fourier Neural Operator, on a Cartesian grid. It mixes channels in Fourier space, where a global convolution is a pointwise product. This is cheap, and exact for periodic problems.
Geo-FNO
An FNO that first maps an irregular domain onto a regular grid through a learned change of coordinates, so that the spectral layers can be used on domains other than a square or a box.
φ-FEM-FNO
The domain itself is an input of the network. A level set marks the grid cells that are inside the domain, so that a single trained network can be used on a whole family of geometries.
U-Net
The input is downsampled while the number of channels is doubled, and then upsampled. At every level, a skip connection joins the encoder and the decoder before each convolution, which preserves fine details.
GINO
Geometry-Informed Neural Operator. It maps an irregular point cloud onto a latent regular grid, applies an FNO on this grid, and maps the result back to the point cloud. Its cost is close to that of an FNO, while it handles the geometry of a mesh.
Pointwise operators
The same map is applied independently at every point, without any mixing across the domain. This is the simplest architecture, and the shared API is validated on it first.
GNN / GNO on a mesh
Message passing on a mesh graph. The choice of neighbourhood determines whether the result is a network on this graph (neighbours are given by the edges of the mesh) or an operator on the domain (neighbours are in a physical ball, summed with the quadrature weights of the mesh). See below.
DeepONet
A branch network reads the source function at a fixed set of sensor points, a trunk network reads the query point, and their combination gives the output of the operator. The encoder and the decoder do not depend on a grid or a point cloud, as they do for the architectures above.
A classical FEM or DG solve can also be seen as an encode/propagate/decode operator. It is presented with the mesh-based solvers, since nothing is learned in it unless a coefficient is (see inverse problems for that case).
4. Graph layers: neural network or neural operator
A single layer is implemented, and the graph layers from the literature correspond to choices of two functions in it. Whether the result depends on the mesh or on the domain is decided by a third choice: the neighbourhood.
Message, aggregation and update
\(m_i = \sum_{j \in N(i)} w_{ij}\, \varphi_\theta(h_i, h_j, x_j - x_i), \qquad h_i' = \gamma_\theta(h_i, m_i)\)
message_kind picks \(\varphi\): GraphSAGE's
\(W h_j\), MoNet's Gaussian bumps on the local coordinates
\((\rho, \text{direction})\), an MLP of the pair (MPNN), or a
kernel \(\kappa_\theta(x_j - x_i)\, h_j\) whose matrix an MLP
returns (GNO). update_kind picks \(\gamma\): an
activation, GIN's \((1+\varepsilon) h + m\), a GRU gate, an MLP of
\([h, m]\), each with a residual connection. MeshGraphNet is the
same layer, with an edge state that is also updated. These choices
do not determine the nature of the layer.
The nature of the layer is determined by the
neighbourhood. With mode="gnn", \(N(i)\)
contains the neighbours along the edges of the mesh, and
\(w_{ij} = 1/d_i\). When the mesh is refined, the edges are twice
shorter, so the same weights act on half the distance: what was
learned depends on the discretisation. With
mode="gno", \(N(i)\) is the ball \(B(x_i, r)\) and
\(w_{ij} = |K_j|\) is the measure of cell \(j\). The sum is then
a Riemann sum of \(\int_{B(x_i, r)} \varphi(\cdot, y)\, dy\),
computed with the quadrature of the mesh (the cell measures, or the
lumped nodal masses, which are the weights of a finite element
assembly). When the mesh is refined, the sum converges to the
integral, so the weights define a map between functions.
The ball only depends on the geometry, so it is computed
beforehand, once, in
MeshData(ball_radii=…), and then stored. The distance
is either Euclidean or geodesic (the sum of edge lengths, which
does not cross holes). The layer only reads the stored ball. This
keeps the layers jittable, allows the mesh to be passed at call
time, and gives the ball layers a continuous version that can be
evaluated at any \(x\).
from scimba_jax.neural_operator.data_for_no.mesh_data import (
MeshData,
)
from scimba_jax.neural_operator.layers.graph_based import (
GraphNetwork,
GraphOperator,
)
# The geometry (mesh edges, balls B(x_i, r)) is computed once
# and stored.
# Cell centres are the nodes, and cell measures the quadrature
# weights.
graph = MeshData(2, mesh, ball_radii=(0.13,))
# "gnn": average over the mesh edges.
# "gno": Riemann sum over the ball, weighted by the cell
# measures.
gnn = GraphNetwork(
2, 1, 1, 16, 4, "gnn", 0.13, "kernel", "mlp", key=key
)
gno = GraphNetwork(
2, 1, 1, 16, 4, "gno", 0.13, "kernel", "mlp", key=key
)
# same (u, mu) call as the other discrete operators
operator = GraphOperator(gno, graph)
projector = NOProjector(operator, (u0_train, uT_train))
_, projector = projector.project(
key, operator, 1000, batch_size=16
)
# Evaluation with the same weights on the nested mesh (4x
# finer),
# with its own balls and measures.
finer = MeshData(2, refined_mesh, ball_radii=(0.13,))
uT_fine = GraphOperator(projector.operator.network, finer)(
u0_fine
)
# gnn kernel: 1.47e-1 on the training mesh, 4.78e-1 on the
# finer one.
# gno kernel: 2.00e-1 -> 1.49e-1; r = 0.2:
# 8.3e-2 -> 4.7e-2.
Numerical test
The test case is a Gaussian advected and diffused on the JET cross-section. The exact operator is a shifted Gaussian convolution, so the errors only come from the networks. Every message is trained in both modes on a mesh of 483 cells, with a radius chosen so that the ball contains about as many neighbours as the edges (4.3 against 3.8). The two modes then have the same cost, and only differ by the definition of the neighbourhood. The same weights are then evaluated on the nested mesh, which is four times finer.
The GNN is accurate on the coarse mesh (MoNet 5.3e-2), but not on the fine one (7.3e-1). The GNO is less accurate on the training mesh (kernel 2.0e-1) and more accurate on the finer mesh (1.5e-1). Indeed, four cells per ball give a crude quadrature, and the finer mesh gives a better approximation of the integral. With nine cells per ball (r = 0.2), the error goes from 8.3e-2 to 4.7e-2, which is under 10 % on a mesh that was not used during training. This behaviour comes from the operator formulation: with a quadrature and a physical radius, the resolution is an input of the network and does not have to be learned.
The same choice is available in the graph U-Net. Pooling is done either by parenthood on a nested hierarchy (gnn), or with a ball of coarse cells around each fine centre (gno). In the second case, the radius doubles at each level, so that the coarse levels handle long-range interactions, as the FNO does in the middle of a GINO.
mode message L2 level 0 L2 level 1 (advected-diffused Gaussian, JET)
gnn kernel 1.47e-01 4.78e-01
gnn monet 5.26e-02 7.28e-01
gnn mp 1.26e-01 4.57e-01
gno kernel 2.00e-01 1.49e-01 r = 0.13, 4 cells per ball
gno mp 2.07e-01 1.36e-01
gno kernel 8.27e-02 4.65e-02 r = 0.20, 9 cells per ball
identity 8.81e-01Trained on level 0 (483 cells), read on level 1 (1932 cells) with the same weights. Relative L2 on 32 test samples.