Skip to content

Add gray-box term-discovery example#1835

Open
oroikono wants to merge 2 commits into
NVIDIA:mainfrom
oroikono:graybox-discovery-example
Open

Add gray-box term-discovery example#1835
oroikono wants to merge 2 commits into
NVIDIA:mainfrom
oroikono:graybox-discovery-example

Conversation

@oroikono

Copy link
Copy Markdown

Add gray-box term-discovery example

Closes #1834

Description

Adds examples/graybox_discovery/, an example that recovers an unknown term in a
PDE (a reaction/closure term) rather than an unknown scalar coefficient.

A 1D reaction-diffusion problem u_t = D u_xx + R(u) is set up with R unknown.
A field network learns u(x, t) and a second network learns R(u), trained
jointly from point observations. The residual is declared in SymPy with R left
as a Function; a PhysicsInformer evaluates it each step (autodiff for the
spatial derivative, a manual torch.autograd.grad for the time derivative since
autodiff spatial gradients only cover x/y/z), so no change to the
residual-assembly code is needed.

The generated data uses a non-polynomial R(u) = 3 sin(pi u) and only observes
u <= 0.75, which leaves R under-determined near u = 1. Adding the known
equilibria R(0) = R(1) = 0 as a direct penalty on the reaction network pins
the closure in the unobserved range.

What's included

  • examples/graybox_discovery/generate_data.py — synthetic data
  • examples/graybox_discovery/graybox_reaction_diffusion.py — training script
  • examples/graybox_discovery/conf/config.yaml — config
  • examples/graybox_discovery/README.md — description, run instructions, and
    the recovered-vs-true R(u) plot
  • examples/graybox_discovery/graybox_reference.py — numpy-only reproduction
    of the mechanism (no PhysicsNeMo/GPU required)
  • examples/graybox_discovery/test_graybox_verifiers.py — checks the prior
    reduces the spread of the recovered closure and improves recovery in the
    unobserved region (per CONTRIBUTING.md, example tests live under
    examples/, not the top-level test/)

No changes to core APIs.

Testing

pytest examples/graybox_discovery/test_graybox_verifiers.py runs the
reference-based checks with no GPU/torch requirement (one additional guarded
test is skipped until the physicsnemo/sym/eq/verifiers.py follow-up PR
lands). The example itself was run end to end in a real PhysicsNeMo
environment via python generate_data.py && python graybox_reaction_diffusion.py; training converges (loss ~0.12 -> ~5e-4 over
20k steps) and the recovered R(u) tracks the true closure well past the
observed range (u_obs = 0.75), which is the effect the physics prior is
meant to produce -- see the plot in the example README.

Checklist

  • Commits signed off (DCO)
  • Apache-2.0 headers on new files
  • Docstrings on new public functions/classes
  • Ran pytest examples/graybox_discovery/test_graybox_verifiers.py
  • Ran ruff check / ruff format --check (clean)
  • Ran test/ci_tests/header_check.py and
    test/ci_tests/check_docstring_coverage.py locally (clean)
  • Ran markdownlint-cli on changed Markdown (clean)
  • Updated CHANGELOG.md

Adds examples/graybox_discovery/, an example that recovers an unknown term
in a PDE (a reaction/closure term) rather than an unknown scalar
coefficient. A field network learns u(x, t) and a second network learns
R(u), trained jointly via a PhysicsInformer residual; a physics prior
(known equilibria R(0) = R(1) = 0) pins the closure in a region the data
does not cover.

- graybox_reaction_diffusion.py / conf/config.yaml: training script and
  config, built on physicsnemo.sym.eq.phy_informer.PhysicsInformer (the
  current physicsnemo.sym API; Solver/Domain/PointwiseConstraint and
  instantiate_arch no longer exist on main).
- generate_data.py: synthetic data generator.
- graybox_reference.py / test_graybox_verifiers.py: a standalone
  numpy+sympy reproduction of the mechanism and its regression tests, so
  CI can validate the core result without GPU/torch. Both live under
  examples/graybox_discovery/ per CONTRIBUTING.md ("example code ... must
  be tested in examples", not the top-level test/).
- README.md: run instructions and the recovered-vs-true R(u) plot from an
  actual training run (loss 0.12 -> 5e-4 over 20k steps; the recovered
  closure tracks the true one well past the observed range u_obs=0.75).

No changes to core APIs. Verified locally against the repo's own tooling:
pytest, ruff (v0.12.5, matching .pre-commit-config.yaml), header_check.py,
check_docstring_coverage.py, and markdownlint-cli (v0.35.0).

Signed-off-by: Orestis Oikonomou <ooikonomou@ethz.ch>
@copy-pr-bot

copy-pr-bot Bot commented Jul 13, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@greptile-apps

greptile-apps Bot commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds a gray-box term-discovery example for a reaction-diffusion PDE. The main changes are:

  • New synthetic data generation and committed example artifacts.
  • New Hydra training script with field and reaction networks.
  • New NumPy/SymPy reference implementation and example tests.
  • New README and changelog entry.

Important Files Changed

Filename Overview
examples/graybox_discovery/graybox_reaction_diffusion.py Adds the main training script; the default data load depends on the caller's working directory.
examples/graybox_discovery/generate_data.py Adds deterministic data generation; the output file is written to the caller's working directory.
examples/graybox_discovery/graybox_reference.py Adds a deterministic NumPy/SymPy reference implementation for the example tests.
examples/graybox_discovery/test_graybox_verifiers.py Adds tests for the reference mechanism and a guarded verifier-helper import check.
examples/graybox_discovery/conf/config.yaml Adds training, scheduler, batch, loss, and data path configuration.
examples/graybox_discovery/README.md Documents the example workflow and expected recovered reaction plot.
CHANGELOG.md Adds a changelog entry for the new gray-box discovery example.

Reviews (1): Last reviewed commit: "Add gray-box term-discovery example" | Re-trigger Greptile

log = PythonLogger(name="graybox_discovery")
log.file_logging()

data = np.load(cfg.data.data_file)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Data Path Depends On CWD

When the training script is launched from the repo root, np.load("graybox_data.npz") looks in the repo root instead of examples/graybox_discovery/, so it misses the bundled data file and can fail before training starts. Resolve the default relative path from the example directory, while still allowing absolute override paths.

ci = rng.choice(X.size, size=4000, replace=False)
x_pde, t_pde = X.ravel()[ci][:, None], Tm.ravel()[ci][:, None]

np.savez(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Generated Data Leaves Example Directory

When this script is run as python examples/graybox_discovery/generate_data.py, it writes graybox_data.npz into the caller's current directory instead of next to the example. A later training run that expects the example-local data can then read stale bundled data or fail to find the newly generated file.

Addresses the Greptile review on NVIDIA#1835: both scripts resolved
graybox_data.npz against the caller's working directory, so running
them from the repo root (rather than inside examples/graybox_discovery/)
wrote the data outside the example and then failed to find it at load
time.

- generate_data.py writes graybox_data.npz next to the script.
- graybox_reaction_diffusion.py resolves a relative cfg.data.data_file
  against the example directory (absolute paths are still honored as-is),
  and writes recovered_R.png there too -- the same working-directory bug
  for the plot output. The anchor is also robust to Hydra's optional
  runtime change of working directory.
- README notes that the data and plot always land in the example
  directory regardless of the launch directory.

Verified by running both from the repo root: generate_data.py writes into
examples/graybox_discovery/ (nothing at CWD), and the training script
loads that data and saves the plot into the example dir.

Signed-off-by: Orestis Oikonomou <ooikonomou@ethz.ch>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Example request: gray-box discovery of an unknown PDE term

1 participant