Fixed path handling to have multiple paths for different stream types. (#50)

Co-authored-by: Christian Lessig <[email protected]>

Author: clessig

src/weathergen/__init__.py

@@ -75,6 +75,7 @@ def train(run_id=None) -> None:
     # directory where input streams are specified
     # cf.streams_directory = './streams_large/'
     cf.streams_directory = "./config/streams/streams_anemoi/"
+    # cf.streams_directory = "./config/streams/streams_mixed/"
     # cf.streams_directory = "./streams_mixed/"
 
     # embed_orientation : 'channels' or 'columns'
@@ -181,10 +182,11 @@ def train(run_id=None) -> None:
     cf.norm_type = "LayerNorm"  #'LayerNorm' #'RMSNorm'
     cf.nn_module = "te"
 
-    cf.data_path = private_cf["data_path"]
-    # "/home/mlx/ai-ml/datasets/stable/"
-    # cf.data_path = '/lus/h2resw01/fws4/lb/project/ai-ml/observations/v1'
-    # cf.data_path = '/leonardo_scratch/large/userexternal/clessig0/obs/v1'
+    # merge private config
+    for k, v in private_cf.items():
+        setattr(cf, k, v)
+    cf.data_path = private_cf["data_path_anemoi"]  # for backward compatibility
+
     cf.start_date = 201301010000
     cf.end_date = 202012310000
     cf.start_date_val = 202101010000

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -31,7 +31,7 @@ class MultiStreamDataSampler(torch.utils.data.IterableDataset):
     ###################################################
     def __init__(
         self,
-        data_path,
+        cf,
         rank,
         num_ranks,
         streams,
@@ -91,9 +91,8 @@ def __init__(
             for fname in stream_info["filenames"]:
                 ds = None
                 if stream_info["type"] == "obs":
-                    c_data_path = "/gpfs/scratch/ehpc01/dop/v1/"
                     ds = ObsDataset(
-                        c_data_path + "/" + fname,
+                        cf.data_path_obs + "/" + fname,
                         start_date,
                         end_date_padded,
                         len_hrs,
@@ -127,12 +126,13 @@ def __init__(
                     stats_offset = 0
 
                 elif stream_info["type"] == "anemoi":
-                    c_data_path = data_path
-                    if "CERRA" in stream_info["name"]:
-                        c_data_path = "/gpfs/scratch/ehpc03/weathergen/"
-
                     ds = AnemoiDataset(
-                        c_data_path + "/" + fname, start_date, end_date, len_hrs, step_hrs, False
+                        cf.data_path_anemoi + "/" + fname,
+                        start_date,
+                        end_date,
+                        len_hrs,
+                        step_hrs,
+                        False,
                     )
                     do = 0
                     geoinfo_idx = [0, 1]

src/weathergen/train/trainer.py

@@ -99,7 +99,7 @@ def evaluate(self, cf, run_id_trained, epoch, run_id_new=False):
         self.init(cf, run_id_trained, epoch, run_id_new, run_mode="evaluate")
 
         self.dataset_val = MultiStreamDataSampler(
-            cf.data_path,
+            cf,
             cf.rank,
             cf.num_ranks,
             cf.streams,
@@ -164,7 +164,7 @@ def evaluate_jac(self, cf, run_id, epoch, mode="row", date=None, obs_id=0, sampl
         self.init(cf, run_id, epoch, run_id_new=True, run_mode="offline")
 
         self.dataset = MultiStreamDataSampler(
-            cf.streams,
+            cf,
             cf.start_date_val,
             cf.end_date_val,
             cf.delta_time,
@@ -298,7 +298,7 @@ def run(self, cf, private_cf, run_id_contd=None, epoch_contd=None, run_id_new=Fa
         self.init(cf, run_id_contd, epoch_contd, run_id_new)
 
         self.dataset = MultiStreamDataSampler(
-            cf.data_path,
+            cf,
             cf.rank,
             cf.num_ranks,
             cf.streams,
@@ -324,7 +324,7 @@ def run(self, cf, private_cf, run_id_contd=None, epoch_contd=None, run_id_new=Fa
             sampling_rate_target=cf.sampling_rate_target,
         )
         self.dataset_val = MultiStreamDataSampler(
-            cf.data_path,
+            cf,
             cf.rank,
             cf.num_ranks,
             cf.streams,

src/weathergen/utils/validation_io.py

@@ -127,7 +127,7 @@ def write_validation(
             ds_source.create_dataset("targets_lens", data=targets_lens_k)
         else:
             rn = rn + f"/{fs}"
-            if source_lens_k.sum() > 0 :
+            if source_lens_k.sum() > 0:
                 ds[f"{rn}/sources"].append(source_k)
             ds[f"{rn}/sources_lens"].append(source_lens_k)
             ds[f"{rn}/preds"].append(preds_k)