fix: error combining 3 or plus inputs in cutout (ecmwf#256)

## Description
ecmwf#249 introduced a bug when there is more than 2 sources. Indeed, when
`combined_mask` is None, it adds an extra dimension to the state.

I also added a test.

## What problem does this change solve?
<!-- Describe if it's a bugfix, new feature, doc update, or breaking
change -->

## What issue or task does this change relate to?
<!-- link to Issue Number -->

##  Additional notes ##
<!-- Include any additional information, caveats, or considerations that
the reviewer should be aware of. -->

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

Author: bdvllrs

src/anemoi/inference/inputs/cutout.py

@@ -12,6 +12,7 @@
 from typing import Iterable
 from typing import List
 from typing import Optional
+from typing import Union
 
 import numpy as np
 
@@ -28,7 +29,7 @@
 def _mask_and_combine_states(
     combined_state: State,
     new_state: State,
-    combined_mask: Optional[np.ndarray],
+    combined_mask: Union[np.ndarray, slice],
     mask: np.ndarray,
     fields: Iterable[str],
 ) -> State:
@@ -119,7 +120,7 @@ def create_input_state(self, *, date: Optional[Date]) -> State:
             combined_state["fields"] = _mask_and_combine_states(
                 combined_state["fields"], new_state["fields"], combined_mask, mask, combined_state["fields"]
             )
-            combined_mask = None
+            combined_mask = slice(0, None)
 
         return combined_state
 
@@ -154,7 +155,7 @@ def load_forcings_state(self, *, variables: List[str], dates: List[Date], curren
             combined_fields = _mask_and_combine_states(
                 combined_fields, new_fields, combined_mask, mask, combined_fields
             )
-            combined_mask = None
+            combined_mask = slice(0, None)
 
         current_state["fields"] |= combined_fields
         return current_state

tests/unit/test_input_cutout.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+from anemoi.inference.inputs.cutout import _mask_and_combine_states
+
+
+def test_mask_and_combine_states():
+    states = [{"a": np.arange(5).astype(float)}, {"a": np.arange(5, 10).astype(float)}, {"a": np.arange(10, 15)}]
+    masks = [np.zeros(5).astype(bool) for _ in states]
+    masks[0][[1, 2]] = True
+    masks[1][[2, 3]] = True
+    masks[2][[2, 4]] = True
+    combined_state = {k: states[0][k][:] for k in states[0]}
+    combined_mask = masks[0]
+    for k in range(1, len(states)):
+        mask = masks[k]
+        new_state = states[k]
+        combined_state = _mask_and_combine_states(combined_state, new_state, combined_mask, mask, ["a"])
+        combined_mask = slice(0, None)
+
+    assert combined_state["a"].shape[0] == 6
+    assert (
+        combined_state["a"]
+        == np.array(
+            [
+                states[0]["a"][1],
+                states[0]["a"][2],
+                states[1]["a"][2],
+                states[1]["a"][3],
+                states[2]["a"][2],
+                states[2]["a"][4],
+            ]
+        )
+    ).all()