Add gray-box term-discovery example#1835
Conversation
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>
Greptile SummaryThis PR adds a gray-box term-discovery example for a reaction-diffusion PDE. The main changes are:
Important Files Changed
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) |
There was a problem hiding this comment.
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( |
There was a problem hiding this comment.
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>
Add gray-box term-discovery example
Closes #1834
Description
Adds
examples/graybox_discovery/, an example that recovers an unknown term in aPDE (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 withRunknown.A field network learns
u(x, t)and a second network learnsR(u), trainedjointly from point observations. The residual is declared in SymPy with
Rleftas a
Function; aPhysicsInformerevaluates it each step (autodiff for thespatial derivative, a manual
torch.autograd.gradfor the time derivative sinceautodiff 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 observesu <= 0.75, which leavesRunder-determined nearu = 1. Adding the knownequilibria
R(0) = R(1) = 0as a direct penalty on the reaction network pinsthe closure in the unobserved range.
What's included
examples/graybox_discovery/generate_data.py— synthetic dataexamples/graybox_discovery/graybox_reaction_diffusion.py— training scriptexamples/graybox_discovery/conf/config.yaml— configexamples/graybox_discovery/README.md— description, run instructions, andthe recovered-vs-true R(u) plot
examples/graybox_discovery/graybox_reference.py— numpy-only reproductionof the mechanism (no PhysicsNeMo/GPU required)
examples/graybox_discovery/test_graybox_verifiers.py— checks the priorreduces the spread of the recovered closure and improves recovery in the
unobserved region (per
CONTRIBUTING.md, example tests live underexamples/, not the top-leveltest/)No changes to core APIs.
Testing
pytest examples/graybox_discovery/test_graybox_verifiers.pyruns thereference-based checks with no GPU/torch requirement (one additional guarded
test is skipped until the
physicsnemo/sym/eq/verifiers.pyfollow-up PRlands). 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 over20k 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 ismeant to produce -- see the plot in the example README.
Checklist
pytest examples/graybox_discovery/test_graybox_verifiers.pyruff check/ruff format --check(clean)test/ci_tests/header_check.pyandtest/ci_tests/check_docstring_coverage.pylocally (clean)markdownlint-clion changed Markdown (clean)CHANGELOG.md