PaperlayerView exact revision on arXiv

Source-linked AI summary

Equation Recast for Canonical Operator Learning Across Parametric PDEs

Qiyun Cheng, Valentin Duruisseaux, Cesar F. Clauser, Md Hossain Sahadath, Huihua Yang, Shaowu Pan, Nathaniel Ferraro, Anima Anandkumar, Wei Ji, Cristina Rea

arXiv:2609.02982v1cs.LGphysics.comp-phphysics.plasm-ph

TL;DR

Parametric operator learning must cover families of inverse operators despite sparse, costly, and uneven simulation data. Equation recast analytically shifts operator variation into effective sources so one canonical operator can be reused. Across PDE examples, it supports broader extrapolation and heterogeneous-data use, while convergence behavior can flag an unreliable continuation path.

  • Problem

    Parametric learning must approximate a family of inverse operators, while costly simulations leave parameter coverage sparse and uneven.

    Source: Results · p. 6, Equation Recast Formulation · p. 18

  • Method

    Equation recast analytically absorbs parameter-induced operator changes into effective sources and reuses a single learned canonical operator.

    Source: Discussion · p. 16

  • Results

    Across parameterized PDE examples, the canonical representation supports broader extrapolation and heterogeneous-data integration, including Tokamak simulations across four geometries.

    Source: Discussion · p. 16

  • Takeaways & Limitations

    The equation-level reformulation can combine with different learned operators to support reusable solvers and shared use of sparse simulation data.

    Source: Discussion · p. 16

    The parameter-induced operator change must be known analytically, and extrapolation to unseen Tokamak geometries was not evaluated.

    Source: Discussion · p. 17, Geometry canonicalization · p. 21

Abstract

from arXiv · show

Learning solution operators across broad parameter ranges can require substantial coverage of both input functions and physical parameters, particularly for purely data-driven parametric models. In addition, the resulting models may fail silently outside the training distribution. We introduce equation recast, which reformulates parametric operator learning as the learning of a single canonical operator. Parameter-induced operator variations are derived analytically from the governing equation and absorbed into effective sources, enabling zero-shot prediction across new parameter regimes. Across multi-parameter, nonlinear, and singular PDE settings, equation recast supports extrapolation, integrates sparse heterogeneous datasets in a shared canonical representation, and uses loss of convergence as an internal warning signal for failure of the recast iteration. In high-fidelity tokamak simulations for nuclear fusion, the framework unifies electron-temperature data across four device geometries through canonical-domain mapping within one jointly trained operator. Equation recast provides a route toward reusable neural PDE solvers combining equation-guided transfer, data efficiency, and monitorable inference.

Results

Equation recast supports extrapolation from a single canonical operator, improves learning from heterogeneous data, and exposes some failures through iteration behavior. Results span parameterized PDEs and four Tokamak geometries, with extrapolation bounded by the learned inverse and the target operator’s regularity.

Results

  • A canonical operator trained at (Pe∗, Da∗) = (4, 2) extrapolates with low error across a broad ADR parameter region, with greatest degradation at high Pe and low Da.The reusable region is robust to the canonical point within the moderate-reaction regime, but deteriorates when the canonical operator is advection dominated or uniformly weak.Source: Results · p. 6, Figure 2 · p. 7
  • Recasting off-canonical data into a shared effective-source representation can enrich underrepresented source regions and improve generalization, suggesting targeted acquisition of data for weakly covered regions.A new parameter configuration contributes a distribution of effective sources, not merely a single point in parameter space.Source: Results · pp. 7–8
  • Recast data at k = 0.05 reduce low-k error several-fold and by nearly an order of magnitude at the most extrapolative values, while reallocating the fixed budget modestly reduces accuracy at large k.The improvement extends beyond the added parameter value, consistent with broader effective-source coverage.Source: Results · p. 7, Figure 3 · p. 9
  • Across the tested portion of k ∈ (k3, k4), recast iterations hit their limit and error rises by many orders despite added data; the Helmholtz resonance makes the inverse singular.The loss of convergence reflects operator structure rather than insufficient training coverage.Source: Results · p. 8, Figure 3 · p. 9
  • With 200 samples at Re∗ = 250, equation recast remains accurate over a substantially broader Reynolds-number range than conditioned baselines, especially below 250; at larger Re, errors concentrate in resolved large-scale modes.The parametric FNO and PINO use the same 200 supervised samples distributed over Re ∈ [200, 300].Source: Results · pp. 10–11, Figure 4 · p. 12
  • Across four included Tokamak geometries, one jointly trained canonical operator predicts electron-temperature increments with 2.2–3.4% mean relative L2 error; extrapolation to unseen devices was not tested.Geometry variation is represented through Jacobian-dependent terms, without device-specific retraining.Source: Results · pp. 13–14, Figure 5 · p. 15

Discussion

