feat(output, plots): Use earthkit plots (ecmwf#293)

## Description
Use `earthkit-plots` for the plotting output

> [!WARNING]
> Requires: ecmwf/earthkit-plots#81

## What problem does this change solve?
Reuse parts of the ECMWF stack for plotting


## Example
Plots go from
<img width="724" height="1231" alt="image"
src="https://github.com/user-attachments/assets/4bac81df-1ad0-4fb5-b209-46bbeb45d69a"
/>
to
<img width="653" height="715" alt="image"
src="https://github.com/user-attachments/assets/c8320d23-ccef-466f-ada6-04fb6457bed1"
/>




***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/***

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: HCookie

pyproject.toml

@@ -76,7 +76,7 @@ optional-dependencies.docs = [
 ]
 
 optional-dependencies.huggingface = [ "huggingface-hub" ]
-optional-dependencies.plot = [ "cartopy", "matplotlib" ]
+optional-dependencies.plot = [ "earthkit-plots" ]
 
 optional-dependencies.plugin = [ "ai-models>=0.7", "tqdm" ]
 optional-dependencies.tests = [ "anemoi-datasets[all]", "anemoi-inference[all]", "hypothesis", "pytest", "pytest-mock" ]

src/anemoi/inference/outputs/plot.py

@@ -11,8 +11,10 @@
 import os
 
 import numpy as np
+from anemoi.utils.grib import units
 
 from anemoi.inference.context import Context
+from anemoi.inference.decorators import main_argument
 from anemoi.inference.types import FloatArray
 from anemoi.inference.types import ProcessorConfig
 from anemoi.inference.types import State
@@ -40,22 +42,26 @@ def fix(lons: FloatArray) -> FloatArray:
 
 
 @output_registry.register("plot")
+@main_argument("path")
 class PlotOutput(Output):
-    """Plot output class."""
+    """Use `earthkit-plots` to plot the outputs."""
 
     def __init__(
         self,
         context: Context,
         path: str,
+        *,
+        variables: list[str] | None = None,
+        mode: str = "subplots",
+        domain: str | list[str] | None = None,
         strftime: str = "%Y%m%d%H%M%S",
-        template: str = "plot_{variable}_{date}.{format}",
-        dpi: int = 300,
+        template: str = "plot_{date}.{format}",
         format: str = "png",
-        variables: list[str] | None = None,
         missing_value: float | None = None,
         post_processors: list[ProcessorConfig] | None = None,
         output_frequency: int | None = None,
         write_initial_state: bool | None = None,
+        **kwargs,
     ) -> None:
         """Initialize the PlotOutput.
 
@@ -67,12 +73,14 @@ def __init__(
             The path to save the plots.
         variables : list, optional
             The list of variables to plot, by default all.
+        mode : str, optional
+            The plotting mode, can be "subplots" or "overlay", by default "subplots".
+        domain : str | list[str] | None, optional
+            The domain/s to plot, by default None.
         strftime : str, optional
             The date format string, by default "%Y%m%d%H%M%S".
         template : str, optional
-            The template for plot filenames, by default "plot_{variable}_{date}.{format}".
-        dpi : int, optional
-            The resolution of the plot, by default 300.
+            The template for plot filenames, by default "plot_{date}.{format}".
         format : str, optional
             The format of the plot, by default "png".
         missing_value : float, optional
@@ -97,8 +105,10 @@ def __init__(
         self.variables = variables
         self.strftime = strftime
         self.template = template
-        self.dpi = dpi
         self.missing_value = missing_value
+        self.domain = domain
+        self.mode = mode
+        self.kwargs = kwargs
 
     def write_step(self, state: State) -> None:
         """Write a step of the state.
@@ -108,50 +118,45 @@ def write_step(self, state: State) -> None:
         state : State
             The state dictionary.
         """
-        import cartopy.crs as ccrs
-        import cartopy.feature as cfeature
-        import matplotlib.pyplot as plt
-        import matplotlib.tri as tri
+        import earthkit.data as ekd
+        import earthkit.plots as ekp
 
         os.makedirs(self.path, exist_ok=True)
 
-        longitudes = state["longitudes"]
+        longitudes = fix(state["longitudes"])
         latitudes = state["latitudes"]
-        triangulation = tri.Triangulation(fix(longitudes), latitudes)
+        date = state["date"]
+        basetime = date - state["step"]
+
+        plotting_fields = []
 
         for name, values in state["fields"].items():
             if self.skip_variable(name):
                 continue
 
-            _, ax = plt.subplots(subplot_kw={"projection": ccrs.PlateCarree()})
-            ax.coastlines()
-            ax.add_feature(cfeature.BORDERS, linestyle=":")
-
-            missing_values = np.isnan(values)
-            missing_value = self.missing_value
-            if missing_value is None:
-                min = np.nanmin(values)
-                missing_value = min - np.abs(min) * 0.001
-
-            values = np.where(missing_values, self.missing_value, values).astype(np.float32)
-
-            _ = ax.tricontourf(triangulation, values, levels=10, transform=ccrs.PlateCarree())
-
-            ax.tricontour(
-                triangulation,
-                values,
-                levels=10,
-                colors="black",
-                linewidths=0.5,
-                transform=ccrs.PlateCarree(),
+            variable = self.context.checkpoint.typed_variables[name]
+            param = variable.param
+
+            plotting_fields.append(
+                ekd.ArrayField(
+                    values,
+                    {
+                        "shortName": param,
+                        "variable_name": param,
+                        "step": state["step"],
+                        "base_datetime": basetime,
+                        "latitudes": latitudes,
+                        "longitudes": longitudes,
+                        "units": units(param),
+                    },
+                )
             )
 
-            date = state["date"].strftime("%Y-%m-%d %H:%M:%S")
-            ax.set_title(f"{name} at {date}")
-
-            date = state["date"].strftime(self.strftime)
-            fname = self.template.format(date=date, variable=name, format=self.format)
-            fname = os.path.join(self.path, fname)
+        fig = ekp.quickplot(
+            ekd.FieldList.from_fields((plotting_fields)), mode=self.mode, domain=self.domain, **self.kwargs
+        )
+        fname = self.template.format(date=date, format=self.format)
+        fname = os.path.join(self.path, fname)
 
-            plt.savefig(fname, dpi=self.dpi, bbox_inches="tight")
-            plt.close()
+        fig.save(fname)
+        del fig