MD-PNOP

This repo includes reproduction code for "MD-PNOP: Equation-recast neural operators for minimal-data extrapolation and PDE solver acceleration", demonstrated on the one-dimensional neutron transport equation.

Q. Cheng, M. H. Sahadath, H. Yang, S. Pan, and W. Ji, "MD-PNOP: Equation-Recast Neural Operators for Minimal-Data Extrapolation and PDE Solver Acceleration," Journal of Computational Physics, 114844 (2026). https://doi.org/10.1016/j.jcp.2026.114844

Preprint: arXiv:2509.01416 — https://arxiv.org/abs/2509.01416

---

Summary

Neural operators can infer PDE solutions orders of magnitude faster than traditional solvers, but they generalize poorly outside the parameter range seen in training and usually require retraining for each new configuration. MD-PNOP (Minimal-Data Parametric Neural Operator Preconditioning) trys to address this limitation. The idea, inspired by perturbation theory, is named as equation recast: a new parameter set $\mathbf{p}'$ is written as a perturbation of the single setting $\mathbf{p}^*$ used for training, $\mathbf{p}' = \mathbf{p}^* + \Delta\mathbf{p}$, and the resulting operator difference is moved to the right-hand side as an additional source term. A neural operator trained at $\mathbf{p}^*$ can then be applied, without retraining, to any new $\mathbf{p}'$ by repeatedly evaluating this modified source.

$$ u^{(k+1)}(x) = \mathcal{L}_{NO}^{-1}(p^*) \left[S(x)-\delta \mathcal{L}(\Delta p)\left[u^{(k)}(x)\right] \right], \quad k=0,1,2,\ldots $$

In the MD-PNOP framework, the operator output is used only as an initial guess for a conventional iterative solver, the final solution is produced entirely by the governing equations — so it is fully physics-constrained and retains full-order accuracy, while the solver converges in fewer iterations. The framework is architecture-agnostic and it is demonstrated here with both a Deep Operator Network (DeepONet) and a Fourier Neural Operator (FNO), on fixed-source, single-group eigenvalue, and multigroup eigenvalue neutron transport problems. A single operator trained on one constant-parameter configuration accelerates problems with heterogeneous, sinusoidal, and discontinuous parameter distributions.

Repository layout

MD-PNOP/
├── common.py                     # operator loaders, normalization constants, GRF source
├── models/                       # the SAME two trained operators are reused everywhere
│   ├── deeponet/  deeponet.py    # architecture
│   │              deeponet.pt    # trained weights
│   └── fno/       fno.py
│                  fno.pt
├── FixedSource/
│   ├── Solver/                   # numerical (reference) + standalone NN + recast-preconditioned solvers
│   ├── sources.npz               # the exact GRF source arrays used in the paper
│   ├── generate_sources.py       # how sources.npz was produced (GRF, seeds 13/1029/618)
│   ├── test1.py                  # Case 1: constant cross sections (= training reference)
│   ├── test2.py                  # Case 2: constant + anisotropic scattering
│   └── test3.py                  # Case 3: sinusoidal cross sections + anisotropic scattering
└── Eigenvalue/
    ├── Single/                   # single-group k-eigenvalue
    │   ├── Solver/               # numerical + SP/CP recast preconditioners
    │   ├── test1.py              # Case 1: constant (= training reference)
    │   ├── test2.py              # Case 2: sinusoidal
    │   └── test3.py              # Case 3: piecewise / discontinuous (multi-slab)
    └── Multi/                    # multigroup (3-group), SP only
        ├── Solver/
        └── test1.py

Each testN.py is self-contained: it defines its cross sections, runs every solver, prints the errors (and $k$ for the eigenvalue cases), and saves the figure. The two operators in models/ are shared by all cases without retraining — this is the minimal-data point of the method. Only the architectures and trained weights are provided; the training data and training scripts are not needed to run the tests.

Reproducibility notes. The fixed-source distributions are Gaussian-random-field realizations (seeds 13/1029/618). A GRF draw over this near-singular kernel depends on the numpy/LAPACK version, so the exact arrays used in the paper (numpy 2.1.3) are shipped in FixedSource/sources.npz and loaded by the tests, making the results independent of the installed numpy version (generate_sources.py shows how they were produced). The eigenvalue problems use deterministic cross sections. For the eigenvalue cases, all solvers (benchmark included) use a uniform 1e-4 convergence criterion; the constrained-preconditioning (CP) step count N_pre is 100, 10 and 1 for cases 1, 2 and 3 (case 3 is a high-dominance-ratio problem where iterating the operator accumulates contamination, so a single step is best).

Setup

pip install -r requirements.txt

A CUDA GPU is used automatically if available, otherwise the code runs on CPU.

Running

cd FixedSource        && python test1.py     # (and test2.py, test3.py)
cd Eigenvalue/Single  && python test1.py     # (and test2.py, test3.py)
cd Eigenvalue/Multi   && python test1.py

Scope and notes on performance

This repository and the paper are intended to introduce and demonstrate the equation-recast method and the MD-PNOP framework, not to ship a SOTA solver. For practical use, the achievable speedup is considerably larger than what these scripts show, through:

1. Warm-up vs. cold start. The wall-clock times reported in the paper are cold-start times: each solver is timed from a fresh process with no warm-up, so they include model loading, the first CUDA kernel launch, and disk I/O, and depend on the machine's I/O. With warm-up (a discarded first call, or amortizing the one-time model load over many solves) the neural-operator solvers are substantially faster. Absolute timings are therefore not comparable across machines; the reproducible quantities are the flux profiles and eigenvalues.

2. Code-level acceleration. The transport sweeps can be vectorized over the discrete ordinates, and group sources computed once per group (as already done in the multigroup solver), giving large speedups with identical results; moving the sweep to a compiled/GPU backend gives more. These do not change accuracy.

3. Relaxation / acceleration schemes. Under-relaxation, Anderson acceleration, and similar schemes can be applied to the recast fixed point and to the source/power iterations to further reduce iteration counts.

Acknowledgement

The scripts in this repository were organized and refactored with the assistance of Claude (Anthropic) AI, and were manually verified by the authors to reproduce the published method and results.