Equation recast reallocates parametric learning to a canonical inverse plus analytical operator differences, with experiments spanning parameter regimes, heterogeneous data, and geometry. Its reach depends on the learned inverse and analytically known operator changes.

  • Equation recast learns one canonical operator, absorbs known operator differences into effective sources, and reuses the canonical inverse in an iterative solve.Source: Discussion · p. 16
  • Single-point training transfers across broader parameter regions, heterogeneous data can enrich effective-source coverage, and recast convergence flags unreliable continuation.The reaction-diffusion example finds that adding off-canonical data can improve predictions over a broader neighboring regime.Source: Discussion · p. 16
  • Harmonic mapping combines geometry and plasma-coefficient variation in one canonical-domain operator trained on data from different devices.Source: Discussion · p. 16
  • The method requires analytically known operator changes; extrapolation depends on canonical-inverse quality and cannot provide a bounded inverse when the target operator is singular.Source: Discussion · p. 17
  • Recasting offers a route to reusable solvers for sparse simulations, but effective-source coverage and targeted sampling strategies remain open.Source: Discussion · p. 17

Method

Equation recast learns a single canonical inverse and handles parameter variation analytically through the governing equation. Its useful range depends on the target operator, the learned inverse, and whether the recast iteration converges.

Equation Recast Formulation

  • Equation recast learns one canonical inverse and analytically transfers parameter variation into effective sources, rather than fitting the full parameter-indexed inverse family.The recast identity is formulated before discretization and can be paired with different discretizations and neural-operator architectures.Source: Equation Recast Formulation · pp. 18–19

Practical Regime and Canonical Selection

Fixed-Point Consistency and Convergence Behavior

  • With exact canonical inversion, any converged fixed point solves the target PDE; with an approximate inverse, the residual depends on inverse error and effective-source deviation.Observed slow convergence or divergence signals a reduced convergence margin, with resonance producing rapidly increasing iteration counts and eventual loss of convergence in the Helmholtz example.Source: Fixed-Point Consistency and… · pp. 19–20

Geometry canonicalization

  • For two-dimensional simply connected domains, harmonic maps send physical geometries to a shared unit disk, with geometric variation entering analytically through Jacobian metrics.The study uses this representation to train across known geometries; extrapolation to an unseen geometry is not evaluated.Source: Geometry canonicalization · pp. 20–21

Code availability

The study provides code and scripts to reproduce its benchmarks and outputs.

Code availability

  • The MIT-licensed repository includes reference solvers, dataset generators, training and inference scripts, plotting scripts, and figure-to-script mappings for all five benchmark problems.The repository is publicly available at https://github.com/QiyunCh/EquationRecast.Source: Code availability · p. 21

S1 Advection–diffusion–reaction equation

The ADR benchmark tests fixed-point equation recast against varying canonical configurations. Intermediate reference configurations retain broad low-error regions, while two extreme configurations have sharply smaller reusable regions.

S1 Advection–diffusion–reaction equation

  • The steady periodic ADR equation is recast around (Pe*, Da*) = (4, 2), keeping the canonical operator fixed while parameter variation enters through the effective source.A one-dimensional FNO represents the canonical inverse, trained on 1,000 canonical samples.Source: S1 Advection–diffusion–reaction equation · p. 26

S1.1 Sensitivity to the canonical configuration

  • The sensitivity scan tests six canonical points over (Pe, Da) ∈ [1, 50]² with unit spacing, using the same source distribution and training-set size and up to 300 iterations.The scan compares the accessible domain below the 1% mean relative L2 error threshold.Source: S1.1 Sensitivity to the canonical… · pp. 26–27
  • 92–96% of the extended domain remains below 1% mean relative L2 error for canonical configurations (2, 2), (2, 4), and (2, 6); less than 10% does for (1, 1) and (25, 2).Figure S.1 compares six canonical configurations over (Pe, Da) ∈ [1, 50]²; (10, 10) also retains a broad reusable region.Source: S1.1 Sensitivity to the canonical… · p. 27, Figure S.1 · p. 27

S2 Reaction–diffusion equation

The reaction–diffusion benchmark examines recasting data from a noncanonical reaction coefficient into a shared canonical representation and describes the evaluation setup.

S2 Reaction–diffusion equation

  • For a target reaction coefficient k, equation recast places coefficient variation in the effective source while retaining the canonical operator.The equation is one-dimensional and steady, with k > 0.Source: S2 Reaction–diffusion equation · p. 28
  • With a total budget of 1,000 samples, one model uses only canonical data at k* = 0.785, while another combines 500 canonical samples with 500 at k = 0.05 recast as canonical-operator data.The high-fidelity solutions for off-canonical samples remain the supervision targets.Source: S2 Reaction–diffusion equation · p. 28
  • Both models are evaluated on 20 test sources across 27 reaction coefficients in [0.01, 10], with accuracy reported as mean relative L2 error.They use the same one-dimensional FNO architecture; recast inference at a target coefficient is iterative.Source: S2 Reaction–diffusion equation · p. 28

S3 Helmholtz equation

The Helmholtz benchmark tests recasting across a wavenumber scan around resonances, using a canonical operator trained at one reference wavenumber. Off-canonical samples are represented with effective sources while retaining their solutions as supervision targets.

