Differentiable Inverse (FDM)

Download .py All tutorials

When the constraint list carries a trainable jno.np.parameter, jno.fdm([...]).solve() returns a differentiable trace node — exactly as fem.solve() does — so it composes straight into jno.core. We recover the unknown amplitude \(s\) of a source \(s\,f_{\text{base}}\) from an observed steady field, letting the parameter's own optimizer drive the fit.

Result

The recovered quantity is a single scalar amplitude (the source \(s\,f_{\text{base}}\) shares the known basis \(f_{\text{base}}\)), so the recovery shows up as the fit, not a spatial field. From the wrong start \(s = 2.5\) the misfit gradient drives \(s \to 1.0000\): the fit residual against the observation collapses to \(\sim\!10^{-8}\) (middle), and \(|s-1|\) falls geometrically each SGD step until it saturates at the iterative forward solver's tolerance floor (\(\sim\!10^{-4}\), right).

The solve is a node inside a jno.core loss

The parameter carries its optimizer; the solve goes straight into the misfit term — the same shape as every FEM inverse tutorial:

s = jno.np.parameter((1,), name="s")
s.optimizer(optax.adam(1e-1))
u = d.unknown(); ui = u.bind(x=x, y=y)

solve = jno.fdm([-ui.d2(x) - ui.d2(y) - s * f_base, u(xb, yb) - 0.0]).solve()   # a trace node
crux  = jno.core([(solve - observed).mse])          # domain inferred from the graph
crux.solve(150)

What to notice

• No adjoint code and no manual gradient loop: crux drives the parameter, and the gradient flows through the solve's implicit custom_root — the same mechanism fem.solve() uses.
• jno.fdm([...]).solve() is still the one entry: with a trainable parameter it is a trace node the crux re-runs each step; without one it returns the solution array eagerly.
• It is a twin experiment — the observation is the forward solve at the true s = 1, so the minimizer is exactly the truth: from a wrong start \(s = 2.5\) the fit recovers \(s \approx 0.999\).

Full script

"""04 - Differentiable inverse through ``jno.fdm`` + ``jno.core``: recover an unknown source amplitude.

When the constraint list carries a trainable ``jno.np.parameter``, ``jno.fdm([...]).solve()`` returns
a differentiable **trace node** (not an array) -- exactly as ``fem.solve()`` does -- so it composes
straight into ``jno.core``. We run a twin experiment: generate a synthetic observation from the
forward solve at the true amplitude ``s = 1``, then recover ``s`` from a deliberately wrong start by
minimising the data misfit, with the parameter's own attached optimizer driving the fit.

    -Delta u = s * f_base,  u = 0 on the boundary.
"""

import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp  # noqa: E402
import numpy as np  # noqa: E402
import optax  # noqa: E402

import jno  # noqa: E402
import jno.jnp_ops as jnn  # noqa: E402

d = jno.Shape.rect(0.0, 0.0, 1.0, 1.0, size=0.08).domain()
x, y, _ = d.variable("interior", split=True)
xb, yb, _ = d.variable("boundary", split=True)
f_base = 2 * np.pi**2 * jnn.sin(np.pi * x) * jnn.sin(np.pi * y)
u = d.unknown()
ui = u.bind(x=x, y=y)

# Synthetic observation: the forward solve at the true amplitude s = 1 (a plain float -> eager array).
observed = jnp.asarray(jno.fdm([-ui.d2(x) - ui.d2(y) - 1.0 * f_base, u(xb, yb) - 0.0]).solve()).reshape(-1)

# Recover s: a trainable parameter with an attached optimizer, driven by crux through the data misfit.
s = jno.np.parameter((1,), name="s")
s.dtype(jnp.float64)
s.initialize(jax.nn.initializers.constant(2.5))  # deliberately wrong start
# A single, well-scaled scalar over a convex (quadratic) misfit: plain gradient descent converges
# straight to the minimum. Adam's per-parameter moment adaptation is counter-productive here — it
# oscillates and can settle at a spurious fixed point away from the true amplitude.
s.optimizer(optax.sgd(1.0))
solve = jno.fdm([-ui.d2(x) - ui.d2(y) - s * f_base, u(xb, yb) - 0.0]).solve()  # a differentiable trace node
crux = jno.core([(solve - observed).mse])  # domain inferred from the graph — no explicit domain= needed
crux.solve(150)

rec = float(np.asarray(crux.eval([s])).reshape(-1)[0])  # the recovered amplitude (do NOT index [0] on a field)
print(f"\nInverse via jno.fdm + jno.core: recovered source amplitude s={rec:.4f}  (true 1.0)")
assert abs(rec - 1.0) < 1e-2, f"did not recover the source amplitude: s={rec:.4f}"