high res data generation and train

.gitignore

@@ -10,4 +10,5 @@ checkpoints
 *.ipynb
 *.pt
 *.pth
-tensordiffeq
+tensordiffeq
+exp

code_test.py

@@ -1,4 +1,10 @@
+import os
+from matplotlib import test
 import yaml
+import numpy as np
+
+import matplotlib.pyplot as plt
+
 
 
 def test_config():
@@ -9,5 +15,21 @@ def test_config():
     print(config['data']['paths'])
 
 
+def test_data():
+    data_path = '../data/NS-T4000.npy'
+    plot_dir = 'exp/debug/plots'
+    os.makedirs(plot_dir, exist_ok=True)
+    data = np.load(data_path)[0]
+    T = data.shape[0]
+    for i in range(T):
+        fig, ax = plt.subplots()
+        im = ax.imshow(data[i])
+        ax.set_title(f'T:{i}')
+        fig.colorbar(im, ax=ax)
+        plot_path = os.path.join(plot_dir, f'T{i}.png')
+        plt.savefig(plot_path)
+
+
 if __name__ == '__main__':
-    test_config()
+    # test_config()
+    test_data()

configs/operator/Re500-05s-1000-FNO.yaml

@@ -0,0 +1,36 @@
+data:
+  paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy', '../data/NS-Re500Part2.npy']
+  Re: 500
+  total_num: 200
+  offset: 0
+  n_samples: 1000
+  t_duration: 0.5
+  data_res: [64, 64, 33]
+  pde_res: [128, 128, 65]
+  shuffle: True
+
+model:
+  layers: [64, 64, 64, 64, 64]
+  modes1: [8, 8, 8, 8]
+  modes2: [8, 8, 8, 8]
+  modes3: [8, 8, 8, 8]
+  fc_dim: 128
+  act: gelu
+  num_pad: 4
+
+train:
+  batchsize: 2
+  epochs: 501
+  milestones: [300]
+  base_lr: 0.001
+  scheduler_gamma: 0.5
+  ic_loss: 0.0
+  f_loss: 0.0
+  xy_loss: 1.0
+  save_step: 100
+
+log:
+  logdir: Re500-05s-1000-FNO
+  entity: hzzheng-pino
+  project: 'PINO-NS'
+  group: 'Re500-05s-1000-FNO'

configs/operator/Re500-05s-1000-PINO.yaml

@@ -0,0 +1,36 @@
+data:
+  paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy', '../data/NS-Re500Part2.npy']
+  Re: 500
+  total_num: 300
+  offset: 0
+  n_samples: 1000
+  t_duration: 0.5
+  data_res: [64, 64, 33]
+  pde_res: [128, 128, 65]
+  shuffle: True
+
+model:
+  layers: [64, 64, 64, 64, 64]
+  modes1: [8, 8, 8, 8]
+  modes2: [8, 8, 8, 8]
+  modes3: [8, 8, 8, 8]
+  fc_dim: 128
+  act: gelu
+  num_pad: 4
+
+train:
+  batchsize: 2
+  epochs: 501
+  milestones: [300]
+  base_lr: 0.001
+  scheduler_gamma: 0.5
+  ic_loss: 1.0
+  f_loss: 1.0
+  xy_loss: 5.0
+  save_step: 100
+
+log:
+  logdir: Re500-05s-1000-PINO
+  entity: hzzheng-pino
+  project: PINO-NS
+  group: Re500-05s-1000-PINO

configs/operator/Re500-1s-FNO.yaml

@@ -0,0 +1,36 @@
+data:
+  paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy', '../data/NS-Re500Part2.npy']
+  Re: 500
+  total_num: 300
+  offset: 0
+  n_samples: 200
+  t_duration: 1.0
+  data_res: [64, 64, 65]
+  pde_res: [128, 128, 129]
+  shuffle: True
+
+model:
+  layers: [64, 64, 64, 64, 64]
+  modes1: [8, 8, 8, 8]
+  modes2: [8, 8, 8, 8]
+  modes3: [8, 8, 8, 8]
+  fc_dim: 128
+  act: gelu
+  num_pad: 4
+
+train:
+  batchsize: 2
+  epochs: 501
+  milestones: [200, 400]
+  base_lr: 0.001
+  scheduler_gamma: 0.5
+  ic_loss: 0.0
+  f_loss: 0.0
+  xy_loss: 1.0
+  save_step: 100
+
+log:
+  logdir: Re500-1s-200-FNO
+  entity: hzzheng-pino
+  project: 'PINO-NS'
+  group: 'Re500-1s-200-FNO'

configs/operator/Re500-4000-FNO.yaml

@@ -0,0 +1,36 @@
+data:
+  paths: ['../data/NS-T4000.npy']
+  Re: 500
+  total_num: 4000
+  offset: 0
+  n_samples: 3200
+  t_duration: 1.0
+  data_res: [64, 64, 65]
+  pde_res: [128, 128, 65]
+  shuffle: True
+
+model:
+  layers: [64, 64, 64, 64, 64]
+  modes1: [8, 8, 8, 8]
+  modes2: [8, 8, 8, 8]
+  modes3: [8, 8, 8, 8]
+  fc_dim: 128
+  act: gelu
+  num_pad: 4
+
+train:
+  batchsize: 2
+  epochs: 501
+  milestones: [300]
+  base_lr: 0.001
+  scheduler_gamma: 0.5
+  ic_loss: 0.0
+  f_loss: 0.0
+  xy_loss: 1.0
+  save_step: 100
+
+log:
+  logdir: Re500-1s-FNO
+  entity: hzzheng-pino
+  project: 'PINO-NS'
+  group: 'Re500-1s-FNO'

