add exps

configs/operator/Re500-1_8-800-FNO-s32.yaml

@@ -0,0 +1,48 @@
+data:
+  name: KF
+  paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']
+  Re: 500
+  offset: 0
+  total_num: 300
+  raw_res: [256, 256, 513]
+  n_data_samples: 100
+  data_res: [32, 32, 129]  # resolution in 1 second
+  pde_res: [32, 32, 129]   # resolution in 1 second
+  a_offset: 0
+  n_a_samples: 250
+  testoffset: 275
+  n_test_samples: 25
+  t_duration: 0.125
+  shuffle: True
+
+model:
+  layers: [64, 64, 64, 64, 64]
+  modes1: [12, 12, 12, 12]
+  modes2: [12, 12, 12, 12]
+  modes3: [12, 12, 12, 12]
+  fc_dim: 128
+  act: gelu
+  pad_ratio: [0.0, 0.125]
+
+train:
+  batchsize: 2
+  start_iter: 0
+  num_iter: 50_001
+  milestones: [20_000, 40_000]
+  base_lr: 0.001
+  scheduler_gamma: 0.5
+  ic_loss: 0.0
+  f_loss: 0.0
+  xy_loss: 1.0
+  save_step: 5000
+  eval_step: 5000
+
+test:
+  batchsize: 1
+  data_res: [64, 64, 257]
+
+log:
+  logdir: Re500-1_8s-800-FNO-s32
+  entity: hzzheng-pino
+  project: PINO-KF-Re500
+  group: Re500-1_8s-800-FNO-s32

configs/operator/Re500-1_8-800-PINO-s.yaml

@@ -6,7 +6,7 @@ data:
   total_num: 300
   raw_res: [256, 256, 513]
   n_data_samples: 100
-  data_res: [32, 32, 129]  # resolution in 1 second
+  data_res: [64, 64, 257]  # resolution in 1 second
   pde_res: [256, 256, 513]   # resolution in 1 second
   a_offset: 0
   n_a_samples: 275
@@ -19,10 +19,10 @@ model:
   layers: [64, 64, 64, 64, 64]
   modes1: [12, 12, 12, 12]
   modes2: [12, 12, 12, 12]
-  modes3: [10, 10, 10, 10]
+  modes3: [12, 12, 12, 12]
   fc_dim: 128
   act: gelu
-  pad_ratio: [0.125, 0.125]
+  pad_ratio: [0.0, 0.125]
 
 train:
   batchsize: 2
@@ -33,7 +33,7 @@ train:
   scheduler_gamma: 0.5
   ic_loss: 10.0
   f_loss: 1.0
-  xy_loss: 50.0
+  xy_loss: 10.0
   save_step: 5000
   eval_step: 5000
 
@@ -42,7 +42,7 @@ test:
   data_res: [256, 256, 513]
 
 log:
-  logdir: Re500-1_8s-800-PINO-s-32
+  logdir: Re500-1_8s-800-PINO-s
   entity: hzzheng-pino
   project: PINO-KF-Re500
-  group: Re500-1_8s-800-PINO-s-32
+  group: Re500-1_8s-800-PINO-s

configs/operator/Re500-1_8-800-PINO-s16.yaml

@@ -0,0 +1,48 @@
+data:
+  name: KF
+  paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']
+  Re: 500
+  offset: 0
+  total_num: 300
+  raw_res: [256, 256, 513]
+  n_data_samples: 100
+  data_res: [16, 16, 65]  # resolution in 1 second
+  pde_res: [256, 256, 513]   # resolution in 1 second
+  a_offset: 0
+  n_a_samples: 275
+  testoffset: 275
+  n_test_samples: 25
+  t_duration: 0.125
+  shuffle: True
+
+model:
+  layers: [64, 64, 64, 64, 64]
+  modes1: [12, 12, 12, 12]
+  modes2: [12, 12, 12, 12]
+  modes3: [12, 12, 12, 12]
+  fc_dim: 128
+  act: gelu
+  pad_ratio: [0.0, 0.125]
+
+train:
+  batchsize: 1
+  start_iter: 0
+  num_iter: 200_001
+  milestones: [20_000, 60_000, 120_000]
+  base_lr: 0.001
+  scheduler_gamma: 0.5
+  ic_loss: 10.0
+  f_loss: 1.0
+  xy_loss: 10.0
+  save_step: 5000
+  eval_step: 5000
+
+test:
+  batchsize: 1
+  data_res: [256, 256, 513]
+
+log:
+  logdir: Re500-1_8s-800-PINO-s-16
+  entity: hzzheng-pino
+  project: PINO-KF-Re500
+  group: Re500-1_8s-800-PINO-s-16

