Inverse: nonlinear reaction coefficient (closed-form)

Download .py Back to chapter

Reaction-coefficient inverse problem with a fixed-target posterior. Recover the scalar k in λ u''(x) + k tanh(u(x)) = f(x) from noisy observations of the forcing f. Because the analytical solution u_exact = sin(πx)/π² is known, we plug it directly into the PDE residual instead of training a surrogate. NUTS samples a fixed-target posterior over k with calibrated, hyperparameter-stable credible intervals — matching Yang et al. 2021's textbook B-PINN HMC result (k ≈ 0.705 ± 0.006 for the small-noise case, our short chain hits k ≈ 0.703 ± 0.043).

If u were unknown (a real PDE solve), the two-stage surrogate pattern via substeps=[([surrogate_constraints], n_train), ([pde], 1)] with .stop_gradient on the surrogate in the PDE residual is the recommended alternative. Tutorial 06 (Inverse FEM Diffusivity) shows the more general case where the forward model is fully numerical (FEM).

Reference

Yang, L., Meng, X., & Karniadakis, G. E. (2021). B-PINNs: Bayesian physics-informed neural networks for forward and inverse PDE problems with noisy data. Journal of Computational Physics, 425, 109913.

Script

"""03 — Bayesian inverse: recover nonlinear reaction coefficient k"""

from pathlib import Path

import blackjax
import jax
import jax.numpy as jnp

import jno

π = jno.np.pi

# ── Physical setup ────────────────────────────────────────────────────────────
λ = 0.01
k_true = 0.7
sigma_obs = 0.005  # noise on observed f(x)

# ── Domain ────────────────────────────────────────────────────────────────────
domain = jno.domain.line(x_range=(-0.7, 0.7), mesh_size=0.05)
x, _ = domain.variable("interior")

# ── Analytical solution + noisy forcing observations under k_true ────────────
u = jno.np.sin(π * x) / π**2  # exact solution
u_xx = -jno.np.sin(π * x)  # analytical 2nd derivative (no FD noise)

# Synthetic noisy forcing observations.  jno.noise.gaussian is redrawn
# each step from the solver's PRNG key — with a fixed global seed the
# realisation is consistent across the whole chain.
f_obs = λ * u_xx + k_true * jno.np.tanh(u) + jno.noise.gaussian(std=sigma_obs)

# ── Bayesian reaction coefficient — NUTS with closed-form forward ────────────
k = jno.np.parameter((1,), key=jax.random.PRNGKey(0), name="k")
k.bayesian(
    blackjax.nuts,
    step_size=1e-2,  # initial guess; window adaptation tunes it
    warmup=200,
    keep=400,
    # adapt=True default — pure Bayesian inference, adaptation is well-defined.
)

# ── Likelihood: PDE-residual scaled by σ ──────────────────────────────────────
residual = (λ * u_xx + k * jno.np.tanh(u) - f_obs) / sigma_obs

# ── Solve — pure Bayesian, no surrogate, no substeps needed ──────────────────
crux = jno.core([residual.mse])
crux.solve(600)

# ── Posterior summary ────────────────────────────────────────────────────────
k_chain = k.posterior_samples
k_mean = float(jnp.mean(k_chain))
k_std = float(jnp.std(k_chain))
k_lo, k_hi = (float(v) for v in jnp.quantile(k_chain, jnp.array([0.05, 0.95])))

print(f"k = {k_mean:.4f} ± {k_std:.4f}")
print(f"   90% CI = [{k_lo:.4f}, {k_hi:.4f}]   truth = {k_true}")

rel_k = abs(k_mean - k_true) / abs(k_true)

results_file = Path(__file__).parent.parent.parent / "tutorial_results.txt"
with open(results_file, "a") as f:
    f.write(
        f"10_bayesian_pinns/03_inverse_reaction_coefficient.py | epochs=600 | "
        f"rel_k={rel_k:.4f} | CI_width={k_hi - k_lo:.4f}\n"
    )

assert rel_k < 0.1, f"posterior-mean k off by {rel_k:.2%}"