/Gaussian-Process-Slover-for-High-Freq-PDE

Solving High Frequency and Multi-Scale PDEs with Gaussian Processes (ICLR 2024)

Primary LanguageJupyter NotebookMIT LicenseMIT

(ICLR 2024) Gaussian-Process-Slover-for-High-Freq-PDE

___

Official implementation of the paper "Solving High Frequency and Multi-Scale PDEs with Gaussian Processes" [OpenReview][Arxiv][知乎] (ICLR 2024), by Shikai Fang*, Madison Cooley*, Da Long*, Shibo Li, Robert Kirby, Shandian Zhe. (*:equal contribution)

Key Idea:

  • model the PDE solution in frequency domain using Gaussian Processes(GP).
  • Efficient Inference by maximizing log-joint prob over structured kernels:

Illustration of the learned solutions of some high-frequency PDEs.

Requirements:

The project is mainly built with Jax 0.4.8 under python 3.10. See detailed info of packages in requirements.txt.

Instructions:

  1. Clone this repository.
  2. Install the required packages by running pip install -r requirements.txt.
  3. Run the solvers for the 1d and 2d PDEs bu running the following scripts in the terminal:
    • 1d PDEs: . run_1d.sh
    • 2d PDEs: . run_2d.sh
  4. Detailed explanations on the PDEs and kernels are provided in the scripts.There are 4 kernels available:
    • "Matern52_Cos_1d"--->GP-HM-Stm
    • "SE_Cos_1d"--->GP-HM-GM
    • "Matern52_1d"--->GP-Matern
    • "SE_1d"--->GP-SE
  5. Hyperparameters of each PDE can be tuned in the .yaml file in the config folder. The current hyperparameters are the best ones we found for each PDEs.
  6. To apply the solver to other PDEs, you can modify the equation_dict variable in the model_GP_solver_1d.py and model_GP_solver_2d.py files.

Citation

Please cite our work if you would like it

@inproceedings{fang2024solving,
      title={Solving High Frequency and Multi-Scale PDEs with Gaussian Processes}, 
      author={Shikai Fang and Madison Cooley and Da Long and Shibo Li and Robert Kirby and Shandian Zhe},
      booktitle={International Conference on Learning Representation},
      year={2024}
}