configs/operator/Re500-1_8-800-PINO-s32.yaml

@@ -0,0 +1,48 @@
+data:
+  name: KF
+  paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']
+  Re: 500
+  offset: 0
+  total_num: 300
+  raw_res: [256, 256, 513]
+  n_data_samples: 100
+  data_res: [32, 32, 129]  # resolution in 1 second
+  pde_res: [256, 256, 513]   # resolution in 1 second
+  a_offset: 0
+  n_a_samples: 275
+  testoffset: 275
+  n_test_samples: 25
+  t_duration: 0.125
+  shuffle: True
+
+model:
+  layers: [64, 64, 64, 64, 64]
+  modes1: [12, 12, 12, 12]
+  modes2: [12, 12, 12, 12]
+  modes3: [12, 12, 12, 12]
+  fc_dim: 128
+  act: gelu
+  pad_ratio: [0.0, 0.125]
+
+train:
+  batchsize: 2
+  start_iter: 0
+  num_iter: 200_001
+  milestones: [20_000, 60_000, 120_000]
+  base_lr: 0.001
+  scheduler_gamma: 0.5
+  ic_loss: 10.0
+  f_loss: 1.0
+  xy_loss: 10.0
+  save_step: 5000
+  eval_step: 5000
+
+test:
+  batchsize: 1
+  data_res: [256, 256, 513]
+
+log:
+  logdir: Re500-1_8s-800-PINO-s-32
+  entity: hzzheng-pino
+  project: PINO-KF-Re500
+  group: Re500-1_8s-800-PINO-s-32

configs/pretrain/Darcy-pretrain.yaml

@@ -6,6 +6,7 @@ data:
   n_sample: 1000
   nx: 421
   sub: 7
+  pde_sub: 2
 
 model:
   layers: [64, 64, 64, 64, 64]
@@ -28,5 +29,6 @@ train:
 log:
   project: 'PINO-Darcy-pretrain'
   group: 'gelu-pino'
+  entity: hzzheng-pino
 
 

train_operator.py

@@ -6,7 +6,7 @@
 
 from solver.random_fields import GaussianRF
 from train_utils import Adam
-from train_utils.datasets import NSLoader, online_loader, DarcyFlow
+from train_utils.datasets import NSLoader, online_loader, DarcyFlow, DarcyCombo
 from train_utils.train_3d import mixed_train
 from train_utils.train_2d import train_2d_operator
 from models import FNO3d, FNO2d
@@ -79,15 +79,23 @@ def train_2d(args, config):
     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
     data_config = config['data']
 
-    dataset = DarcyFlow(data_config['datapath'],
-                        nx=data_config['nx'], sub=data_config['sub'],
-                        offset=data_config['offset'], num=data_config['n_sample'])
+    # dataset = DarcyFlow(data_config['datapath'],
+                        # nx=data_config['nx'], sub=data_config['sub'],
+                        # offset=data_config['offset'], num=data_config['n_sample'])
+
+    dataset = DarcyCombo(datapath=data_config['datapath'], 
+                         nx=data_config['nx'], 
+                         sub=data_config['sub'], 
+                         pde_sub=data_config['pde_sub'], 
+                         num=data_config['n_samples'], 
+                         offset=data_config['offset'])
     train_loader = DataLoader(dataset, batch_size=config['train']['batchsize'], shuffle=True)
     model = FNO2d(modes1=config['model']['modes1'],
                   modes2=config['model']['modes2'],
                   fc_dim=config['model']['fc_dim'],
                   layers=config['model']['layers'],
-                  act=config['model']['act']).to(device)
+                  act=config['model']['act'], 
+                  pad_ratio=config['model']['pad_ratio']).to(device)
     # Load from checkpoint
     if 'ckpt' in config['train']:
         ckpt_path = config['train']['ckpt']

train_utils/datasets.py

@@ -457,10 +457,7 @@ def __init__(self,
                  nx, sub,
                  offset=0,
                  num=1):
