# ProbHardE2E-PDEs

This repository contains code and notebooks to reproduce ProbHardE2E PDE benchmarks using VarianceNO as the base unconstrained model from https://github.com/amazon-science/operator-probconserv, which is based on the Fourier Neural Operator from https://github.com/zongyi-li/fourier_neural_operator. The methods are applied to both linear and nonlinear PDEs, enforcing constraints such as initial conditions, mass conservation, and total variation denoising.

## Directory Structure

```
probharde2e-pdes/
├── 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:

```bash
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 Constraint)**:
  - `ProbHardE2E-TVD.ipynb`

- **Table 3 (Nonlinear Porous Medium Equation Constraints)**:
  - `ProbhardE2E_Nonlinear_PME.ipynb`

To run these notebooks:

```bash
jupyter notebook
```
