make rollout_iterations and lead_time_hours params

geoarches/configs/dataloader/era5.yaml

@@ -1,12 +1,13 @@
 dataset: 
   _target_: geoarches.dataloaders.era5.Era5Forecast
   path: data/era5_240/full/
-  lead_time_hours: 24 # mixed
+  lead_time_hours: 24
+  multistep: ${oc.select:module.rollout_iterations,1}
   norm_scheme: pangu
   load_prev: True
 
 validation_args:
-  multistep: 2
+  multistep: ${oc.select:module.validation.rollout_iterations,1}
 
 test_args:
-  multistep: 2
+  multistep: ${module.inference.rollout_iterations}

geoarches/configs/module/archesweather.yaml

@@ -11,6 +11,11 @@ module:
   use_graphcast_coeffs: True
   use_prev: True
   loss_delta_normalization: True
+  lead_time_hours: ${dataloader.dataset.lead_time_hours}
+  rollout_iterations: 1
+
+inference:
+  rollout_iterations: 2 # number of rollouts/multistep
 
 backbone:
   # default backbone

geoarches/configs/module/archesweathergen.yaml

@@ -33,6 +33,14 @@ module:
   num_warmup_steps: 5000
   num_training_steps: ${max_steps}
 
+  # for logging metrics
+  lead_time_hours: ${dataloader.dataset.lead_time_hours}
+  rollout_iterations: 1 # For training, number of rollouts
+
+validation:
+  num_members: 5
+  rollout_iterations: 2 # number of rollouts
+
 inference:
   num_steps: 25
   num_members: 10

geoarches/lightning_modules/diffusion.py

@@ -53,6 +53,7 @@ def __init__(
         num_cycles=0.5,
         learn_residual=False,
         sd3_timestep_sampling=True,
+        lead_time_hours=24,
         **kwargs,
     ):
         """
@@ -93,9 +94,11 @@ def __init__(
 
         # set up metrics
         save_memory = cfg.inference.num_members > 10
-        val_kwargs = dict(lead_time_hours=24, rollout_iterations=1, save_memory=save_memory)
+        val_kwargs = dict(
+            lead_time_hours=lead_time_hours, rollout_iterations=1, save_memory=save_memory
+        )
         test_kwargs = dict(
-            lead_time_hours=24,
+            lead_time_hours=lead_time_hours,
             rollout_iterations=cfg.inference.rollout_iterations,
             save_memory=save_memory,
         )
@@ -380,8 +383,8 @@ def sample_rollout(
 
     def validation_step(self, batch, batch_nb):
         # for the validation, we make some generations and log them
-        val_num_members = 5
-        val_rollout_iterations = 2
+        val_num_members = self.cfg.validation.num_steps
+        val_rollout_iterations = self.cfg.validation.rollout_iterations
         samples = [
             self.sample_rollout(
                 batch,

geoarches/lightning_modules/forecast.py

@@ -39,6 +39,7 @@ def __init__(
         add_input_state=False,
         save_test_outputs=False,
         use_weatherbench_lat_coeffs=True,
+        lead_time_hours=24,
         rollout_iterations=1,
         test_filename_suffix="",
         **kwargs,
@@ -96,7 +97,7 @@ def __init__(
             compute_lat_weights_fn=compute_lat_weights_weatherbench
             if use_weatherbench_lat_coeffs
             else compute_lat_weights,
-            lead_time_hours=24,
+            lead_time_hours=lead_time_hours,
             rollout_iterations=rollout_iterations,
         )