# PDE-PFN: Prior-Data Fitted Neural PDE Solver Under Arbitrary Mesh

We provide codes for **family of 2D CDR equations**, **shallow water equations (SWE)**, **compressible Navier-Stokes equations (CNSE)**, and **Darcy Flow** experiments.

In this document, we will describe the procedures of the experiments:
* Environment setup, 
* Data generation or download
* Model pre-training(fine-tuning) and testing 



## Setup Requirements

You can set up the environment with following command:

```
conda env update -f environment.yaml
```

---


## Pre-training on Family of 2D CDR Equations



### Set Directory
First, set the top-level folder we provided as the working directory.

### Data Generation (if needed)
1. Run the Python file as shown below to create the directory for storing the dataset.
```
python data_gen/utils/dataset_util_2D.py
```
2. Execute below shell file to generate the seen coefficient dataset. The logs will be stored in the `data_gen/Generation_logs` directory.
```
bash data_gen/Generation_shell/2D_cdr/generate_2D_cdrrr.sh
```

3. Execute below shell file to generate the PINN prior seen coefficient dataset. The logs will be stored in the `data_gen/Generation_logs` directory.
```
bash data_gen/Generation_shell/2D_cdr/generate_2D_cdrrr_PINN.sh
```

4. Execute below shell file to generate the unseen coefficient dataset. The logs will be stored in the `data_gen/Generation_logs` directory.
```
bash data_gen/Generation_shell/2D_cdr/generate_2D_cdrrr_unseen_inter.sh
bash data_gen/Generation_shell/2D_cdr/generate_2D_cdrrr_unseen_extra.sh
```
Before running this `.sh` file, ensure that your GPU is available. Modify the file accordingly based on your hardware configuration.


### Pre-train Model on the Family of 2D CDR Equations
To pre-train PDE-PFN on the famliy of 2D CDR equations, execute the shell file as follows. 

```
bash shell/2D_cdr/total_2d_cdr.sh
```
Before running this `.sh` file, ensure that your GPU is available. Modify the file accordingly based on your hardware configuration.

If there is no problem, 
1. The logs will be stored in `logs/PDE_PFN/2D_cdr`. You can check both the training and testing results inside the log file.
2. the directory `parameter` is generated to save model parameter.
3. The trained model's parameters will be saved inside the `parameter/PDE_PFN/2D_cdr` folder


For pre-training using PINN prior, please use the corresponding shell files:
```
bash shell/2D_cdr/total_2d_cdr_PINN.sh
```

**❗ we provide the pretrained weight with codes ❗**


## Different PDE Fine-tuning Experiments

### Set Directory
First, set the top-level folder we provided as the working directory.

### Data Download
Run the Python file as shown below to download the datasets.
```
python data_gen/Download_direct.py
```
Then you will get the preprocessed dataset.

### Finetuning Model
:warning: Before fine-tuning the model, you must pretrain the model on 2D CDR equation first. :warning:

To fine-tune our model, run the corresponding shell file. 
We will use the SWE dataset as the basis for our explanation.

```
bash shell/SWE/total_SWE.sh
```
Before running this `.sh` file, ensure that your GPU is available. Modify the file accordingly based on your hardware configuration.

If there is no problem, 
1. The logs will be stored in `logs/PDE_PFN/SWE`. You can check both the training and testing results inside the log file.
2. the directory `parameter` is generated to save net parameter.
3. The trained model's parameters will be saved inside the `parameter/PDE_PFN/SWE` folder.