Introduce geoinfo_channels for anemoi datasets (#1760)

* implemented

* remove eval in interface

* lint

* incoporate requested changes

* fix imports

* Introduced geoinfo_channel parameter for anemoi dataset

---------

Co-authored-by: moritzhauschulz <[email protected]>

Author: clessig

src/weathergen/datasets/data_reader_anemoi.py

@@ -112,8 +112,16 @@ def __init__(
         # get target channel weights from stream config
         self.target_channel_weights = self.parse_target_channel_weights()
 
-        self.geoinfo_channels = []
-        self.geoinfo_idx = []
+        # select/filter requested geoinfo channels (can be any variable, not just constant-in-time)
+        self.geoinfo_idx = self.select_geoinfo_channels(ds0)
+        self.geoinfo_channels = [ds.variables[i] for i in self.geoinfo_idx]
+        # set geoinfo normalization statistics
+        if len(self.geoinfo_idx) > 0:
+            self.mean_geoinfo = ds.statistics["mean"][self.geoinfo_idx]
+            self.stdev_geoinfo = ds.statistics["stdev"][self.geoinfo_idx]
+        else:
+            self.mean_geoinfo = np.zeros(0)
+            self.stdev_geoinfo = np.ones(0)
 
         ds_name = stream_info["name"]
         _logger.info(f"{ds_name}: source channels: {self.source_channels}")
@@ -177,12 +185,13 @@ def _get(self, idx: TIndex, channels_idx: list[int]) -> ReaderData:
                 num_data_fields=len(channels_idx), num_geo_fields=len(self.geoinfo_idx)
             )
 
+        # coords-first representation and collapse multiple steps
+        data = data.transpose([0, 2, 1]).reshape((data.shape[0] * data.shape[2], -1))
+
+        # extract geoinfo channels (can be time-varying, so read from dataset)
+        geoinfos = data[:, list(self.geoinfo_idx)]
         # extract channels
-        data = (
-            data[:, list(channels_idx)]
-            .transpose([0, 2, 1])
-            .reshape((data.shape[0] * data.shape[2], -1))
-        )
+        data = data[:, list(channels_idx)]
 
         # construct lat/lon coords
         latlon = np.concatenate(
@@ -195,9 +204,6 @@ def _get(self, idx: TIndex, channels_idx: list[int]) -> ReaderData:
         # repeat latlon len(t_idxs) times
         coords = np.vstack((latlon,) * len(t_idxs))
 
-        # empty geoinfos for anemoi
-        geoinfos = np.zeros((len(data), 0), dtype=data.dtype)
-
         # date time matching #data points of data
         # Assuming a fixed frequency for the dataset
         datetimes = np.repeat(self.ds.dates[didx_start:didx_end], len(data) // len(t_idxs))
@@ -254,6 +260,40 @@ def select_channels(self, ds0: anemoi_datasets, ch_type: str) -> NDArray[np.int6
 
         return np.array(chs_idx, dtype=np.int64)
 
+    def select_geoinfo_channels(self, ds0: anemoi_datasets) -> NDArray[np.int64]:
+        """
+        Select geoinfo channels (can be any variable, not just constant-in-time)
+
+        Parameters
+        ----------
+        ds0 :
+            raw anemoi dataset with available channels
+
+        Returns
+        -------
+        NDArray of channel indices for geoinfo variables
+
+        """
+
+        geoinfo_channels = self.stream_info.get("geoinfo_channels", [])
+
+        if len(geoinfo_channels) == 0:
+            return np.array([], dtype=np.int64)
+
+        # Select channels that match the geoinfo list (exact match required)
+        chs_idx = np.sort(
+            [ds0.name_to_index[k] for k in ds0.typed_variables.keys() if k in geoinfo_channels]
+        )
+
+        if len(chs_idx) == 0 and len(geoinfo_channels) > 0:
+            stream_name = self.stream_info["name"]
+            _logger.warning(
+                f"No matching geoinfo channels found for {stream_name}. "
+                f"Requested: {geoinfo_channels}"
+            )
+
+        return np.array(chs_idx, dtype=np.int64)
+
 
 def _clip_lat(lats: NDArray) -> NDArray[np.float32]:
     """

src/weathergen/datasets/data_reader_base.py

@@ -610,7 +610,9 @@ def normalize_geoinfos(self, geoinfos: NDArray[DType]) -> NDArray[DType]:
 
         assert geoinfos.shape[-1] == len(self.geoinfo_idx), "incorrect number of geoinfo channels"
         for i, _ in enumerate(self.geoinfo_idx):
-            geoinfos[..., i] = (geoinfos[..., i] - self.mean_geoinfo[i]) / self.stdev_geoinfo[i]
+            # for constant fields, just center the data (resulting in 0s after subtracting mean)
+            stdev = 1.0 if np.isclose(self.stdev_geoinfo[i], 0) else self.stdev_geoinfo[i]
+            geoinfos[..., i] = (geoinfos[..., i] - self.mean_geoinfo[i]) / stdev
 
         return geoinfos