Skip to content

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

Description

@oroikono

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

What I'd like to add

An example that recovers the unknown functional form of a term in a PDE, not
just an unknown scalar coefficient. The setup is a closure/source term whose
shape isn't known — the standard situation in subgrid/turbulence and reaction
modeling.

Why it fits the existing design

Inverse PINNs in physicsnemo.sym already recover unknown scalar coefficients by
detaching known fields and letting a small network produce the unknown value. The
same graph machinery covers an unknown term: I declare the term as a Function
in the SymPy residual and provide it with a second network that maps the state to
the term value. The residual assembly (SympyToTorch / Add.make_args) and
autodiff are unchanged. So this needs no core change — it's an example that shows
a capability the framework already supports but doesn't demonstrate.

Sketch:

class GrayBoxReactionDiffusion(PDE):
    def __init__(self, D=0.02):
        x, t = Symbol("x"), Symbol("t")
        u = Function("u")(x, t)
        R = Function("R")(x, t)            # supplied by a network node
        self.equations = {"reaction_diffusion": u.diff(t) - D*u.diff(x, 2) - R}

r_net = instantiate_arch(input_keys=[Key("u")], output_keys=[Key("R")], ...)
nodes = eq.make_nodes() + [u_net.make_node("u_network"), r_net.make_node("reaction_network")]

Physics priors (e.g. known equilibria R(0)=R(1)=0) go in as ordinary
PointwiseConstraints on the reaction network. In the example they matter: the
data only covers part of the state range, so the term is under-determined without
them.

Scope

I'd start with a single self-contained example under examples/graybox_discovery/
(data generator, training script, config, README) plus a small test. If there's
interest, a natural follow-up is a thin helper layer — a discover_terms= option
on PDE.make_nodes that generalizes the existing freeze_terms indexing, and a
few reusable physics-constraint classes — but I'd keep that out of the first PR.

Prior art

Learnable-term-in-PINN with symbolic post-processing is established (e.g.
Kiyani et al. 2023, arXiv:2305.10706; UPINN, ICML 2023; AI-Aristotle, 2024;
PIKAN gray-box, 2025). The example just makes the pattern reproducible inside
PhysicsNeMo with the physics-prior constraints attached.

Logistics

Happy to implement and maintain it, follow the style guide (Apache headers,
docstrings, tests), and sign off commits under the DCO. Opening this first to
check for overlap and get feedback on where the example should live.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions