## Installation

Ensure that you have Python 3.10 or later installed on your system.
Our code is GPU-only.
We highly recommend using the most recent versions of JAX and JAX-lib, along with compatible CUDA and cuDNN versions.
The code has been tested and confirmed to work with the following versions:

- JAX 0.4.35
- CUDA 12.4
- cuDNN 8.9

You can install the latest versions of JAX and JAX-lib with the following commands:
```
pip3 install -U pip
pip3 install --upgrade jax jaxlib
```

Install the package with the following commands:

``` 
pip install .
```

## Quickstart

We use [Weights & Biases](https://wandb.ai/site) to log and monitor training metrics. 
Please ensure you have Weights & Biases installed and properly set up with your account before proceeding. 
You can follow the installation guide provided [here](https://docs.wandb.ai/quickstart).

To illustrate how to use our code, we will use the advection equation as an example. 
First, navigate to the advection directory within the `examples` folder:

``` 
cd jaxpi/examples/allen_cahn
``` 
To train the model, run the following command:
```
python3 main.py 
```

To customize your experiment configuration, you may want to specify a different config file as follows:

```
python3 main.py --config=configs/soap.py 
```


Our code automatically supports multi-GPU execution. 
You can specify the GPUs you want to use with the `CUDA_VISIBLE_DEVICES` environment variable. For example, to use the first two GPUs (0 and 1), use the following command:

```
CUDA_VISIBLE_DEVICES=0,1 python3 main.py
```

**Note on Memory Usage**: Different models and examples may require varying amounts of GPU memory. 
If you encounter an out-of-memory error, you can decrease the batch size using the `--config.batch_size_per_device` option.

To evaluate the model's performance, you can switch to evaluation mode with the following command:

```
python3 main.py --config.mode=eval
```


**Warning:** The data of 2D PDEs is too large to be uploaded. Please test 1D PDEs.

## Results


### Comparison of Optimizer Performance

*Comparison of optimizer performance obtained by training PINNs with Adam and SOAP, respectively, across various PDEs, following the training pipeline. The evaluation metric is the relative $L^2$ error over the entire spatio-temporal domain.*

| Benchmark | Adam | SOAP |
|-----------|------|------|
| Wave | ${5.15 \times 10^{-5}}$ | **${8.05 \times 10^{-6}}$** |
| Burgers | ${8.20 \times 10^{-5}}$ | **${4.03 \times 10^{-5}}$** |
| Allen-Cahn | ${2.24 \times 10^{-5}}$ | **${3.48 \times 10^{-6}}$** |  
| Korteweg–De Vries | ${7.04 \times 10^{-4}}$ | **${3.40 \times 10^{-4}}$** |
| Kuramoto-Sivashinsky | ${7.48 \times 10^{-2}}$ | **${3.86 \times 10^{-2}}$** |
| Grey-Scott | ${3.61 \times 10^{-3}}$ | **${1.18 \times 10^{-4}}$** |
| Ginzburg-Landau | ${1.49 \times 10^{-2}}$ | **${4.78 \times 10^{-3}}$** |
| Lid-driven cavity | ${3.24 \times 10^{-1}}$ | **${3.99 \times 10^{-2}}$** |
| Kolmogorov flow | ${2.04 \times 10^{-1}}$ | **${6.99 \times 10^{-2}}$** |
| Rayleigh-Taylor instability | ${7.32 \times 10^{-2}}$ | **${5.22 \times 10^{-3}}$** |


## Animations 

### Kolmogorov flow

![Kolmogorov flow](animations/kf.gif)


### Rayleigh-Taylor instability

![Rayleigh-Taylor instability](animations/rt.gif)



