Skip to content

Latest commit

 

History

History
242 lines (202 loc) · 9.42 KB

README.md

File metadata and controls

242 lines (202 loc) · 9.42 KB
Raw

LOD: Learnable Orthogonal Decomposition🔥


The code repository for the paper Learnable Orthogonal Decomposition for Non-Regressive Prediction for PDE.

Table of Contents📖

  1. Introduction📖
  2. Datasets📚
  3. POD preprocess🌊
  4. LOD Training🤗
  5. Evaluation🌟
  6. LOD Inference🌊
  7. Ablation Study🥛
  8. References

Introduction📖
Comparison of Auto-Regressive (FNO) and Non-Regressive (LOD) Methods on Advection Dataset.
Red: Ground truth, Green: FNO, Blue: LOD.

Understanding spatio-temporal data is a central challenge in the field of deep learning, particularly in solving Partial Differential Equations (PDEs). Existing approaches, such as Transformers and Neural Operators, have successfully mapped input conditions to PDE solutions but often face challenges due to their auto-regressive nature, which leads to increased computational costs and error accumulation over time. In this paper, we introduce a novel approach called 🔥Learnable Orthogonal Decomposition (LOD)🤗, inspired by the classical Proper Orthogonal Decomposition (POD) technique and enhanced by deep learning. LOD effectively decouples temporal and spatial information in PDE data, simplifying the learning process for neural networks. By focusing on the core architecture of LOD, the method is designed to maximize efficiency while maintaining high accuracy. Spatial bases are initialized with POD-generated components, which are set as learnable parameters within the model. The deep learning model then predicts the temporal coefficients in a single inference step, enabling non-regressive prediction of the entire time series. Experiments demonstrate that LOD accurately captures the dynamics of complex physical systems, achieving both high accuracy and low computational cost across various PDE benchmark datasets. While LOD may show higher error in certain datasets compared to traditional auto-regressive models, it consistently outperforms them in long-term prediction stability, particularly in capturing complex physical dynamics over extended time series, as exemplified by the Burgers equation.

Datasets📚

Before training, you must download the PDEBench dataset.
You can download data through above link, also can utilize this code.

Here, below is folder structure we recommend.

data2
├── PDEBench                  
│   ├── 1D
│       ├── Advection
│       ├── Burgers
│       ├── CFD
│       ├── ReactionDiffusion
│       ├── diffusion-sorption
│   ├── 2D           
│       ├── shallow-water

POD preprocess🌊

def POD(T, N_eigen):

    # Eigenvalue problem
    U = T @ T.T

    if (U==U.T).all(): # symmetric
      D, V = np.linalg.eigh(U)
    else:
      print('Not symmetric')
      D, V = np.linalg.eig(U)

    del U
    
    # Sorting eigenvalues and eigenvectors
    indices = D.argsort()[::-1]
    D = D[indices]
    V = V[:, indices]
    
    # Calculating cumulative energy ratio
    cumulative_energy_ratio = np.cumsum(D) / np.sum(D)
    #print(cumulative_energy_ratio >= 1 - epsilon)
    
    # Finding the number of eigenvalues to satisfy the energy threshold
    # n = np.argmax(cumulative_energy_ratio >= 1 - epsilon) + 1 # You can use threshold...
    n = N_eigen # hyperparameter
    
    # Normalizing eigenvectors
    EV = V[:, :n] / np.sqrt(D[:n])
    
    # Calculating the projection matrix
    phi = EV.T @ T
    
    # Reconstructing T
    Tr = T @ phi.T

    return Tr, phi, cumulative_energy_ratio # coeff, bases

This is the core code that performs the POD!
With the above function, you can orthogonalize the time series data and generate coefficients and bases that have the shape of N_eigen.

1D-PDE

Use the make_1D_POD yaml files.

dataset:
    root_path: '/data2/PDEBench/1D'
    save_path: '/data2/PDEBench/POD/' # We recommend
    data_path: ['1D_diff-sorp_NA_NA.h5'] # change data
    N_eigen: 64 # change hyperparameter

Then, implement below code.

python POD_1D_process.py

If you want to preprocess about CFD,

python POD_1D_CFD_process.py

Shallow-water

python POD_2D_process.py

LOD Training🤗

All our experiments were conducted using an NVIDIA H100 80GB GPU with Pytorch 2.2.0.

You can see the training config files.
We provided Advection, Burgers, Diffusion-Reaction, Diffusion-Sorption, CFD, and Shallow-Water.

1D-PDE

  • Advection, Burgers, Diffusion-Reaction, and Diffusion-Sorption
python LOD_1D.py --pde [choose ...advection, burgers, reaction, sorption...]
  • 1D-CFD
python LOD_CFD.py

Shallow-water

python LOD_2D.py

Evaluation🌟

Efficiency

Model Inference_Time VRAM # of parameters
FNO 43s 34.83MB 43137
LOD-small 3s 34.83MB 43325
LOD 7s 51.96MB 4451194

In 1D-PDE dataset, LOD-small is about 14.3x faster than FNO.
LOD have a higher VRAM usage than FNO, but about 6.1x faster inference speed.

Performance (Advection)

Model Advection beta 0.1 Advection beta 0.4 Advection beta 1.0 Advection beta 4.0
POD(Train) 0.001737 0.001668 0.001874 0.002480
FNO 0.009380 0.01261 0.009883 0.005690
PINN 0.7800 0.9200 0.4000 0.6600
TranSolver 0.003981 0.09200 0.2333 0.01509
OFormer 0.004631 0.005718 0.007607 0.01281
LOD-small 0.004515 0.01425 0.01027 0.07331
LOD 0.003422 0.004887 0.003890 0.005382

LOD Inference🌊

Code

python 1D_visualization.py --pde [choose ...advection, burgers, reaction, sorption...]
python 2D_visualization.py

We provided some checkpoints.
You can easily implement our code..!

checkpoint
├── lod_Advection_beta0.1.pt
├── lod_Burgers_Nu1.0.pt
├── lod_ReactionDiffusion_Nu0.5_Rho1.0.pt
├── lod_ReactionDiffusion_Nu5.0_Rho10.0.pt
└── lod-small_shallow-water_NA.pt

Results

LOD results
"Advection case1 - LOD” "Advection case2 - LOD" "Advection case3 - LOD"
"Burgers case1 - LOD” "Burgers case2 - LOD" "Burgers case3 - LOD"
"Burgers case4 - LOD” "Diffusion-Reaction case1 - LOD" "Diffusion-Reaction case2 - LOD"
  • Red: Ground Truth
  • Blue: Prediction
LOD-small results
"Shallow-Water case1 - GT” "Shallow-Water case1 - LOD-small" "Shallow-Water case2 - GT" "Shallow-Water case2 - LOD-small"

Ablation Studies🥛

For ablation studies, we have summarized the results in the ablation folder.

References