ProbHardE2E

This repository contains code and notebooks to reproduce ProbHardE2E PDE benchmarks. The methods are applied to both linear and nonlinear PDEs, enforcing hard constraints, e.g., initial conditions, mass conservation, and total variation diminishing (TVD).

Directory Structure

├── FNO_HardC.ipynb
├── FNO_OrthoC.ipynb
├── ProbHardE2E_Linear.ipynb
├── ProbHardE2E_Soft.ipynb
├── ProbHardE2E_TVD.ipynb
├── ProbhardE2E_Nonlinear_PME.ipynb
├── plot_timing.ipynb
├── models/
│   ├── FNO1d.py, FNO2d.py, UncertainNO.py, DiverseFNO2d.py
├── utils.py
├── datasets.py
├── probconserv.py
├── nonlinear_projection.py
├── commands.sh
├── requirements.txt
└── results/

Installation

To set up the environment, clone the repository and install the required packages:

pip install -r requirements.txt

Reproducing Experimental Results

The following notebooks can be used to reproduce the tables presented in the paper:

• Table 1 (Linear Constraints and Other Baselines):

  • FNO_HardC.ipynb
  • FNO_OrthoC.ipynb
  • ProbHardE2E_Linear.ipynb
  • ProbHardE2E_Soft.ipynb

• Table 2 (Total Variation Diminishing (TVD) Constraint):

  • ProbHardE2E_TVD.ipynb

• Table 3 (Nonlinear Porous Medium Equation Constraints):

  • ProbhardE2E_Nonlinear_PME.ipynb

To run these notebooks:

jupyter notebook

Sources

This repo contains modified versions of the code found in the following repos:

https://github.com/zongyi-li/fourier_neural_operator: For implementation of the Fourier Neural Operator (FNO) (MIT license)

https://github.com/amazon-science/operator-probconserv: For implementation of Variance-NO (Apache 2.0 license)

https://github.com/amazon-science/probconserv: For implementation of ProbConserv (Apache 2.0 license)

Citation

If you use this code, or our work, please cite:

@inproceedings{utkarsh2026_probharde2e,
    title={End-to-end probabilistic framework for learning with hard constraints},
    author={Utkarsh ., Maddix, D.C., Ma, R., Mahoney, M.W., Wang, Y.},
    booktitle={The Fourteenth International Conference on Learning Representations},
    year={2026},
    url={https://openreview.net/forum?id=teE4pl9ftK}
}

Security

See CONTRIBUTING for more information.

License

This project is licensed under the Apache-2.0 License.