feat: Mars with Nested models (ecmwf#236)

## Description
The mars input and by association other inputs cannot be used with
nested models.
See ecmwf#235 for context.

In addition to the changes in this PR, some extra manual patches of the
checkpoint are required.


This enables an input of 
```yaml
cutout:
	lam_0:
		mars:
			grid: '0.01/0.01'
			area: '2/-2/-2/2'
		mask: 'source0/trimedge_mask'
	global: 'mars'
```

Author: HCookie

src/anemoi/inference/checkpoint.py

@@ -612,13 +612,18 @@ def mars_requests(
                 if use_grib_paramid and "param" in r:
                     r["param"] = [shortname_to_paramid(_) for _ in r["param"]]
 
-                # Simplyfie the request
-
-                for k, v in r.items():
+                # Simplify the request
 
+                for k in list(r.keys()):
+                    v = r[k]
                     if len(v) == 1:
-                        r[k] = v[0]
+                        v = v[0]
 
+                    # Remove empty values for when tree is not fully defined
+                    if v == "-":
+                        r.pop(k)
+                        continue
+                    r[k] = v
                 result.append(r)
 
         return result

src/anemoi/inference/inputs/cutout.py

@@ -24,6 +24,36 @@
 LOG = logging.getLogger(__name__)
 
 
+def _mask_and_nest_state(
+    state: State,
+    _state: State,
+    mask: np.ndarray,
+) -> State:
+    """Mask and nest the state with the given mask.
+
+    Parameters
+    ----------
+    state : State
+        The state to be masked and nested.
+    _state : State
+        The state to be masked and nested with.
+    mask : np.ndarray
+        The mask to be applied.
+
+    Returns
+    -------
+    State
+        The masked and nested state.
+    """
+    for field, values in state["fields"].items():
+        state["fields"][field] = np.concatenate([values, _state["fields"][field][..., mask]], axis=-1)
+
+    state["latitudes"] = np.concatenate([state["latitudes"], _state["latitudes"][..., mask]], axis=-1)
+    state["longitudes"] = np.concatenate([state["longitudes"], _state["longitudes"][..., mask]], axis=-1)
+
+    return state
+
+
 @input_registry.register("cutout")
 class Cutout(Input):
     """Combines one or more LAMs into a global source using cutouts."""
@@ -43,7 +73,10 @@ def __init__(self, context, **sources: dict[str, dict]):
         self.sources: dict[str, Input] = {}
         self.masks: dict[str, np.ndarray] = {}
         for src, cfg in sources.items():
-            mask = cfg.pop("mask", f"{src}/cutout_mask")
+            if isinstance(cfg, str):
+                mask = f"{src}/cutout_mask"
+            else:
+                mask = cfg.pop("mask", f"{src}/cutout_mask")
             self.sources[src] = create_input(context, cfg)
             self.masks[src] = self.sources[src].checkpoint.load_supporting_array(mask)
 
@@ -72,10 +105,7 @@ def create_input_state(self, *, date: Optional[Date]) -> State:
             mask = self.masks[source]
             _state = self.sources[source].create_input_state(date=date)
 
-            state["latitudes"] = np.concatenate([state["latitudes"], _state["latitudes"][..., mask]], axis=-1)
-            state["longitudes"] = np.concatenate([state["longitudes"], _state["longitudes"][..., mask]], axis=-1)
-            for field, values in state["fields"].items():
-                state["fields"][field] = np.concatenate([values, _state["fields"][field][..., mask]], axis=-1)
+            state = _mask_and_nest_state(state, _state, mask)
 
         return state
 

src/anemoi/inference/inputs/mars.py

@@ -124,6 +124,9 @@ def postproc(
     if grid_is_valid(grid):
         pproc["grid"] = grid
 
+    if isinstance(area, str):
+        area = [float(x) for x in area.split("/")]
+
     if area_is_valid(area):
         pproc["area"] = rounded_area(area)
 
@@ -219,7 +222,7 @@ def __init__(
         patches : Optional[List[Tuple[Dict[str, Any], Dict[str, Any]]]]
             Optional list of patches for the input.
         **kwargs : Any
-            Additional keyword arguments.
+            Additional keyword to pass to the request to MARS.
         """
         super().__init__(context, namer=namer)
         self.kwargs = kwargs
@@ -282,12 +285,12 @@ def retrieve(self, variables: List[str], dates: List[Date]) -> Any:
 
         kwargs = self.kwargs.copy()
         kwargs.setdefault("expver", "0001")
+        kwargs.setdefault("grid", self.checkpoint.grid)
+        kwargs.setdefault("area", self.checkpoint.area)
 
         return retrieve(
             requests,
-            self.checkpoint.grid,
-            self.checkpoint.area,
-            self.patch,
+            patch=self.patch,
             **kwargs,
         )
 

src/anemoi/inference/metadata.py

@@ -596,7 +596,7 @@ def mars_requests(self, *, variables: List[str]) -> Iterator[DataRequest]:
 
             mars = self.variables_metadata[variable]["mars"].copy()
 
-            for k in ("date", "hdate", "time"):
+            for k in ("date", "hdate", "time", "valid_datetime", "variable"):
                 mars.pop(k, None)
 
             yield mars