fix: resolve critical bugs for Windows compatibility and core functionality

.gitignore

@@ -177,5 +177,10 @@ debug/
 .DS_Store
 
 # neural net, sciml artefacts
-.pt
-.hdf5
+*.pt
+*.h5
+*.pickle
+
+# data analysis
+*.csv
+*.pdf

pdebench/models/analyse_result_forward.py

@@ -176,6 +176,7 @@ def main():
     index1, index2, index3 = [], [], []
     for _j, fl in enumerate(files):
         with Path(fl).open("rb") as f:
+            fl = str(fl)
             title = fl.split("\\")[-1][:-7].split("_")
             if title[0] == "1D":
                 if title[1] == "CFD":
@@ -234,8 +235,12 @@ def main():
 
     fig, ax = plt.subplots(figsize=(8, 6))
     for i in range(num_models):
-        pos = x - 0.3 + 0.5 / (num_models - 1) * i
-        ax.bar(pos, data[data.index.isin([models[i]], level=2)]["MSE"], width)
+        if num_models == 1:
+            pos = x
+            ax.bar(pos, data[data.index.isin([models[i]], level=2)]["MSE"], width)
+        else:
+            pos = x - 0.3 + 0.5 / (num_models - 1) * i
+            ax.bar(pos, data[data.index.isin([models[i]], level=2)]["MSE"], width)
 
     ax.set_xticks(x)
     ax.set_xticklabels(pdes, fontsize=30)

pdebench/models/config/config.yaml

@@ -9,28 +9,29 @@ hydra:
     dir: .
 
 args:
-  model_name: "FNO"
+  model_name: "Unet"
   if_training: False
   continue_training: False
-  num_workers: 2
+  num_workers: 0
   batch_size: 5
   initial_step: 10
   t_train: 101
   model_update: 10
+  data_path: "../data/"
   filename: "2D_diff-react_NA_NA"
   single_file: False
   reduced_resolution: 1
   reduced_resolution_t: 1
   reduced_batch: 1
-  epochs: 500
+  epochs: 1
   learning_rate: 1.e-3
   scheduler_step: 100
   scheduler_gamma: 0.5
   #Unet
   in_channels: 2
   out_channels: 2
   ar_mode: True
-  pushforward: True
+  pushforward: False
   unroll_step: 20
   #FNO
   num_channels: 2

pdebench/models/train_models_forward.py

@@ -156,7 +156,7 @@
 logger = logging.getLogger(__name__)
 
 
-@hydra.main(version_base="1.2", config_path="config", config_name="config_rdb")
+@hydra.main(version_base="1.2", config_path="config", config_name="config")
 def main(cfg: DictConfig):
     if cfg.args.model_name == "FNO":
         from pdebench.models.fno.train import run_training as run_training_FNO

pdebench/models/unet/train.py

@@ -88,16 +88,16 @@ def run_training(
 
         train_data = UNetDatasetMult(
             flnm,
-            reduced_resolution=reduced_resolution,
-            reduced_resolution_t=reduced_resolution_t,
-            reduced_batch=reduced_batch,
+            # reduced_resolution=reduced_resolution,
+            # reduced_resolution_t=reduced_resolution_t,
+            # reduced_batch=reduced_batch,
             saved_folder=base_path,
         )
         val_data = UNetDatasetMult(
-            flnm,
-            reduced_resolution=reduced_resolution,
-            reduced_resolution_t=reduced_resolution_t,
-            reduced_batch=reduced_batch,
+            filename=flnm,
+            # reduced_resolution=reduced_resolution,
+            # reduced_resolution_t=reduced_resolution_t,
+            # reduced_batch=reduced_batch,
             if_test=True,
             saved_folder=base_path,
         )
@@ -191,7 +191,8 @@ def run_training(
             mode="Unet",
             initial_step=initial_step,
         )
-        pickle.dump(errs, Path.open(model_name + ".pickle", "wb"))
+        pickle_path = Path(model_name + ".pickle")
+        pickle.dump(errs, pickle_path.open("wb"))
 
         return
 
@@ -236,7 +237,7 @@ def run_training(
 
                 if training_type in ["autoregressive"]:
                     # Initialize the prediction tensor
-                    pred = yy_tensor[..., :initial_step, :]
+                    pred = yy[..., :initial_step, :].to(device)
 
                     # Extract shape of the input tensor for reshaping (i.e. stacking the
                     # time and channels dimension together)
@@ -312,19 +313,19 @@ def run_training(
                     loss.backward()
                     optimizer.step()
 
-            if training_type in ["single"]:
-                x = xx[..., 0, :]
-                y = yy[..., t_train - 1 : t_train, :]
-                pred = model(x.permute([0, 2, 1])).permute([0, 2, 1])
-                _batch = yy.size(0)
-                loss += loss_fn(pred.reshape(_batch, -1), y.reshape(_batch, -1))
+                if training_type in ["single"]:
+                    x = xx[..., 0, :]
+                    y = yy[..., t_train - 1 : t_train, :]
+                    pred = model(x.permute([0, 2, 1])).permute([0, 2, 1])
+                    _batch = yy.size(0)
+                    loss += loss_fn(pred.reshape(_batch, -1), y.reshape(_batch, -1))
 
-                train_l2_step += loss.item()
-                train_l2_full += loss.item()
+                    train_l2_step += loss.item()
+                    train_l2_full += loss.item()
 
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
 
             if ep % model_update == 0:
                 val_l2_step = 0
@@ -370,15 +371,15 @@ def run_training(
                                 _pred.reshape(_batch, -1), _yy.reshape(_batch, -1)
                             ).item()
 
-                    if training_type in ["single"]:
-                        x = xx[..., 0, :]
-                        y = yy[..., t_train - 1 : t_train, :]
-                        pred = model(x.permute([0, 2, 1])).permute([0, 2, 1])
-                        _batch = yy.size(0)
-                        loss += loss_fn(pred.reshape(_batch, -1), y.reshape(_batch, -1))
+                        if training_type in ["single"]:
+                            x = xx[..., 0, :]
+                            y = yy[..., t_train - 1 : t_train, :]
+                            pred = model(x.permute([0, 2, 1])).permute([0, 2, 1])
+                            _batch = yy.size(0)
+                            loss += loss_fn(pred.reshape(_batch, -1), y.reshape(_batch, -1))
 
-                        val_l2_step += loss.item()
-                        val_l2_full += loss.item()
+                            val_l2_step += loss.item()
+                            val_l2_full += loss.item()
 
                     if val_l2_full < loss_val_min:
                         loss_val_min = val_l2_full