Inverse problems — Scimba Jax

1. Unknown coefficient

\(-\Delta u + c(x)\, u = f\)

The problem is posed on a disk, where u is observed at a few points and c is unknown. Compared to the forward problem, the only change is to replace the fixed coefficient with a small network wrapped in an ApproximationSpace. The forward FEM solve, the mesh and the basis stay the same.

import jax
from scimba_jax.linear_approximation.solvers import LinearSolve
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.approximation_spaces.dg_approximation_spaces import (
    DGEllipticApproximationSpace,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.physical_models.classical_weakform.laplacian_inverse_linear_problem import (
    LaplacianReactionWeakFormLearnableCoeff,
)
# u observed at a few points, c unknown (mesh, model and variables from the FEM page).
# c(x) = softplus(MLP(x)) is an ApproximationSpace, and softplus keeps c positive.
def positive(*args):
    return jax.nn.softplus(args[0])
coefficient = ApproximationSpace(
    {"x": 2},
    [(MLP(in_size=2, out_size=1, hidden_sizes=[16, 16], key=key), "scalar", None)],
    post_processing=positive,
)
model.add_weak_form(
    "main", LaplacianReactionWeakFormLearnableCoeff(dim=2, c=coefficient, f=f_source)
)
# The forward solve is the same linear system as on the FEM page.
scheme = EllipticFEscheme(model, variables)
space = DGEllipticApproximationSpace(
    dims={"x": 2, "dofsl": 1},
    list_assemblers=[scheme],
    model_type="x_dofsl",
    newton_kwargs={"solver": LinearSolve(tol=1e-7)},
)
# 89 DOFs, 40 observations, 20 ENG epochs: loss 5.7e-07, relative L2 error
# on c of 8.0e-02, against 0.36 for the best constant c.

2. Data loss and natural gradient

Here, the only loss of the Projector compares the output of the solve with the observations, and there is no PDE residual. The forward solve is computed inside this loss, so the optimization differentiates through the linear solve. The map from c to u is smoothing, so plain gradient descent performs poorly. The natural gradient optimizer (ENG, see PINNs) rescales the gradient with the metric of the parametrization, and it is needed in practice for this class of problems to converge.

from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.abstract_physical_model import AbstractPhysicalModel
from scimba_jax.physical_models.data_residuals import CollocDataResidual
# Data loss: compares u_h with the measurements at the observation points.
class DataFittingModel(AbstractPhysicalModel):
    def __init__(self, domain, data):
        super().__init__(main_domain=domain)
        self.data_residuals["data"] = CollocDataResidual(
            size=1, model_type="x_dofsl", data=data
        )
model = DataFittingModel(domain, data=(points, observations))
sampler = TensorizedSampler(
    [DomainSampler(domain)], bc=False, data_samplers=model.data_residuals
)
# ENG is needed here: c -> u is smoothing, so with plain gradient descent,
# the loss decreases while c stays far from the exact coefficient.
projector = Projector(model, space, sampler, optimizer="ENG", matrix_regularization=1e-6)
key, projector = projector.project(key, space, n_epochs=20, n_dl=len(points))
# `matrix_regularization` is required: the Gram matrix is rank deficient (some
# directions are not seen by the data), so Cholesky gives nan without it.

The same approach applies to the other discretizations (DG, finite volumes or a collocation basis), and to several coefficients at once, as long as each coefficient is wrapped in the same way and can be identified from the available data.