configs/operator/Re500-FNO.yaml

@@ -1,42 +1,36 @@
 data:
-  datapath: '../data/NS-Re500Part0.npy'
+  paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy']
   Re: 500
-  total_num: 4000
+  total_num: 200
   offset: 0
-  n_sample: 400
-  time_interval: 0.5
-  nx: 64
-  nt: 64
-  sub: 1
-  sub_t: 1
+  n_samples: 700
+  t_duration: 0.5
+  data_res: [64, 64, 33]
+  pde_res: [128, 128, 65]
   shuffle: True
-  S2: 128
-  T2: 65
 
 model:
   layers: [64, 64, 64, 64, 64]
   modes1: [8, 8, 8, 8]
   modes2: [8, 8, 8, 8]
   modes3: [8, 8, 8, 8]
   fc_dim: 128
+  act: gelu
+  num_pad: 4
 
 train:
-  batchsize: 1
-  epochs: 40000
-  milestones: [10000, 20000, 30000]
+  batchsize: 2
+  epochs: 501
+  milestones: [300]
   base_lr: 0.001
   scheduler_gamma: 0.5
-  ic_loss: 5.0
-  f_loss: 1.0
+  ic_loss: 0.0
+  f_loss: 0.0
   xy_loss: 1.0
-  save_dir: 'Re500-Pretrain500'
-  save_name: 'PINO-pretrain-Re500-1s-400.pt'
-  data_iter: 1
-  eqn_iter: 1
+  save_step: 100
 
 log:
+  logdir: Re500-05s-FNO
   entity: hzzheng-pino
-  project: 'PINO-pretrain-ICLR'
-  group: 'Re500-05s-4C0'
-
-
+  project: 'PINO-NS'
+  group: 'Re500-05s-FNO'

configs/operator/Re500-PINO.yaml

@@ -3,8 +3,8 @@ data:
   Re: 500
   total_num: 200
   offset: 0
-  n_sample: 700
-  time_duration: 0.5
+  n_samples: 700
+  t_duration: 0.5
   data_res: [64, 64, 33]
   pde_res: [128, 128, 65]
   shuffle: True
@@ -16,23 +16,23 @@ model:
   modes3: [8, 8, 8, 8]
   fc_dim: 128
   act: gelu
-  num_pad: 5
+  num_pad: 4
 
 train:
-  batchsize: 1
+  batchsize: 2
   epochs: 501
   milestones: [300]
   base_lr: 0.001
   scheduler_gamma: 0.5
-  ic_loss: 5.0
-  f_loss: 1.0
-  xy_loss: 1.0
+  ic_loss: 1.0
+  f_loss: 5.0
+  xy_loss: 5.0
   save_step: 100
 
 log:
-  logdir: Re500-05s-PINO
+  logdir: Re500-05s-PINO-pde5
   entity: hzzheng-pino
-  project: 'PINO-pretrain-ICLR'
-  group: 'Re500-05s-4C0'
+  project: 'PINO-NS'
+  group: 'Re500-05s-PINO'
 
 

generate_data.py

@@ -0,0 +1,61 @@
+import math
+import numpy as np
+import os
+from tqdm import tqdm
+
+import torch
+from solver.random_fields import GaussianRF
+from solver.kolmogorov_flow import KolmogorovFlow2d
+from timeit import default_timer
+import argparse
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--re', type=float, default=40.0)
+    parser.add_argument('--sub_x', type=int, default=4)
+    parser.add_argument('--T', type=int, default=2000)
+    parser.add_argument('--outdir', type=str, default='data')
+    parser.add_argument('--t_res', type=int, default=256)
+    args = parser.parse_args()
+    save_dir = args.outdir
+    os.makedirs(save_dir, exist_ok=True)
+    device = torch.device('cuda:0')
+    s = 1024
+    sub = args.sub_x
+
+    n = 4   # forcing
+    Re = args.re
+
+    T_in = 100.0
+    T = args.T
+    t = args.t_res
+    dt = 1.0 / t
+
+    GRF = GaussianRF(2, s, 2 * math.pi, alpha=2.5, tau=7, device=device)
+    u0 = GRF.sample(1)
+
+    NS = KolmogorovFlow2d(u0, Re, n)
+    NS.advance(T_in, delta_t=1e-3)
+
+    sol = np.zeros((T, t + 1, s // sub, s // sub))
+    sol_ini = NS.vorticity().squeeze(0).cpu().numpy()[::sub, ::sub]
+    pbar = tqdm(range(T))
+    for i in pbar:
+        sol[i, 0, :, :] = sol_ini
+        for j in range(t):
+            t1 = default_timer()
+            NS.advance(dt, delta_t=1e-3)
+            sol[i, j + 1, :, :] = NS.vorticity().squeeze(0).cpu().numpy()[::sub, ::sub]
+            t2 = default_timer()
+        pbar.set_description(
+            (
+                f'{i}, time cost: {t2-t1}'
+            )
+        )
+        sol_ini = sol[i, -1, :, :]
+
+    save_path = os.path.join(save_dir, f'NS-Re{Re}_T{t}.npy')
+    # np.save('NS_fine_Re500_S512_s64_T500_t128.npy', sol)
+    np.save(save_path, sol)