diff --git a/applications/README.md b/applications/README.md new file mode 100644 index 00000000..c4b35239 --- /dev/null +++ b/applications/README.md @@ -0,0 +1,22 @@ +# Applications + +Research papers built **on** `mononet` as a library. Distinct from +[`benchmarks/`](../benchmarks/) (which reproduces the `mononet` paper's own +Tables 1 & 2) — applications are *downstream* work that uses the published +layers to do new science. Nothing here ships in the PyPI wheel. + +Each paper is a self-contained subpackage with its own `README.md` (abstract + +headline), `RUNBOOK.md` (exact reproduction commands), a Markdown manuscript +under `paper/`, and an executed notebook rendered into the Sphinx docs. + +## Papers + +| Dir | Paper | Status | +|---|---|---| +| [`pinn/`](pinn/) | Structure-Preserving PINNs — hard monotonicity as a PDE admissibility prior (inverse conservation-law / traffic flagship) | In progress (Paper 1) | + +Planned follow-ups (own specs/plans): high-dim HJB (Paper 2), Fokker–Planck +(Paper 3), eikonal (Paper 4), arbitrage-free surfaces (Paper 5, regression — +not a PINN). + +Design: [`docs/superpowers/specs/2026-07-12-applications-structure-preserving-pinns-design.md`](../docs/superpowers/specs/2026-07-12-applications-structure-preserving-pinns-design.md). diff --git a/applications/__init__.py b/applications/__init__.py new file mode 100644 index 00000000..feae9e02 --- /dev/null +++ b/applications/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Research applications built on top of the ``mononet`` package. + +This top-level area hosts downstream research papers that *depend on* +``mononet`` as a library. It is distinct from ``benchmarks/`` (which reproduces +the ``mononet`` paper's own results) and is **not** shipped in the PyPI wheel. +""" + +from __future__ import annotations diff --git a/applications/_common/__init__.py b/applications/_common/__init__.py new file mode 100644 index 00000000..a40de0f4 --- /dev/null +++ b/applications/_common/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared, framework-free helpers reused across application papers.""" + +from __future__ import annotations diff --git a/applications/_common/metrics_io.py b/applications/_common/metrics_io.py new file mode 100644 index 00000000..19e989dd --- /dev/null +++ b/applications/_common/metrics_io.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +"""JSON read/write for application result artifacts. + +Result artifacts are small JSON mappings written to ``results/`` and consumed +by report/notebook code. Mirrors the conventions used in ``benchmarks/``. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + + +def write_result(path: str | Path, obj: dict[str, Any]) -> None: + """Write `obj` as sorted, pretty-printed JSON to `path`. + + Parent directories are created if missing. + + :param path: Destination file path. + :param obj: JSON-serializable result mapping. + """ + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(json.dumps(obj, indent=2, sort_keys=True) + "\n") + + +def read_result(path: str | Path) -> dict[str, Any]: + """Read a JSON result mapping from `path`. + + :param path: Source file path. + :returns: The parsed mapping. + """ + data: dict[str, Any] = json.loads(Path(path).read_text()) + return data diff --git a/applications/_common/plot_theme.py b/applications/_common/plot_theme.py new file mode 100644 index 00000000..66d775bd --- /dev/null +++ b/applications/_common/plot_theme.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared matplotlib theme for application figures. + +Kept as a plain rcParams mapping so importing it pulls in no plotting backend; +callers apply it via ``matplotlib.rcParams.update(theme_rc())``. +""" + +from __future__ import annotations + +from typing import Any + + +def theme_rc() -> dict[str, Any]: + """Return the shared matplotlib rcParams overrides for figures. + + :returns: A mapping of rcParam keys to values, suitable for + `matplotlib.rcParams.update`. + """ + return { + "figure.dpi": 120, + "axes.grid": True, + "grid.alpha": 0.3, + "font.size": 10, + "figure.autolayout": True, + } diff --git a/applications/_common/seeding.py b/applications/_common/seeding.py new file mode 100644 index 00000000..794e0da3 --- /dev/null +++ b/applications/_common/seeding.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic RNG helpers shared across applications. + +All randomness in the applications area flows through these helpers so that +point sampling and observation generation are reproducible and identical +across backends. +""" + +from __future__ import annotations + +import numpy as np + + +def rng(seed: int) -> np.random.Generator: + """Return a NumPy generator seeded deterministically. + + :param seed: Non-negative integer seed. + :returns: A `numpy.random.Generator` (PCG64) seeded with `seed`. + """ + return np.random.default_rng(seed) + + +def split_seeds(seed: int, n: int) -> list[int]: + """Derive `n` independent integer seeds from a base `seed`. + + Uses `numpy.random.SeedSequence` so the child seeds are well separated and + reproducible. + + :param seed: Base seed. + :param n: Number of child seeds to produce. + :returns: A list of `n` deterministic non-negative integer seeds. + """ + state = np.random.SeedSequence(seed).generate_state(n) + return [int(s) for s in state] diff --git a/applications/pinn/README.md b/applications/pinn/README.md new file mode 100644 index 00000000..77b11faa --- /dev/null +++ b/applications/pinn/README.md @@ -0,0 +1,16 @@ +# Structure-Preserving PINNs (Paper 1) + +Hard, expressive monotonicity (via `mononet`) as a PDE **admissibility prior**: +a Physics-Informed Neural Network whose architecture guarantees a +Total-Variation-Diminishing, entropy-admissible, oscillation-free solution **by +construction**, where soft-penalty PINNs only approximate it. + +- **Forward mechanism tier:** Burgers-Riemann, Burgers smooth→shock, linear + advection, LWR — validated against exact / TVD finite-volume ground truth. +- **Deep flagship (inverse):** traffic state estimation — reconstruct a + monotone-front density field from sparse, noisy observations. + +Scope: the **monotone-solution class** (single-front / queue / Riemann +scenarios). Abstract and headline result land here once experiments run +(manuscript in [`paper/paper.md`](paper/paper.md); reproduction in +[`RUNBOOK.md`](RUNBOOK.md)). diff --git a/applications/pinn/RUNBOOK.md b/applications/pinn/RUNBOOK.md new file mode 100644 index 00000000..718e8647 --- /dev/null +++ b/applications/pinn/RUNBOOK.md @@ -0,0 +1,30 @@ +# RUNBOOK — Structure-Preserving PINNs (Paper 1) + +Exact commands to regenerate every figure and table. Filled in as the +implementation lands (see the plan: +`docs/superpowers/plans/2026-07-12-applications-structure-preserving-pinns-paper1.md`). + +## Environment + +Application deps are opt-in uv groups (`pinn` includes shared `applications`); +backends are package extras. Sync the app env before running anything here: + +```bash +uv sync --extra jax --group pinn # CPU JAX + optax + matplotlib (dev/tests) +uv sync --extra jax-gpu --group pinn # GPU JAX (heavy search); needs working CUDA +uv sync --extra all-cpu --group pinn # both backends (cross-backend equivalence) + +uv run pytest tests/applications -q # unit + smoke suite +``` + +- **Development / unit + smoke tests:** `default` devcontainer (`all-cpu`, both + backends) or a single-backend env (the other backend's tests skip). +- **Heavy Optuna search + sweeps:** `gpu-jax` devcontainer (primary GPU backend). +- **Cross-backend equivalence:** needs both backends → an `all-cpu` sync. + +## Reproduction (to be completed) + +- [ ] Optuna search (equal budget per method) — `experiments/search.py` +- [ ] Forward mechanism-tier sweep — `experiments/sweep.py` +- [ ] Inverse flagship sparsity × noise sweep — `experiments/sweep.py` +- [ ] Fill manuscript results + render notebook diff --git a/applications/pinn/__init__.py b/applications/pinn/__init__.py new file mode 100644 index 00000000..ac79329a --- /dev/null +++ b/applications/pinn/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Structure-Preserving PINNs (Paper 1). + +Hard monotonicity as a PDE admissibility prior. This package must import +without pulling in any ML backend: access backend code explicitly via +``applications.pinn.models.jax`` / ``applications.pinn.models.torch``. +""" + +from __future__ import annotations diff --git a/applications/pinn/configs/.gitkeep b/applications/pinn/configs/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/applications/pinn/core/__init__.py b/applications/pinn/core/__init__.py new file mode 100644 index 00000000..545e26be --- /dev/null +++ b/applications/pinn/core/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Backend-agnostic science for the PINN application (pure NumPy). + +Single source of truth: problem definitions, exact solutions, the reference +solver, sampling, metrics, and plotting. Imports **no** ML framework. +""" + +from __future__ import annotations diff --git a/applications/pinn/core/admissibility.py b/applications/pinn/core/admissibility.py new file mode 100644 index 00000000..4839f286 --- /dev/null +++ b/applications/pinn/core/admissibility.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Admissibility specification and violation metric. + +The *admissibility condition* of a structure-preserving PINN is a monotonicity +(and, for later papers, convexity) statement about the solution field. This +module holds the declarative spec and the non-negative violation functional that +is the paper's headline metric: it is ``0`` by construction for the hard-monotone +model and ``> 0`` for unconstrained / soft-penalty baselines. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import numpy.typing as npt + + +@dataclass(frozen=True, slots=True) +class AdmissibilitySpec: + """Declarative admissibility condition for a problem's solution field. + + :param mask: Monotonicity sign per input axis — ``+1`` non-decreasing, + ``-1`` non-increasing, ``0`` unconstrained (mirrors + `mononet.MonotonicityMask`). + :param convex_axes: Input axes in which the solution is convex. + :param concave_axes: Input axes in which the solution is concave. + """ + + mask: tuple[int, ...] + convex_axes: tuple[int, ...] = () + concave_axes: tuple[int, ...] = () + + +def violation( + field: npt.NDArray[np.floating], + *, + axis: int, + sign: int, +) -> float: + """Total wrong-sign variation of ``field`` along ``axis``. + + For a non-increasing target (``sign == -1``) every *upward* step is a + violation; the returned value is the summed magnitude of those steps — i.e. + the total mass by which the field fails to be monotone. It is ``0`` iff the + field is monotone in the required direction along ``axis``. + + :param field: Solution field sampled on a grid. + :param axis: Axis along which monotonicity is required. + :param sign: ``-1`` non-increasing, ``+1`` non-decreasing, ``0`` + unconstrained (always ``0``). + :returns: Non-negative total wrong-sign variation. + """ + if sign == 0: + return 0.0 + diffs = np.diff(field, axis=axis) + wrong = np.maximum(diffs, 0.0) if sign == -1 else np.maximum(-diffs, 0.0) + return float(wrong.sum()) diff --git a/applications/pinn/core/exact.py b/applications/pinn/core/exact.py new file mode 100644 index 00000000..aa9d0ccd --- /dev/null +++ b/applications/pinn/core/exact.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Closed-form entropy solutions for scalar conservation laws. + +These are the exact ground truth used to validate the reference solver and to +score PINNs on the forward mechanism tier. All functions are elementwise over +NumPy arrays and handle ``t == 0`` (the initial step) without division issues. + +Conventions: the Riemann discontinuity sits at ``x0`` (default ``0``); ``x`` and +``t`` broadcast against each other. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import numpy.typing as npt + +if TYPE_CHECKING: + from collections.abc import Callable + +Array = npt.NDArray[np.floating] + + +def _riemann_step(x: Array, x0: float, left: float, right: float) -> Array: + """Return the initial Riemann step: ``left`` for ``x < x0`` else ``right``.""" + return np.where(x < x0, left, right) + + +def burgers_riemann( + x: Array, + t: Array | float, + u_l: float, + u_r: float, + *, + x0: float = 0.0, +) -> Array: + """Entropy solution of Burgers' equation for Riemann data. + + Flux ``f(u) = u^2 / 2`` (convex). For ``u_l > u_r`` a shock travels at the + Rankine-Hugoniot speed ``s = (u_l + u_r) / 2``; for ``u_l < u_r`` a + rarefaction fan connects the states; equal states are constant. + + :param x: Spatial coordinates. + :param t: Time(s), broadcast against ``x``; ``t == 0`` gives the initial step. + :param u_l: Left state. + :param u_r: Right state. + :param x0: Location of the initial discontinuity. + :returns: ``u(x, t)`` on the entropy solution. + """ + x = np.asarray(x, dtype=float) + t_arr = np.asarray(t, dtype=float) + t_safe = np.where(t_arr > 0.0, t_arr, 1.0) + xi = (x - x0) / t_safe + step = _riemann_step(x, x0, u_l, u_r) + if u_l > u_r: # shock + s = 0.5 * (u_l + u_r) + moving = np.where(xi < s, u_l, u_r) + elif u_l < u_r: # rarefaction: u = xi clamped to [u_l, u_r] + moving = np.clip(xi, u_l, u_r) + else: + moving = np.full_like(x, u_l) + return np.where(t_arr > 0.0, moving, step) + + +def advection( + x: Array, + t: Array | float, + a: float, + u0: Callable[[Array], Array], +) -> Array: + """Exact solution of linear advection ``u_t + a u_x = 0``. + + The initial profile is transported unchanged: ``u(x, t) = u0(x - a t)``. + + :param x: Spatial coordinates. + :param t: Time(s), broadcast against ``x``. + :param a: Constant advection speed. + :param u0: Initial-condition callable. + :returns: ``u(x, t)``. + """ + x = np.asarray(x, dtype=float) + t_arr = np.asarray(t, dtype=float) + return u0(x - a * t_arr) + + +def greenshields_flux(rho: Array, v_max: float, rho_max: float) -> Array: + """Greenshields LWR flux ``Q(rho) = v_max * rho * (1 - rho / rho_max)``.""" + return v_max * rho * (1.0 - rho / rho_max) + + +def greenshields_flux_prime(rho: Array, v_max: float, rho_max: float) -> Array: + """Characteristic speed ``Q'(rho) = v_max * (1 - 2 rho / rho_max)``.""" + return v_max * (1.0 - 2.0 * rho / rho_max) + + +def lwr_riemann( + x: Array, + t: Array | float, + rho_l: float, + rho_r: float, + *, + v_max: float = 1.0, + rho_max: float = 1.0, + x0: float = 0.0, +) -> Array: + """Entropy solution of the LWR traffic model (Greenshields, concave flux). + + For the concave flux, ``rho_l < rho_r`` gives a shock (a forming queue) + travelling at ``s = (Q(rho_r) - Q(rho_l)) / (rho_r - rho_l)``; ``rho_l > + rho_r`` gives a rarefaction (a dissolving queue). Both profiles are monotone + in ``x``. + + :param x: Spatial coordinates. + :param t: Time(s), broadcast against ``x``; ``t == 0`` gives the initial step. + :param rho_l: Left density state. + :param rho_r: Right density state. + :param v_max: Free-flow speed. + :param rho_max: Jam density. + :param x0: Location of the initial discontinuity. + :returns: ``rho(x, t)`` on the entropy solution. + """ + x = np.asarray(x, dtype=float) + t_arr = np.asarray(t, dtype=float) + t_safe = np.where(t_arr > 0.0, t_arr, 1.0) + xi = (x - x0) / t_safe + step = _riemann_step(x, x0, rho_l, rho_r) + if rho_l < rho_r: # shock + q_l = greenshields_flux(np.asarray(rho_l), v_max, rho_max) + q_r = greenshields_flux(np.asarray(rho_r), v_max, rho_max) + s = float((q_r - q_l) / (rho_r - rho_l)) + moving = np.where(xi < s, rho_l, rho_r) + elif rho_l > rho_r: # rarefaction + lam_l = float(greenshields_flux_prime(np.asarray(rho_l), v_max, rho_max)) + lam_r = float(greenshields_flux_prime(np.asarray(rho_r), v_max, rho_max)) + fan = 0.5 * rho_max * (1.0 - xi / v_max) # invert Q'(rho) = xi + moving = np.where(xi <= lam_l, rho_l, np.where(xi >= lam_r, rho_r, fan)) + else: + moving = np.full_like(x, rho_l) + return np.where(t_arr > 0.0, moving, step) + + +def breaking_time(du0_dx_min: float) -> float: + """Shock breaking time ``t_b = -1 / min(u0')`` for Burgers. + + :param du0_dx_min: The minimum of the initial-derivative ``u0'`` (negative + for shock-forming data). + :returns: The finite breaking time. + :raises ValueError: If ``du0_dx_min >= 0`` (no shock forms). + """ + if du0_dx_min >= 0.0: + raise ValueError("no shock forms when min(u0') >= 0") + return -1.0 / du0_dx_min + + +def burgers_characteristic( + x: Array, + t: float, + u0: Callable[[Array], Array], + *, + iters: int = 500, + tol: float = 1e-13, +) -> Array: + """Pre-shock Burgers solution via the method of characteristics. + + Solves the implicit relation ``u = u0(x - u t)`` by fixed-point iteration, + which converges for ``t < t_b`` (before any shock forms). Use the reference + solver for ``t >= t_b``. + + :param x: Spatial coordinates. + :param t: Time (scalar), must be strictly before the breaking time. + :param u0: Initial-condition callable. + :param iters: Maximum fixed-point iterations. + :param tol: Convergence tolerance on successive iterates. + :returns: ``u(x, t)``. + """ + x = np.asarray(x, dtype=float) + u = u0(x) + for _ in range(iters): + u_next = u0(x - u * t) + if float(np.max(np.abs(u_next - u))) < tol: + return u_next + u = u_next + return u diff --git a/applications/pinn/core/metrics.py b/applications/pinn/core/metrics.py new file mode 100644 index 00000000..2eaccc65 --- /dev/null +++ b/applications/pinn/core/metrics.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Error and admissibility metrics for scoring PINN solutions. + +All metrics operate on NumPy arrays sampled on the evaluation grid and are used +identically across backends and methods. +""" + +from __future__ import annotations + +import numpy as np +import numpy.typing as npt + +Array = npt.NDArray[np.floating] + + +def l1(pred: Array, ref: Array, *, dx: float = 1.0) -> float: + """Discrete L1 error ``sum|pred - ref| * dx``.""" + return float(np.abs(pred - ref).sum() * dx) + + +def l2(pred: Array, ref: Array, *, dx: float = 1.0) -> float: + """Discrete L2 error ``sqrt(sum(pred - ref)^2 * dx)``.""" + return float(np.sqrt((np.square(pred - ref)).sum() * dx)) + + +def tv(profile: Array) -> float: + """Total variation of a 1-D profile, ``sum|diff|``.""" + return float(np.abs(np.diff(profile)).sum()) + + +def tv_curve(field: Array) -> Array: + """Total variation at each time slice of a ``(nt, nx)`` field.""" + per_slice: Array = np.abs(np.diff(field, axis=1)).sum(axis=1) + return per_slice + + +def overshoot(pred: Array, ref: Array) -> float: + """Excursion of ``pred`` beyond the range of ``ref`` (the Gibbs overshoot). + + Zero when ``pred`` stays within ``[min(ref), max(ref)]``; otherwise the + largest amount by which it exceeds that range on either side. + + :param pred: Predicted values. + :param ref: Reference values. + :returns: Non-negative overshoot magnitude. + """ + above = float(np.max(pred) - np.max(ref)) + below = float(np.min(ref) - np.min(pred)) + return max(above, below, 0.0) + + +def shock_position(profile: Array, x_values: Array, level: float) -> float: + """Locate the ``level`` crossing of a monotone profile (the front position). + + :param profile: 1-D solution values on ``x_values``. + :param x_values: Spatial grid. + :param level: Crossing level (e.g. the mean of the two states). + :returns: The x-coordinate nearest the level crossing. + """ + idx = int(np.argmin(np.abs(profile - level))) + return float(x_values[idx]) + + +def shock_position_error( + pred: Array, ref: Array, x_values: Array, *, level: float +) -> float: + """Absolute error in the front position between ``pred`` and ``ref``.""" + return abs( + shock_position(pred, x_values, level) - shock_position(ref, x_values, level) + ) + + +def mass_error(pred: Array, ref: Array, *, dx: float = 1.0) -> float: + """Absolute difference in total mass ``|sum(pred) - sum(ref)| * dx``.""" + return float(abs(pred.sum() - ref.sum()) * dx) + + +def reconstruction_error(field: Array, ref: Array, *, dx: float = 1.0) -> float: + """L2 error over a whole ``(nt, nx)`` field (inverse-flagship headline).""" + return l2(field.ravel(), ref.ravel(), dx=dx) diff --git a/applications/pinn/core/plotting.py b/applications/pinn/core/plotting.py new file mode 100644 index 00000000..0144c67d --- /dev/null +++ b/applications/pinn/core/plotting.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Matplotlib figures for the PINN paper. + +Figures are built with the object-oriented API (no ``pyplot`` global state, no +GUI backend), so the functions are pure and safe to call in tests and notebooks. +Callers save or embed the returned :class:`~matplotlib.figure.Figure`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import numpy.typing as npt +from matplotlib.figure import Figure + +if TYPE_CHECKING: + from collections.abc import Mapping + +Array = npt.NDArray[np.floating] + + +def profiles( + x_values: Array, + series: Mapping[str, Array], + *, + title: str = "", +) -> Figure: + """Plot several named solution profiles against ``x``. + + :param x_values: Shared spatial axis. + :param series: Mapping of label to profile values. + :param title: Optional axes title. + :returns: The figure. + """ + fig = Figure() + ax = fig.subplots() + for label, values in series.items(): + ax.plot(x_values, values, label=label) + ax.set_xlabel("x") + ax.set_ylabel("u") + if title: + ax.set_title(title) + ax.legend() + return fig + + +def tv_curves( + t_values: Array, + curves: Mapping[str, Array], + *, + title: str = "", +) -> Figure: + """Plot total-variation-vs-time curves for several methods. + + :param t_values: Time axis. + :param curves: Mapping of label to TV(t) values. + :param title: Optional axes title. + :returns: The figure. + """ + fig = Figure() + ax = fig.subplots() + for label, values in curves.items(): + ax.plot(t_values, values, label=label) + ax.set_xlabel("t") + ax.set_ylabel("TV(t)") + if title: + ax.set_title(title) + ax.legend() + return fig + + +def field_heatmap( + field: Array, + x_values: Array, + t_values: Array, + *, + title: str = "", +) -> Figure: + """Render a space-time field as a heatmap. + + :param field: Field of shape ``(len(t_values), len(x_values))``. + :param x_values: Spatial axis. + :param t_values: Temporal axis. + :param title: Optional axes title. + :returns: The figure. + """ + fig = Figure() + ax = fig.subplots() + extent = ( + float(x_values[0]), + float(x_values[-1]), + float(t_values[0]), + float(t_values[-1]), + ) + image = ax.imshow( + field, origin="lower", aspect="auto", extent=extent, cmap="viridis" + ) + fig.colorbar(image, ax=ax) + ax.set_xlabel("x") + ax.set_ylabel("t") + if title: + ax.set_title(title) + return fig diff --git a/applications/pinn/core/problems/__init__.py b/applications/pinn/core/problems/__init__.py new file mode 100644 index 00000000..749d4f77 --- /dev/null +++ b/applications/pinn/core/problems/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Problem registry: one plug-in module per PDE family. + +Importing this package registers the built-in problems (conservation laws). +Follow-up papers add ``hjb`` / ``fokker_planck`` / ``eikonal`` modules here. +""" + +from __future__ import annotations + +from applications.pinn.core.problems import conservation # noqa: F401 (registers) +from applications.pinn.core.problems.base import ( + Problem, + available, + get, + register, +) + +__all__ = ["Problem", "available", "get", "register"] diff --git a/applications/pinn/core/problems/base.py b/applications/pinn/core/problems/base.py new file mode 100644 index 00000000..b5484279 --- /dev/null +++ b/applications/pinn/core/problems/base.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The ``Problem`` protocol and the problem registry. + +Each PDE family is a plug-in module that defines one or more `Problem` +implementations and registers them by name. Downstream code (models, trainers, +experiments) discovers problems through :func:`get` / :func:`available` and never +imports the concrete classes directly. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol, TypeVar, cast, runtime_checkable + +import numpy as np +import numpy.typing as npt + +if TYPE_CHECKING: + from collections.abc import Callable + + from applications.pinn.core.admissibility import AdmissibilitySpec + +Array = npt.NDArray[np.floating] +_T = TypeVar("_T") + + +@runtime_checkable +class Problem(Protocol): + """A scalar PDE problem on a 1-D spatial domain over time. + + Concrete problems are constructed with their parameters (flux constants, + Riemann states, …) and expose a uniform interface so the framework can build + residuals, enforce admissibility, and score against ground truth. + """ + + #: Registry key (class-level identifier, e.g. ``"burgers"``). + key: str + + @property + def domain(self) -> tuple[tuple[float, float], tuple[float, float]]: + """Return ``((x_min, x_max), (t_min, t_max))``.""" + + def admissibility(self) -> AdmissibilitySpec: + """Return the monotonicity/convexity spec for the solution field.""" + + def flux(self, u: Array) -> Array: + """Evaluate the flux ``f(u)``.""" + + def flux_prime(self, u: Array) -> Array: + """Evaluate the flux derivative ``f'(u)`` (characteristic speed).""" + + def initial(self, x: Array) -> Array: + """Evaluate the initial condition ``u(x, 0)`` (forward mode).""" + + def ground_truth(self, x: Array, t: Array) -> Array | None: + """Return the reference solution on ``(x, t)`` if available, else None.""" + + +_REGISTRY: dict[str, type] = {} + + +def register(key: str) -> Callable[[type[_T]], type[_T]]: + """Class decorator registering a `Problem` implementation under ``key``. + + Generic in the decorated class so the decorated symbol keeps its concrete + type; the registry surfaces entries as `Problem` (the contract implementers + are expected to satisfy structurally). + + :param key: Unique registry name. + :returns: The decorator. + :raises KeyError: If ``key`` is already registered. + """ + + def decorate(cls: type[_T]) -> type[_T]: + if key in _REGISTRY: + raise KeyError(f"problem {key!r} already registered") + _REGISTRY[key] = cls + return cls + + return decorate + + +def get(key: str) -> type[Problem]: + """Return the registered `Problem` class for ``key``. + + :param key: Registry name. + :returns: The problem class. + :raises KeyError: If ``key`` is not registered. + """ + if key not in _REGISTRY: + raise KeyError(f"unknown problem {key!r}; available: {available()}") + return cast("type[Problem]", _REGISTRY[key]) + + +def available() -> list[str]: + """Return the sorted list of registered problem keys.""" + return sorted(_REGISTRY) diff --git a/applications/pinn/core/problems/conservation.py b/applications/pinn/core/problems/conservation.py new file mode 100644 index 00000000..88e1f8c1 --- /dev/null +++ b/applications/pinn/core/problems/conservation.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Scalar conservation-law problems: Burgers, linear advection, LWR traffic. + +Each problem is registered and exposes the flux, the admissibility spec (the +monotone direction of the entropy solution in ``x``), the initial condition, and +a ground-truth evaluator. All ground truth is closed-form except the smooth Burgers +case, which forms a shock and uses the Godunov reference past the breaking time. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import ClassVar + +import numpy as np +import numpy.typing as npt + +from applications.pinn.core import exact, reference_solver +from applications.pinn.core.admissibility import AdmissibilitySpec +from applications.pinn.core.problems.base import register + +Array = npt.NDArray[np.floating] + + +def _mono_sign(left: float, right: float) -> int: + """Return the monotone direction in ``x`` of a left>right / left right else 1 + + +@register("burgers_riemann") +@dataclass(frozen=True, slots=True) +class BurgersRiemann: + """Burgers' equation with Riemann initial data (pure shock or rarefaction).""" + + key: ClassVar[str] = "burgers_riemann" + u_l: float = 1.0 + u_r: float = 0.0 + x0: float = 0.0 + x_range: tuple[float, float] = (-2.0, 3.0) + t_max: float = 1.5 + + @property + def domain(self) -> tuple[tuple[float, float], tuple[float, float]]: + """Return the space-time domain.""" + return (self.x_range, (0.0, self.t_max)) + + @property + def sonic(self) -> float: + """Sonic point of the convex Burgers flux.""" + return 0.0 + + def admissibility(self) -> AdmissibilitySpec: + """Monotone in ``x`` in the direction set by the Riemann states.""" + return AdmissibilitySpec(mask=(_mono_sign(self.u_l, self.u_r), 0)) + + def flux(self, u: Array) -> Array: + """Burgers flux ``u^2 / 2`` (backend-polymorphic).""" + return 0.5 * u**2 + + def flux_prime(self, u: Array) -> Array: + """Characteristic speed ``u`` (backend-polymorphic).""" + return u + + def initial(self, x: Array) -> Array: + """Return the initial Riemann step.""" + return exact.burgers_riemann(x, 0.0, self.u_l, self.u_r, x0=self.x0) + + def ground_truth(self, x: Array, t: Array) -> Array: + """Exact entropy solution.""" + return exact.burgers_riemann(x, t, self.u_l, self.u_r, x0=self.x0) + + +@register("advection") +@dataclass(frozen=True, slots=True) +class LinearAdvection: + """Linear advection of a monotone-decreasing front at constant speed ``a``.""" + + key: ClassVar[str] = "advection" + a: float = 1.0 + steepness: float = 3.0 + x0: float = 0.0 + x_range: tuple[float, float] = (-3.0, 4.0) + t_max: float = 1.5 + + @property + def domain(self) -> tuple[tuple[float, float], tuple[float, float]]: + """Return the space-time domain.""" + return (self.x_range, (0.0, self.t_max)) + + def admissibility(self) -> AdmissibilitySpec: + """Return the spec: the transported front is non-increasing in ``x``.""" + return AdmissibilitySpec(mask=(-1, 0)) + + def _u0(self, x: Array) -> Array: + profile = 0.5 * (1.0 - np.tanh(self.steepness * (x - self.x0))) + return np.asarray(profile, dtype=float) + + def flux(self, u: Array) -> Array: + """Linear flux ``a * u`` (backend-polymorphic).""" + result: Array = self.a * u + return result + + def flux_prime(self, u: Array) -> Array: + """Constant characteristic speed ``a`` (backend-polymorphic broadcast).""" + result: Array = self.a + 0.0 * u + return result + + def initial(self, x: Array) -> Array: + """Monotone-decreasing initial front.""" + return self._u0(np.asarray(x, dtype=float)) + + def ground_truth(self, x: Array, t: Array) -> Array: + """Exact transported profile ``u0(x - a t)``.""" + return exact.advection(x, t, self.a, self._u0) + + +@register("lwr_riemann") +@dataclass(frozen=True, slots=True) +class LwrRiemann: + """LWR traffic (Greenshields flux) with Riemann density data.""" + + key: ClassVar[str] = "lwr_riemann" + rho_l: float = 0.2 + rho_r: float = 0.8 + v_max: float = 1.0 + rho_max: float = 1.0 + x0: float = 0.0 + x_range: tuple[float, float] = (-2.0, 2.0) + t_max: float = 1.5 + + @property + def domain(self) -> tuple[tuple[float, float], tuple[float, float]]: + """Return the space-time domain.""" + return (self.x_range, (0.0, self.t_max)) + + @property + def sonic(self) -> float: + """Sonic point of the concave Greenshields flux.""" + return self.rho_max / 2.0 + + def admissibility(self) -> AdmissibilitySpec: + """Monotone in ``x`` in the direction set by the density states.""" + return AdmissibilitySpec(mask=(_mono_sign(self.rho_l, self.rho_r), 0)) + + def flux(self, u: Array) -> Array: + """Greenshields flux (backend-polymorphic).""" + return exact.greenshields_flux(u, self.v_max, self.rho_max) + + def flux_prime(self, u: Array) -> Array: + """Greenshields characteristic speed (backend-polymorphic).""" + return exact.greenshields_flux_prime(u, self.v_max, self.rho_max) + + def initial(self, x: Array) -> Array: + """Return the initial Riemann density step.""" + return exact.lwr_riemann( + x, + 0.0, + self.rho_l, + self.rho_r, + v_max=self.v_max, + rho_max=self.rho_max, + x0=self.x0, + ) + + def ground_truth(self, x: Array, t: Array) -> Array: + """Exact entropy solution.""" + return exact.lwr_riemann( + x, + t, + self.rho_l, + self.rho_r, + v_max=self.v_max, + rho_max=self.rho_max, + x0=self.x0, + ) + + +@register("burgers_smooth") +@dataclass(slots=True) +class BurgersSmoothShock: + """Burgers with smooth monotone-decreasing data that steepens into a shock. + + Exact pre-breaking (method of characteristics); past the breaking time the + ground truth is the Godunov reference, interpolated. The reference field is + built once and cached. + """ + + key: ClassVar[str] = "burgers_smooth" + steepness: float = 1.0 + x_range: tuple[float, float] = (-4.0, 4.0) + t_max: float = 2.0 + ref_nx: int = 1000 + ref_nt: int = 400 + _cache: dict[str, Array] = field(default_factory=dict, compare=False, repr=False) + + @property + def domain(self) -> tuple[tuple[float, float], tuple[float, float]]: + """Return the space-time domain.""" + return (self.x_range, (0.0, self.t_max)) + + @property + def sonic(self) -> float: + """Sonic point of the convex Burgers flux.""" + return 0.0 + + @property + def breaking_time(self) -> float: + """Breaking time ``t_b = -1 / min(u0')`` for the chosen initial data.""" + # u0 = -tanh(k x): min slope is -k at x = 0. + return 1.0 / self.steepness + + def admissibility(self) -> AdmissibilitySpec: + """Monotone non-increasing in ``x``.""" + return AdmissibilitySpec(mask=(-1, 0)) + + def _u0(self, x: Array) -> Array: + return -np.tanh(self.steepness * np.asarray(x, dtype=float)) + + def flux(self, u: Array) -> Array: + """Burgers flux ``u^2 / 2`` (backend-polymorphic).""" + return 0.5 * u**2 + + def flux_prime(self, u: Array) -> Array: + """Characteristic speed ``u`` (backend-polymorphic).""" + return u + + def initial(self, x: Array) -> Array: + """Smooth monotone-decreasing initial profile.""" + return self._u0(x) + + def _reference(self) -> tuple[Array, Array, Array]: + if "field" not in self._cache: + (x_lo, x_hi), (_t0, t_hi) = self.domain + edges = np.linspace(x_lo, x_hi, self.ref_nx + 1) + xv = 0.5 * (edges[:-1] + edges[1:]) + tv = np.linspace(0.0, t_hi, self.ref_nt) + fld = reference_solver.godunov( + self._u0(xv), + xv, + list(tv), + flux=self.flux, + flux_prime=self.flux_prime, + sonic=self.sonic, + ) + self._cache["x"], self._cache["t"], self._cache["field"] = xv, tv, fld + return self._cache["x"], self._cache["t"], self._cache["field"] + + def ground_truth(self, x: Array, t: Array) -> Array: + """Characteristic solution pre-breaking; Godunov reference afterwards.""" + xv, tv, fld = self._reference() + return reference_solver.interpolate(fld, xv, tv, x, t) diff --git a/applications/pinn/core/reference_solver.py b/applications/pinn/core/reference_solver.py new file mode 100644 index 00000000..1dab4676 --- /dev/null +++ b/applications/pinn/core/reference_solver.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +"""A TVD finite-volume (Godunov) reference solver for scalar conservation laws. + +This is the numerical ground truth: a first-order Godunov scheme that converges +to the entropy solution and is total-variation-diminishing by construction. It +handles both convex (Burgers) and concave (LWR) fluxes through the exact Godunov +interface flux, which needs only the flux and its single sonic point ``u*`` +(where ``f'(u*) = 0``). Outflow (zero-gradient) boundaries; CFL-limited stepping. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import numpy.typing as npt + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + +Array = npt.NDArray[np.floating] + + +def godunov_flux( + u_left: Array, + u_right: Array, + flux: Callable[[Array], Array], + sonic: float, +) -> Array: + """Exact Godunov numerical flux at cell interfaces. + + For ``u_left <= u_right`` the flux is the minimum of ``f`` over + ``[u_left, u_right]``; otherwise the maximum over ``[u_right, u_left]``. With + a single sonic point ``u*`` these extrema are attained at the endpoints or at + ``u*`` when it lies inside the interval — exact for convex or concave ``f``. + + :param u_left: Left interface states. + :param u_right: Right interface states. + :param flux: Flux function ``f``. + :param sonic: The point ``u*`` where ``f'(u*) = 0``. + :returns: Godunov flux per interface. + """ + f_l = flux(u_left) + f_r = flux(u_right) + f_s = flux(np.full_like(u_left, sonic)) + lo = np.minimum(u_left, u_right) + hi = np.maximum(u_left, u_right) + sonic_inside = (lo <= sonic) & (sonic <= hi) + minimum = np.minimum(f_l, f_r) + maximum = np.maximum(f_l, f_r) + min_flux = np.where(sonic_inside, np.minimum(minimum, f_s), minimum) + max_flux = np.where(sonic_inside, np.maximum(maximum, f_s), maximum) + return np.where(u_left <= u_right, min_flux, max_flux) + + +def godunov( + u0_values: Array, + x_centers: Array, + t_eval: Sequence[float], + *, + flux: Callable[[Array], Array], + flux_prime: Callable[[Array], Array], + sonic: float, + cfl: float = 0.4, +) -> Array: + """Evolve initial data with the Godunov scheme, snapshotting at ``t_eval``. + + :param u0_values: Initial cell-centre values. + :param x_centers: Uniform cell-centre coordinates. + :param t_eval: Non-negative, increasing times at which to snapshot. + :param flux: Flux ``f``. + :param flux_prime: Flux derivative ``f'`` (for the CFL condition). + :param sonic: Sonic point ``u*`` (``f'(u*) = 0``). + :param cfl: CFL number in ``(0, 1)``. + :returns: Array of shape ``(len(t_eval), len(x_centers))``. + """ + dx = float(x_centers[1] - x_centers[0]) + u = np.array(u0_values, dtype=float) + snapshots: list[Array] = [] + current = 0.0 + for target in t_eval: + while current < target - 1e-15: + speed = float(np.max(np.abs(flux_prime(u)))) + 1e-12 + dt = min(cfl * dx / speed, target - current) + extended = np.concatenate([u[:1], u, u[-1:]]) + interface = godunov_flux(extended[:-1], extended[1:], flux, sonic) + u = u - (dt / dx) * (interface[1:] - interface[:-1]) + current += dt + snapshots.append(u.copy()) + return np.asarray(snapshots) + + +def interpolate( + field: Array, + x_values: Array, + t_values: Array, + x_query: Array, + t_query: Array, +) -> Array: + """Bilinearly interpolate a reference ``field`` at scattered query points. + + :param field: Field of shape ``(len(t_values), len(x_values))``. + :param x_values: Ascending spatial grid axis. + :param t_values: Ascending temporal grid axis. + :param x_query: Query x-coordinates. + :param t_query: Query t-coordinates (same shape as ``x_query``). + :returns: Interpolated values, shape of ``x_query``. + """ + xq = np.asarray(x_query, dtype=float) + tq = np.asarray(t_query, dtype=float) + ix = np.clip(np.searchsorted(x_values, xq) - 1, 0, len(x_values) - 2) + it = np.clip(np.searchsorted(t_values, tq) - 1, 0, len(t_values) - 2) + x0, x1 = x_values[ix], x_values[ix + 1] + t0, t1 = t_values[it], t_values[it + 1] + wx = np.clip((xq - x0) / (x1 - x0), 0.0, 1.0) + wt = np.clip((tq - t0) / (t1 - t0), 0.0, 1.0) + top = field[it, ix] * (1 - wx) + field[it, ix + 1] * wx + bot = field[it + 1, ix] * (1 - wx) + field[it + 1, ix + 1] * wx + return top * (1 - wt) + bot * wt diff --git a/applications/pinn/core/sampling.py b/applications/pinn/core/sampling.py new file mode 100644 index 00000000..793edb2a --- /dev/null +++ b/applications/pinn/core/sampling.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic point sampling for PINN training and evaluation. + +All point sets are generated once in NumPy from a seed so JAX and PyTorch train +on identical inputs. Provides interior collocation, initial/boundary points, a +structured evaluation grid, and — for the inverse flagship — a sparse, noisy +observation sampler drawn from a reference field. +""" + +from __future__ import annotations + +import numpy as np +import numpy.typing as npt + +from applications._common import seeding + +Array = npt.NDArray[np.floating] +Domain = tuple[tuple[float, float], tuple[float, float]] + + +def _uniform(gen: np.random.Generator, lo: float, hi: float, n: int) -> Array: + return lo + (hi - lo) * gen.random(n) + + +def _latin(gen: np.random.Generator, lo: float, hi: float, n: int) -> Array: + """One-dimensional Latin-hypercube sample of ``n`` points in ``[lo, hi]``.""" + edges = np.arange(n) + jitter = gen.random(n) + unit = (edges + jitter) / n + return lo + (hi - lo) * gen.permutation(unit) + + +def collocation( + domain: Domain, n: int, *, seed: int, strategy: str = "uniform" +) -> Array: + """Sample ``n`` interior collocation points ``(x, t)``. + + :param domain: ``((x_min, x_max), (t_min, t_max))``. + :param n: Number of points. + :param seed: RNG seed. + :param strategy: ``"uniform"`` or ``"lhs"`` (Latin hypercube). + :returns: Array of shape ``(n, 2)`` with columns ``[x, t]``. + :raises ValueError: If ``strategy`` is unknown. + """ + gen = seeding.rng(seed) + (x_lo, x_hi), (t_lo, t_hi) = domain + if strategy == "uniform": + x = _uniform(gen, x_lo, x_hi, n) + t = _uniform(gen, t_lo, t_hi, n) + elif strategy == "lhs": + x = _latin(gen, x_lo, x_hi, n) + t = _latin(gen, t_lo, t_hi, n) + else: + raise ValueError(f"unknown strategy {strategy!r}") + return np.column_stack([x, t]) + + +def initial_points(domain: Domain, n: int, *, seed: int) -> Array: + """Sample ``n`` points on the initial line ``t = t_min``. + + :param domain: ``((x_min, x_max), (t_min, t_max))``. + :param n: Number of points. + :param seed: RNG seed. + :returns: Array of shape ``(n, 2)`` with ``t == t_min``. + """ + gen = seeding.rng(seed) + (x_lo, x_hi), (t_lo, _t_hi) = domain + x = _uniform(gen, x_lo, x_hi, n) + return np.column_stack([x, np.full(n, t_lo)]) + + +def boundary_points(domain: Domain, n: int, *, seed: int) -> Array: + """Sample ``n`` points split across the two spatial boundaries. + + :param domain: ``((x_min, x_max), (t_min, t_max))``. + :param n: Total number of points (half per boundary). + :param seed: RNG seed. + :returns: Array of shape ``(n, 2)`` with ``x`` at ``x_min`` or ``x_max``. + """ + gen = seeding.rng(seed) + (x_lo, x_hi), (t_lo, t_hi) = domain + n_left = n // 2 + t_left = _uniform(gen, t_lo, t_hi, n_left) + t_right = _uniform(gen, t_lo, t_hi, n - n_left) + left = np.column_stack([np.full(n_left, x_lo), t_left]) + right = np.column_stack([np.full(n - n_left, x_hi), t_right]) + return np.vstack([left, right]) + + +def eval_grid(domain: Domain, nx: int, nt: int) -> tuple[Array, Array]: + """Return the structured evaluation grid axes ``(x_values, t_values)``. + + :param domain: ``((x_min, x_max), (t_min, t_max))``. + :param nx: Number of spatial samples. + :param nt: Number of temporal samples. + :returns: ``(x_values, t_values)`` of shapes ``(nx,)`` and ``(nt,)``. + """ + (x_lo, x_hi), (t_lo, t_hi) = domain + return np.linspace(x_lo, x_hi, nx), np.linspace(t_lo, t_hi, nt) + + +def observations( + field: Array, + x_values: Array, + t_values: Array, + *, + n_obs: int, + noise_std: float, + seed: int, +) -> tuple[Array, Array]: + """Draw sparse, noisy observations from a reference field. + + Emulates scattered probe / loop-detector data: ``n_obs`` grid cells are + sampled without replacement and their values perturbed by Gaussian noise. + + :param field: Reference field of shape ``(len(t_values), len(x_values))``. + :param x_values: Spatial grid axis. + :param t_values: Temporal grid axis. + :param n_obs: Number of observations. + :param noise_std: Standard deviation of additive Gaussian noise. + :param seed: RNG seed. + :returns: ``(coords, values)`` with ``coords`` of shape ``(n_obs, 2)`` + (columns ``[x, t]``) and ``values`` of shape ``(n_obs,)``. + """ + gen = seeding.rng(seed) + nt, nx = field.shape + flat = gen.choice(nt * nx, size=n_obs, replace=False) + ti, xi = np.divmod(flat, nx) + coords = np.column_stack([x_values[xi], t_values[ti]]) + values = field[ti, xi] + noise_std * gen.standard_normal(n_obs) + return coords, values diff --git a/applications/pinn/experiments/__init__.py b/applications/pinn/experiments/__init__.py new file mode 100644 index 00000000..ebe73dff --- /dev/null +++ b/applications/pinn/experiments/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CLI entry points: single runs, Optuna search, and the full sweep.""" + +from __future__ import annotations diff --git a/applications/pinn/experiments/config.py b/applications/pinn/experiments/config.py new file mode 100644 index 00000000..c35d52a6 --- /dev/null +++ b/applications/pinn/experiments/config.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Configuration for a single PINN experiment run. + +A ``RunConfig`` fully determines one ``(problem, method, backend, seed)`` run: +architecture, point counts, optimisation, loss weights, and the evaluation grid. +It is JSON-round-trippable (stdlib dataclasses) for reproducibility. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from applications.pinn.models.protocol import Backend, Method, ModelConfig + + +@dataclass(frozen=True, slots=True) +class RunConfig: + """Everything needed to run and score one experiment. + + :param problem: Registry key of the problem. + :param method: One of the four methods. + :param backend: ``"torch"`` or ``"jax"``. + :param seed: Seed for model init and point sampling. + :param model: Architecture configuration. + :param n_collocation: Interior collocation points. + :param n_ic: Initial-condition points (forward tier). + :param n_bc: Boundary points (forward tier). + :param steps: Optimisation steps. + :param lr: Adam learning rate. + :param residual_weight: Residual loss weight. + :param ic_weight: IC loss weight. + :param bc_weight: BC loss weight. + :param soft_penalty: Monotonicity-penalty weight for the ``soft`` method + (ignored otherwise; the whole point is that only ``soft`` uses it). + :param eval_nx: Evaluation-grid spatial resolution. + :param eval_nt: Evaluation-grid temporal resolution. + """ + + problem: str + method: Method + backend: Backend + seed: int = 0 + model: ModelConfig = field(default_factory=ModelConfig) + n_collocation: int = 2000 + n_ic: int = 256 + n_bc: int = 128 + steps: int = 2000 + lr: float = 1e-3 + residual_weight: float = 1.0 + ic_weight: float = 10.0 + bc_weight: float = 1.0 + soft_penalty: float = 1.0 + eval_nx: int = 200 + eval_nt: int = 50 diff --git a/applications/pinn/experiments/run.py b/applications/pinn/experiments/run.py new file mode 100644 index 00000000..5974e2be --- /dev/null +++ b/applications/pinn/experiments/run.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run and score a single PINN experiment. + +Orchestrates: build model (backend + method) -> assemble training data -> +train -> predict on the evaluation grid -> score against ground truth. Returns a +plain-dict artifact (JSON-serialisable) with the headline metrics, including the +admissibility violation that is ``0`` by construction for the hard-monotone model. +""" + +from __future__ import annotations + +import argparse +from typing import TYPE_CHECKING, Any + +import numpy as np +import numpy.typing as npt + +from applications.pinn.core import metrics, sampling +from applications.pinn.core.admissibility import violation +from applications.pinn.core.problems import get +from applications.pinn.models.protocol import ModelConfig, build +from applications.pinn.training.losses import LossWeights, TrainingData + +if TYPE_CHECKING: + from applications.pinn.experiments.config import RunConfig + +Array = npt.NDArray[np.floating] + + +def _ground_truth(problem: object, x: Array, t: Array) -> Array: + """Evaluate a problem's ground truth, requiring it to be available.""" + value = problem.ground_truth(x, t) # type: ignore[attr-defined] + if value is None: + raise ValueError("problem has no ground truth for scoring") + result: Array = value + return result + + +def _predict(model: object, coords: Array, backend: str) -> Array: + """Evaluate a trained model on ``coords`` (N, 2), returning a NumPy vector.""" + x = coords[:, 0:1] + t = coords[:, 1:2] + if backend == "torch": + import torch + + xt = torch.as_tensor(x, dtype=torch.float32) + tt = torch.as_tensor(t, dtype=torch.float32) + return model(xt, tt).detach().numpy().ravel() # type: ignore[operator,no-any-return] + import jax.numpy as jnp + + out = model(jnp.asarray(x), jnp.asarray(t)) # type: ignore[operator] + return np.asarray(out).ravel() + + +def _train( + problem: object, model: object, data: TrainingData, cfg: RunConfig +) -> object: + weights = LossWeights( + residual=cfg.residual_weight, + ic=cfg.ic_weight, + bc=cfg.bc_weight, + mono=cfg.soft_penalty if cfg.method == "soft" else 0.0, + ) + sign_x = int(problem.admissibility().mask[0]) # type: ignore[attr-defined] + kwargs: dict[str, Any] = { + "weights": weights, + "sign_x": sign_x, + "lr": cfg.lr, + "steps": cfg.steps, + } + if cfg.backend == "torch": + from applications.pinn.training import torch_trainer + + trained, _ = torch_trainer.train(problem, model, data, **kwargs) # type: ignore[arg-type] + return trained + from applications.pinn.training import jax_trainer + + trained_jax, _ = jax_trainer.train(problem, model, data, **kwargs) # type: ignore[arg-type] + return trained_jax + + +def run_one(cfg: RunConfig) -> dict[str, Any]: + """Run one experiment and return its metrics artifact. + + :param cfg: The fully-specified run configuration. + :returns: A JSON-serialisable dict of configuration + headline metrics. + """ + problem = get(cfg.problem)() + domain = problem.domain + model = build(problem, cfg.model, cfg.method, cfg.backend) + + collocation = sampling.collocation(domain, cfg.n_collocation, seed=cfg.seed) + ic_pts = sampling.initial_points(domain, cfg.n_ic, seed=cfg.seed + 1) + ic_vals = problem.initial(ic_pts[:, 0]) + bc_pts = sampling.boundary_points(domain, cfg.n_bc, seed=cfg.seed + 2) + bc_vals = _ground_truth(problem, bc_pts[:, 0], bc_pts[:, 1]) + data = TrainingData( + collocation=collocation, ic=(ic_pts, ic_vals), bc=(bc_pts, bc_vals) + ) + + trained = _train(problem, model, data, cfg) + + x_values, t_values = sampling.eval_grid(domain, cfg.eval_nx, cfg.eval_nt) + grid_x, grid_t = np.meshgrid(x_values, t_values) + coords = np.column_stack([grid_x.ravel(), grid_t.ravel()]) + pred = _predict(trained, coords, cfg.backend).reshape(cfg.eval_nt, cfg.eval_nx) + ref = _ground_truth(problem, grid_x.ravel(), grid_t.ravel()).reshape( + cfg.eval_nt, cfg.eval_nx + ) + + sign_x = int(problem.admissibility().mask[0]) + dx = float(x_values[1] - x_values[0]) + viol = max(violation(pred[i], axis=0, sign=sign_x) for i in range(cfg.eval_nt)) + over = max(metrics.overshoot(pred[i], ref[i]) for i in range(cfg.eval_nt)) + return { + "problem": cfg.problem, + "method": cfg.method, + "backend": cfg.backend, + "seed": cfg.seed, + "l1": metrics.l1(pred, ref, dx=dx), + "l2": metrics.l2(pred, ref, dx=dx), + "admissibility_violation": viol, + "overshoot": over, + } + + +def main() -> None: + """CLI: run one experiment and print its metrics artifact as JSON.""" + import json + + from applications.pinn.experiments.config import RunConfig + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("problem") + parser.add_argument("method") + parser.add_argument("backend", choices=["torch", "jax"]) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--steps", type=int, default=2000) + args = parser.parse_args() + cfg = RunConfig( + problem=args.problem, + method=args.method, + backend=args.backend, + seed=args.seed, + steps=args.steps, + model=ModelConfig(), + ) + print(json.dumps(run_one(cfg), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/applications/pinn/experiments/search.py b/applications/pinn/experiments/search.py new file mode 100644 index 00000000..ace13739 --- /dev/null +++ b/applications/pinn/experiments/search.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optuna hyperparameter search with an equal budget across methods. + +The headline claim ("hard beats soft") is only credible if every method is tuned +equally hard. So all methods share the **identical search space and trial +budget**; the soft baseline additionally searches its penalty weight (it must be +tuned, not fixed). The objective is the L2 error of the run against ground truth. +Best configs are frozen to ``configs/`` for the sweep to consume. +""" + +from __future__ import annotations + +import json +from dataclasses import replace +from pathlib import Path +from typing import TYPE_CHECKING + +import optuna + +from applications.pinn.experiments.config import RunConfig +from applications.pinn.experiments.run import run_one +from applications.pinn.models.protocol import ModelConfig + +if TYPE_CHECKING: + from applications.pinn.models.protocol import Backend, Method + +optuna.logging.set_verbosity(optuna.logging.WARNING) + + +def _config_from_trial(trial: optuna.Trial, base: RunConfig) -> RunConfig: + """Build a RunConfig from a trial's suggested hyperparameters.""" + lr = trial.suggest_float("lr", 1e-4, 1e-2, log=True) + width = trial.suggest_categorical("width", [16, 32, 64]) + ic_weight = trial.suggest_float("ic_weight", 1.0, 100.0, log=True) + soft_penalty = ( + trial.suggest_float("soft_penalty", 1e-2, 1e1, log=True) + if base.method == "soft" + else base.soft_penalty + ) + return replace( + base, + lr=lr, + ic_weight=ic_weight, + soft_penalty=soft_penalty, + model=replace(base.model, width=width), + ) + + +def search( + problem: str, + method: Method, + backend: Backend, + *, + n_trials: int, + template: RunConfig | None = None, + seed: int = 0, +) -> dict[str, float]: + """Tune hyperparameters for one ``(problem, method, backend)`` cell. + + :param problem: Problem registry key. + :param method: Method to tune (identical space across methods; soft adds its + penalty weight). + :param backend: Backend to tune on. + :param n_trials: Trial budget (identical across methods). + :param template: Base config (steps, point counts, eval grid). Defaults to a + moderate config for the given problem/method/backend. + :param seed: Sampler seed. + :returns: The best hyperparameters found. + """ + base = template or RunConfig(problem=problem, method=method, backend=backend) + base = replace(base, problem=problem, method=method, backend=backend) + + def objective(trial: optuna.Trial) -> float: + cfg = _config_from_trial(trial, base) + return float(run_one(cfg)["l2"]) + + study = optuna.create_study( + direction="minimize", sampler=optuna.samplers.TPESampler(seed=seed) + ) + study.optimize(objective, n_trials=n_trials) + best: dict[str, float] = dict(study.best_params) + return best + + +def freeze( + problem: str, + method: Method, + backend: Backend, + best: dict[str, float], + *, + configs_dir: str | Path, +) -> Path: + """Write a tuned config to ``configs/__.json``. + + :param problem: Problem key. + :param method: Method. + :param backend: Backend. + :param best: Best hyperparameters from :func:`search`. + :param configs_dir: Directory to write into. + :returns: The path written. + """ + payload = { + "problem": problem, + "method": method, + "backend": backend, + "best_params": best, + } + path = Path(configs_dir) / f"{problem}_{method}_{backend}.json" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + return path + + +def tuned_config( + problem: str, method: Method, backend: Backend, path: Path +) -> RunConfig: + """Reconstruct a RunConfig from a frozen tuned-config file. + + :param problem: Problem key. + :param method: Method. + :param backend: Backend. + :param path: Path to the frozen JSON. + :returns: A RunConfig with the tuned hyperparameters applied. + """ + best = json.loads(Path(path).read_text())["best_params"] + base = RunConfig(problem=problem, method=method, backend=backend) + return replace( + base, + lr=float(best.get("lr", base.lr)), + ic_weight=float(best.get("ic_weight", base.ic_weight)), + soft_penalty=float(best.get("soft_penalty", base.soft_penalty)), + model=ModelConfig(width=int(best.get("width", base.model.width))), + ) diff --git a/applications/pinn/experiments/sweep.py b/applications/pinn/experiments/sweep.py new file mode 100644 index 00000000..bad6ef91 --- /dev/null +++ b/applications/pinn/experiments/sweep.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run the full experiment matrix and collect metrics. + +Enumerates ``problem x method x backend x seed`` (consuming tuned configs when +available) and returns one metrics artifact per cell. The heavy execution is a +``slow``, RUNBOOK-documented step; this module just composes ``run_one``. +""" + +from __future__ import annotations + +from dataclasses import replace +from itertools import product +from typing import TYPE_CHECKING, Any + +from applications.pinn.experiments.config import RunConfig +from applications.pinn.experiments.run import run_one + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + + from applications.pinn.models.protocol import Backend, Method + + +def build_matrix( + problems: Sequence[str], + methods: Sequence[Method], + backends: Sequence[Backend], + seeds: Sequence[int], + *, + template: RunConfig | None = None, +) -> list[RunConfig]: + """Enumerate the experiment matrix as a list of RunConfigs. + + :param problems: Problem keys. + :param methods: Methods. + :param backends: Backends. + :param seeds: Seeds. + :param template: Base config for shared settings (steps, grids, ...). + :returns: One RunConfig per cell. + """ + base = template or RunConfig( + problem=problems[0], method=methods[0], backend=backends[0] + ) + return [ + replace(base, problem=p, method=m, backend=b, seed=s) + for p, m, b, s in product(problems, methods, backends, seeds) + ] + + +def run_matrix(configs: Iterable[RunConfig]) -> list[dict[str, Any]]: + """Run every config and collect its metrics artifact. + + :param configs: RunConfigs to execute. + :returns: One artifact dict per config. + """ + return [run_one(cfg) for cfg in configs] diff --git a/applications/pinn/models/__init__.py b/applications/pinn/models/__init__.py new file mode 100644 index 00000000..c05d2520 --- /dev/null +++ b/applications/pinn/models/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-backend model builders (JAX, PyTorch) behind one protocol.""" + +from __future__ import annotations diff --git a/applications/pinn/models/jax/__init__.py b/applications/pinn/models/jax/__init__.py new file mode 100644 index 00000000..0450c4f7 --- /dev/null +++ b/applications/pinn/models/jax/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +"""JAX (Flax NNX) model builders. Importing this module requires JAX.""" + +from __future__ import annotations diff --git a/applications/pinn/models/jax/builders.py b/applications/pinn/models/jax/builders.py new file mode 100644 index 00000000..d1ed8e1f --- /dev/null +++ b/applications/pinn/models/jax/builders.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +"""JAX (Flax NNX) model builders for the four PINN methods. + +Mirrors the PyTorch builders: a callable ``u(x, t)`` on column vectors, with the +monotone-in-``x`` / free-in-``t`` structure from embedding ``t`` through an +unconstrained MLP and feeding ``[x, h(t)]`` to a stack monotone in all inputs. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx + +from mononet.core.types import ActivationSpec, MonotonicityMask +from mononet.jax import MonoInput, MonoLinear, MonoResidual + +if TYPE_CHECKING: + from applications.pinn.models.protocol import Method, ModelConfig + +_ACTS = {"tanh": jnp.tanh, "elu": jax.nn.elu, "softplus": jax.nn.softplus} + + +class _PlainMLP(nnx.Module): + """Unconstrained two-hidden-layer MLP.""" + + def __init__( + self, in_dim: int, width: int, out_dim: int, activation: str, *, rngs: nnx.Rngs + ) -> None: + self.l1 = nnx.Linear(in_dim, width, rngs=rngs) + self.l2 = nnx.Linear(width, width, rngs=rngs) + self.l3 = nnx.Linear(width, out_dim, rngs=rngs) + self.act = _ACTS[activation] + + def __call__(self, x: jax.Array) -> jax.Array: + """Apply the MLP.""" + x = self.act(self.l1(x)) + x = self.act(self.l2(x)) + return self.l3(x) + + +class _ClampedLinear(nnx.Module): + """Linear layer with non-negative weights (inexpressive monotone map).""" + + def __init__(self, in_features: int, out_features: int, *, rngs: nnx.Rngs) -> None: + key = rngs.params() + scale = (2.0 / (in_features + out_features)) ** 0.5 + self.w = nnx.Param(jax.random.normal(key, (in_features, out_features)) * scale) + self.b = nnx.Param(jnp.zeros((out_features,))) + + def __call__(self, x: jax.Array) -> jax.Array: + """Apply ``x @ |W| + b``.""" + return x @ jnp.abs(self.w.get_value()) + self.b.get_value() + + +class VanillaMLP(nnx.Module): + """Unconstrained MLP over ``[x, t]`` (also the ``soft`` architecture).""" + + def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs) -> None: + """Build the unconstrained MLP.""" + self.net = _PlainMLP(2, cfg.width, 1, cfg.plain_activation, rngs=rngs) + + def __call__(self, x: jax.Array, t: jax.Array) -> jax.Array: + """Return ``u(x, t)``.""" + return self.net(jnp.concatenate([x, t], axis=-1)) + + +class WeightClipMono(nnx.Module): + """Monotone-in-``x`` net via non-negative weights (inexpressive baseline).""" + + def __init__(self, cfg: ModelConfig, sign_x: int, *, rngs: nnx.Rngs) -> None: + """Build the clamped-weight monotone stack and the ``t`` embedding.""" + self.sign_x = float(sign_x) + self.t_embed = _PlainMLP( + 1, cfg.t_embed_width, cfg.t_embed_dim, cfg.plain_activation, rngs=rngs + ) + in_dim = 1 + cfg.t_embed_dim + self.c1 = _ClampedLinear(in_dim, cfg.width, rngs=rngs) + self.c2 = _ClampedLinear(cfg.width, cfg.width, rngs=rngs) + self.c3 = _ClampedLinear(cfg.width, 1, rngs=rngs) + + def __call__(self, x: jax.Array, t: jax.Array) -> jax.Array: + """Return ``u(x, t)`` (monotone in ``x`` via non-negative weights).""" + z = jnp.concatenate([self.sign_x * x, self.t_embed(t)], axis=-1) + z = jax.nn.softplus(self.c1(z)) + z = jax.nn.softplus(self.c2(z)) + return self.c3(z) + + +class HardMonoField(nnx.Module): + """Expressive hard-monotone-in-``x`` field built from ``mononet`` layers.""" + + def __init__(self, cfg: ModelConfig, sign_x: int, *, rngs: nnx.Rngs) -> None: + """Build the mononet monotone stack and the unconstrained ``t`` embedding.""" + self.t_embed = _PlainMLP( + 1, cfg.t_embed_width, cfg.t_embed_dim, cfg.plain_activation, rngs=rngs + ) + in_dim = 1 + cfg.t_embed_dim + mask = MonotonicityMask( + np.array([sign_x, *([1] * cfg.t_embed_dim)], dtype=np.int8) + ) + self.mono_input = MonoInput(mask) + act = ActivationSpec(cfg.mono_activation) # type: ignore[arg-type] + self.n_blocks = cfg.n_blocks + # nnx tracks Module-valued attributes; a plain list is not a pytree node, + # so register each block under its own attribute name. + self.block0 = MonoResidual( + in_dim, + cfg.width, + mode=cfg.mode, + activation=act, + rngs=rngs, # type: ignore[arg-type] + ) + for i in range(1, cfg.n_blocks): + setattr( + self, + f"block{i}", + MonoResidual( + cfg.width, cfg.width, mode=cfg.mode, activation=act, rngs=rngs + ), # type: ignore[arg-type] + ) + self.head = MonoLinear( + cfg.width, 1, mode=cfg.mode, activation="identity", rngs=rngs + ) # type: ignore[arg-type] + + def __call__(self, x: jax.Array, t: jax.Array) -> jax.Array: + """Return ``u(x, t)`` (monotone in ``x`` by construction).""" + z = jnp.concatenate([x, self.t_embed(t)], axis=-1) + z = self.mono_input(z) + for i in range(self.n_blocks): + z = getattr(self, f"block{i}")(z) + return self.head(z) + + +def build_jax(problem: object, cfg: ModelConfig, method: Method) -> nnx.Module: + """Build a JAX (Flax NNX) model for ``method``. + + :param problem: A registered `Problem` (its admissibility mask sets ``sign_x``). + :param cfg: Architecture configuration. + :param method: One of ``vanilla`` / ``soft`` / ``weight_clip`` / ``hard_monotone``. + :returns: A Flax NNX module callable as ``u(x, t)``. + :raises ValueError: If ``method`` is unknown. + """ + rngs = nnx.Rngs(cfg.seed) + sign_x = int(problem.admissibility().mask[0]) # type: ignore[attr-defined] + if method in ("vanilla", "soft"): + return VanillaMLP(cfg, rngs=rngs) + if method == "weight_clip": + return WeightClipMono(cfg, sign_x, rngs=rngs) + if method == "hard_monotone": + return HardMonoField(cfg, sign_x, rngs=rngs) + raise ValueError(f"unknown method {method!r}") diff --git a/applications/pinn/models/protocol.py b/applications/pinn/models/protocol.py new file mode 100644 index 00000000..34284473 --- /dev/null +++ b/applications/pinn/models/protocol.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Backend-agnostic model configuration and the build dispatcher. + +A PINN model is a callable ``u(x, t)`` mapping two column vectors to one. The +four *methods* differ only in architecture (the loss is identical across them, +except the soft baseline's penalty, which lives in the trainer): + +- ``vanilla`` — unconstrained MLP. +- ``soft`` — same unconstrained MLP; monotonicity via a loss penalty. +- ``weight_clip`` — inexpressive hard-monotone net (non-negative weights). +- ``hard_monotone`` — expressive hard-monotone net built from ``mononet`` layers. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + from mononet.core.config import Mode + +Method = Literal["vanilla", "soft", "weight_clip", "hard_monotone"] +Backend = Literal["torch", "jax"] + + +@dataclass(frozen=True, slots=True) +class ModelConfig: + """Architecture hyperparameters shared by all methods and backends. + + Kept deliberately shallow (a few ``MonoResidual`` blocks): the repo's + depth-vs-scale finding is that width, not depth, is the capacity lever. + + :param width: Hidden width of the monotone / plain stack. + :param n_blocks: Number of residual blocks (``~4`` layers total with the + head). + :param t_embed_dim: Dimension of the unconstrained ``t`` embedding that makes + the field free (non-monotone) in ``t``. + :param t_embed_width: Hidden width of the ``t`` embedding MLP. + :param mono_activation: Smooth activation for the monotone stack. + :param plain_activation: Activation for unconstrained sub-networks. + :param mode: ``mononet`` construction mode (``"absolute"`` or ``"switch"``). + :param seed: Seed for parameter initialisation. + """ + + width: int = 32 + n_blocks: int = 2 + t_embed_dim: int = 8 + t_embed_width: int = 32 + mono_activation: str = "softplus" + plain_activation: str = "tanh" + mode: Mode = "absolute" + seed: int = 0 + + +def build( + problem: object, cfg: ModelConfig, method: Method, backend: Backend +) -> object: + """Dispatch to the backend-specific model builder. + + :param problem: A registered `Problem` (supplies the ``x`` monotonicity sign). + :param cfg: Architecture configuration. + :param method: One of the four methods. + :param backend: ``"torch"`` or ``"jax"``. + :returns: A backend-native model callable ``u(x, t)``. + :raises NotImplementedError: If the backend builder is not yet available. + """ + if backend == "torch": + from applications.pinn.models.torch import builders as torch_builders + + return torch_builders.build_torch(problem, cfg, method) + if backend == "jax": + from applications.pinn.models.jax import builders as jax_builders + + return jax_builders.build_jax(problem, cfg, method) + raise NotImplementedError(f"backend {backend!r} builder not yet implemented") diff --git a/applications/pinn/models/torch/__init__.py b/applications/pinn/models/torch/__init__.py new file mode 100644 index 00000000..38d50f62 --- /dev/null +++ b/applications/pinn/models/torch/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PyTorch model builders. Importing this module requires PyTorch.""" + +from __future__ import annotations diff --git a/applications/pinn/models/torch/builders.py b/applications/pinn/models/torch/builders.py new file mode 100644 index 00000000..f53d878b --- /dev/null +++ b/applications/pinn/models/torch/builders.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PyTorch model builders for the four PINN methods. + +Every model is a callable ``u(x, t)`` on column vectors. The monotone-in-``x`` / +free-in-``t`` structure is achieved by embedding ``t`` through an *unconstrained* +MLP and feeding ``[x, h(t)]`` to a stack that is monotone-increasing in all +inputs: ``x`` gets its sign via the mask, while ``monotone(arbitrary(t))`` is +arbitrary in ``t``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import torch +from torch import nn + +from mononet.core.types import ActivationSpec, MonotonicityMask +from mononet.torch import MonoInput, MonoLinear, MonoResidual + +if TYPE_CHECKING: + from applications.pinn.models.protocol import Method, ModelConfig + + +def _plain_mlp(in_dim: int, width: int, out_dim: int, activation: str) -> nn.Module: + act = {"tanh": nn.Tanh, "elu": nn.ELU, "softplus": nn.Softplus}[activation] + return nn.Sequential( + nn.Linear(in_dim, width), + act(), + nn.Linear(width, width), + act(), + nn.Linear(width, out_dim), + ) + + +class _ClampedLinear(nn.Module): + """Linear layer with non-negative weights (inexpressive monotone map).""" + + def __init__(self, in_features: int, out_features: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(out_features, in_features)) + self.bias = nn.Parameter(torch.zeros(out_features)) + nn.init.xavier_uniform_(self.weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply ``x @ |W|^T + b`` (weights forced non-negative).""" + out: torch.Tensor = torch.nn.functional.linear(x, self.weight.abs(), self.bias) + return out + + +class VanillaMLP(nn.Module): + """Unconstrained MLP over ``[x, t]`` (also the ``soft`` architecture).""" + + def __init__(self, cfg: ModelConfig) -> None: + """Build the unconstrained MLP.""" + super().__init__() + self.net = _plain_mlp(2, cfg.width, 1, cfg.plain_activation) + + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Return ``u(x, t)``.""" + out: torch.Tensor = self.net(torch.cat([x, t], dim=-1)) + return out + + +class WeightClipMono(nn.Module): + """Monotone-in-``x`` net via non-negative weights (inexpressive baseline).""" + + def __init__(self, cfg: ModelConfig, sign_x: int) -> None: + """Build the clamped-weight monotone stack and the ``t`` embedding.""" + super().__init__() + self.sign_x = float(sign_x) + self.t_embed = _plain_mlp( + 1, cfg.t_embed_width, cfg.t_embed_dim, cfg.plain_activation + ) + in_dim = 1 + cfg.t_embed_dim + self.stack = nn.Sequential( + _ClampedLinear(in_dim, cfg.width), + nn.Softplus(), + _ClampedLinear(cfg.width, cfg.width), + nn.Softplus(), + _ClampedLinear(cfg.width, 1), + ) + + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Return ``u(x, t)`` (monotone in ``x`` via non-negative weights).""" + z = torch.cat([self.sign_x * x, self.t_embed(t)], dim=-1) + out: torch.Tensor = self.stack(z) + return out + + +class HardMonoField(nn.Module): + """Expressive hard-monotone-in-``x`` field built from ``mononet`` layers.""" + + def __init__(self, cfg: ModelConfig, sign_x: int) -> None: + """Build the mononet monotone stack and the unconstrained ``t`` embedding.""" + super().__init__() + self.t_embed = _plain_mlp( + 1, cfg.t_embed_width, cfg.t_embed_dim, cfg.plain_activation + ) + in_dim = 1 + cfg.t_embed_dim + mask = MonotonicityMask( + np.array([sign_x, *([1] * cfg.t_embed_dim)], dtype=np.int8) + ) + self.mono_input = MonoInput(mask) + act = ActivationSpec(cfg.mono_activation) # type: ignore[arg-type] + blocks: list[nn.Module] = [ + MonoResidual(in_dim, cfg.width, mode=cfg.mode, activation=act) # type: ignore[arg-type] + ] + blocks.extend( + MonoResidual(cfg.width, cfg.width, mode=cfg.mode, activation=act) # type: ignore[arg-type] + for _ in range(cfg.n_blocks - 1) + ) + blocks.append(MonoLinear(cfg.width, 1, mode=cfg.mode, activation="identity")) # type: ignore[arg-type] + self.mono = nn.Sequential(*blocks) + + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Return ``u(x, t)`` (monotone in ``x`` by construction).""" + z = torch.cat([x, self.t_embed(t)], dim=-1) + out: torch.Tensor = self.mono(self.mono_input(z)) + return out + + +def build_torch(problem: object, cfg: ModelConfig, method: Method) -> nn.Module: + """Build a torch model for ``method`` given a problem's monotonicity sign. + + :param problem: A registered `Problem` (its admissibility mask sets ``sign_x``). + :param cfg: Architecture configuration. + :param method: One of ``vanilla`` / ``soft`` / ``weight_clip`` / ``hard_monotone``. + :returns: A torch ``nn.Module`` with ``forward(x, t)``. + :raises ValueError: If ``method`` is unknown. + """ + torch.manual_seed(cfg.seed) + sign_x = int(problem.admissibility().mask[0]) # type: ignore[attr-defined] + if method in ("vanilla", "soft"): + return VanillaMLP(cfg) + if method == "weight_clip": + return WeightClipMono(cfg, sign_x) + if method == "hard_monotone": + return HardMonoField(cfg, sign_x) + raise ValueError(f"unknown method {method!r}") diff --git a/applications/pinn/notebooks/.gitkeep b/applications/pinn/notebooks/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/applications/pinn/paper/figures/.gitkeep b/applications/pinn/paper/figures/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/applications/pinn/paper/paper.md b/applications/pinn/paper/paper.md new file mode 100644 index 00000000..46581bab --- /dev/null +++ b/applications/pinn/paper/paper.md @@ -0,0 +1,255 @@ +# Structure-Preserving Physics-Informed Neural Networks: Hard Monotonicity as a Total-Variation-Diminishing Prior for Conservation Laws + +**Authors:** Davor Runje (and collaborators — TBD) +**Status:** Draft scaffold — narrative complete; numerical results are placeholders (`[[…]]`) pending the experimental runs (see `RUNBOOK.md`). + +> **Scaffold note (remove before submission).** This manuscript is written +> *before* the experiments per the project's manuscript-first workflow. Every +> claim about *what will be shown* is committed here so the implementation only +> has to fill in numbers. Placeholders are written as `[[NAME]]` and each has a +> matching planned experiment. If a result contradicts a claim below, the claim +> is revised to match the evidence — not the other way around. + +## Abstract + +Physics-Informed Neural Networks (PINNs) fail on hyperbolic conservation laws +with shocks: as unconstrained continuous approximators they exhibit the Gibbs +phenomenon, producing spurious oscillations that violate the total-variation- +diminishing (TVD) property and the entropy condition. The standard remedy — +adding entropy/total-variation *penalty* terms to the loss — turns training into +an unstable multi-objective problem and offers no guarantee: the network can +still oscillate. We instead move the physics into the *architecture*. Using an +expressive hard-monotonicity construction (`mononet`), we build PINNs whose +output is monotone in the spatial coordinate **by construction**, so that within +the monotone-solution class the reconstruction is TVD and entropy-admissible +with **zero** achievable oscillation — a structural guarantee, not a penalty. +We first validate the mechanism on forward scalar conservation laws against exact +and TVD finite-volume ground truth, where the admissibility violation is +`[[FWD-VIOLATION-HARD]]` for our method versus `[[FWD-VIOLATION-SOFT]]` for a +tuned soft-penalty baseline. We then apply it where a PINN genuinely beats a +classical solver — the *inverse* problem of traffic state estimation from sparse, +noisy observations — and show reconstruction error `[[INV-ERROR-HARD]]` versus +`[[INV-ERROR-SOFT]]` at `[[OBS-SPARSITY]]` observation coverage, with the +guarantee holding across backends (JAX and PyTorch agree to `[[XBACKEND-TOL]]`). +Code and the trained models are released as part of the `mononet` package. + +## 1. Introduction + +Hyperbolic conservation laws, `u_t + f(u)_x = 0`, model transport-dominated +phenomena and develop discontinuous *shocks* in finite time even from smooth +initial data. Standard PINNs approximate the solution with an unconstrained +network trained on the PDE residual plus initial/boundary terms. Near a shock +this fails characteristically: the network overshoots and oscillates (the Gibbs +phenomenon), the total variation of the profile increases, and the recovered +weak solution can violate the entropy condition — i.e. it is *unphysical*. + +The prevailing fix is *soft*: add penalty terms that discourage total-variation +growth or entropy violation. This has two structural defects. (i) It converts +training into a multi-objective problem in which the physics residual, the data, +and the admissibility penalty compete; the penalty weight is a fragile +hyperparameter. (ii) It offers **no guarantee** — a soft penalty is minimized, +not enforced, so the trained network can and does still oscillate wherever the +optimizer trades penalty for residual. + +**This paper moves the admissibility condition from the loss into the +architecture.** For a scalar conservation law whose entropy solution is monotone +in space, TVD and the (Oleinik) entropy condition are *equivalent to a +monotonicity statement about the solution*. We therefore parameterize the PINN +with a network that is monotone in the spatial coordinate **by construction**, +using the expressive constrained-monotonic construction of `mononet` +[@runje2023constrained; @sartor2025advancing] — expressive because, unlike naive +weight-clipping or convex-only monotone nets, its convex/concave activation split +is a universal approximator of monotone functions. A monotone profile cannot +oscillate; hence within the monotone-solution class the reconstruction is TVD +and entropy-admissible with zero achievable oscillation, independent of the +optimizer. + +**Contributions.** +1. A **structure-preserving PINN framework**: an *admissibility abstraction* that + maps a PDE to a monotonicity/convexity constraint on the solution field, and a + *problem registry* over a backend-agnostic core with JAX and PyTorch trainers. +2. A **theorem** (Section 4) that, within the monotone-solution class, a + spatially-monotone network is TVD with total variation fixed by the boundary + states, and satisfies the Oleinik entropy condition — by construction. +3. **Mechanism validation** on forward scalar conservation laws against exact and + TVD finite-volume references, isolating the architectural effect from sampling + and tuning (Section 6.1). +4. The **flagship result**: on *inverse* traffic state estimation from sparse, + noisy data — a regime where classical finite-volume solvers do not apply + because the full initial condition is unknown — the hard-monotone PINN + reconstructs the shock/queue front without oscillation and dominates + unconstrained and soft-penalty PINNs as observations become sparse and noisy + (Section 6.2). +5. A **cross-backend** guarantee: identical results (to tolerance) from JAX and + PyTorch, dogfooding `mononet`'s cross-backend equivalence. + +**Scope.** We restrict to the *monotone-solution class* of scalar conservation +laws — problems whose entropy solution is monotone in `x` for all `t` (monotone +initial data producing a shock; Riemann problems; single-front queue formation). +This class is exactly where the guarantee is unconditional. Non-monotone data +(e.g. `u_0 = -\sin\pi x`), rarefactions, N-waves, and multi-front congestion are +out of scope and discussed in Section 7. + +## 2. Related work + +**PINNs for conservation laws.** [@raissi2019pinn] introduced PINNs; their +failure at shocks is well documented, and remedies add entropy/viscosity or +total-variation *penalties* to the loss (e.g. weak-form entropy-stable PINNs +[@wepinn]) or use derivative/shape penalties with adaptive weighting [@dcpinn]. +These are *soft*: admissibility is minimized, not guaranteed. We enforce it +architecturally. + +**Traffic state estimation with PINNs.** [@shi2021physics] estimate traffic state +from sparse data with a PINN informed by second-order traffic models, using soft +physics losses. This is the closest prior art to our flagship; it exhibits the +oscillation-at-shock failure a hard-monotone architecture removes. + +**Shape-constrained networks.** Hard monotonicity/convexity by architecture has a +long lineage — non-negative-weight + convex-activation nets and Input Convex +Neural Networks [@amos2017icnn] — but these are the *inexpressive convex-only* +regime and pay a documented accuracy penalty. `mononet` +[@runje2023constrained; @sartor2025advancing] is expressive (convex/concave +activation split, a universal approximator of monotone functions), and HardNet +[@hardnet2024] shows hard-constrained nets retain universal approximation — +refuting the "hard constraints kill expressiveness" premise of soft-penalty work. +To our knowledge no prior PINN embeds *expressive* hard monotonicity in the +architecture for conservation laws. + +**High-dimensional neural PDE solvers.** Deep BSDE [@han2018solving] and the Deep +Galerkin Method [@sirignano2018dgm] target high-dimensional PDEs with +unconstrained networks and impose no shape structure. Structure-preserving +extensions to those settings (HJB, Fokker–Planck, eikonal) are the subject of +follow-up papers. + +## 3. Method + +### 3.1 Structure-preserving PINN + +For a scalar conservation law `u_t + f(u)_x = 0` on `[a,b] × [0,T]`, we +parameterize the solution field `u_θ(x, t)` by a network built from `mononet` +layers with a monotonicity mask that constrains `u_θ` to be non-increasing in `x` +(sign chosen per scenario) and unconstrained in `t`. The network is a shallow +stack of `MonoResidual` blocks (≈4 layers; depth beyond a few layers does not +help this construction, so width is the capacity lever). + +### 3.2 Admissibility abstraction and problem registry + +Each PDE is a plug-in exposing: the residual, an `AdmissibilitySpec` (the +monotonicity mask + convex/concave axes), the initial/boundary conditions or the +observation operator (inverse mode), and an exact-or-reference ground truth. The +admissibility-violation functional — the total wrong-sign spatial-derivative mass +— is `0` for the hard-monotone model by construction and positive otherwise; it +is the paper's headline metric. The framework is backend-agnostic (pure-NumPy +core) with thin JAX (Flax NNX) and PyTorch trainers. + +### 3.3 Losses and baselines + +The training loss is standard PINN — residual + initial/boundary (forward) or +residual + data-fit (inverse) — **with no entropy/TV penalty**; the architecture +supplies admissibility, so every method shares an identical loss and differs only +in the network. Baselines: **vanilla** (unconstrained MLP), **soft** +(unconstrained + TV/entropy penalty, penalty weight tuned), **weight-clip** +(inexpressive non-negative-weight monotone net), and **hard-monotone** (ours). + +## 4. Theory + +Let `u_θ(·, t)` be continuous and non-increasing in `x` on `[a, b]`. + +**Total variation.** `TV(u_θ(·,t)) = ∫_a^b |∂_x u_θ| dx = u_θ(a,t) − u_θ(b,t)`. +With bounded boundary states the total variation is fixed by those states; the +network cannot create new local extrema, so it is TVD (non-increasing in `t` when +the boundary states are non-increasing in `t`) and exhibits **zero** overshoot. + +**Entropy condition.** For a convex flux `f`, the Oleinik condition across a shock +requires `u_L > s > u_R`, i.e. `∂_x u ≤ 0` at the front. Constraining the +hypothesis space to non-increasing `u_θ` restricts it to the entropy-admissible +set; entropy-violating expansion shocks are unrepresentable. + +**Scope of the guarantee.** The statements hold within the monotone-solution +class. A continuous monotone network represents a shock as a steep but continuous +ramp (zero overshoot), not a true discontinuity; Section 6 quantifies the ramp +width versus grid/observation resolution. *(Precise statements and proofs to be +finalized here; cross-check with the Lean formalization track where applicable.)* + +## 5. Experimental design + +**Forward mechanism tier (Section 6.1).** Burgers-Riemann, Burgers smooth→shock, +linear advection, LWR (Greenshields flux). Ground truth: closed-form entropy +solutions and a TVD Godunov reference. All four methods, both backends. Metrics: +L¹/L² error, admissibility violation, TV(t), near-shock overshoot, shock +speed/position error, mass conservation. + +**Inverse flagship (Section 6.2).** LWR traffic state estimation: reconstruct a +monotone-front density field from sparse, noisy scattered observations (probe / +loop-detector geometry), no full initial condition. Sweep over observation +sparsity and noise. Metrics: reconstruction L¹/L² vs sparsity/noise, admissibility +violation, front-position error. + +**Protocol.** Hyperparameters are tuned per method with Optuna under an +**identical search space and trial budget** (the soft baseline's penalty weight is +searched, not fixed), so comparisons are not confounded by tuning. Sampling and +observation masks are seeded and identical across methods. Reported numbers are +multi-seed. + +## 6. Results + +> All values below are placeholders pending the runs; each has a matching planned +> experiment in `RUNBOOK.md`. + +### 6.1 Forward mechanism tier + +`[[TABLE-forward-tier]]` — per problem × method: L¹/L² error, admissibility +violation, overshoot, shock-speed error. *Expected:* hard-monotone violation and +overshoot ≈ 0; vanilla/soft show Gibbs overshoot `[[OVERSHOOT-SOFT]]`. + +`[[FIG-tv-curve]]` — TV(t) for each method on Burgers-Riemann. *Expected:* +flat/non-increasing for hard-monotone; bumps for baselines. + +`[[FIG-profiles]]` — solution profiles at `t = [[T-SNAP]]` near the shock. +*Expected:* clean monotone ramp (ours) vs ringing (baselines). + +### 6.2 Inverse flagship — traffic state estimation + +`[[FIG-inverse-sweep]]` — reconstruction L² vs observation sparsity and noise, all +methods. *Expected:* hard-monotone degrades gracefully; soft/vanilla degrade and +oscillate as data thins. + +`[[TABLE-inverse]]` — reconstruction error, admissibility violation, front-position +error at representative (sparsity, noise) operating points. + +`[[FIG-inverse-field]]` — reconstructed `ρ(x,t)` vs reference field with +observations overlaid, at `[[OBS-SPARSITY]]`. + +### 6.3 Cross-backend equivalence + +`[[TABLE-xbackend]]` — JAX vs PyTorch hard-monotone agreement (max abs difference +`[[XBACKEND-TOL]]`) on identical points. + +## 7. Discussion and limitations + +- **Monotone-solution class.** The guarantee is unconditional *within* this class + and empty outside it: non-monotone initial data, rarefactions, N-waves, and + multi-front congestion cannot be represented by a globally monotone field. This + is a genuine restriction, stated up front, not a caveat discovered late. +- **Continuous ramp, not a true jump.** The shock is a steep continuous ramp; + Section 6 reports its width. This is adequate (and oscillation-free) for the + targeted applications but is not a sharp-interface method. +- **Forward tier is a mechanism check, not a solver competitor.** A TVD + finite-volume scheme dominates forward 1-D problems; the forward tier exists to + validate the constraint against exact ground truth. Relevance comes from the + inverse setting, where no such solver applies. +- **Outlook.** The same primitive extends to other PDEs whose admissibility is a + monotonicity/convexity statement — high-dimensional HJB (with the domain- + restriction and expressiveness caveats of the de-risking analysis), + Fokker–Planck (valid-density-by-construction), and eikonal (causality) — each a + follow-up paper. + +## Acknowledgments / Reproducibility + +Built on `mononet` (`pip install mononet`); all experiments reproduce via +`applications/pinn/RUNBOOK.md`. Code, configs, and trained-model artifacts are in +the repository. + +## References + +See `references.bib`. diff --git a/applications/pinn/paper/references.bib b/applications/pinn/paper/references.bib new file mode 100644 index 00000000..383bbc54 --- /dev/null +++ b/applications/pinn/paper/references.bib @@ -0,0 +1,119 @@ +% References for "Structure-Preserving PINNs" (Paper 1). +% Entries marked `note = {VERIFY}` were surfaced during the 2026 research pass +% and need a citation check before submission (see spec §9). + +@article{raissi2019pinn, + author = {Raissi, Maziar and Perdikaris, Paris and Karniadakis, George E.}, + title = {Physics-informed neural networks: A deep learning framework for + solving forward and inverse problems involving nonlinear partial + differential equations}, + journal = {Journal of Computational Physics}, + volume = {378}, + pages = {686--707}, + year = {2019}, + doi = {10.1016/j.jcp.2018.10.045} +} + +@article{han2018solving, + author = {Han, Jiequn and Jentzen, Arnulf and E, Weinan}, + title = {Solving high-dimensional partial differential equations using + deep learning}, + journal = {Proceedings of the National Academy of Sciences}, + volume = {115}, + number = {34}, + pages = {8505--8510}, + year = {2018}, + doi = {10.1073/pnas.1718942115} +} + +@article{sirignano2018dgm, + author = {Sirignano, Justin and Spiliopoulos, Konstantinos}, + title = {{DGM}: A deep learning algorithm for solving partial differential + equations}, + journal = {Journal of Computational Physics}, + volume = {375}, + pages = {1339--1364}, + year = {2018}, + doi = {10.1016/j.jcp.2018.08.029} +} + +@book{leveque2002finite, + author = {LeVeque, Randall J.}, + title = {Finite Volume Methods for Hyperbolic Problems}, + publisher = {Cambridge University Press}, + year = {2002}, + doi = {10.1017/CBO9780511791253} +} + +@inproceedings{runje2023constrained, + author = {Runje, Davor and Shankaranarayana, Sharath M.}, + title = {Constrained Monotonic Neural Networks}, + booktitle = {Proceedings of the 40th International Conference on Machine + Learning (ICML)}, + year = {2023}, + eprint = {2205.11775}, + archivePrefix = {arXiv} +} + +@inproceedings{sartor2025advancing, + author = {Sartor, Davide and others}, + title = {Advancing Constrained Monotonic Neural Networks}, + booktitle = {Proceedings of the 42nd International Conference on Machine + Learning (ICML)}, + year = {2025}, + eprint = {2505.02537}, + archivePrefix = {arXiv}, + note = {VERIFY author list} +} + +@article{shi2021physics, + author = {Shi, Rongye and Mo, Zhaobin and Di, Xuan}, + title = {Physics-informed deep learning for traffic state estimation: A + hybrid paradigm informed by second-order traffic models}, + journal = {Proceedings of the AAAI Conference on Artificial Intelligence}, + year = {2021}, + note = {VERIFY volume/pages} +} + +@inproceedings{amos2017icnn, + author = {Amos, Brandon and Xu, Lei and Kolter, J. Zico}, + title = {Input Convex Neural Networks}, + booktitle = {Proceedings of the 34th International Conference on Machine + Learning (ICML)}, + year = {2017}, + eprint = {1609.07152}, + archivePrefix = {arXiv} +} + +@misc{hardnet2024, + title = {{HardNet}: Hard-Constrained Neural Networks with Universal + Approximation Guarantees}, + year = {2024}, + eprint = {2410.10807}, + archivePrefix = {arXiv}, + note = {VERIFY authors/venue} +} + +@misc{lighthill1955kinematic, + author = {Lighthill, M. J. and Whitham, G. B.}, + title = {On kinematic waves. II. A theory of traffic flow on long crowded + roads}, + year = {1955}, + note = {LWR model; Proc. R. Soc. London A 229(1178):317--345} +} + +% --- Soft-penalty / shape-constrained PINN prior art (VERIFY before citing) --- + +@misc{wepinn, + title = {Weak-form entropy-stable PINNs for hyperbolic conservation laws + (WE-PINN)}, + year = {2025}, + note = {VERIFY exact title/authors/venue (spec §9)} +} + +@misc{dcpinn, + title = {Derivative-constrained PINNs with self-adaptive loss balancing + (DC-PINN)}, + year = {2026}, + note = {VERIFY exact title/authors/venue (spec §9)} +} diff --git a/applications/pinn/results/.gitkeep b/applications/pinn/results/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/applications/pinn/studies/.gitkeep b/applications/pinn/studies/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/applications/pinn/training/__init__.py b/applications/pinn/training/__init__.py new file mode 100644 index 00000000..bc7ec983 --- /dev/null +++ b/applications/pinn/training/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Loss specs and per-backend training loops.""" + +from __future__ import annotations diff --git a/applications/pinn/training/jax_trainer.py b/applications/pinn/training/jax_trainer.py new file mode 100644 index 00000000..84fef11b --- /dev/null +++ b/applications/pinn/training/jax_trainer.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +"""JAX (optax) training loop for the PINN methods. + +The model is split into a pure ``params`` pytree via ``nnx.split`` so that +``jax.grad`` cleanly provides (a) the input derivatives ``u_x, u_t`` that build +the PDE residual and (b) the parameter gradients of the total loss. First-order +scalar conservation laws need only first derivatives, so no Hessian is taken. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import jax +import jax.numpy as jnp +import optax +from flax import nnx + +if TYPE_CHECKING: + from applications.pinn.training.losses import LossWeights, TrainingData + + +def train( + problem: object, + model: nnx.Module, + data: TrainingData, + *, + weights: LossWeights, + sign_x: int, + lr: float = 1e-3, + steps: int = 200, +) -> tuple[nnx.Module, list[float]]: + """Train ``model`` with optax Adam on the PINN loss; return it and the history. + + :param problem: Registered problem (supplies ``flux_prime`` for the residual). + :param model: A Flax NNX model callable as ``u(x, t)``. + :param data: Collocation and (IC/BC or observation) point sets. + :param weights: Loss-term weights. + :param sign_x: Desired monotonicity sign in ``x`` (for the soft penalty). + :param lr: Adam learning rate. + :param steps: Number of optimisation steps. + :returns: ``(trained_model, loss_history)``. + """ + graphdef, params, nontrain = nnx.split(model, nnx.Param, ...) # type: ignore[misc] + + xc = jnp.asarray(data.collocation[:, 0]) + tc = jnp.asarray(data.collocation[:, 1]) + supervised = { + name: (jnp.asarray(c), jnp.asarray(v)) + for name, pair in (("ic", data.ic), ("bc", data.bc), ("obs", data.obs)) + if pair is not None + for c, v in [pair] + } + term_weight = {"ic": weights.ic, "bc": weights.bc, "obs": weights.data} + + def apply(p: Any, x: jax.Array, t: jax.Array) -> jax.Array: + out: jax.Array = nnx.merge(graphdef, p, nontrain)(x, t) + return out + + def scalar_u(p: Any, x: jax.Array, t: jax.Array) -> jax.Array: + return apply(p, x.reshape(1, 1), t.reshape(1, 1))[0, 0] + + du_dx = jax.grad(scalar_u, argnums=1) + du_dt = jax.grad(scalar_u, argnums=2) + + def loss_fn(p: Any) -> jax.Array: + u = jax.vmap(scalar_u, in_axes=(None, 0, 0))(p, xc, tc) + u_x = jax.vmap(du_dx, in_axes=(None, 0, 0))(p, xc, tc) + u_t = jax.vmap(du_dt, in_axes=(None, 0, 0))(p, xc, tc) + residual = u_t + problem.flux_prime(u) * u_x # type: ignore[attr-defined] + loss = weights.residual * jnp.mean(residual**2) + if weights.mono > 0.0: + loss += weights.mono * jnp.mean(jnp.maximum(-sign_x * u_x, 0.0) ** 2) + for name, (coords, values) in supervised.items(): + pred = apply(p, coords[:, 0:1], coords[:, 1:2]).ravel() + loss += term_weight[name] * jnp.mean((pred - values) ** 2) + return loss + + optimizer = optax.adam(lr) + opt_state = optimizer.init(params) + + @jax.jit + def step(p: Any, o: Any) -> tuple[Any, Any, jax.Array]: + loss, grads = jax.value_and_grad(loss_fn)(p) + updates, o = optimizer.update(grads, o, p) + return optax.apply_updates(p, updates), o, loss + + history: list[float] = [] + for _ in range(steps): + params, opt_state, loss = step(params, opt_state) + history.append(float(loss)) + + return nnx.merge(graphdef, params, nontrain), history diff --git a/applications/pinn/training/losses.py b/applications/pinn/training/losses.py new file mode 100644 index 00000000..680c3139 --- /dev/null +++ b/applications/pinn/training/losses.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Backend-agnostic loss configuration and training-data container. + +The PDE residual and its derivatives are backend-specific (they use autodiff), +so they live in the per-backend trainers. What is shared — and identical across +methods and backends — is the *weighting* of the loss terms and the *data* the +loss is evaluated on. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import numpy.typing as npt + +Array = npt.NDArray[np.floating] +Supervised = tuple[Array, Array] # (coords (N, 2) [x, t], values (N,)) + + +@dataclass(frozen=True, slots=True) +class LossWeights: + """Relative weights of the PINN loss terms. + + :param residual: PDE-residual term. + :param ic: Initial-condition term (forward tier). + :param bc: Boundary-condition term (forward tier). + :param data: Observation data-fit term (inverse tier). + :param mono: Soft monotonicity-penalty weight — **nonzero only for the + ``soft`` baseline**; the architecture supplies monotonicity otherwise. + """ + + residual: float = 1.0 + ic: float = 1.0 + bc: float = 1.0 + data: float = 1.0 + mono: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class TrainingData: + """Point sets a PINN trains on (NumPy; each backend converts to its tensors). + + :param collocation: Interior points ``(N, 2)`` where the residual is enforced. + :param ic: Optional ``(coords, values)`` on the initial line (forward tier). + :param bc: Optional ``(coords, values)`` on the boundaries (forward tier). + :param obs: Optional ``(coords, values)`` sparse observations (inverse tier). + """ + + collocation: Array + ic: Supervised | None = None + bc: Supervised | None = None + obs: Supervised | None = None diff --git a/applications/pinn/training/torch_trainer.py b/applications/pinn/training/torch_trainer.py new file mode 100644 index 00000000..a47cccd4 --- /dev/null +++ b/applications/pinn/training/torch_trainer.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PyTorch (Adam) training loop for the PINN methods. + +Mirrors the JAX trainer. Input derivatives ``u_x, u_t`` for the PDE residual come +from ``torch.autograd.grad`` with ``create_graph=True`` so the residual stays in +the graph and its parameter gradients flow. First-order scalar conservation laws +need only first derivatives. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import nn + + from applications.pinn.training.losses import LossWeights, TrainingData + + +def _col(values: object) -> torch.Tensor: + """Return a float32 column tensor from a NumPy array.""" + return torch.as_tensor(values, dtype=torch.float32).reshape(-1, 1) + + +def train( + problem: object, + model: nn.Module, + data: TrainingData, + *, + weights: LossWeights, + sign_x: int, + lr: float = 1e-3, + steps: int = 200, +) -> tuple[nn.Module, list[float]]: + """Train ``model`` with Adam on the PINN loss; return it and the loss history. + + :param problem: Registered problem (supplies ``flux_prime`` for the residual). + :param model: A torch model callable as ``u(x, t)``. + :param data: Collocation and (IC/BC or observation) point sets. + :param weights: Loss-term weights. + :param sign_x: Desired monotonicity sign in ``x`` (for the soft penalty). + :param lr: Adam learning rate. + :param steps: Number of optimisation steps. + :returns: ``(trained_model, loss_history)``. + """ + xc = _col(data.collocation[:, 0]).requires_grad_(True) + tc = _col(data.collocation[:, 1]).requires_grad_(True) + supervised = [ + (_col(pair[0][:, 0]), _col(pair[0][:, 1]), _col(pair[1]), w) + for pair, w in ( + (data.ic, weights.ic), + (data.bc, weights.bc), + (data.obs, weights.data), + ) + if pair is not None + ] + + optimizer = torch.optim.Adam(model.parameters(), lr=lr) + history: list[float] = [] + for _ in range(steps): + optimizer.zero_grad() + u = model(xc, tc) + (u_x,) = torch.autograd.grad(u.sum(), xc, create_graph=True) + (u_t,) = torch.autograd.grad(u.sum(), tc, create_graph=True) + residual = u_t + problem.flux_prime(u) * u_x # type: ignore[attr-defined] + loss = weights.residual * (residual**2).mean() + if weights.mono > 0.0: + loss = loss + weights.mono * torch.relu(-sign_x * u_x).pow(2).mean() + for cx, ct, cv, w in supervised: + loss = loss + w * (model(cx, ct) - cv).pow(2).mean() + loss.backward() + optimizer.step() + history.append(float(loss.detach())) + + return model, history diff --git a/docs/superpowers/plans/2026-07-12-applications-structure-preserving-pinns-paper1.md b/docs/superpowers/plans/2026-07-12-applications-structure-preserving-pinns-paper1.md new file mode 100644 index 00000000..66c28822 --- /dev/null +++ b/docs/superpowers/plans/2026-07-12-applications-structure-preserving-pinns-paper1.md @@ -0,0 +1,203 @@ +# Structure-Preserving PINNs — Paper 1 Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Spec:** [`docs/superpowers/specs/2026-07-12-applications-structure-preserving-pinns-design.md`](../specs/2026-07-12-applications-structure-preserving-pinns-design.md). Read it first; this plan implements **Paper 1 only**. + +**Goal:** Stand up the `applications/` area and the `applications/pinn/` package: the framework (admissibility abstraction, problem registry, backend-agnostic NumPy `core/`, per-backend JAX+Torch trainers, Optuna HP search), the **forward conservation-law mechanism tier** (Burgers-Riemann, Burgers smooth→shock, linear advection, LWR), and the **deep inverse flagship** (traffic state estimation from sparse noisy data). Produce a Markdown manuscript, an executed notebook, and a RUNBOOK — with the manuscript **scaffolded before any implementation**. + +**Architecture:** `applications/` is repo-only, out of the `mononet` wheel. `core/` is pure NumPy (single source of truth); framework code lives only in `models/{jax,torch}` and `training/{jax,torch}_trainer.py`. Each PDE is a plug-in `core/problems/*.py` module behind one `Problem` protocol. Both tiers share the field `u_θ(x,t)` (shallow `MonoResidual`, ≈4 layers) and differ only in the loss. + +**Tech Stack:** Python 3.11+, NumPy, JAX + Flax NNX + optax, PyTorch, Optuna (v4.9.0, already present), matplotlib, pytest, Sphinx/myst-nb. `mononet.jax` / `mononet.torch` layers (Sub-project A, locked API). + +## Global Constraints + +- **No `mononet` package/kernel/layer change.** Consume the locked public API only (`MonoLinear`/`MonoDense`, `MonoResidual`, `MonoInput`, `MonotonicityMask`, `MonoConfig`, `MonoResidualConfig`). If a genuine gap appears, STOP and raise it — do not patch the wheel from here. +- **Lazy backend imports preserved.** `import mononet` must not import torch/jax/keras. Backend code imports `from mononet.jax …` / `from mononet.torch …` locally. +- **`core/` imports no ML framework** — NumPy only. This is what makes JAX and Torch numbers comparable and is asserted by a test. +- **Monotone-solution class only** (spec §4 non-goals). Every problem module documents its scope; non-monotone data is out. +- **Determinism:** all point/observation generation is seeded in NumPy and identical across backends and methods. +- Line length 88 (ruff); **strict mypy** — `applications/` is added to the mypy `files` list and must pass. MyST field-list docstrings on all public functions/classes. +- Backends optional in CI: use `pytest.importorskip`; select with `MONONET_TEST_BACKEND`. Heavy training/search marked `slow` and excluded from default CI. +- **Branch:** `spec/applications-structure-preserving-pinns` (already checked out). Commit per task. **Never commit to `main`.** Pre-commit must pass (no `--no-verify`); locale-only errors → prefix `LC_ALL=C.UTF-8 LANG=C.UTF-8`. +- **Devcontainer strategy** (only `default` has both backends; `gpu-jax`=JAX-only, `gpu-torch`=torch-only): + - **Phases 0–4 (dev): `default`** (CPU, `MONONET_EXTRAS=all-cpu` → both backends). Writing/unit-testing both backends and the cross-backend equivalence test require both installed in one venv; smoke/unit tests are cheap on CPU. + - **Phases 5–6 (heavy Optuna search + sweeps): `gpu-jax` is the primary GPU backend** (JAX + `jit` for the repeated residual/Hessian). Run the full search/sweep in JAX there. + - **Cross-backend result = equivalence test + reproducing the JAX-tuned config on Torch** (forward-equivalence + a confirmation run) — cheap, no second HP search. Do the Torch confirmation on `default` (CPU) or a short `gpu-torch` session; do **not** attempt a combined torch+jax GPU venv (CUDA 13 vs 12.9 wheel conflicts). +- Tests live under `tests/applications/pinn/` mirroring the package. +- **Dependency strategy (interim + follow-up).** Application deps are opt-in uv + **dependency-groups** in the root `pyproject.toml`, mirroring `bench`: + `applications` (shared, e.g. matplotlib) and `pinn` (`{include-group = + "applications"}` + optax). Backends stay package extras; `dev` stays lean. Work + on the PINN app with `uv sync --extra jax --group pinn` (or `jax-gpu`). + Application tests `importorskip` matplotlib/optax **and** the backend, so the + default dev `uv run pytest` skips them cleanly. **Agreed follow-up:** migrate + each application to an *independent* uv project (own `pyproject`/`uv.lock`/`.venv`, + `mononet` via editable path source) so apps can't conflict — see the + `applications-dependency-strategy` memory. Single-backend venvs mean the + cross-backend equivalence test (3.2) runs only in an all-cpu (both-backends) sync. + +--- + +## Phase 0 — Area scaffold & package skeleton + +### Task 0.1: `applications/` area + `pinn/` package skeleton + +**Files:** +- Create: `applications/README.md`, `applications/__init__.py`, `applications/_common/{__init__,seeding,metrics_io,plot_theme}.py` +- Create: `applications/pinn/{__init__.py,README.md,RUNBOOK.md}` and package dirs `core/`, `core/problems/`, `models/{jax,torch}/`, `training/`, `configs/`, `experiments/`, `results/`, `notebooks/`, `paper/figures/`, `studies/` +- Modify: `pyproject.toml` (add `applications` to mypy `files`; add optax + flax to a dev/pinn group only if not already present — confirm before adding) +- Test: `tests/applications/pinn/test_skeleton.py` + +**Interfaces:** +- `_common/seeding.py`: `rng(seed: int) -> numpy.random.Generator`, `split_seeds(seed, n) -> list[int]`. +- `_common/metrics_io.py`: `write_result(path, obj: dict)` / `read_result(path) -> dict` (JSON). + +**Steps:** +- [ ] **Step 1 (RED):** `test_skeleton.py` asserts (a) `applications.pinn` importable without importing torch/jax (`sys.modules` check), (b) `seeding.rng(0)` reproducible, (c) `metrics_io` round-trips a dict. +- [ ] **Step 2 (GREEN):** create the tree + minimal module bodies with MyST docstrings. +- [ ] **Step 3:** README/RUNBOOK stubs; `applications/README.md` explains applications vs benchmarks and indexes the papers. +- [ ] **Step 4:** wire mypy/ruff; run `uv run ruff check`, `uv run mypy`, `uv run pytest tests/applications -q`. +- [ ] **Step 5:** commit `feat(pinn): applications/ area + pinn package skeleton`. + +> **REVIEW CHECKPOINT 0** — confirm dependency additions are acceptable and the lazy-import test passes. + +--- + +## Phase 1 — Markdown paper scaffold (NO implementation) — spec §11 + +Write everything that does not depend on results. **No code, no numbers — placeholders only.** + +### Task 1.1: `references.bib` +- [ ] Gather citations from spec §9/§9a: Raissi 2019; Han–Jentzen–E 2018; Sirignano–Spiliopoulos 2018; LeVeque 2002; Runje 2023 & Sartor 2025; Shi et al. 2021 (PINN-TSE); WE-PINN, DC-PINN; HardNet 2024; ICNN (Amos 2017). Mark unverified 2026 arXiv IDs `note = {VERIFY}`. +- [ ] commit `docs(pinn-paper): references.bib`. + +### Task 1.2: `paper.md` scaffold +Write in full except results (which get explicit placeholders): +- [ ] Title, author, abstract (headline number → `[[TVD-VIOLATION-NUMBER]]`). +- [ ] **1 Introduction** — soft-vs-hard constraints; the named PINN failure (Gibbs at shocks); thesis; contributions. +- [ ] **2 Related work** — grounded in §9/§9a: soft-penalty conservation/shape PINNs + PINN-TSE (Shi 2021), hard-but-inexpressive nets (Dugas/ICNN), expressiveness (HardNet, mononet). +- [ ] **3 Method** — admissibility abstraction; registry; `mononet` field, mask `x→−1`,`t→0`; shallow `MonoResidual`; loss terms; cross-backend. +- [ ] **4 Theory** — TVD/entropy-by-construction for the monotone-solution class (precise, with the class restriction); continuous-ramp shock representation. +- [ ] **5 Experimental design** — benchmarks (forward + inverse), baselines, metrics (§6), Optuna equal-budget (§6a), sparsity×noise sweep, cross-backend equivalence. +- [ ] **6 Results** — subsection headers + empty tables/figures with placeholders (`[[TABLE-forward-tier]]`, `[[FIG-inverse-sweep]]`, `[[FIG-tv-curve]]`, …), one-line captions saying what each will show. +- [ ] **7 Discussion / limitations** — monotone-class scope; forward tier is a mechanism check; pointer to Papers 2–5. +- [ ] Verify no orphan placeholder lacks a planned experiment; commit `docs(pinn-paper): manuscript scaffold with results placeholders`. + +> **REVIEW CHECKPOINT 1** — human reads the scaffolded manuscript; everything except numbers reads as a finished paper. Adjust framing here before writing code. + +--- + +## Phase 2 — `core/` backend-agnostic science (TDD, NumPy only) + +### Task 2.1: `Problem` protocol + registry +**Files:** `core/problems/base.py`, `core/problems/__init__.py`; Test `test_registry.py` +- [ ] RED: registry register/get/available; `get("nope")` raises. GREEN + commit. + +### Task 2.2: `admissibility.py` +**Files:** `core/admissibility.py`; Test `test_admissibility.py` +- [ ] RED: monotone-decreasing field → violation 0; a bump → violation = bump rise. GREEN + commit. + +### Task 2.3: `exact.py` — closed-form entropy solutions +**Files:** `core/exact.py`; Test `test_exact.py` +- Interfaces: `burgers_riemann(x,t,uL,uR)` (shock, `s=(uL+uR)/2`); `burgers_smooth_shock(x,t,u0)` (characteristics pre-`t_b`, R-H post; `t_b=-1/min(u0')`); `advection(x,t,a,u0)`; `lwr_riemann(x,t,ρL,ρR,flux)`. +- [ ] RED: Rankine–Hugoniot speed, self-similarity, monotonicity preserved, sampled values. GREEN + commit. + +### Task 2.4: `reference_solver.py` — TVD finite-volume (Godunov) +**Files:** `core/reference_solver.py`; Test `test_reference_solver.py` +- [ ] RED: converges to `burgers_riemann` under refinement (L¹↓ below tol); TV non-increasing; correct shock speed. GREEN + commit. + +### Task 2.5: `sampling.py` — deterministic point sets + sparse observations +**Files:** `core/sampling.py`; Test `test_sampling.py` +- Interfaces: `collocation`, `initial_points`, `boundary_points`, `eval_grid`, `observations(reference_field, grid, n_obs, noise_std, seed)`. +- [ ] RED: same seed → identical; obs count exact; noise reproducible; coords in domain. GREEN + commit. + +### Task 2.6: `conservation.py` problems (forward + inverse mode) +**Files:** `core/problems/conservation.py`; Test `test_conservation.py` +- `Burgers`, `LinearAdvection`, `LWR` (Greenshields `Q(ρ)=v_max ρ(1−ρ/ρ_max)`; document choice). +- [ ] RED: residual of exact solution ≈ 0 on interior (finite-diff); mask matches known monotone direction. GREEN + commit. + +### Task 2.7: `metrics.py` + `plotting.py` +**Files:** `core/metrics.py`, `core/plotting.py`; Test `test_metrics.py` +- `l1,l2,tv,tv_curve,overshoot,shock_position_error,mass_error,reconstruction_error`; plotting returns Figures (no I/O). +- [ ] RED: metrics on analytic inputs (monotone TV = endpoint diff; zero overshoot). GREEN + commit. + +> **REVIEW CHECKPOINT 2** — `core/` complete, framework-free (assert no torch/jax import), validated against closed forms. This is the ruler. + +--- + +## Phase 3 — Models & trainers (per-backend, TDD) + +### Task 3.1: model protocol + builders (both backends) +**Files:** `models/protocol.py`, `models/jax/*.py`, `models/torch/*.py`; Tests `test_builders_{jax,torch}.py` (`importorskip`) +- `build(problem, cfg, method)` for `method ∈ {vanilla, soft, weight_clip, hard_monotone}`; hard = shallow `MonoResidual` (≈4 layers) with the mask; baselines matched in depth/width. +- [ ] RED: forward-pass shape; `hard_monotone` monotone in `x` on a grid (finite-diff sign) per backend. GREEN + commit each. + +### Task 3.2: cross-backend forward-equivalence test +**Files:** `test_cross_backend.py` +- [ ] With identical ported weights, JAX vs Torch `hard_monotone` agree within tol; skip unless both importable. commit. + +### Task 3.3: losses + trainers +**Files:** `training/losses.py`, `training/jax_trainer.py`, `training/torch_trainer.py`; Tests `test_trainer_{jax,torch}.py` +- `losses.py` term specs (residual, IC, BC, data-fit). Trainers: `train(problem, method, cfg, points, backend)`; residual via `jax.grad/hessian` / `autograd.grad(create_graph=True)`. +- [ ] RED (smoke): few-step train on Burgers-Riemann reduces loss; `hard_monotone` violation 0 throughout. Full runs `slow`. GREEN + commit. + +> **REVIEW CHECKPOINT 3** — smoke train runs on both backends; monotone model provably non-oscillatory. + +--- + +## Phase 4 — Experiments, Optuna search, configs + +### Task 4.1: `experiments/run.py` (Typer CLI); Test `test_run_smoke.py` +- [ ] `run(problem, method, backend, seed, config)` → results JSON. RED: smoke run → valid artifact. GREEN + commit. + +### Task 4.2: `experiments/search.py` — Optuna, equal budget; Test `test_search_smoke.py` +- [ ] Reuse `benchmarks/_common/search.py` pattern. **Identical search space + `n_trials` across all methods**; soft penalty weight IS searched. Freeze best config per (problem, method, backend) to `configs/`. RED: 2-trial smoke → frozen config. GREEN + commit. + +### Task 4.3: `experiments/sweep.py` + configs; Test `test_sweep_smoke.py` +- [ ] Enumerate matrix (problem × method × backend × seed × obs-sparsity/noise) consuming tuned configs. RED: smoke 2-cell sweep. GREEN + commit. + +> **REVIEW CHECKPOINT 4** — machinery works at smoke scale; RUNBOOK updated with search+sweep commands. + +--- + +## Phase 5 — Forward mechanism tier (results) — `gpu-jax` + +### Task 5.1: run forward tier + references +- [ ] `slow`: Optuna search (equal budget) + multi-seed sweep for all four forward problems, all methods, both backends where feasible. Cache TVD references to `results//reference.npz` (committed). +- [ ] Assert headline: `hard_monotone` violation = 0, overshoot ≈ 0; baselines show Gibbs. commit results. + +--- + +## Phase 6 — Inverse flagship (traffic state estimation) (results) — `gpu-jax` + +### Task 6.1: inverse problem wiring +- [ ] Define monotone-front scenarios (queue behind bottleneck / red signal / incident — resolve spec §10; document). Observation model. RED: inverse loss uses observations, no full IC. GREEN + commit. + +### Task 6.2: sparsity × noise sweep + results +- [ ] `slow`: reconstruct from sparse noisy observations across the sparsity×noise grid, all methods, both backends. Reference = high-res Godunov field. +- [ ] Headline curves: reconstruction L¹/L² vs sparsity/noise (hard-monotone holds; soft/vanilla degrade+oscillate); violation → 0; front-position error; cross-backend equivalence. commit results. + +> **REVIEW CHECKPOINT 6** — the headline result exists and matches the scaffolded claims (or the claims get revised to match reality — evidence over assertion). + +--- + +## Phase 7 — Fill manuscript, notebook, docs + +### Task 7.1: fill results into `paper.md` +- [ ] Replace every `[[…]]` placeholder with real tables/numbers/figures. If a result contradicts a scaffolded claim, revise the claim. commit. + +### Task 7.2: notebook + Sphinx integration +- [ ] Executed notebook (loads committed results/references — no heavy training at build time). Add "Applications" to the Sphinx toctree. RED: `./tools/build-docs.sh` renders it. GREEN + commit. + +### Task 7.3: finalize RUNBOOK + README headline +- [ ] RUNBOOK reproduces every figure/table (search → sweep → fill → notebook). `pinn/README.md` gets abstract + headline number. commit. + +> **REVIEW CHECKPOINT 7 (final)** — full suite green (excluding `slow` in CI), `ruff`/`mypy` clean, docs build, manuscript complete. Then use superpowers:finishing-a-development-branch. + +--- + +## Out of scope (registry-ready, not implemented here) +- HJB (Paper 2, with §9a caveats), Fokker–Planck (Paper 3), eikonal (Paper 4): only the registry must admit them. +- Arbitrage-free surfaces (Paper 5): separate, non-PINN, own brainstorm. +- Adaptive resampling (RAR): optional `slow` ablation only, not the headline. diff --git a/docs/superpowers/specs/2026-07-12-applications-structure-preserving-pinns-design.md b/docs/superpowers/specs/2026-07-12-applications-structure-preserving-pinns-design.md new file mode 100644 index 00000000..b3af0aef --- /dev/null +++ b/docs/superpowers/specs/2026-07-12-applications-structure-preserving-pinns-design.md @@ -0,0 +1,448 @@ +# Applications program — Structure-Preserving PINNs (Paper 1 + program map) + +**Date:** 2026-07-12 +**Author:** Davor Runje +**Status:** Draft (brainstorming output); pending user review. +**Parent spec:** [`2026-05-21-mononet-package-design.md`](2026-05-21-mononet-package-design.md) +**Depends on:** [Sub-project A](2026-06-27-A-core-algorithm-and-backends-design.md) (locked public layer API: `MonoLinear`, `MonoResidual`, `MonoInput`, `MonotonicityMask`, `MonoConfig`) +**Primary references:** +- Runje & Shankaranarayana, *Constrained Monotonic Neural Networks*, ICML 2023 — +- Sartor et al., *Advancing Constrained Monotonic Neural Networks*, ICML 2025 — +- Raissi, Perdikaris & Karniadakis, *Physics-Informed Neural Networks*, JCP 2019 +- LeVeque, *Finite Volume Methods for Hyperbolic Problems*, 2002 (entropy/TVD, Oleinik condition) +- Shi et al., *Physics-informed deep learning for traffic state estimation*, 2021 (soft-constrained PINN-TSE — the baseline to beat) + +## 1. Context and the relevance problem + +This spec introduces a new top-level area, `applications/`, for downstream +research papers built **on** `mononet` as a library. It is distinct from +`benchmarks/`, which reproduces the `mononet` paper's own tables. Nothing under +`applications/` ships in the PyPI wheel. + +The seed idea proposed enforcing monotonicity **in the network architecture** so +that a PINN's solution is Total-Variation-Diminishing (TVD) and entropy-admissible +*by construction*, rather than via soft loss penalties that only approximate those +properties and turn training into an unstable multi-objective problem. + +Two findings from brainstorming reshaped the scope: + +1. **A forward 1-D problem alone is not relevant.** A TVD finite-volume solver + dominates it on speed and accuracy. Relevance requires problems where a PINN is + genuinely the right tool: high dimension (grids fail) or inverse/data- + assimilation (a solver needs the full IC; a PINN does not). +2. **The monotonicity idea generalizes beyond hyperbolic PDEs.** The same + architectural primitive enforces the *admissibility condition* of several PDE + families, each of which is a monotonicity or convexity statement. This is a + framework, not a single application. + +A prior-art pass (deep-research, 2026-07-12; recovered from the run journal after +the auto-synthesis stage degraded — see §9) confirmed the gap and sharpened it: + +- Conservation-law PINNs that target entropy/TV (WE-PINN) and shape-constrained + PINNs (DC-PINN), and PINN traffic-state estimators (Shi et al. 2021), enforce + these properties via **soft penalties**, not architecture. No published PINN + embeds hard monotonicity in the *architecture* for these PDEs. +- Existing *hard* shape-constrained pricing networks (Dugas et al. 2001; + Chataigner–Crépey "Deep Local Volatility" 2020; ICNN) are the **inexpressive + convex-only special case** — `mononet`'s degenerate `s = (m, 0, 0)` branch + (digest Corollary 3) — and pay a documented ~10× accuracy penalty. `mononet`'s + convex/concave activation split is precisely the escape from that trade-off. +- HardNet (arXiv:2410.10807) and `mononet` establish that hard-constrained nets + can be universal approximators, refuting the "hard constraints kill + expressiveness" position that soft-penalty papers (e.g. Ackerer et al.) rely on. + +## 2. Thesis and contribution + +**Thesis.** A single architectural primitive — `mononet`'s hard, *expressive* +monotonicity/convexity — is a structure-preserving inductive bias that enforces +the admissibility condition of many PDEs **by construction**, across +forward/inverse and low/high dimension. Because each admissibility condition is a +monotonicity or convexity statement, one mechanism covers them all. + +| PDE family | Admissibility condition | `mononet` constraint | +|---|---|---| +| Scalar conservation laws (Burgers, LWR, advection) | TVD + Oleinik entropy | monotone in `x` | +| HJB / stochastic control | comparison principle; admissible value function | monotone in state + convex | +| Fokker–Planck | valid density (non-negative, normalized) | monotone CDF | +| Eikonal | causality | monotone arrival time | + +**Contribution of Paper 1 (framework + inverse-conservation-law flagship):** +1. The **admissibility abstraction** — a small interface mapping each PDE to + (a) a `MonotonicityMask`/convexity spec, (b) an admissibility-violation metric, + (c) an exact-or-reference ground truth. +2. A **problem registry** so each PDE is a plug-in module over a backend-agnostic + core, with thin per-backend (JAX + PyTorch) trainers. +3. **Mechanism validation on the forward conservation-law tier** (exact ground + truth): hard-monotone PINN vs. vanilla / soft-penalty / weight-clipping + baselines and a TVD finite-volume + closed-form reference, cross-backend. The + credibility anchor — the constraint yields TVD/entropy to machine precision on + an exactly-solvable case *before* we deploy it on the harder inverse problem. +4. **The deep flagship — inverse scalar conservation laws / traffic state + estimation.** Reconstruct a monotone-front density field `ρ(x, t)` (a queue / + shock forming behind a bottleneck, incident, or signal — LWR traffic) from + **sparse, noisy observations** (scattered probe-vehicle / loop-detector data), + with hard monotonicity guaranteeing a **TVD, entropy-admissible, oscillation- + free reconstruction by construction**. This is the headline: an *inverse* + setting where PINNs genuinely beat classical solvers (a solver needs the full + initial + boundary data; the PINN assimilates scattered points), attacking the + most-cited named PINN failure — spurious oscillation at shocks — see §2a. + +### 2a. Why this shape, and why it is impactful + +Citation reality (Semantic Scholar, 2026-07-12): the PINN field is enormous +(Raissi et al. 2019 ≈ 18,166 cites) and therefore crowded — incremental soft-loss +variants vanish. What gets *noticed* is a **provable guarantee against a named +failure mode**, a result **where the PINN genuinely beats the classical +alternative**, and **reusable released code** (`mononet` is `pip install`-able). + +The inverse-conservation flagship hits all three. The named failure is concrete: +PINNs (and PINN-based traffic state estimators, e.g. Shi et al. 2021) produce +spurious Gibbs oscillations near shocks/queues; a monotone-by-construction field +*cannot* oscillate. The PINN-beats-solver case is genuine: classical finite-volume +schemes cannot run without the full initial/boundary conditions, whereas the +inverse problem provides only sparse scattered observations — exactly the data- +assimilation regime PINNs own. And unlike the HJB alternative we de-risked (see +§9a), the monotonicity guarantee here is **unconditional within a cleanly-stated +problem class** (the monotone-solution class: single-front / queue / Riemann +scenarios), not a domain-restricted half-space caveat. + +Hence: **one deep flagship, not a spread of shallow vignettes.** Breadth (HJB, +Fokker–Planck, eikonal) is delivered by the follow-up papers (§3), each at proper +depth. + +## 3. Program map (paper series) + +This spec fully designs **Paper 1** and records the map for the follow-ups. Each +follow-up gets its own brainstorm → spec → plan, reusing the framework. + +- **Paper 1 — Framework + inverse-conservation flagship** (this spec): primitive + + admissibility abstraction + registry + forward conservation-law mechanism tier + (full) + deep inverse traffic-state-estimation result. Audience: scientific ML / + numerical PDE / intelligent transport systems. +- **Paper 2 — High-dim HJB / stochastic control** (follow-up): admissible value + functions by construction. **Carries the de-risking caveats (§9a):** monotonicity + is domain-restricted (positive-orthant / non-negative state), the convex-HJB-by- + construction lane is partly occupied (Liu et al. 2023), so the defensible + contribution is the *conjunction* — joint monotone+convex, expressive beyond + ICNN, high-dimensional, on the HJB residual, with a correct derived policy. + Benchmark on naturally non-negative-state control (energy storage, inventory, + epidemic) plus LQ (convex) and Merton (concave, 1-D validation). +- **Paper 3 — Fokker–Planck**: valid-density-by-construction; eliminates the + negative-density failure of FP-PINNs; bridges to Sub-project D (flows). +- **Paper 4 — Eikonal / seismic traveltime**: causal-by-construction; geophysics + inverse problems; beats established eikonal-PINN baselines. +- **Paper 5 — Expressive arbitrage-free surfaces** (separate, **not a PINN**): + shape-constrained *regression* of price/IV surfaces. Different method and + audience (quant finance). Sketched in §8; own brainstorm later. + +## 4. Scope of Paper 1 + +### Goals +- Establish `applications/` (README, `_common`, conventions) as the home for + research papers built on `mononet`. +- Ship the framework: admissibility abstraction, problem registry, backend- + agnostic `core/`, per-backend trainers (JAX Flax NNX + PyTorch), Optuna search. +- Fully treat the **forward conservation-law mechanism tier** (exact ground + truth): Burgers-Riemann, Burgers smooth→shock, linear advection, LWR traffic. +- Deliver the **deep inverse flagship — traffic state estimation**: reconstruct a + monotone-front density field from sparse, noisy observations with a TVD/entropy- + admissible, oscillation-free result by construction, benchmarked against a + high-resolution TVD reference and against vanilla / soft-penalty PINN-TSE + baselines under varying observation sparsity and noise. +- Executed notebook rendered into the Sphinx "Applications" nav. + +### Non-goals +- No training loops, dataset loaders, or PINN code in the `mononet` wheel. + Everything here is application-local. +- **Monotone-solution class only.** Both tiers restrict to problems whose entropy + solution is monotone in `x` for all `t` (monotone initial data → shock; Riemann + problems; single-front queue formation). Stated up front as the problem class. + Explicitly out of scope and documented: non-monotone data (`u₀ = −sin πx`), + rarefactions, N-waves, multi-front congestion. A continuous monotone network + represents a shock as a steep-but-continuous ramp (zero overshoot), not a true + discontinuity. (This class restriction is *unconditional within the class* — + contrast the domain-restricted HJB monotonicity, §9a.) +- **No shallow vignettes.** HJB (Paper 2), Fokker–Planck (Paper 3) and eikonal + (Paper 4) are *not* touched in Paper 1 beyond the registry being designed to + admit them. Impact comes from depth on the inverse flagship, not breadth across + thin demos (§2a). +- The forward conservation-law tier is a **mechanism check**, not a relevance + claim: a finite-volume solver dominates forward 1-D scalar problems. Its role is + exact-ground-truth validation of the constraint, nothing more. Relevance comes + from the inverse flagship, where no solver competes. +- No adaptive/residual-based resampling (RAR) in the headline; optional, clearly + flagged ablation only (it would help every method and confound the architecture + comparison). + +## 5. Architecture + +### 5.1 Folder layout + +``` +applications/ +├── README.md # what applications are; index; contrast with benchmarks/ +├── _common/ # shared, minimal until a 2nd app needs it +│ ├── __init__.py +│ ├── seeding.py # deterministic RNG helpers +│ ├── metrics_io.py # JSON results read/write (reuse benchmarks conventions) +│ └── plot_theme.py # shared matplotlib theme +├── pinn/ # Paper 1: framework + forward tier + inverse flagship +│ ├── README.md # abstract + headline result +│ ├── RUNBOOK.md # exact commands to regenerate every figure/table +│ ├── paper/ # Markdown paper scaffold (written before implementation, §11) +│ │ ├── paper.md # full manuscript; results tables/figures filled in post-implementation +│ │ ├── figures/ # generated figures (committed; regenerable via RUNBOOK) +│ │ └── references.bib # citations gathered during brainstorming/research +│ ├── __init__.py +│ ├── core/ # backend-agnostic science (pure NumPy) — single source of truth +│ │ ├── admissibility.py # AdmissibilitySpec: mask, convexity/concavity flags, violation metric +│ │ ├── problems/ # problem registry (one plug-in module per PDE) +│ │ │ ├── __init__.py # registry: name -> Problem +│ │ │ ├── base.py # Problem protocol (residual, admissibility, IC/BC or observations, ground truth) +│ │ │ └── conservation.py # Burgers, linear advection, LWR — forward tier + inverse mode (Paper 1, full) +│ │ │ # hjb.py / fokker_planck.py / eikonal.py added by Papers 2 / 3 / 4 (registry ready; not in Paper 1) +│ │ ├── exact.py # closed-form entropy solutions (Riemann; characteristics + R-H) +│ │ ├── reference_solver.py # TVD finite-volume (Godunov) — forward ground truth + inverse reference field +│ │ ├── sampling.py # deterministic collocation / IC / BC / eval sets; sparse-noisy observation sampler +│ │ ├── metrics.py # L1/L2, admissibility violation, TV(t), overshoot, shock speed/position, mass +│ │ └── plotting.py # profiles, TV(t) curves, error heatmaps, observation overlays +│ ├── models/ +│ │ ├── protocol.py # PINNModel protocol: build(problem, cfg) -> callable u(x, t) +│ │ ├── jax/ # mononet.jax hard-monotone + vanilla + soft + weight-clip builders +│ │ └── torch/ # mononet.torch equivalents +│ ├── training/ +│ │ ├── losses.py # residual / IC / BC / data-fit term specs (backend-agnostic) +│ │ ├── jax_trainer.py # jax.grad/hessian residual; optax; jit +│ │ └── torch_trainer.py # autograd.grad(create_graph=True); torch optim +│ ├── configs/ # one JSON per (problem × method × backend × seed × obs-sparsity/noise) +│ ├── experiments/ +│ │ ├── run.py # CLI: one (problem, method, backend, seed) -> results artifact +│ │ ├── search.py # Optuna HP search (Typer CLI; mirrors benchmarks/_common/search.py) +│ │ └── sweep.py # full matrix, using per-method tuned configs +│ ├── results/ # committed metrics JSON + small figures; /reference.npz +│ ├── notebooks/ +│ │ └── structure-preserving-pinn.ipynb +│ └── studies/ # follow-up papers extend here (hjb/, fokker_planck/, eikonal/) +└── arbitrage/ # Paper 5 (separate; regression, not PINN) — sketch only +``` + +Key invariants: +- `core/` is pure NumPy and imports **no** framework. It is the single source of + truth that makes JAX and Torch numbers provably comparable. +- Framework code lives **only** in `models/{jax,torch}` and + `training/{jax,torch}_trainer.py`. Baselines are alternate model builders behind + one protocol. +- Adding a PDE = adding one registered `problems/*.py` module; no infrastructure + change. This is what makes the follow-up papers incremental. + +### 5.2 The admissibility abstraction + +`core/admissibility.py` defines, per problem: +- `mask: MonotonicityMask` — sign per input axis (e.g. `x → −1`, `t → 0`). +- `convex_axes` / `concave_axes` — which inputs the solution is convex/concave in + (drives `mononet`'s convex/concave activation split). +- `violation(u_grid) -> float` — a non-negative admissibility-violation measure + (e.g. total positive part of `∂u/∂x` for a non-increasing target). The headline + claim is that this is **0 by construction** for the hard-monotone model and + **> 0** for soft/vanilla baselines. + +### 5.3 Model construction + +The solution field `u_θ(x, t)` is built from `mononet` layers with mask `x → −1` +(non-increasing; or `+1` per scenario), `t → 0`. Both tiers use the same field; +they differ only in the loss (data term present in the inverse tier). + +**Architecture: `MonoResidual` blocks, shallow (≈4 layers total).** Use +`MonoResidual` blocks rather than a plain stack of `MonoLinear` layers, kept +deliberately shallow — on the order of 4 layers total. This follows the repo's own +depth-vs-scale finding that additional depth does not improve accuracy for this +construction (see `2026-07-05-deep-residual-accuracy-design.md`, +`2026-07-03-deep-monotonic-residual-design.md`, and the loan size-ladder +experiment); width/scale, not depth, is the lever. A shallow residual model also +keeps higher-order input derivatives (needed for the PDE residual) cheap and +well-conditioned. Depth is a config knob so the "depth doesn't help here either" +check can be reproduced as a small ablation, not re-litigated. + +Cross-backend: JAX (Flax NNX) and PyTorch, using the locked Sub-project A API. +Baselines (same across tiers): **vanilla** (unconstrained MLP), **soft** +(unconstrained + TV/entropy penalty — the strawman), **weight-clip** (non-negative- +weight + monotone-activation, the inexpressive hard baseline), **hard-monotone** +(`mononet`, proposed). The inverse tier is additionally positioned against +published PINN-TSE (traffic state estimation) practice, which is unconstrained/soft. + +### 5.4 Data generation + +PINN "data" is generated **deterministically from a seed** at run time (no large +committed files): +- **Forward tier:** interior collocation, IC points `(x, 0)`, BC points, dense + evaluation grid. +- **Inverse flagship:** interior collocation (for the residual) + a **sparse, + noisy observation set** — scattered `(x_k, t_k, ρ_k + ε)` points emulating + probe-vehicle / loop-detector data — and a dense evaluation grid for scoring. + **No** full IC/BC is given; that is the inverse setting. A sweep over observation + sparsity and noise level is the core experiment. + +- Points generated once in NumPy (`core/sampling.py`); each backend converts the + **same** arrays to its tensor type — JAX and Torch train on identical sets. +- Sampling and observation masks are **identical across all methods**, recorded in + config. The architecture, not the data, must do the work. +- Ground truth: closed-form entropy solution where available; otherwise a + high-resolution cached TVD finite-volume reference field + (`results//reference.npz`, small, committed, regenerable). Observations + are subsampled from this reference field + seeded noise. No `datasets/` directory. + +### 5.5 Constraint enforcement + +- **Primary:** soft loss terms (residual + data-fit for the inverse tier; residual + + IC + BC for the forward tier), **no** TV/entropy penalty — the architecture + supplies admissibility. Clean apples-to-apples: every method shares the loss and + differs only in architecture. +- **Ablation:** hard IC via an output ansatz `u = g + φ·N` where a monotonicity- + preserving construction exists (forward tier), for the stronger "no soft + constraints" story. + +## 6. Metrics + +Forward tier (vs. exact/reference): +- L¹ / L² error; **admissibility violation** (§5.2, headline: 0 for hard-monotone, + > 0 for baselines); TV(t) evolution and near-shock overshoot (Gibbs); shock + speed/position error; mass-conservation error. + +Inverse flagship (vs. reference field, over the sparsity × noise sweep): +- **Reconstruction L¹ / L² error** as a function of observation sparsity and noise + — the headline curves (where unconstrained/soft PINNs degrade and oscillate near + the front, and the hard-monotone model holds); +- **admissibility violation** (near-shock oscillation / positive-`∂ρ/∂x` mass) → 0 + by construction vs. > 0 for baselines; +- shock/queue front position error; robustness across seeds and observation masks. + +Cross-backend equivalence: JAX and Torch hard-monotone runs on identical points +agree within a documented tolerance — a secondary result dogfooding `mononet`'s +cross-backend guarantee. + +### 6a. Hyperparameter search and fair-comparison budget + +HP search uses **Optuna** (already a repo dependency, v4.9.0), reusing the +Phase-2a pattern (`benchmarks/_common/search.py` + a Typer CLI, `n_trials` budget, +seeded samplers). **Every method — vanilla, soft, weight-clip, hard-monotone — +gets the identical search space where applicable and the identical trial budget.** +This is load-bearing for the headline claim: "hard beats soft" is only credible if +the soft baseline was tuned at least as hard as the proposed model (the soft +penalty weight in particular must be searched, not fixed). Tuned per-method configs +are frozen to JSON and consumed by `sweep.py`; the search itself is a `slow`-marked, +RUNBOOK-documented step, not part of CI. + +## 7. Testing, docs, dependencies + +- **Testing.** `core/` unit-tested against closed forms (exact entropy solution; FV + solver on a Riemann problem with known speed) to fixed tolerance; `sampling.py` + determinism tests (same seed → identical arrays); a fast smoke-train per backend + in CI under `-m "not slow"`; full training runs marked `slow`. Backends selected + as elsewhere in the repo (`importorskip`). +- **Docs.** The notebook executes via myst-nb into a new Sphinx "Applications" nav + section. Sphinx + myst-nb is the source of truth (not MkDocs). +- **Dependencies.** optax, torch optim, and the FV solver are **application-local + dev dependencies**, never added to the `mononet` wheel. Optuna (v4.9.0) is already + a repo dev dependency — reuse it, do not add another HP-search library. Preserve + lazy backend imports. +- **Compute / devcontainer.** Develop on `default` (CPU, all backends — needed for + both-backend code + the equivalence test). Run the heavy Optuna search and sweeps + on **`gpu-jax`** (primary GPU backend; JAX `jit` for the repeated residual/Hessian + — CPU is too slow for the search). The cross-backend result is the equivalence + test plus reproducing the JAX-tuned config on Torch, not a second full search. No + combined torch+jax GPU venv (CUDA-wheel conflicts). + +## 8. Paper 5 sketch — arbitrage-free surfaces (separate, not a PINN) + +Recorded for completeness; own brainstorm later. Shape-constrained **regression** +(not PDE-residual): expressive exact monotone+convex fitting of price/IV surfaces +vs. soft-penalty prior art (Ackerer et al.) and inexpressive hard prior art +(Dugas/Chataigner–Crépey, Bernstein-QP). Model-free core constraint: call price +`C(K)` decreasing and convex in **strike** `K`, monotone in maturity `T`. Spot-space +`delta ∈ [0,1]` / `gamma ≥ 0` is a **model-dependent** variant (Black–Scholes / 1-D +diffusion), to be flagged as such — strike-space monotonicity/convexity is the +model-free theorem, spot-space Greeks are not (see §9 rigor correction). Metric of +record: arbitrage-violation count → 0 by construction vs. residual violations for +soft methods. + +## 9. Research provenance and rigor corrections + +The deep-research pass (2026-07-12) verified claims adversarially (22 confirmed, 3 +killed). Corrections folded into this spec: + +1. **Do not equate prior hard nets with `mononet`.** The Dugas-2001 / ICNN / + Chataigner–Crépey construction is the convex-only inexpressive special case, not + `mononet`'s expressive activation split. Cite them as prior lineage; the novelty + is expressiveness. +2. **Strike-space vs spot-space (Paper 5).** `C(K)` monotone-decreasing and convex + in strike is a model-free no-arbitrage theorem. `delta ∈ [0,1]`, `gamma ≥ 0`, + `vega ≥ 0` in spot are model-dependent (hold under Black–Scholes / 1-D diffusion; + can fail under jumps / correlated stochastic vol). Keep the two separate. +3. **Expressiveness is defensible.** HardNet + `mononet` universal-approximation + results refute "hard constraints kill expressiveness"; that refutation is the + framing lever, not a claim to hedge. + +The workflow's auto-synthesis output was placeholder-degraded; findings were +recovered from `journal.jsonl` (19 source extractions + 75 verification votes) and +manually verified. Several supporting citations were recent (2026) arXiv IDs not +independently re-fetched; the load-bearing prior art (Dugas 2001, Ackerer, +Chataigner–Crépey 2020, Cohen–Reisinger–Wang, ICNN, HardNet) is well established. + +### 9a. HJB de-risking findings (2026-07-12) — recorded for Paper 2 + +Two focused research passes on the HJB option (before it was demoted to a follow-up) +surfaced findings that **Paper 2 must inherit**: + +1. **Monotonicity is domain-restricted for HJB.** Convexity and monotonicity are + independent properties. The clean high-D HJB value functions are symmetric bowls + (LQ, Black–Scholes–Barenblatt `∝ ‖x‖²`): globally convex but monotone only on a + half-space / positive orthant. A nondegenerate convex quadratic is never globally + monotone. Merton is genuinely monotone + concave but effectively 1-D in state. + The popular Han–Jentzen–E d=100 "LQG" benchmark (`g = ln((1+‖x‖²)/2)`) is + **neither convex nor concave nor monotone** — a trap. ⇒ Paper 2 must use naturally + non-negative-state problems (energy storage, inventory, epidemic; or positive- + orthant pricing) and state the domain explicitly. +2. **The convex-HJB-by-construction lane is partly occupied.** Liu et al. 2023 + (arXiv:2309.09953) already use an ICNN-style convex net + a viscosity-solution + theorem (convex-only, 1-D). ICNN value/Q-functions (Amos 2017), convex control + (Chen–Shi–Zhang 2019), ICNN Lyapunov certificates (Manek–Kolter 2019), ISNN + (joint monotone+convex, 2025, mechanics-only), and Bokanowski et al. 2026 + (structure-in-loss, not architecture) all border the idea. ⇒ Paper 2's defensible + contribution is the **conjunction**: joint monotone+convex, *expressive beyond + ICNN*, high-dimensional, on the HJB residual, with a correct derived policy — and + it must explicitly distinguish those five works. + +These are exactly why the flagship moved to inverse conservation laws, whose +monotonicity guarantee is unconditional within its problem class. + +## 10. Open questions for the implementation plan + +- **Inverse-flagship problem design.** LWR flux choice (Greenshields + `Q(ρ) = v_max ρ(1 − ρ/ρ_max)`); which monotone-front scenarios (queue behind a + bottleneck / red signal / incident); observation model (probe vs. loop-detector + geometry); the sparsity × noise grid for the headline sweep. +- Whether a hard IC ansatz is meaningful in the inverse setting (likely forward-tier + only, since the inverse problem has no full IC). +- LWR flux domain and BC handling for the forward tier. +- Tolerance thresholds for cross-backend equivalence on trained models. + +## 11. Deliverable sequencing — paper scaffold first + +Per the user's direction, work proceeds **manuscript-first**: + +1. **Paper scaffold (Markdown, no results).** Write `pinn/paper/paper.md` as far as + possible *before* implementation: title, abstract, introduction, related work + (grounded in the §9/§9a research), the method (framework + admissibility + abstraction + registry), the theory (TVD/entropy-by-construction statements and + the monotone-solution-class scope), experimental design (benchmarks, baselines, + metrics, sweeps), and results/discussion sections with **explicit placeholders** + for numbers, tables, and figures. Gather `references.bib` from the research + already done. This front-loads all reasoning that does not depend on runs and + makes the implementation a matter of *filling in* pre-specified tables/figures. +2. **Implementation.** Build `core/` → models → trainers → experiments, TDD, per the + plan produced by `writing-plans`. +3. **Fill results.** Execute the sweep, populate the placeholders, render the + notebook into docs, finalize the manuscript. + +LaTeX conversion is deferred (Markdown → LaTeX is mechanical). The plan from +`writing-plans` will make step 1 (scaffold) its first phase. diff --git a/pyproject.toml b/pyproject.toml index ea3baae7..59dd4b9f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,19 @@ dev = [ "nest-asyncio==1.6.0", "hypothesis>=6.115", ] + +# Application research code (applications/) is out of the wheel and has its own +# opt-in groups, mirroring `bench`. Shared app deps live in `applications`; +# per-paper groups include it and add their specifics. Work on an application +# with e.g. `uv sync --extra jax --group pinn`. +applications = [ + "matplotlib>=3.8", +] +pinn = [ + {include-group = "applications"}, + "optax>=0.2", + "optuna>=4.0", +] docs = [ "sphinx==9.0.4; python_version < '3.12'", "sphinx==9.1.0; python_version >= '3.12'", @@ -131,7 +144,7 @@ include = ["mononet"] include = ["mononet"] [tool.mypy] -files = ["mononet", "tests", "benchmarks"] +files = ["mononet", "tests", "benchmarks", "applications"] strict = true strict_bytes = true local_partial_types = true @@ -165,6 +178,8 @@ module = [ "mononet.jax.layers", "mononet.keras.layers", "benchmarks._common.model_builder", + "applications.pinn.models.torch.builders", + "applications.pinn.models.jax.builders", ] disallow_subclassing_any = false @@ -193,6 +208,7 @@ include = [ "docs/**/*.py", "tools/**/*.py", "benchmarks/**/*.py", + "applications/**/*.py", "pyproject.toml", ] exclude = ["docs/docs_src"] @@ -242,6 +258,10 @@ ignore = [ "T20", # flake8-print (tools use print for user feedback) ] +"applications/**/experiments/*.py" = [ + "T201", # CLI entry points print their JSON artifact to stdout +] + "tests/**/*.py" = [ "D100", # missing docstring in public module "D101", diff --git a/tests/applications/__init__.py b/tests/applications/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/applications/pinn/__init__.py b/tests/applications/pinn/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/applications/pinn/core/__init__.py b/tests/applications/pinn/core/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/applications/pinn/core/test_admissibility.py b/tests/applications/pinn/core/test_admissibility.py new file mode 100644 index 00000000..7dc218a6 --- /dev/null +++ b/tests/applications/pinn/core/test_admissibility.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the admissibility spec and violation metric.""" + +from __future__ import annotations + +import numpy as np + +from applications.pinn.core.admissibility import AdmissibilitySpec, violation + + +def test_monotone_decreasing_field_has_zero_violation() -> None: + """A strictly decreasing profile violates a non-increasing target by 0.""" + x = np.linspace(0.0, 1.0, 50) + field = 1.0 - x # decreasing in x (axis 0) + assert violation(field, axis=0, sign=-1) == 0.0 + + +def test_bump_violation_equals_its_rise() -> None: + """For a non-increasing target, violation = total upward jump (the rise).""" + field = np.array([3.0, 2.0, 2.5, 1.0]) # one upward step of +0.5 + assert violation(field, axis=0, sign=-1) == 0.5 + + +def test_increasing_target_sign_is_symmetric() -> None: + """For a non-decreasing target, downward steps are the violations.""" + field = np.array([0.0, 1.0, 0.7, 2.0]) # one downward step of -0.3 + assert np.isclose(violation(field, axis=0, sign=1), 0.3) + + +def test_unconstrained_axis_never_violates() -> None: + """A sign of 0 (unconstrained) yields zero violation regardless.""" + field = np.array([[0.0, 5.0], [9.0, -3.0]]) + assert violation(field, axis=1, sign=0) == 0.0 + + +def test_violation_along_named_axis_of_2d_field() -> None: + """Violation is measured along the requested axis of a 2-D field.""" + # rows decreasing in axis 0; columns flat -> zero violation on axis 0 + field = np.array([[3.0, 3.0], [2.0, 2.0], [1.0, 1.0]]) + assert violation(field, axis=0, sign=-1) == 0.0 + + +def test_spec_fields() -> None: + """AdmissibilitySpec stores mask and convex/concave axes.""" + spec = AdmissibilitySpec(mask=(-1, 0), convex_axes=(), concave_axes=()) + assert spec.mask == (-1, 0) + assert spec.convex_axes == () diff --git a/tests/applications/pinn/core/test_conservation.py b/tests/applications/pinn/core/test_conservation.py new file mode 100644 index 00000000..13ba11b4 --- /dev/null +++ b/tests/applications/pinn/core/test_conservation.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the conservation-law problems and their registration.""" + +from __future__ import annotations + +import numpy as np + +from applications.pinn.core import problems +from applications.pinn.core.problems import conservation + + +def test_all_conservation_problems_registered() -> None: + """The four forward-tier problems are discoverable via the registry.""" + for key in ("burgers_riemann", "burgers_smooth", "advection", "lwr_riemann"): + assert key in problems.available() + assert problems.get(key).key == key + + +def test_burgers_riemann_mask_matches_monotone_direction() -> None: + """A decreasing Riemann (u_l>u_r) is non-increasing in x; increasing flips.""" + dec = conservation.BurgersRiemann(u_l=1.0, u_r=0.0) + inc = conservation.BurgersRiemann(u_l=0.0, u_r=1.0) + assert dec.admissibility().mask == (-1, 0) + assert inc.admissibility().mask == (1, 0) + + +def test_lwr_mask_matches_density_direction() -> None: + """Forming queue (rho_l None: + """Flux and characteristic speed match the analytic forms.""" + b = conservation.BurgersRiemann() + u = np.array([-1.0, 0.0, 2.0]) + assert np.allclose(b.flux(u), 0.5 * u**2) + assert np.allclose(b.flux_prime(u), u) + adv = conservation.LinearAdvection(a=0.7) + assert np.allclose(adv.flux_prime(u), 0.7) + + +def test_advection_residual_is_small() -> None: + """The exact advection solution satisfies u_t + a u_x = 0 (finite diff).""" + adv = conservation.LinearAdvection(a=0.8) + x = np.linspace(-2.0, 3.0, 400) + t, dt, dx = 0.5, 1e-4, x[1] - x[0] + t_plus = np.full_like(x, t + dt) + t_minus = np.full_like(x, t - dt) + u_t = (adv.ground_truth(x, t_plus) - adv.ground_truth(x, t_minus)) / (2 * dt) + f = adv.flux(adv.ground_truth(x, np.full_like(x, t))) + f_x = np.gradient(f, dx) + interior = slice(5, -5) + assert np.max(np.abs((u_t + f_x)[interior])) < 1e-2 + + +def test_burgers_smooth_ground_truth_matches_characteristic_pre_breaking() -> None: + """Before the breaking time, the Godunov ground truth ~ the characteristic.""" + p = conservation.BurgersSmoothShock(steepness=1.0) # t_b = 1.0 + x = np.linspace(-4.0, 4.0, 300) + t = 0.5 + gt = p.ground_truth(x, np.full_like(x, t)) + char = exact_characteristic(p, x, t) + assert np.max(np.abs(gt - char)) < 2e-2 + # and it is monotone non-increasing in x + assert np.all(np.diff(gt) <= 1e-6) + + +def exact_characteristic( + p: conservation.BurgersSmoothShock, x: np.ndarray, t: float +) -> np.ndarray: + """Return the pre-breaking characteristic solution for the smooth problem.""" + from applications.pinn.core import exact + + return exact.burgers_characteristic(x, t, p._u0) + + +def test_burgers_smooth_develops_shock_and_stays_monotone() -> None: + """Past the breaking time the profile is a monotone (non-oscillating) shock.""" + p = conservation.BurgersSmoothShock(steepness=1.0) + x = np.linspace(-4.0, 4.0, 400) + gt = p.ground_truth(x, np.full_like(x, 1.8)) # t > t_b = 1.0 + assert np.all(np.diff(gt) <= 1e-6) # monotone, no overshoot diff --git a/tests/applications/pinn/core/test_core_is_framework_free.py b/tests/applications/pinn/core/test_core_is_framework_free.py new file mode 100644 index 00000000..3551b898 --- /dev/null +++ b/tests/applications/pinn/core/test_core_is_framework_free.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Invariant: the pure-NumPy core imports no ML framework.""" + +from __future__ import annotations + +import subprocess +import sys + + +def test_core_imports_no_backend() -> None: + """Importing the whole core (incl. the problem registry) pulls in no backend. + + Run in a fresh interpreter so the check is independent of test order. + """ + code = ( + "import sys\n" + "import applications.pinn.core.problems\n" + "import applications.pinn.core.metrics\n" + "import applications.pinn.core.sampling\n" + "import applications.pinn.core.reference_solver\n" + "assert 'torch' not in sys.modules, 'torch imported by core'\n" + "assert 'jax' not in sys.modules, 'jax imported by core'\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr diff --git a/tests/applications/pinn/core/test_exact.py b/tests/applications/pinn/core/test_exact.py new file mode 100644 index 00000000..e81f6274 --- /dev/null +++ b/tests/applications/pinn/core/test_exact.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for closed-form entropy solutions.""" + +from __future__ import annotations + +import numpy as np + +from applications.pinn.core import exact + + +def test_burgers_shock_speed_and_monotonicity() -> None: + """u_l>u_r gives a shock at s=(u_l+u_r)/2, monotone decreasing in x.""" + u_l, u_r, t = 1.0, 0.0, 2.0 + s = 0.5 * (u_l + u_r) + x = np.linspace(-2.0, 3.0, 401) + u = exact.burgers_riemann(x, t, u_l, u_r) + # left of the shock -> u_l, right -> u_r + assert np.all(u[x < s * t - 0.05] == u_l) + assert np.all(u[x > s * t + 0.05] == u_r) + # monotone non-increasing + assert np.all(np.diff(u) <= 1e-12) + + +def test_burgers_rarefaction_is_self_similar() -> None: + """u_l u_l) & (xi < u_r) + assert np.allclose(u[inside], xi[inside]) + assert np.all(np.diff(u) >= -1e-12) + + +def test_burgers_initial_step_at_t0() -> None: + """At t=0 the solution is the Riemann step.""" + x = np.array([-1.0, -0.1, 0.1, 1.0]) + u = exact.burgers_riemann(x, 0.0, 2.0, 0.5) + assert np.array_equal(u, np.array([2.0, 2.0, 0.5, 0.5])) + + +def test_advection_translates_profile() -> None: + """Linear advection transports the initial profile at speed a.""" + a, t = 0.75, 2.0 + + def u0(x: np.ndarray) -> np.ndarray: + out: np.ndarray = np.exp(-(x**2)) + return out + + x = np.linspace(-3.0, 5.0, 200) + u = exact.advection(x, t, a, u0) + assert np.allclose(u, u0(x - a * t)) + + +def test_lwr_shock_speed_rankine_hugoniot() -> None: + """LWR shock (rho_l s * t + 0.05] == rho_r) + # density increasing in x (forming queue) + assert np.all(np.diff(rho) >= -1e-12) + + +def test_lwr_rarefaction_monotone_decreasing() -> None: + """LWR rarefaction (rho_l>rho_r) is a monotone-decreasing fan.""" + rho_l, rho_r, t = 0.8, 0.2, 1.0 + x = np.linspace(-2.0, 2.0, 401) + rho = exact.lwr_riemann(x, t, rho_l, rho_r) + assert np.all(np.diff(rho) <= 1e-12) + assert np.isclose(rho[0], rho_l) + assert np.isclose(rho[-1], rho_r) + + +def test_breaking_time() -> None: + """t_b = -1/min(u0'); raises when no shock forms.""" + assert np.isclose(exact.breaking_time(-2.0), 0.5) + with np.testing.assert_raises(ValueError): + exact.breaking_time(1.0) + + +def test_characteristic_satisfies_implicit_relation() -> None: + """Pre-shock solution satisfies u = u0(x - u t) to tolerance.""" + + def u0(x: np.ndarray) -> np.ndarray: + out: np.ndarray = -np.arctan(x) # decreasing; min slope -1 -> t_b=1 + return out + + t = 0.5 # < t_b = 1 + x = np.linspace(-4.0, 4.0, 200) + u = exact.burgers_characteristic(x, t, u0) + assert np.max(np.abs(u - u0(x - u * t))) < 1e-10 diff --git a/tests/applications/pinn/core/test_metrics.py b/tests/applications/pinn/core/test_metrics.py new file mode 100644 index 00000000..0ed59ef9 --- /dev/null +++ b/tests/applications/pinn/core/test_metrics.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for scoring metrics and plotting.""" + +from __future__ import annotations + +import numpy as np + +from applications.pinn.core import metrics + + +def test_l1_l2() -> None: + """L1/L2 errors match closed-form values.""" + a = np.array([0.0, 1.0, 2.0]) + b = np.array([0.0, 0.0, 0.0]) + assert metrics.l1(a, b) == 3.0 + assert np.isclose(metrics.l2(a, b), np.sqrt(5.0)) + assert metrics.l1(a, b, dx=0.5) == 1.5 + + +def test_tv_of_monotone_is_endpoint_difference() -> None: + """TV of a monotone profile equals |first - last|.""" + prof = np.array([3.0, 2.0, 1.0, 0.5]) + assert np.isclose(metrics.tv(prof), 2.5) + + +def test_tv_of_bump_exceeds_monotone() -> None: + """A non-monotone profile has larger TV than its monotone envelope.""" + bump = np.array([3.0, 2.0, 2.6, 1.0]) + assert metrics.tv(bump) > 2.0 + + +def test_tv_curve_shape_and_values() -> None: + """tv_curve returns one TV per time slice.""" + field = np.array([[3.0, 2.0, 1.0], [2.0, 1.0, 0.0]]) + curve = metrics.tv_curve(field) + assert curve.shape == (2,) + assert np.allclose(curve, [2.0, 2.0]) + + +def test_overshoot_zero_when_within_range() -> None: + """A profile inside the reference range has zero overshoot.""" + ref = np.array([1.0, 0.5, 0.0]) + pred = np.array([0.9, 0.4, 0.1]) + assert metrics.overshoot(pred, ref) == 0.0 + + +def test_overshoot_measures_excursion() -> None: + """Overshoot equals the largest excursion beyond the reference range.""" + ref = np.array([1.0, 0.0]) + pred = np.array([1.2, -0.1]) # 0.2 above, 0.1 below -> 0.2 + assert np.isclose(metrics.overshoot(pred, ref), 0.2) + + +def test_shock_position_error() -> None: + """Front-position error tracks a shifted level crossing.""" + x = np.linspace(0.0, 1.0, 101) + ref = 1.0 - x # crosses 0.5 at x = 0.5 + pred = 1.1 - x # crosses 0.5 at x = 0.6 + err = metrics.shock_position_error(pred, ref, x, level=0.5) + assert np.isclose(err, 0.1, atol=0.02) + + +def test_mass_error() -> None: + """Mass error is the absolute difference of integrals.""" + a = np.array([1.0, 1.0, 1.0]) + b = np.array([1.0, 0.0, 1.0]) + assert np.isclose(metrics.mass_error(a, b, dx=0.5), 0.5) + + +def test_reconstruction_error_over_field() -> None: + """Reconstruction error is the L2 over the flattened field.""" + field = np.zeros((2, 3)) + ref = np.ones((2, 3)) + assert np.isclose(metrics.reconstruction_error(field, ref), np.sqrt(6.0)) diff --git a/tests/applications/pinn/core/test_plotting.py b/tests/applications/pinn/core/test_plotting.py new file mode 100644 index 00000000..d8231d65 --- /dev/null +++ b/tests/applications/pinn/core/test_plotting.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for plotting helpers (require matplotlib; skipped without it).""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("matplotlib") + +from matplotlib.figure import Figure + +from applications.pinn.core import plotting + + +def test_plotting_returns_figures() -> None: + """Plot helpers return Figure objects with the expected content.""" + x = np.linspace(0.0, 1.0, 20) + t = np.linspace(0.0, 1.0, 10) + fig1 = plotting.profiles(x, {"exact": x, "pred": x**2}, title="p") + fig2 = plotting.tv_curves(t, {"hard": np.ones_like(t)}) + fig3 = plotting.field_heatmap(np.outer(t, x), x, t) + assert isinstance(fig1, Figure) + assert isinstance(fig2, Figure) + assert isinstance(fig3, Figure) + assert len(fig1.axes[0].get_lines()) == 2 diff --git a/tests/applications/pinn/core/test_reference_solver.py b/tests/applications/pinn/core/test_reference_solver.py new file mode 100644 index 00000000..2093538d --- /dev/null +++ b/tests/applications/pinn/core/test_reference_solver.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the Godunov reference solver.""" + +from __future__ import annotations + +from itertools import pairwise + +import numpy as np + +from applications.pinn.core import exact +from applications.pinn.core.reference_solver import godunov + + +def _burgers_flux(u: np.ndarray) -> np.ndarray: + return 0.5 * u**2 + + +def _burgers_flux_prime(u: np.ndarray) -> np.ndarray: + return u + + +def _centers(a: float, b: float, nx: int) -> np.ndarray: + edges = np.linspace(a, b, nx + 1) + return 0.5 * (edges[:-1] + edges[1:]) + + +def test_converges_to_burgers_riemann_shock() -> None: + """L1 error to the exact shock decreases under grid refinement.""" + a, b, t = -2.0, 3.0, 1.5 + u_l, u_r = 1.0, 0.0 + errors = [] + for nx in (200, 400, 800): + x = _centers(a, b, nx) + u0 = exact.burgers_riemann(x, 0.0, u_l, u_r) + field = godunov( + u0, x, [t], flux=_burgers_flux, flux_prime=_burgers_flux_prime, sonic=0.0 + ) + ref = exact.burgers_riemann(x, t, u_l, u_r) + dx = (b - a) / nx + errors.append(float(np.abs(field[0] - ref).sum() * dx)) + assert errors[0] > errors[1] > errors[2] + assert errors[-1] < 2e-2 + + +def test_shock_speed_is_correct() -> None: + """The recovered shock sits at s*t within one cell.""" + a, b, t, nx = -2.0, 3.0, 1.5, 800 + u_l, u_r = 1.0, 0.0 + s = 0.5 * (u_l + u_r) + x = _centers(a, b, nx) + u0 = exact.burgers_riemann(x, 0.0, u_l, u_r) + field = godunov( + u0, x, [t], flux=_burgers_flux, flux_prime=_burgers_flux_prime, sonic=0.0 + )[0] + # midpoint crossing of (u_l+u_r)/2 + crossing = x[np.argmin(np.abs(field - 0.5 * (u_l + u_r)))] + assert abs(crossing - s * t) < 2 * (b - a) / nx + + +def test_total_variation_non_increasing() -> None: + """TV(t) is non-increasing (TVD), here constant for monotone data.""" + a, b, nx = -2.0, 3.0, 400 + x = _centers(a, b, nx) + u0 = exact.burgers_riemann(x, 0.0, 1.0, 0.0) + ts = [0.0, 0.5, 1.0, 1.5] + field = godunov( + u0, x, ts, flux=_burgers_flux, flux_prime=_burgers_flux_prime, sonic=0.0 + ) + tvs = [float(np.abs(np.diff(field[i])).sum()) for i in range(len(ts))] + for earlier, later in pairwise(tvs): + assert later <= earlier + 1e-9 + + +def test_converges_to_lwr_riemann() -> None: + """Godunov with the concave Greenshields flux matches the LWR shock.""" + a, b, t, nx = -2.0, 2.0, 1.0, 800 + rho_l, rho_r = 0.2, 0.8 + v_max, rho_max = 1.0, 1.0 + + def flux(r: np.ndarray) -> np.ndarray: + return exact.greenshields_flux(r, v_max, rho_max) + + def flux_prime(r: np.ndarray) -> np.ndarray: + return exact.greenshields_flux_prime(r, v_max, rho_max) + + x = _centers(a, b, nx) + rho0 = exact.lwr_riemann(x, 0.0, rho_l, rho_r, v_max=v_max, rho_max=rho_max) + field = godunov( + rho0, x, [t], flux=flux, flux_prime=flux_prime, sonic=rho_max / 2.0 + )[0] + ref = exact.lwr_riemann(x, t, rho_l, rho_r, v_max=v_max, rho_max=rho_max) + dx = (b - a) / nx + assert float(np.abs(field - ref).sum() * dx) < 2e-2 diff --git a/tests/applications/pinn/core/test_registry.py b/tests/applications/pinn/core/test_registry.py new file mode 100644 index 00000000..4de48844 --- /dev/null +++ b/tests/applications/pinn/core/test_registry.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the problem registry.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from applications.pinn.core.admissibility import AdmissibilitySpec +from applications.pinn.core.problems import base + +if TYPE_CHECKING: + from collections.abc import Iterator + + from applications.pinn.core.problems.base import Array + + +class _StubProblem: + """Minimal `Problem`-conforming stub for registry tests.""" + + key = "stub" + + @property + def domain(self) -> tuple[tuple[float, float], tuple[float, float]]: + return ((0.0, 1.0), (0.0, 1.0)) + + def admissibility(self) -> AdmissibilitySpec: + return AdmissibilitySpec(mask=(-1, 0)) + + def flux(self, u: Array) -> Array: + return u + + def flux_prime(self, u: Array) -> Array: + return np.ones_like(u) + + def initial(self, x: Array) -> Array: + return x + + def ground_truth(self, x: Array, t: Array) -> Array | None: + return None + + +@pytest.fixture(autouse=True) +def _isolate_registry() -> Iterator[None]: + """Snapshot and restore the module registry around each test.""" + saved = dict(base._REGISTRY) + yield + base._REGISTRY.clear() + base._REGISTRY.update(saved) + + +def test_register_get_available_roundtrip() -> None: + """A registered class is retrievable and listed.""" + + @base.register("dummy") + class _Dummy(_StubProblem): + key = "dummy" + + assert base.get("dummy") is _Dummy + assert "dummy" in base.available() + + +def test_get_unknown_raises() -> None: + """Requesting an unregistered key raises KeyError.""" + with pytest.raises(KeyError): + base.get("does-not-exist") + + +def test_duplicate_registration_raises() -> None: + """Re-registering the same key raises KeyError.""" + + @base.register("dup") + class _A(_StubProblem): + key = "dup" + + with pytest.raises(KeyError): + + @base.register("dup") + class _B(_StubProblem): + key = "dup" + + +def test_available_is_sorted() -> None: + """`available` returns keys in sorted order.""" + + @base.register("zeta") + class _Z(_StubProblem): + key = "zeta" + + @base.register("alpha") + class _A(_StubProblem): + key = "alpha" + + keys = base.available() + assert keys == sorted(keys) diff --git a/tests/applications/pinn/core/test_sampling.py b/tests/applications/pinn/core/test_sampling.py new file mode 100644 index 00000000..165d57f1 --- /dev/null +++ b/tests/applications/pinn/core/test_sampling.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for deterministic sampling and the observation sampler.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from applications.pinn.core import sampling + +DOMAIN = ((-1.0, 2.0), (0.0, 1.5)) + + +@pytest.mark.parametrize("strategy", ["uniform", "lhs"]) +def test_collocation_deterministic_and_in_domain(strategy: str) -> None: + """Same seed -> identical points; all points lie in the domain.""" + a = sampling.collocation(DOMAIN, 256, seed=0, strategy=strategy) + b = sampling.collocation(DOMAIN, 256, seed=0, strategy=strategy) + c = sampling.collocation(DOMAIN, 256, seed=1, strategy=strategy) + assert np.array_equal(a, b) + assert not np.array_equal(a, c) + assert a.shape == (256, 2) + (x_lo, x_hi), (t_lo, t_hi) = DOMAIN + assert np.all((a[:, 0] >= x_lo) & (a[:, 0] <= x_hi)) + assert np.all((a[:, 1] >= t_lo) & (a[:, 1] <= t_hi)) + + +def test_unknown_strategy_raises() -> None: + """An unknown sampling strategy is rejected.""" + with pytest.raises(ValueError, match="unknown strategy"): + sampling.collocation(DOMAIN, 8, seed=0, strategy="sobol") + + +def test_initial_points_on_initial_line() -> None: + """Initial points all have t == t_min.""" + pts = sampling.initial_points(DOMAIN, 64, seed=3) + assert pts.shape == (64, 2) + assert np.all(pts[:, 1] == DOMAIN[1][0]) + + +def test_boundary_points_on_boundaries() -> None: + """Boundary points lie on x_min or x_max, split evenly.""" + pts = sampling.boundary_points(DOMAIN, 40, seed=2) + assert pts.shape == (40, 2) + on_edge = np.isclose(pts[:, 0], DOMAIN[0][0]) | np.isclose(pts[:, 0], DOMAIN[0][1]) + assert np.all(on_edge) + assert np.isclose(pts[:, 0], DOMAIN[0][0]).sum() == 20 + + +def test_eval_grid_axes() -> None: + """Eval grid axes span the domain with the requested counts.""" + x_values, t_values = sampling.eval_grid(DOMAIN, 50, 30) + assert x_values.shape == (50,) + assert t_values.shape == (30,) + assert x_values[0] == DOMAIN[0][0] + assert t_values[-1] == DOMAIN[1][1] + + +def test_observations_noiseless_match_field() -> None: + """With zero noise, observations equal the sampled field values.""" + x_values, t_values = sampling.eval_grid(DOMAIN, 20, 10) + field = np.outer(t_values, x_values) # arbitrary reference field + coords, values = sampling.observations( + field, x_values, t_values, n_obs=15, noise_std=0.0, seed=7 + ) + assert coords.shape == (15, 2) + assert values.shape == (15,) + # each observed value must appear in the field grid + for (x, t), v in zip(coords, values, strict=True): + i = int(np.argmin(np.abs(t_values - t))) + j = int(np.argmin(np.abs(x_values - x))) + assert np.isclose(v, field[i, j]) + + +def test_observations_deterministic_and_noise_reproducible() -> None: + """Same seed reproduces coordinates and noisy values exactly.""" + x_values, t_values = sampling.eval_grid(DOMAIN, 20, 10) + field = np.outer(t_values, x_values) + c1, v1 = sampling.observations( + field, x_values, t_values, n_obs=15, noise_std=0.1, seed=7 + ) + c2, v2 = sampling.observations( + field, x_values, t_values, n_obs=15, noise_std=0.1, seed=7 + ) + assert np.array_equal(c1, c2) + assert np.array_equal(v1, v2) diff --git a/tests/applications/pinn/experiments/__init__.py b/tests/applications/pinn/experiments/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/applications/pinn/experiments/test_run_smoke.py b/tests/applications/pinn/experiments/test_run_smoke.py new file mode 100644 index 00000000..3dd74d4c --- /dev/null +++ b/tests/applications/pinn/experiments/test_run_smoke.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Smoke test for the single-run experiment orchestrator.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("jax") +pytest.importorskip("flax") +pytest.importorskip("optax") + +import numpy as np + +from applications.pinn.experiments.config import RunConfig +from applications.pinn.experiments.run import run_one +from applications.pinn.models.protocol import ModelConfig + +_TINY = ModelConfig(width=8, n_blocks=2, t_embed_dim=4, seed=0) + + +def _cfg(method: str) -> RunConfig: + return RunConfig( + problem="burgers_riemann", + method=method, # type: ignore[arg-type] + backend="jax", + seed=0, + model=_TINY, + n_collocation=128, + n_ic=32, + n_bc=16, + steps=15, + eval_nx=40, + eval_nt=3, + ) + + +def test_run_one_produces_valid_artifact() -> None: + """A tiny run returns an artifact with finite headline metrics.""" + art = run_one(_cfg("hard_monotone")) + assert art["problem"] == "burgers_riemann" + assert art["method"] == "hard_monotone" + for key in ("l1", "l2", "admissibility_violation", "overshoot"): + assert np.isfinite(art[key]) + # hard-monotone is admissible by construction + assert art["admissibility_violation"] < 1e-4 + + +def test_soft_baseline_runs() -> None: + """The soft baseline runs end-to-end (penalty path exercised).""" + art = run_one(_cfg("soft")) + assert np.isfinite(art["l2"]) diff --git a/tests/applications/pinn/experiments/test_search_sweep_smoke.py b/tests/applications/pinn/experiments/test_search_sweep_smoke.py new file mode 100644 index 00000000..9e269051 --- /dev/null +++ b/tests/applications/pinn/experiments/test_search_sweep_smoke.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Smoke tests for the Optuna search and the sweep composer.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +pytest.importorskip("jax") +pytest.importorskip("flax") +pytest.importorskip("optax") +pytest.importorskip("optuna") + +from applications.pinn.experiments import search, sweep +from applications.pinn.experiments.config import RunConfig +from applications.pinn.models.protocol import ModelConfig + +if TYPE_CHECKING: + from pathlib import Path + +_TEMPLATE = RunConfig( + problem="burgers_riemann", + method="hard_monotone", + backend="jax", + model=ModelConfig(width=8, n_blocks=2, t_embed_dim=4), + n_collocation=96, + n_ic=24, + n_bc=12, + steps=8, + eval_nx=30, + eval_nt=3, +) + + +def test_search_returns_best_params_and_freezes(tmp_path: Path) -> None: + """A 2-trial search yields best params and freezes a reloadable config.""" + best = search.search( + "burgers_riemann", + "hard_monotone", + "jax", + n_trials=2, + template=_TEMPLATE, + seed=0, + ) + assert "lr" in best + assert "width" in best + path = search.freeze( + "burgers_riemann", "hard_monotone", "jax", best, configs_dir=tmp_path + ) + assert path.exists() + cfg = search.tuned_config("burgers_riemann", "hard_monotone", "jax", path) + assert cfg.model.width in (16, 32, 64) + + +def test_soft_search_includes_penalty() -> None: + """The soft method's search space includes its penalty weight.""" + soft_template = _TEMPLATE.__class__( + problem="burgers_riemann", + method="soft", + backend="jax", + model=_TEMPLATE.model, + n_collocation=96, + n_ic=24, + n_bc=12, + steps=8, + eval_nx=30, + eval_nt=3, + ) + best = search.search( + "burgers_riemann", "soft", "jax", n_trials=2, template=soft_template, seed=0 + ) + assert "soft_penalty" in best + + +def test_sweep_runs_small_matrix() -> None: + """The sweep composes run_one over a 2-cell matrix.""" + configs = sweep.build_matrix( + ["burgers_riemann"], + ["hard_monotone", "vanilla"], + ["jax"], + [0], + template=_TEMPLATE, + ) + assert len(configs) == 2 + results = sweep.run_matrix(configs) + assert len(results) == 2 + assert {r["method"] for r in results} == {"hard_monotone", "vanilla"} diff --git a/tests/applications/pinn/models/__init__.py b/tests/applications/pinn/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/applications/pinn/models/test_builders_jax.py b/tests/applications/pinn/models/test_builders_jax.py new file mode 100644 index 00000000..ccceec45 --- /dev/null +++ b/tests/applications/pinn/models/test_builders_jax.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the JAX (Flax NNX) model builders (monotonicity by construction).""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("jax") +pytest.importorskip("flax") + +import jax.numpy as jnp + +from applications.pinn.core.problems import conservation +from applications.pinn.models.jax import builders +from applications.pinn.models.protocol import Method, ModelConfig + +CFG = ModelConfig(width=16, n_blocks=2, t_embed_dim=4, seed=0) +METHODS: list[Method] = ["vanilla", "soft", "weight_clip", "hard_monotone"] + + +def _grid(n: int = 80) -> tuple[jnp.ndarray, jnp.ndarray]: + x = jnp.linspace(-1.5, 2.0, n).reshape(-1, 1) + t = jnp.full((n, 1), 0.3) + return x, t + + +@pytest.mark.parametrize("method", METHODS) +def test_forward_shape(method: Method) -> None: + """Every method returns a column vector of the right shape.""" + model = builders.build_jax(conservation.BurgersRiemann(), CFG, method) + x, t = _grid() + assert model(x, t).shape == (80, 1) + + +def test_hard_monotone_decreasing_in_x() -> None: + """A decreasing problem (u_l>u_r) yields a non-increasing field in x.""" + model = builders.build_jax( + conservation.BurgersRiemann(u_l=1.0, u_r=0.0), CFG, "hard_monotone" + ) + x, t = _grid() + u = np.asarray(model(x, t)).ravel() + assert np.all(np.diff(u) <= 1e-4) + + +def test_hard_monotone_increasing_in_x() -> None: + """A forming-queue LWR problem (rho_l= -1e-4) + + +def test_weight_clip_is_monotone_in_x() -> None: + """The weight-clip baseline is also monotone in x by construction.""" + model = builders.build_jax( + conservation.BurgersRiemann(u_l=1.0, u_r=0.0), CFG, "weight_clip" + ) + x, t = _grid() + u = np.asarray(model(x, t)).ravel() + assert np.all(np.diff(u) <= 1e-4) + + +def test_hard_monotone_depends_on_t() -> None: + """The field genuinely varies with t (free, not constant in t).""" + model = builders.build_jax(conservation.BurgersRiemann(), CFG, "hard_monotone") + x = jnp.linspace(-1.0, 1.0, 40).reshape(-1, 1) + u0 = np.asarray(model(x, jnp.zeros_like(x))) + u1 = np.asarray(model(x, jnp.ones_like(x))) + assert float(np.abs(u0 - u1).max()) > 1e-4 diff --git a/tests/applications/pinn/models/test_builders_torch.py b/tests/applications/pinn/models/test_builders_torch.py new file mode 100644 index 00000000..9c10ce30 --- /dev/null +++ b/tests/applications/pinn/models/test_builders_torch.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the PyTorch model builders (monotonicity by construction).""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("torch") + +import torch + +from applications.pinn.core.problems import conservation +from applications.pinn.models.protocol import Method, ModelConfig +from applications.pinn.models.torch import builders + +CFG = ModelConfig(width=16, n_blocks=2, t_embed_dim=4, seed=0) +METHODS: list[Method] = ["vanilla", "soft", "weight_clip", "hard_monotone"] + + +def _grid(n: int = 80) -> tuple[torch.Tensor, torch.Tensor]: + x = torch.linspace(-1.5, 2.0, n, dtype=torch.float32).reshape(-1, 1) + t = torch.full((n, 1), 0.3, dtype=torch.float32) + return x, t + + +@pytest.mark.parametrize("method", METHODS) +def test_forward_shape(method: Method) -> None: + """Every method returns a column vector of the right shape.""" + problem = conservation.BurgersRiemann() + model = builders.build_torch(problem, CFG, method) + x, t = _grid() + out = model(x, t) + assert out.shape == (80, 1) + + +def test_hard_monotone_decreasing_in_x() -> None: + """A decreasing problem (u_l>u_r) yields a non-increasing field in x.""" + problem = conservation.BurgersRiemann(u_l=1.0, u_r=0.0) # sign_x = -1 + model = builders.build_torch(problem, CFG, "hard_monotone") + x, t = _grid() + u = model(x, t).detach().numpy().ravel() + assert np.all(np.diff(u) <= 1e-4) + + +def test_hard_monotone_increasing_in_x() -> None: + """A forming-queue LWR problem (rho_l= -1e-4) + + +def test_weight_clip_is_monotone_in_x() -> None: + """The weight-clip baseline is also monotone in x by construction.""" + problem = conservation.BurgersRiemann(u_l=1.0, u_r=0.0) # sign_x = -1 + model = builders.build_torch(problem, CFG, "weight_clip") + x, t = _grid() + u = model(x, t).detach().numpy().ravel() + assert np.all(np.diff(u) <= 1e-4) + + +def test_hard_monotone_depends_on_t() -> None: + """The field genuinely varies with t (free, not constant/separable-trivial).""" + problem = conservation.BurgersRiemann() + model = builders.build_torch(problem, CFG, "hard_monotone") + x = torch.linspace(-1.0, 1.0, 40, dtype=torch.float32).reshape(-1, 1) + u0 = model(x, torch.zeros_like(x)) + u1 = model(x, torch.ones_like(x)) + assert float((u0 - u1).abs().max().detach()) > 1e-4 diff --git a/tests/applications/pinn/models/test_cross_backend.py b/tests/applications/pinn/models/test_cross_backend.py new file mode 100644 index 00000000..0dbb3051 --- /dev/null +++ b/tests/applications/pinn/models/test_cross_backend.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cross-backend parity for the PINN models. + +Exact *bitwise* cross-backend equivalence of the underlying layers is mononet's +responsibility and is covered by ``tests/equivalence/``. Here we assert the +*application-level* parity that matters: with identical configuration and inputs, +both backends' hard-monotone field satisfies the same by-construction constraint +(monotone in ``x``) and produces a finite field of the same shape. Numeric +agreement of *trained* runs across backends is an empirical result reported in +the experiment sweep, not a unit test. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("torch") +pytest.importorskip("jax") +pytest.importorskip("flax") + +import jax.numpy as jnp +import torch + +from applications.pinn.core.admissibility import violation +from applications.pinn.core.problems import conservation +from applications.pinn.models.jax import builders as jax_builders +from applications.pinn.models.protocol import ModelConfig +from applications.pinn.models.torch import builders as torch_builders + +CFG = ModelConfig(width=16, n_blocks=2, t_embed_dim=4, seed=0) + + +def test_both_backends_hard_monotone_parity() -> None: + """Both backends' hard-monotone field is monotone in x on identical inputs.""" + problem = conservation.BurgersRiemann(u_l=1.0, u_r=0.0) # sign_x = -1 + x_np = np.linspace(-1.5, 2.0, 80, dtype=np.float32) + t_val = 0.4 + + torch_model = torch_builders.build_torch(problem, CFG, "hard_monotone") + xt = torch.as_tensor(x_np).reshape(-1, 1) + tt = torch.full_like(xt, t_val) + u_torch = torch_model(xt, tt).detach().numpy().ravel() + + jax_model = jax_builders.build_jax(problem, CFG, "hard_monotone") + xj = jnp.asarray(x_np).reshape(-1, 1) + tj = jnp.full_like(xj, t_val) + u_jax = np.asarray(jax_model(xj, tj)).ravel() + + assert u_torch.shape == u_jax.shape == (80,) + assert np.all(np.isfinite(u_torch)) + assert np.all(np.isfinite(u_jax)) + # the shared by-construction guarantee holds in both backends + assert violation(u_torch, axis=0, sign=-1) < 1e-4 + assert violation(u_jax, axis=0, sign=-1) < 1e-4 + + +def test_both_backends_respect_increasing_sign() -> None: + """A forming-queue LWR problem is non-decreasing in x in both backends.""" + problem = conservation.LwrRiemann(rho_l=0.2, rho_r=0.8) # sign_x = +1 + x_np = np.linspace(-1.5, 1.5, 60, dtype=np.float32) + + ut = ( + torch_builders.build_torch(problem, CFG, "hard_monotone")( + torch.as_tensor(x_np).reshape(-1, 1), + torch.full((60, 1), 0.3), + ) + .detach() + .numpy() + .ravel() + ) + uj = np.asarray( + jax_builders.build_jax(problem, CFG, "hard_monotone")( + jnp.asarray(x_np).reshape(-1, 1), jnp.full((60, 1), 0.3) + ) + ).ravel() + + assert violation(ut, axis=0, sign=1) < 1e-4 + assert violation(uj, axis=0, sign=1) < 1e-4 diff --git a/tests/applications/pinn/test_skeleton.py b/tests/applications/pinn/test_skeleton.py new file mode 100644 index 00000000..122ed235 --- /dev/null +++ b/tests/applications/pinn/test_skeleton.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Skeleton smoke tests for the applications.pinn package.""" + +from __future__ import annotations + +import subprocess +import sys +from typing import TYPE_CHECKING + +import numpy as np + +from applications._common import metrics_io, seeding + +if TYPE_CHECKING: + from pathlib import Path + + +def test_import_pinn_does_not_import_backends() -> None: + """Importing `applications.pinn` must not eagerly import torch or jax. + + Run in a fresh interpreter so the check is independent of test order + (other tests in the suite may already have imported a backend). + """ + code = ( + "import sys, applications.pinn\n" + "assert 'torch' not in sys.modules, 'torch was imported'\n" + "assert 'jax' not in sys.modules, 'jax was imported'\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + +def test_rng_is_reproducible() -> None: + """`seeding.rng` is deterministic per seed and varies across seeds.""" + a = seeding.rng(0).standard_normal(5) + b = seeding.rng(0).standard_normal(5) + c = seeding.rng(1).standard_normal(5) + assert np.array_equal(a, b) + assert not np.array_equal(a, c) + + +def test_split_seeds_are_deterministic_and_distinct() -> None: + """`split_seeds` yields the requested count, deterministically.""" + s1 = seeding.split_seeds(0, 4) + s2 = seeding.split_seeds(0, 4) + assert s1 == s2 + assert len(s1) == 4 + assert len(set(s1)) == 4 + + +def test_metrics_io_roundtrips(tmp_path: Path) -> None: + """`write_result`/`read_result` round-trip a JSON mapping exactly.""" + obj = {"metric": "l2", "value": 1.5, "seeds": [0, 1, 2]} + path = tmp_path / "nested" / "result.json" + metrics_io.write_result(path, obj) + assert metrics_io.read_result(path) == obj diff --git a/tests/applications/pinn/training/__init__.py b/tests/applications/pinn/training/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/applications/pinn/training/test_trainer_jax.py b/tests/applications/pinn/training/test_trainer_jax.py new file mode 100644 index 00000000..91642762 --- /dev/null +++ b/tests/applications/pinn/training/test_trainer_jax.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Smoke tests for the JAX PINN trainer.""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("jax") +pytest.importorskip("flax") +pytest.importorskip("optax") + +import jax.numpy as jnp + +from applications.pinn.core import sampling +from applications.pinn.core.admissibility import violation +from applications.pinn.core.problems import conservation +from applications.pinn.models.jax import builders +from applications.pinn.models.protocol import ModelConfig +from applications.pinn.training import jax_trainer +from applications.pinn.training.losses import LossWeights, TrainingData + +CFG = ModelConfig(width=16, n_blocks=2, t_embed_dim=4, seed=0) + + +def _burgers_forward_data() -> tuple[conservation.BurgersRiemann, TrainingData]: + problem = conservation.BurgersRiemann(u_l=1.0, u_r=0.0) + domain = problem.domain + coll = sampling.collocation(domain, 256, seed=0) + ic_pts = sampling.initial_points(domain, 64, seed=1) + ic_vals = problem.initial(ic_pts[:, 0]) + return problem, TrainingData(collocation=coll, ic=(ic_pts, ic_vals)) + + +def test_training_reduces_loss() -> None: + """A few optax steps reduce the PINN loss (residual autodiff works).""" + problem, data = _burgers_forward_data() + model = builders.build_jax(problem, CFG, "hard_monotone") + _, history = jax_trainer.train( + problem, model, data, weights=LossWeights(), sign_x=-1, lr=3e-3, steps=60 + ) + assert history[-1] < history[0] + assert np.isfinite(history[-1]) + + +def test_hard_monotone_stays_admissible_after_training() -> None: + """The trained hard-monotone field is still non-increasing in x (violation 0).""" + problem, data = _burgers_forward_data() + model = builders.build_jax(problem, CFG, "hard_monotone") + trained, _ = jax_trainer.train( + problem, model, data, weights=LossWeights(), sign_x=-1, lr=3e-3, steps=40 + ) + x_values, _ = sampling.eval_grid(problem.domain, 60, 3) + x = jnp.asarray(x_values).reshape(-1, 1) + u = np.asarray(trained(x, jnp.full_like(x, 0.5))).ravel() + assert violation(u, axis=0, sign=-1) < 1e-5 + + +def test_vanilla_also_trains() -> None: + """The unconstrained baseline also runs and reduces its loss.""" + problem, data = _burgers_forward_data() + model = builders.build_jax(problem, CFG, "vanilla") + _, history = jax_trainer.train( + problem, model, data, weights=LossWeights(), sign_x=-1, lr=3e-3, steps=40 + ) + assert history[-1] < history[0] diff --git a/tests/applications/pinn/training/test_trainer_torch.py b/tests/applications/pinn/training/test_trainer_torch.py new file mode 100644 index 00000000..73bd97e3 --- /dev/null +++ b/tests/applications/pinn/training/test_trainer_torch.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Smoke tests for the PyTorch PINN trainer.""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("torch") + +import torch + +from applications.pinn.core import sampling +from applications.pinn.core.admissibility import violation +from applications.pinn.core.problems import conservation +from applications.pinn.models.protocol import ModelConfig +from applications.pinn.models.torch import builders +from applications.pinn.training import torch_trainer +from applications.pinn.training.losses import LossWeights, TrainingData + +CFG = ModelConfig(width=16, n_blocks=2, t_embed_dim=4, seed=0) + + +def _burgers_forward_data() -> tuple[conservation.BurgersRiemann, TrainingData]: + problem = conservation.BurgersRiemann(u_l=1.0, u_r=0.0) + domain = problem.domain + coll = sampling.collocation(domain, 256, seed=0) + ic_pts = sampling.initial_points(domain, 64, seed=1) + ic_vals = problem.initial(ic_pts[:, 0]) + return problem, TrainingData(collocation=coll, ic=(ic_pts, ic_vals)) + + +def test_training_reduces_loss() -> None: + """A few Adam steps reduce the PINN loss (residual autodiff works).""" + problem, data = _burgers_forward_data() + model = builders.build_torch(problem, CFG, "hard_monotone") + _, history = torch_trainer.train( + problem, model, data, weights=LossWeights(), sign_x=-1, lr=3e-3, steps=60 + ) + assert history[-1] < history[0] + assert np.isfinite(history[-1]) + + +def test_hard_monotone_stays_admissible_after_training() -> None: + """The trained hard-monotone field is still non-increasing in x (violation 0).""" + problem, data = _burgers_forward_data() + model = builders.build_torch(problem, CFG, "hard_monotone") + trained, _ = torch_trainer.train( + problem, model, data, weights=LossWeights(), sign_x=-1, lr=3e-3, steps=40 + ) + x_values, _ = sampling.eval_grid(problem.domain, 60, 3) + x = torch.as_tensor(x_values, dtype=torch.float32).reshape(-1, 1) + t = torch.full_like(x, 0.5) + u = trained(x, t).detach().numpy().ravel() + assert violation(u, axis=0, sign=-1) < 1e-5 + + +def test_vanilla_also_trains() -> None: + """The unconstrained baseline also runs and reduces its loss.""" + problem, data = _burgers_forward_data() + model = builders.build_torch(problem, CFG, "vanilla") + _, history = torch_trainer.train( + problem, model, data, weights=LossWeights(), sign_x=-1, lr=3e-3, steps=40 + ) + assert history[-1] < history[0] diff --git a/uv.lock b/uv.lock index 719c805b..2f8034c3 100644 --- a/uv.lock +++ b/uv.lock @@ -1920,6 +1920,9 @@ torch-gpu = [ ] [package.dev-dependencies] +applications = [ + { name = "matplotlib" }, +] bench = [ { name = "matplotlib" }, { name = "optuna" }, @@ -1964,6 +1967,11 @@ lint = [ { name = "types-pyyaml" }, { name = "types-setuptools" }, ] +pinn = [ + { name = "matplotlib" }, + { name = "optax" }, + { name = "optuna" }, +] [package.metadata] requires-dist = [ @@ -1987,6 +1995,7 @@ requires-dist = [ provides-extras = ["torch", "torch-cpu", "jax", "keras", "all", "all-cpu", "torch-gpu", "jax-gpu", "keras-gpu"] [package.metadata.requires-dev] +applications = [{ name = "matplotlib", specifier = ">=3.8" }] bench = [ { name = "matplotlib", specifier = "==3.11.0" }, { name = "optuna", specifier = "==4.9.0" }, @@ -2031,6 +2040,11 @@ lint = [ { name = "types-pyyaml", specifier = "==6.0.12.20260518" }, { name = "types-setuptools", specifier = "==83.0.0.20260706" }, ] +pinn = [ + { name = "matplotlib", specifier = ">=3.8" }, + { name = "optax", specifier = ">=0.2" }, + { name = "optuna", specifier = ">=4.0" }, +] [[package]] name = "mpmath"