Develop ssl

config/jepa_multi_data_pipeline.yml

@@ -24,51 +24,32 @@ stages:
   # Stage 1: JEPA pre-training with EMA teacher + deep SSL
   - name: pretrain
     command: train
-    base-config: ./config/config_jepa_multi_data_all_years.yml
-    config: ./config/config_jepa_multi_data_all_years.yml
-    # base-config: ./config/config_jepa_self_flow_multi_data.yml
-    # config: ./config/config_jepa_self_flow_multi_data.yml
+    base-config: config/jepa/space_jepa/space_jepa_and_reconstruction_pretraining_era5_o96_obs_o256_1979_2015.yml
+    config: config/jepa/space_jepa/space_jepa_and_reconstruction_pretraining_era5_o96_obs_o256_1979_2015.yml
     options: []
     chain: 2
     nodes: 2
 
-  # Stage 1: JEPA pre-training with EMA teacher + deep SSL
+  # Stage 2: JEPA pre-training with EMA teacher + deep SSL
   - name: pretrain-ft
-    command: train
-    config: ./config/config_jepa_multi_data_all_years_ft.yml
+    command: train-continue
+    config: config/jepa/space_jepa/space_jepa_and_reconstruction_pretraining_era5_o96_obs_o256_2016_2022.yml
     options: []
-    chain: 2
+    chain: 1
     nodes: 2
 
-  # # Stage 1.1: JEPA pre-training with EMA teacher + deep SSL
-  # - name: pretrain-cooldown
-  #   command: train-continue
-  #   from: pretrain
-  #   config: ./config/config_jepa_self_flow_cooldown.yml
-  #   options: []
-  #   chain: 1
-  #   nodes: 2
-
   # Stage 3: Forecasting finetuning (freezes encoder)
   - name: finetune-forecast
     command: train-continue
     from: pretrain-ft
-    config: ./config/config_jepa_multi_data_ft_forecast_all_years.yml
-    chain: 4
+    config: config/jepa/space_jepa/space_jepa_frozen_forecast_era5_o96_obs_o256.yml
+    chain: 2
     nodes: 2
 
-  # Stage 3: Forecasting finetuning (freezes encoder)
-  - name: finetune-forecast-ft
+  # Stage 3: Forecasting finetuning (unfreezes encoder)
+  - name: finetune-forecast-unfrozen
     command: train-continue
-    from: finetune-forecast
-    config: ./config/config_jepa_multi_data_ft_forecast.yml
-    chain: 3
+    from: pretrain-ft
+    config: config/jepa/space_jepa/space_jepa_unfrozen_forecast_era5_o96_obs_o256.yml
+    chain: 2
     nodes: 2
-
-  # # Stage 3: Forecasting finetuning (unfreezes encoder)
-  # - name: finetune-forecast-2
-  #   command: train-continue
-  #   from: finetune-forecast
-  #   config: ./config/config_jepa_multi_data_ft_forecast_unfreeze_encoder.yml
-  #   chain: 2
-  #   nodes: 2

config/streams/jepa_forecast_multi_data_all_years/avhrr.yml

@@ -8,7 +8,7 @@ METOP_ABC_AVHRR_IASI:
   forcing: True
   loss_weight: 1.0
   token_size: 512
-  tokenize_spacetime: True
+  tokenize_spacetime: False
   embed:
     net: transformer
     num_tokens: 1

config/streams/jepa_forecast_multi_data_od/avhrr.yml

@@ -8,7 +8,7 @@ METOP_ABC_AVHRR_IASI:
   forcing: True
   loss_weight: 1.0
   token_size: 512
-  tokenize_spacetime: True
+  tokenize_spacetime: False
   embed:
     net: transformer
     num_tokens: 1

packages/common/src/weathergen/common/config.py

