Inverse: Recover a Diffusivity Field (differentiable FEM)
An inverse problem: recover the entire spatially-varying diffusivity field \(k(x)\) in \(-\nabla\!\cdot(k\nabla u)=f\) — a hidden high-conductivity inclusion — from the measured response \(u\), by differentiating the FEM solve end to end. This is the FEM flavour of parameter-field identification / diffuse tomography.
A nodal field parameter + smoothness prior
jno.np.parameter(phi) is a trainable P1 field (one DOF per node) on the trial space;
fem.solve() is the differentiable forward; field inversion is ill-posed, so an H1-seminorm
prior k.regularize("h1seminorm") keeps it stable. crux fits the data plus the prior:
k = jno.np.parameter(phi, name="k")
fem = jno.fem([k * (ui.x * vi.x + ui.y * vi.y) - f * vi, u(xb, yb) - 0.0], quad_degree=3)
crux = jno.core([(fem.solve() - u_obs).mse, 1e-3 * k.regularize("h1seminorm").mean], domain=obs)
crux.solve(500)
rec = np.asarray(crux.eval([k])).reshape(-1) # the recovered field (do not index [0])
What to notice
jno.np.parameter(phi)makes the unknown a full nodal field, not a scalar.fem.solve()is differentiable, so the gradient reaches every node ofk.k.regularize(...)(h1seminorm/l2/tv/nonneg/bounded) is the FE-exact prior; the inclusion is recovered to rel-\(L^2 \approx 6\times10^{-2}\) (the prior smooths its peak slightly).
Full script
"""09 - Inverse problem: recover a hidden diffusivity field k(x) through a differentiable FEM solve.
Forward: -div(k(x) grad u) = f, u = 0 on the boundary, with unknown k(x) > 0.
Given the measured response ``u`` to a known source, recover the *entire* nodal diffusivity
field ``k(x)`` -- a hidden high-conductivity inclusion -- by differentiating the FEM solve end
to end. ``k = jno.np.parameter(phi)`` is a trainable P1 field on the trial space;
``fem.solve()`` is the differentiable forward solve; ``crux`` minimises the data misfit plus an
H1-seminorm smoothness prior ``k.regularize("h1seminorm")`` (field inversion is ill-posed
without regularisation). This is the FEM flavour of parameter-field identification / tomography.
"""
import jax
import jax.numpy as jnp
import numpy as np
import optax
from shapely.geometry import box
import jno
d = jno.domain(box(0.0, 0.0, 1.0, 1.0), mesh_size=0.1)
u, phi = d.fem_symbols()
xi, yi, _ = d.variable("interior", split=True)
xb, yb, _ = d.variable("boundary", split=True)
ui, vi = u.bind(x=xi, y=yi), phi.bind(x=xi, y=yi)
f = 30.0 * (xi * (1 - xi) + yi * (1 - yi)) # strong source so u is sensitive to k
nodes = np.asarray(d.built_mesh.points)[:, :2]
k_true = 1.0 + 0.8 * np.exp(
-((nodes[:, 0] - 0.5) ** 2 + (nodes[:, 1] - 0.5) ** 2) / (2 * 0.12**2)
) # background + inclusion
# One parametric assembly: generate clean full-field data by evaluating it at the true k ...
k = jno.np.parameter(phi, name="k")
fem = jno.fem([k * (ui.x * vi.x + ui.y * vi.y) - f * vi, u(xb, yb) - 0.0], quad_degree=3)
A_true, b = fem.operator.evaluate({"k": jnp.asarray(k_true)})
A_true = A_true.todense() if hasattr(A_true, "todense") else jnp.asarray(A_true) # operator is BCOO
u_obs = jnp.linalg.solve(A_true, jnp.asarray(b).reshape(-1))
# ... then recover k(x) from u_obs through the differentiable solve + an H1 smoothness prior.
k.dtype(jnp.float64)
k.initialize(jax.nn.initializers.constant(1.0)) # start from a uniform field
k.optimizer(optax.adam(2e-2))
crux = jno.core(
[(fem.solve() - u_obs).mse, 1e-3 * k.regularize("h1seminorm").mean],
domain=jno.domain.from_array({"_": np.zeros((1, 1))}),
)
crux.solve(500)
rec = np.asarray(crux.eval([k])).reshape(-1) # the recovered nodal field (do NOT index [0])
rel = float(np.linalg.norm(rec - k_true) / np.linalg.norm(k_true))
print(
f"\nInverse diffusivity field: nodes={k_true.shape[0]} k(x) rel_L2={rel:.3e} peak rec/true={rec.max():.3f}/{k_true.max():.3f}"
)
assert rel < 0.1