fix: boundary forcings (#388)

## Description
Fixes #387 
## What problem does this change solve?

Boundary forcings have been broken for the umptieth time. 
We fix it and also introduce an new integration test to ensure they wont
be broken so easily anymore.

##  Additional notes ##
We also introduce support for tests with checks on multiple outputs.


By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

Author: dietervdb-meteo

src/anemoi/inference/forcings.py

@@ -281,7 +281,7 @@ def load_forcings_array(self, dates: list[Date], current_state: State) -> FloatA
             The loaded forcings as a numpy array.
         """
         data = self._state_to_numpy(
-            self.input.load_forcings_state(variables=self.variables, dates=dates, current_state=current_state),
+            self.input.load_forcings_state(dates=dates, current_state=current_state),
             self.variables,
             dates,
         )

src/anemoi/inference/runner.py

@@ -332,7 +332,7 @@ def create_dynamic_forcings_inputs(self, input_state: State) -> list[Forcings]:
 
     def create_boundary_forcings_inputs(self, input_state: State) -> list[Forcings]:
 
-        if not self.checkpoint.has_supporting_array("boundary"):
+        if not self.checkpoint.has_supporting_array("output_mask"):
             return []
 
         result = []

src/anemoi/inference/runners/default.py

@@ -287,7 +287,7 @@ def create_boundary_forcings_input(self) -> Input:
         Input
             The created boundary forcings input.
         """
-        variables = self.variables.retrieved_boundary_forcings_variables()
+        variables = self.variables.retrieved_prognostic_variables()
         config = self._input_forcings("boundary_forcings", "-boundary", "forcings", "input") if variables else "empty"
         input = create_input(self, config, variables=variables, purpose="boundary_forcings")
         LOG.info("Boundary forcings input: %s", input)

src/anemoi/inference/testing/checks.py

@@ -197,6 +197,7 @@ def check_cutout_with_xarray(
         for var in checkpoint.prognostic_variables:
             assert var in ds.data_vars, f"Variable {var} not found in output file."
             ref_idx = ref_ds.name_to_index[var]
+            # loop through time dimension
             for data in ds[var]:
                 assert np.allclose(
                     data.values, ref_ds[0, ref_idx, 0, :]
@@ -206,6 +207,88 @@ def check_cutout_with_xarray(
         raise NotImplementedError("Reference file check is not implemented yet.")
 
 
+@testing_registry.register("check_boundary_forcings_with_xarray")
+def check_boundary_forcings_with_xarray(
+    *,
+    file: Path,
+    checkpoint: "Checkpoint",
+    reference_dataset={},
+    reference_file=None,
+    **kwargs,
+) -> None:
+    LOG.info(f"Checking boundary forcings: {file}")
+
+    # get boundary mask from checkpoint
+    supporting_arrays = checkpoint.supporting_arrays
+    LOG.info(f"Supporting arrays in checkpoint: {supporting_arrays.keys()}")
+    if "output_mask" not in supporting_arrays:
+        LOG.warning("Boundary forcings check is trivial. Consider removing from test config.")
+        return
+    else:
+        boundary_mask = ~supporting_arrays["output_mask"]
+
+    import numpy as np
+    import xarray as xr
+
+    ds = xr.open_dataset(file)
+
+    # check if boundary mask compatible with output
+    n_grid = len(ds["latitude"].values)
+    n_mask = len(boundary_mask)
+    assert (
+        n_grid == n_mask
+    ), f"Number of grid points ({n_grid}) does not match size of output mask in checkpoint ({n_mask})."
+    dates = ds["time"].astype("datetime64[s]").values
+    freq = dates[1] - dates[0]
+    if reference_dataset:
+        from anemoi.datasets import open_dataset
+
+        ref_ds = open_dataset(**reference_dataset, start=dates[0])
+        ref_freq = np.timedelta64(ref_ds.frequency)
+        ref_dates = ref_ds.dates.astype("datetime64[s]")
+        step = freq // ref_freq
+
+        # make sure all dates needed are present and we will step through them consistently
+        assert set(dates[:-1]).issubset(ref_dates), f"Reference dataset is missing dates {set(dates) - set(ref_dates)}"
+        assert step == freq / ref_freq, f"Frequency mismatch between output ({freq}) and reference ({ref_freq})"
+        LOG.info(f"Inference output has a timestep that is {step} times that of the reference dataset.")
+
+        if ref_ds.shape[2] != 1:
+            raise NotImplementedError("Support for ensembles is not implemented yet.")
+        ref_values = ref_ds[:, :, 0, :]
+
+        # make sure we have the reference dataset on the output grid
+        lats = ref_ds.latitudes
+        lons = ref_ds.longitudes
+        if "grid_indices" in supporting_arrays:
+            LOG.info("Using grid indices for boundary forcings check.")
+            grid_indices = supporting_arrays["grid_indices"]
+            ref_values = ref_values[:, :, grid_indices]
+            lats = lats[grid_indices]
+            lons = lons[grid_indices]
+        assert np.allclose(lats, ds.latitude.values), "Latitudes don't match between output and reference."
+        assert np.allclose(lons, ds.longitude.values), "Longitudes don't match between output and reference."
+
+        # check boundary forcings
+        # each inference step takes us from input i to output i
+        # boundary forcings are applied to output i in the creation of input i+1
+        # the current mock inference model simply passes the input, so output i+1 == input i+1
+        # the boundary forcing applied on output i (ref dataset at i) appear thus directly in output i+1
+        for var in checkpoint.prognostic_variables:
+            assert var in ds.data_vars, f"Variable {var} not found in output file."
+            ref_idx = ref_ds.name_to_index[var]
+            for i in range(len(dates) - 1):
+                out = ds[var].isel(time=i + 1).values
+                forcing = ref_values[i * step, ref_idx]
+                assert np.allclose(
+                    out[boundary_mask], forcing[boundary_mask]
+                ), f"Boundary forcing for variable {var} does not match reference data at {ref_dates[i*step]}."
+
+    elif reference_file:
+        # check against a reference file, implement when needed
+        raise NotImplementedError("Reference file check is not implemented yet.")
+
+
 @testing_registry.register("check_file_exist")
 def check_file_exist(*, file: Path, **kwargs) -> None:
     LOG.info(f"Checking file exists: {file}")

tests/integration/rmi-lam/config.yaml

@@ -1,14 +1,26 @@
 - name: dataset-in-netcdf-out
   input: null  # anemoi-datasets will download the zarr.zip at runtime
-  output: output.nc
+  output:
+    - lam_output.nc
+    - full_output.nc
   checks:
     - check_cutout_with_xarray:
+        file: ${output:0}
         mask: 'lam_0'
         reference_date: 2020-01-02T00:00:00
         reference_dataset:
           dataset: ${s3:}/aifs-ea-an-oper-0001-mars-o48-2020-6h.zarr.zip
           area: [70, -55, 10, 70]
+    - check_boundary_forcings_with_xarray:
+        file: ${output:1}
+        reference_dataset:
+          dataset:
+            cutout:
+              - dataset: ${s3:}/aifs-ea-an-oper-0001-mars-o48-2020-6h.zarr.zip
+                area: [70, -55, 10, 70]
+              - dataset: ${s3:}/aifs-ea-an-oper-0001-mars-o32-2020-6h.zarr.zip
     - check_with_xarray:
+        file: ${output:0}
         check_accum: tp
         check_nans: true
   inference_config:
@@ -22,11 +34,13 @@
               area: [70, -55, 10, 70]
             - dataset: ${s3:}/aifs-ea-an-oper-0001-mars-o32-2020-6h.zarr.zip
     output:
-      netcdf:
-        post_processors:
-          - extract_mask:
-              mask: lam_0/cutout_mask
-              as_slice: true
-        path: ${output:}
+      tee:
+      - netcdf:
+          post_processors:
+            - extract_mask:
+                mask: lam_0/cutout_mask
+                as_slice: true
+          path: ${output:0}
+      - netcdf: ${output:1}
     post_processors:
       - accumulate_from_start_of_forecast

tests/integration/test_integration.py

@@ -69,17 +69,22 @@ def _markers(config: DictConfig):
 
 class Setup(NamedTuple):
     config: OmegaConf
-    output: Path
+    output: list[Path]
 
 
 @pytest.fixture(params=MODEL_CONFIGS)
 def test_setup(request, get_test_data: GetTestData, tmp_path: Path) -> Setup:
     model, config = request.param
     input = config.input
-    output = tmp_path / config.output
+    output = config.output
     inference_config = config.inference_config
     s3_path = f"anemoi-integration-tests/inference/{model}"
 
+    # set output path(s)
+    if not isinstance(output, (list, ListConfig)):
+        output = [output]
+    output = [tmp_path / file_name for file_name in output]
+
     # download input file(s)
     if not input:
         input = []
@@ -116,7 +121,7 @@ def load_array(name):
 
     # substitute inference config with real paths
     OmegaConf.register_new_resolver("input", lambda i=0: str(input_data[i]), replace=True)
-    OmegaConf.register_new_resolver("output", lambda: str(output), replace=True)
+    OmegaConf.register_new_resolver("output", lambda i=0: str(output[i]), replace=True)
     OmegaConf.register_new_resolver("checkpoint", lambda: str(checkpoint_path), replace=True)
     OmegaConf.register_new_resolver("s3", lambda: str(f"{TEST_DATA_URL}{s3_path}"), replace=True)
     OmegaConf.register_new_resolver("sys.prefix", lambda: sys.prefix, replace=True)
@@ -144,7 +149,8 @@ def test_integration(test_setup: Setup, tmp_path: Path) -> None:
     runner = create_runner(config)
     runner.execute()
 
-    assert (test_setup.output).exists(), "Output file was not created."
+    for file in test_setup.output:
+        assert file.exists(), f"Output file was not created: {file}."
 
     checkpoint_output_variables = _typed_variables_output(runner._checkpoint)
     LOG.info(f"Checkpoint output variables: {checkpoint_output_variables}")
@@ -159,9 +165,10 @@ def test_integration(test_setup: Setup, tmp_path: Path) -> None:
             VariableFromMarsVocabulary(var, {"param": var}) for var in expected_variables_config
         ] or checkpoint_output_variables
 
+        file = kwargs.pop("file", test_setup.output[0])
         testing_registry.create(
             check,
-            file=test_setup.output,
+            file=file,
             expected_variables=expected_variables,
             checkpoint=runner._checkpoint,
             **kwargs,