@@ -422,6 +422,9 @@ def load_merge_configs(
     else:
         base_config = load_run_config(from_run_id, mini_epoch, None)
         from_run_id = get_run_id_from_config(base_config)
+        with open_dict(base_config):
+            # one-shot action key: must be set per stage, not inherited from the previous run
+            base_config.pop("reset_modules", None)
     with open_dict(base_config):
         base_config.from_run_id = from_run_id
     # use OmegaConf.unsafe_merge if too slow

src/weathergen/datasets/data_reader_obs.py

@@ -48,7 +48,7 @@ def __init__(
         # To read idx convert to a string, format e.g.: 197001010000
         base_date_str = dt_obj.strftime("%Y%m%d%H%M")
         self.hrly_index = self.z[f"idx_{base_date_str}_1"]
-        self.colnames = self.data.attrs["colnames"]
+        self.colnames = list(self.data.attrs["colnames"])
 
         data_colnames = [col for col in self.colnames if "obsvalue" in col]
         data_idx = [i for i, col in enumerate(self.colnames) if "obsvalue" in col]
@@ -74,7 +74,20 @@ def __init__(
         self.target_idx = np.array(self.target_idx, dtype=np.int64)
 
         # determine idx for coords and geoinfos
-        self.coords_idx = [self.colnames.index("lat"), self.colnames.index("lon")]
+        coords_channels = stream_info.get("coords_channels", ["lat", "lon"])
+        assert len(coords_channels) == 2, (
+            f"{stream_info['name']}: 'coords_channels' must be a list of exactly two "
+            f"names [lat, lon], got {coords_channels!r}."
+        )
+        lat_name, lon_name = coords_channels
+        for name in (lat_name, lon_name):
+            n = self.colnames.count(name)
+            assert n == 1, (
+                f"{stream_info['name']}: coordinate column not found in {self.filename}. "
+                f"Looked for '{lat_name}'/'{lon_name}'; available colnames: {self.colnames}. "
+                f"Set 'coords_channels' in the stream config to match data."
+            )
+        self.coords_idx = [self.colnames.index(lat_name), self.colnames.index(lon_name)]
 
         # geoinfo channels
         sname = stream_info["name"]
@@ -200,7 +213,7 @@ def _setup_sample_index(self) -> None:
                 self.indices_start = np.append(
                     self.indices_start,
                     np.ones(
-                        (diff_in_hours_end - self.hrly_index.shape[0] - 1) // step_hrs, dtype=int
+                        (diff_in_hours_end - (self.hrly_index.shape[0] - 1)) // step_hrs, dtype=int
                     )
                     * self.indices_start[-1],
                 )
@@ -209,7 +222,8 @@ def _setup_sample_index(self) -> None:
                     self.indices_end,
                     np.ones(
                         # add (len_hrs + 1) since above we also have diff_in_hours_start + len_hrs
-                        (diff_in_hours_end - self.hrly_index.shape[0] + (len_hrs + 1)) // step_hrs,
+                        (diff_in_hours_end - (self.hrly_index.shape[0] - 1) + (len_hrs + 1))
+                        // step_hrs,
                         dtype=int,
                     )
                     * self.indices_end[-1],

src/weathergen/datasets/masking.py

@@ -774,7 +774,7 @@ def apply_noise_to_data(
                     noise_level_per_cell = torch.zeros(
                         geoinfos.shape[0], dtype=geoinfos.dtype, device=geoinfos.device
                     )
-                    rd.geoinfos = torch.cat([geoinfos,noise_level_per_cell.unsqueeze(1)], dim=-1)
+                    rd.geoinfos = torch.cat([geoinfos, noise_level_per_cell.unsqueeze(1)], dim=-1)
                     input_data_geoinfo.append(rd)
                 input_data = input_data_geoinfo
             return input_data
@@ -824,7 +824,7 @@ def apply_noise_to_data(
                     noise_level_per_cell = torch.where(
                         point_noise_mask, noise_level_per_cell, noise_level_t
                     )
-                rd.geoinfos = torch.cat([rd.geoinfos,noise_level_per_cell.unsqueeze(1)], dim=-1)
+                rd.geoinfos = torch.cat([rd.geoinfos, noise_level_per_cell.unsqueeze(1)], dim=-1)
 
             noised.append(rd)
         return noised

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -131,6 +131,7 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage):
             self.teacher_time_offset = 0
 
         self.batch_size = get_batch_size_from_config(mode_cfg)
+        self.num_workers = cf.data_loading.num_workers
         self.shuffle = mode_cfg.shuffle
 
         self.len_timedelta = mode_cfg.time_window_len
@@ -194,8 +195,13 @@ def check_samples(self, fsm: int):
 
         # streamlined calculation of length
         epoch_len = self.samples_per_mini_epoch
-        # adjust len to split loading across all workers and ensure it is multiple of batch_size
-        self.len = ((epoch_len // self.world_size) // self.batch_size) * self.batch_size
+        # adjust len to split loading across all workers and ensure it is multiple of batch_size;
+        # also account for num_workers so per-worker slice is a multiple of batch_size,
+        # preventing the range-loop in __iter__ from yielding extra batches via ceiling division
+        effective_workers = max(1, self.num_workers)
+        self.len = ((epoch_len // self.world_size) // (self.batch_size * effective_workers)) * (
+            self.batch_size * effective_workers
+        )
 
         n_duplicates = self.len * self.world_size - available_samples
         if not self.repeat_data:
@@ -732,7 +738,7 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
                     input_data,
                     source_masks.metadata[sidx],
                     is_student=True,
-                    add_geoinfo_noise="noise_time" in stream_info.get("geoinfo_channels",[]),
+                    add_geoinfo_noise="noise_time" in stream_info.get("geoinfo_channels", []),
                 )
 
                 sdata = self._build_stream_data(
@@ -761,8 +767,10 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
 
                 # Apply self-flow noise to teacher data (handled by masker)
                 input_data_target = self.masker.apply_noise_to_data(
-                    input_data_target_orig, target_masks.metadata[tidx], is_student=False,
-                    add_geoinfo_noise="noise_time" in stream_info.get("geoinfo_channels",[]),
+                    input_data_target_orig,
+                    target_masks.metadata[tidx],
+                    is_student=False,
+                    add_geoinfo_noise="noise_time" in stream_info.get("geoinfo_channels", []),
                 )
 
                 sdata = self._build_stream_data(

src/weathergen/datasets/stream_data.py

@@ -462,7 +462,7 @@ def spoof(healpix_level: int, datetime, geoinfo_size, num_channels) -> IOReaderD
     coords = np.stack([lats.deg, lons.deg], axis=-1, dtype=np.float32)
     # spoof two tokens to avoid unnecessary computational load
     coords = coords[np.random.choice(coords.shape[0], size=2, replace=False)]
-    
+
     geoinfos = np.zeros((coords.shape[0], geoinfo_size), dtype=np.float32)
 
     data = np.zeros((coords.shape[0], num_channels), dtype=np.float32)

src/weathergen/datasets/tokenizer_utils.py

@@ -91,9 +91,9 @@ def encode_times_target(times, time_win) -> torch.tensor:
     dt = pd.to_datetime(times)
     dt_win = pd.to_datetime(time_win)
     # for target only provide local time
-    dt_delta = torch.tensor(
-        np.atleast_1d((dt - dt_win[0]).seconds), dtype=torch.float32
-    ).unsqueeze(1)
+    dt_delta = torch.tensor(np.atleast_1d((dt - dt_win[0]).seconds), dtype=torch.float32).unsqueeze(
+        1
+    )
     time_tensor = torch.cat(
         (
             dt_delta,
@@ -543,4 +543,11 @@ def get_target_coords_local(
     zi = 99
     a[..., (geoinfo_offset + zi) :] = target_coords[..., (geoinfo_offset + 2) :]
 
+    # Careful when merging develop-ssl into develop.
+    # This is not to be merged in to develop.
+    a[..., 98] = np.sin(coords[:, 0])
+    a[..., 97] = np.cos(coords[:, 0])
+    a[..., 96] = np.sin(coords[:, 1])
+    a[..., 95] = np.cos(coords[:, 1])
+
     return a

src/weathergen/model/ema.py

@@ -30,10 +30,10 @@ def __init__(
         self.rampup_ratio = rampup_ratio
         self.ema_model = empty_model
         self.is_model_sharded = is_model_sharded
-        self.batch_size = 1
-        # Build a name → param map once
+        # Build a name → param map once; used to precompute the update pairs.
         self.src_params = dict(self.original_model.named_parameters())
-
+        self._ema_update_params: list[torch.nn.Parameter] = []
+        self._src_update_params: list[torch.nn.Parameter] = []
         self.reset()
 
     @torch.no_grad()
@@ -56,13 +56,32 @@ def reset(self):
         if needs_strip:
             maybe_sharded_sd = {k.removeprefix("module."): v for k, v in maybe_sharded_sd.items()}
         mkeys, ukeys = self.ema_model.load_state_dict(maybe_sharded_sd, strict=False, assign=False)
+        self._rebuild_update_pairs()
         self.ema_model.eval()
 
-    def requires_grad_(self, flag: bool):
-        for p in self.ema_model.parameters():
-            p.requires_grad = flag
+    def _resolve_src_param(self, name: str):
+        p_src = self.src_params.get(name, None)
+        if p_src is None:
+            p_src = self.src_params.get("module." + name, None)
+        return p_src
+
+    def _rebuild_update_pairs(self):
+        """Precompute EMA/source parameter pairs once so update() stays branch-light."""
+        ema_params = []
+        src_params = []
+        for name, p_ema in self.ema_model.named_parameters():
+            p_src = self._resolve_src_param(name)
+            if p_src is None:
+                raise AssertionError(
+                    f"{name}: All parameters of the EMA model must be in the base model."
+                )
+            ema_params.append(p_ema)
+            src_params.append(p_src)
 
-    def get_current_beta(self, cur_step: int) -> float:
+        self._ema_update_params = ema_params
+        self._src_update_params = src_params
+
+    def get_current_beta(self, cur_step: int, batch_size: int) -> float:
         """
         Get current EMA beta value for monitoring.
 
@@ -78,30 +97,20 @@ def get_current_beta(self, cur_step: int) -> float:
         halflife_steps = self.halflife_steps
         if self.rampup_ratio is not None:
             halflife_steps = min(halflife_steps, cur_step * self.rampup_ratio)
-        beta = 0.5 ** (self.batch_size / max(halflife_steps, 1e-6))
+        beta = 0.5 ** (batch_size / max(halflife_steps, 1e-6))
         return beta
 
     @torch.no_grad()
-    def update(self, cur_step, batch_size):
+    def update(self, cur_step: int, batch_size: int):
         # ensure model remains sharded
         if self.is_model_sharded:
             self.ema_model.reshard()
-        # determine correct interpolation params
-        self.batch_size = batch_size
-        beta = self.get_current_beta(cur_step)
 
-        for name, p_ema in self.ema_model.named_parameters():
-            p_src = self.src_params.get(name, None)
-            # Due to DDP only being applied only to the student the names may missmatch
-            # Thus, we check for the alternate naming scheme
-            p_src = self.src_params.get("module." + name, None) if p_src is None else p_src
-            if "identity" in name.lower():
-                continue
-            if p_src is None:
-                # EMA-only param or intentionally excluded
-                assert False, f"{name}: All parameters of the EMA model must be in the base model."
+        # determine correct interpolation params
+        beta = self.get_current_beta(cur_step, batch_size)
 
-            p_ema.lerp_(p_src, 1.0 - beta)
+        torch._foreach_mul_(self._ema_update_params, beta)
+        torch._foreach_add_(self._ema_update_params, self._src_update_params, alpha=1.0 - beta)
 
     @torch.no_grad()
     def forward_eval(self, *args, **kwargs):