-        if sub == 1:
-            self.S = nx
-        else:
-            self.S = int(nx // sub) + 1
+        self.S = int(nx // sub) + 1 if sub > 1 else nx
         data = scipy.io.loadmat(datapath)
         a = data['coeff']
         u = data['sol']
@@ -482,14 +479,17 @@ def __init__(self,
                  nx, sub,
                  offset=0,
                  num=1):
-        if sub == 1:
-            self.S = nx
-        else:
-            self.S = int(nx // sub) + 1
+        self.S = int(nx // sub) + 1 if sub > 1 else nx
         data = scipy.io.loadmat(datapath)
         a = data['coeff']
         self.a = torch.tensor(a[offset: offset + num, ::sub, ::sub], dtype=torch.float)
         self.mesh = torch2dgrid(self.S, self.S)
+        data = scipy.io.loadmat(datapath)
+        a = data['coeff']
+        u = data['sol']
+        self.a = torch.tensor(a[offset: offset + num, ::sub, ::sub], dtype=torch.float)
+        self.u = torch.tensor(u[offset: offset + num, ::sub, ::sub], dtype=torch.float)
+        self.mesh = torch2dgrid(self.S, self.S)
 
     def __len__(self):
         return self.a.shape[0]
@@ -499,8 +499,36 @@ def __getitem__(self, item):
         return torch.cat([fa.unsqueeze(2), self.mesh], dim=2) 
 
 
+class DarcyCombo(Dataset):
+    def __init__(self, 
+                 datapath, 
+                 nx, 
+                 sub, pde_sub, 
+                 num=1000, offset=0) -> None:
+        super().__init__()
+        self.S = int(nx // sub) + 1 if sub > 1 else nx
+        self.pde_S = int(nx // pde_sub) + 1 if sub > 1 else nx
+        data = scipy.io.loadmat(datapath)
+        a = data['coeff']
+        u = data['sol']
+        self.a = torch.tensor(a[offset: offset + num, ::sub, ::sub], dtype=torch.float)
+        self.u = torch.tensor(u[offset: offset + num, ::sub, ::sub], dtype=torch.float)
+        self.mesh = torch2dgrid(self.S, self.S)
+        self.pde_a = torch.tensor(a[offset: offset + num, ::pde_sub, ::pde_sub], dtype=torch.float)
+        self.pde_mesh = torch2dgrid(self.pde_S, self.pde_S)
+
+    def __len__(self):
+        return self.a.shape[0]
+
+    def __getitem__(self, item):
+        fa = self.a[item]
+        pde_a = self.pde_a[item]
+        data_ic = torch.cat([fa.unsqueeze(2), self.mesh], dim=2)
+        pde_ic = torch.cat([pde_a.unsqueeze(2), self.pde_mesh], dim=2)
+        return data_ic, self.u[item], pde_ic
+
 '''
-dataset class for loading initial conditions
+dataset class for loading initial conditions for Komogrov flow
 '''
 class KFaDataset(Dataset):
     def __init__(self, paths, 

train_utils/train_2d.py

@@ -57,20 +57,24 @@ def train_2d_operator(model,
     mesh = train_loader.dataset.mesh
     mollifier = torch.sin(np.pi * mesh[..., 0]) * torch.sin(np.pi * mesh[..., 1]) * 0.001
     mollifier = mollifier.to(rank)
+    pde_mesh = train_loader.dataset.pde_mesh
+    pde_mol = torch.sin(np.pi * pde_mesh[..., 0]) * torch.sin(np.pi * pde_mesh[..., 1]) * 0.001
+    pde_mol = pde_mol.to(rank)
     for e in pbar:
         loss_dict = {'train_loss': 0.0,
                      'data_loss': 0.0,
                      'f_loss': 0.0,
                      'test_error': 0.0}
-        for x, y in train_loader:
-            x, y = x.to(rank), y.to(rank)
+        for data_ic, u, pde_ic in train_loader:
+            data_ic, u, pde_ic = data_ic.to(rank), u.to(rank), pde_ic.to(rank)
 
             optimizer.zero_grad()
 
-            pred = model(x).reshape(y.shape)
-            pred = pred * mollifier
-
-            data_loss = myloss(pred, y)
+            # data loss
+            if data_weight > 0:
+                pred = model(data_ic).squeeze(dim=-1)
+                pred = pred * mollifier
+                data_loss = myloss(pred, y)
 
             a = x[..., 0]
             f_loss = darcy_loss(pred, a)