S3 Helmholtz equation

  • At the canonical wavenumber k∗= 0.785, Model 1 uses 1000 training samples; Model 2 replaces half with off-canonical samples.The supplied setup passage truncates before stating the off-canonical training wavenumber.Source: S3 Helmholtz equation · p. 29
  • For an off-canonical pair, the effective source is the canonical-operator input while the high-fidelity solution remains the supervision target.This places canonical and off-canonical samples in one learning problem.Source: S3 Helmholtz equation · p. 29
  • Inference uses Aitken’s ∆2 relaxation and evaluates 20 test sources across k ∈[0.630, 1.200], including neighborhoods of resonances k2, k3, and k4.Reference solutions use the same finite-difference solver.Source: S3 Helmholtz equation · p. 29

S4 Two-dimensional steady Navier–Stokes equation

The Navier–Stokes benchmark tests equation recast over Reynolds numbers using a canonical FNO and compares iterative updates and training-data strategies. Redistributing heterogeneous samples into the canonical representation yields the strongest data-budget scaling, while Aitken relaxation is the robust default among tested updates.

S4 Two-dimensional steady Navier–Stokes equation

  • The benchmark uses a 128 × 128 periodic grid, a canonical dataset of 200 samples at Re∗= 250, and conditioned baselines trained on 200 samples over Re ∈[200, 300].All models use the same FNO backbone; PINO adds a governing-equation residual to its objective.Source: S4 Two-dimensional steady… · p. 30

S4.1 Inference behavior and iterative-solver ablation

  • Iteration counts are smallest near Re∗= 250, where the recast correction vanishes, and increase toward the scan edges as larger fixed-point corrections are required.Iteration count measures numerical effort, not prediction accuracy.Source: S4.1 Inference behavior and… · p. 31
  • Aitken relaxation converges in all 300 cases; Anderson mixing fails in 21, and fixed under-relaxation fails in 53 with considerably larger final prediction error.These results motivate Aitken as the default update for this benchmark.Source: S4.1 Inference behavior and… · p. 31
  • Equation recast costs approximately 7.8 ms per sample near Re∗ and 13.6 ms toward the tested extremes, versus 20.1–43.7 ms for the reference solver and 0.99 ms for one neural-operator pass.The comparison matches prediction accuracy; recast adds iterative overhead over direct inference but remains faster than the reference solver.Source: S4.1 Inference behavior and… · p. 33

S4.2 Data-budget and heterogeneous-data ablation

  • Redistributed training recasts equal-size sample blocks from Re ∈{50, 100, 200, 300, 400} into the shared canonical representation.The comparison includes canonical recast, full-range parametric FNO, PINO, and redistributed recast; the redistributed model is limited at N = 200, with 40 samples per Reynolds number.Source: S4.2 Data-budget and… · pp. 33–34
  • At N = 1500, the redistributed equation-recast model reaches 4.0% mean relative L2 error and has the lowest error from N = 500 onward.It shows the strongest improvement as the data budget increases.Source: S4.2 Data-budget and… · p. 34
  • As data availability rises, all model families improve; the redistributed recast gains most and reduces extrapolation error over most of the Reynolds-number range.The PINO maintains low PDE residuals across budgets, while its reference-based L2 error improves more modestly.Source: S4.2 Data-budget and… · p. 35

S5 Multi-device Tokamak MHD simulation

The Tokamak benchmark maps four device geometries to a shared canonical domain and incorporates coefficient and geometry variation through an effective source. The same recast inputs work with both FNO and LocalNO architectures, with LocalNO-L attaining the lowest reported increment errors.

S5 Multi-device Tokamak MHD simulation

  • The benchmark models one-step electron-temperature updates from high-fidelity simulations across four device geometries, with other simulated physics entering through time- and space-dependent coefficients and sources.Advective transport is not included; the increment target is more sensitive to learned evolution than the background-dominated full temperature field.Source: S5 Multi-device Tokamak MHD simulation · p. 36
  • Coefficient variation and geometry-induced metric variation enter through the same canonical effective-source representation.Alcator C-Mod is the reference geometry; harmonic maps and their Jacobians are computed once per device and reused in preprocessing and inference.Source: S5 Multi-device Tokamak MHD simulation · p. 36
  • The benchmark combines 4,208 one-step pairs from all four geometries, so it tests multi-device unification rather than extrapolation to an unseen geometry.Fifteen percent of pairs are held out for validation, and fields are mapped to a 256 × 256 canonical disk.Source: S5 Multi-device Tokamak MHD simulation · pp. 36–37

S5.1 Neural-operator architecture ablation

  • LocalNO-L has the lowest mean relative L2 errors across all four device geometries, between 2.2% and 3.4%, while both FNO and LocalNO achieve low error with the same recast inputs.The result uses the more discriminating temperature-increment target; full-temperature errors are sub-percent but compress architectural differences.Source: S5.1 Neural-operator architecture… · p. 37, Table S.3 · p. 37, Figure S.5 · p. 38
Loading 2609.02982v1…