# Reproducing Experiments
This guide provides a step-by-step process for reproducing all experiments described in the submission.

**Hardware Specifications**:  All experiments were run on a single compute node with: Dual AMD EPYC 9124 CPUs (32 total CPU cores), 768 GB of DDR5 4800 MHz memory, and 8 NVIDIA RTX 6000 Ada Generation GPUs.  

**Environment Setup**:  Throughout all experiments, we use Python 3.11.2, CUDA 12.4.1, and OpenSSL 1.1.1k FIPS (25 Mar 2021). All other required Python packages match those specified by [PDEArena](https://pdearena.github.io/pdearena/). To install the same package used in our work, we have provided the `requirements.txt` file. 
1. Install `cliffordlayers` package by following the instructions in [PDEArena](https://pdearena.github.io/pdearena/).
2. Then, modify the `requirements.txt` file to point to the correct path of `cliffordlayers` package.
3. You can install the packages by running the following command:
  ```bash
  pip install -r requirements.txt
  ```
We also recommend installing compatible OpenSSL to avoid potential errors.  

## Data Generation
Due to the size limit of submission systems and the anonymous requirements, we do not provide the whole generated dataset. Instead, we provide the seeded data generation pipeline to generate the data. 

### Data Configurations
We provide all data configurations for all considered resolutions. As mentioned in our paper, applying the numerical solver during the training phase would be extremely time-consuming; therefore, we generate the low-resolution solution before the training. Here, we describe the file structure to help locate the data configurations and the scripts for generating the whole dataset. 
```
pdemoe/
├── pdedatagen/ 
│   └── pdemoe_configs/
│       ├── navierstokes2dsmoke_xcoarse.yaml
│       ├── navierstokes2dsmoke_coarse.yaml
│       ├── navierstokes2dsmoke_medium.yaml
│       ├── navierstokes2dsmoke_fine.yaml
│       └── navierstokes2dsmoke_high_res.yaml 
├── pdemoe_data/ 
├── SCRIPTS_DATA_GENERATION.sh 
├── LICENSE
└── .gitignore
```

* `pdedatagen/pdemoe_configs/` includes all configuration files we use in this project. 
* The training data will be generated and saved to `pdemoe_data/`.
* The script file `SCRIPTS_DATA_GENERATION.sh` are used to generate the training data.   

### Generate Data
We present here the step-by-step introduction to generate the XCoarse-resolution dataset. 
1. Investigate the file `SCRIPTS_DATA_GENERATION.sh`. Set up the environment parameters (e.g. activate the virtual environment, export OpenSSL path, load required modules). Minimum $64$ GB memory is suggested. 
2. Setup the following parameters: `NUM_SAMPLES` (we use $100$ in our paper), `CONFIG_PATH` (put the absolute path of the configuration file `navierstokes2dsmoke_xcoarse.yaml`), and `DATA_PATH`(put the absolute path of the folder `pdemoe_data/xcoarse/`). 
3. In the command line, run the following command:
    ```bash
    sbatch SCRIPTS_DATA_GENERATION.sh
    ```
    Then it will create a `pde-datagen` task on the cluster. 

By mannually updating the provided `SCRIPTS_DATA_GENERATION.sh` script, it can create multiple parallel runs on the cluster machine to accelerate the generation process. The runtime warning `RuntimeWarning: divide by zero encountered in divide
  inv_k2 = np.where(k2 > 0, 1.0 / k2, 0.0)` will be raised but it doesn't affect the results. 

## Training
After the data generation process ends, the folder structure looks like this:
```
pdemoe/
├── pdedatagen/ 
│   └── pdemoe_configs/
│       ├── navierstokes2dsmoke_xcoarse.yaml
│       ├── navierstokes2dsmoke_coarse.yaml
│       ├── navierstokes2dsmoke_medium.yaml
│       ├── navierstokes2dsmoke_fine.yaml
│       └── navierstokes2dsmoke_high_res.yaml 
├── pdemoe_data/ 
│   ├── xcoarse/
│   ├── coarse/
│   ├── medium/
│   ├── fine/ 
│   └── gt/
├── SCRIPTS_DATA_GENERATION.sh 
├── SCRIPTS_BASELINES.sh 
├── LICENSE
└── .gitignore
```

### Reproducing Baselines 
We have included the script to reproduce all baselines used in this paper. Here is the step-by-step introduction to repeat the baseline of ResNet-128 model:
1. Investigate the file `SCRIPTS_BASELINES.sh`. Change `MODEL="<FILL IN MODEL NAME>"` to, for example, `MODEL="ResNet-128".
2. Specify the data directory in `--data.data_dir="<PUT YOUR DATA DIR>"`.
3. Edit `pdemoe_train_scripts/train_hybrid.py` to use the desired logger and hyper-parameters. 
4. Run `SCRIPTS_BASELINES.sh`. 

Then, it will create a `pde-baselines` task on the cluster machine. 

### Reproducing PDEMoE
Similar to the baseline, we have provided the script to reproduce the PDEMoE model. 
1. Investigate the file `SCRIPTS_PDEMOE.sh`. Change `MODEL="<FILL IN MODEL NAME>"` to, for example, `MODEL="ResNet-128"`. This parameter is used to specify the base model used in the MF-MoE framework (including the residual approximation and the pure neural operator).  
2. Specify the data directory in `--data.data_dir="<PUT YOUR DATA DIR>"`.
3. Edit `pdemoe_train_scripts/train_hybrid.py` to use the desired logger and hyper-parameters. 
4. Run `SCRIPTS_PDEMOE.sh`. 

Then, it will create a `pde-pdemoe` task on the cluster machine. 

---

The following content provides additional information for reproducing other experiments described in the submission.

## 1. Physics-Aware Features Construction

In this experiment we control two senario: (1) the physical feature "viscosity $\mu$" is included in the input of the gating network; (2)  the physical feature is not included in the input of the gating network.

We fix the same gating network structure and fix the training parameters across both scenarios to ensure a fair comparison. By comparing the performance of these two settings, we can evaluate the impact of incorporating physical domain knowledge into the model's routing decisions.

### Gating Neural Network Architecture

The gating neural network employs a multi-layer perceptron (MLP) architecture implemented through the `GatingModel` class. It consists of three fully connected layers: the first layer transforms the high-dimensional input (with dimensions 4×3×128×128) to a hidden representation, followed by a second hidden layer of the same dimension, both activated by ReLU functions. A notable design feature is the optional inclusion of physical parameters, specifically viscosity, as an additional input feature to the final layer. When the `use_fidelity` flag is enabled, the viscosity value is concatenated with the hidden representation before passing through the final linear layer, which outputs logits corresponding to each expert in the mixture. During training, these logits are processed through a Gumbel-Softmax function with a temperature parameter to produce differentiable discrete selections, enabling the model to perform hard routing while maintaining gradient flow through the backpropagation process.

### Training and Evaluation

We use four pre-trained hybrid neural network using different base neural operators to validate the effect of physics-aware gating network. This operation is used to avoid the additional noises introduced by the neural operator updating. The training data is described in the paper. We use the test error in the last 30 logged steps during training to evaluate the performance.

## 2. Gating Mechanism

In this experiment, we systematically analyze how the constraint parameter $c$ influences the routing probabilities in our gating neural network. By testing three distinct values ($c=0,10,100$), we observe significant variations in the network's routing behavior. Additionally, we carefully calibrate the hyper-parameter $\lambda$, which governs the time penalty strength, to accentuate the performance differences between experts. This fine-tuning ensures that the trade-offs between computational efficiency and solution accuracy are clearly demonstrated across our experimental conditions.

## 3. Data Configuration

The Navier-Stokes 2D smoke simulation is configured with the following parameters:

### Temporal Settings

- Time range: $t \in [18.0, 60.0]$
- Number of timesteps: 8-512 (at minimum dt)
- Initial timesteps to ignore: 0

### Spatial Settings

- Domain size: $L_x = 32.0, L_y = 32.0$
- Spatial resolution: $16 \times 16$  ~ $256\times 256$ grid points
- Components:
  - Scalar components: 1
  - Vector components: 1

### Physical Parameters

- Viscosity ($\nu$): Manually defined in the generation process.
- Buoyancy force in vertical direction: Manually defined in the generation process.
- Correction strength: 1.0
- Force strength: 1.0
- Force frequency: 4

We also provide a generation instruction to generate the data set used in our experiment.
