Improve plotter

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "cmocean (>=4.0.3,<5.0.0)",
     "pyyaml (>=6.0.3,<7.0.0)",
     "bitsea @ git+https://github.com/inogs/bit.sea.git",
+    "rioxarray (>=0.19.0,<1.0.0)",
     "tabulate (>=0.10.0,<1.0.0)",
 ]
 

src/medunda/plots/maps.py

@@ -17,57 +17,123 @@ def configure_parser(subparsers):
     plotting_maps.add_argument(
         "--time",
         type=date_from_str,
-        required=True,
+        required=False,
         help="Date to plot the map for (format: YYYY-MM-DD)",
     )
+
     plotting_maps.add_argument(
-        "--latitude",
-        type=float,
-        required=False,
-        help="Latitude of the area to plot",
+        "--aggregation-dimension",
+        type=str,
+        help="Dimension along which to aggregate the data"
+        "By default, no aggregation is performed and the data is plotted as is."
+        "If the data has more than 2 dimensions, an aggregation method must be specified using '--aggregation-method'",
     )
     plotting_maps.add_argument(
-        "--longitude",
-        type=float,
-        required=False,
-        help="Longitude of the area to plot",
+        "--aggregation-method",
+        type=str,
+        choices=["mean", "max", "min"],
+        default="mean",
+        help="Method to use for aggregating data along the specified dimension",
     )
 
 
-def plotting_maps(data: "xr.DataArray", metadata: dict, time):
-    selected_time = pd.to_datetime(time)
+def plotting_maps(
+    data: "xr.DataArray",
+    metadata: dict,
+    time,
+    aggregation_dimension,
+    aggregation_method,
+    output_dir,
+    show_plot,
+):
+    # If data is from a raster file, we always want to
+    # select the first band for plotting, because we do not
+    # generate files with multiple bands.
+    if "band" in data.dims:
+        data_slice = data.isel(band=0)
+    else:
+        data_slice = data
 
-    data["time"] = pd.to_datetime(data["time"].values)
+    if time is None:
+        actual_time = None
+    else:
+        selected_time = pd.to_datetime(time)
+
+        # Convert the time coordinate to datetime if it's not already
+        data["time"] = pd.to_datetime(data["time"].values)
 
-    try:
-        data_slice = data.sel(time=selected_time)
-    except Exception:
         try:
             data_slice = data.sel(time=selected_time, method="nearest")
         except Exception as e:
             raise ValueError(f"Could not select time {selected_time}: {e}")
 
-    non_spatial_dims = [
-        dim
-        for dim in data_slice.dims
-        if dim not in ["lat", "latitude", "lon", "longitude"]
-    ]
-    if non_spatial_dims:
-        data_slice = data_slice.mean(dim=non_spatial_dims)
+        n_dims = len(data_slice.dims)
+
+        if n_dims == 2 and aggregation_dimension is not None:
+            LOGGER.warning(
+                "Data has 2 dimensions, but an aggregation dimension was specified. "
+                "Aggregation will be skipped."
+            )
+
+        if n_dims > 2:
+            if aggregation_dimension is None:
+                raise ValueError(
+                    "Data has more than 2 dimensions, an aggregation dimension must be specified."
+                    " Please provide an aggregation dimension using the '--aggregation-dimension'"
+                )
 
-    actual_time = pd.to_datetime(data_slice["time"].values)
+            if aggregation_dimension not in data_slice.dims:
+                raise ValueError(
+                    f"Aggregation dimension '{aggregation_dimension}' not found in data "
+                    f"dimensions: {data_slice.dims}"
+                )
+
+            LOGGER.info(
+                f"Applying {aggregation_method} aggregation along dimension "
+                f"'{aggregation_dimension}'"
+            )
+
+            data_slice = getattr(data_slice, aggregation_method)(
+                dim=aggregation_dimension
+            )
+
+        if n_dims != 2:
+            raise ValueError(
+                f"Data has {n_dims} dimensions after aggregation, expected 2 dimensions for plotting"
+            )
+
+        if (
+            "latitude" not in data_slice.dims
+            or "longitude" not in data_slice.dims
+        ):
+            raise ValueError(
+                "Data must have 'latitude' and 'longitude' dimensions for 2D plotting."
+            )
+        else:
+            LOGGER.debug("Data has 2 dimensions, proceeding with plotting.")
+
+        actual_time = pd.to_datetime(data_slice["time"].values)
 
     fig, ax = plt.subplots(figsize=(10, 5))
     cmap = metadata.get("cmap", "viridis")
-
     im = ax.pcolormesh(data_slice, cmap=cmap)
 
-    title_str = f"{metadata['label']} at {actual_time.strftime('%Y-%m-%d')}"
+    if actual_time is not None:
+        title_str = (
+            f"{metadata['label']} at {actual_time.strftime('%Y-%m-%d')}"
+        )
+    else:
+        title_str = metadata["label"]
+
     ax.set_title(title_str)
     ax.set_xlabel("Longitude")
     ax.set_ylabel("Latitude")
 
     cbar = fig.colorbar(im, ax=ax, orientation="vertical")
     cbar.set_label(f"{metadata['label']} ({metadata.get('unit', '')})")
 
-    plt.show()
+    if output_dir:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        plt.savefig(output_dir / f"{metadata['label']}_map.png")
+    if show_plot:
+        plt.show()

src/medunda/plots/timeseries.py

@@ -14,41 +14,52 @@ def configure_parser(subparsers):
         help="Plots timeseries for a specific variable over a defined period of time",
     )
     plotting_timeseries.add_argument(
-        "--start-time",
+        "--start-date",
         type=date_from_str,
-        required=True,
+        required=False,
         help="Start date of the period to plot (format: YYYY-MM-DD)",
     )
     plotting_timeseries.add_argument(
-        "--end-time",
+        "--end-date",
         type=date_from_str,
-        required=True,
+        required=False,
         help="End date of the period to plot (format: YYYY-MM-DD)",
     )
 
 
 def plotting_timeseries(
-    data: "xr.DataArray", metadata: dict, start_time, end_time
+    data: "xr.DataArray",
+    metadata: dict,
+    start_date,
+    end_date,
+    output_dir,
+    show_plot,
 ):
-    start = start_time
-    end = end_time
-
     if "time" not in data.dims:
         raise ValueError(
             "This dataset has no 'time' dimension therefore cannot plot time series."
         )
 
-    # Aggregate spatial dims by mean over lat and lon if they exist
+    if start_date is None and end_date is None:
+        start_date = data["time"].values[0]
+        end_date = data["time"].values[-1]
+    else:
+        if start_date is not None:
+            start_date = start_date
+        if end_date is not None:
+            end_date = end_date
+
     spatial_dims = [
         dim
         for dim in ["lat", "latitude", "lon", "longitude"]
         if dim in data.dims
     ]
+
     data_mean = data.mean(dim=spatial_dims)
 
-    ts = data_mean.sel(time=slice(start, end))
+    ts = data_mean.sel(time=slice(start_date, end_date))
 
-    fig, ax = plt.subplots(figsize=(12, 6))
+    _, ax = plt.subplots(figsize=(12, 6))
     ax.plot(ts.time, ts)
 
     ylabel = f"{metadata['label']}"
@@ -59,4 +70,8 @@ def plotting_timeseries(
     ax.set_xlabel("Time")
     ax.set_ylabel(ylabel)
 
-    plt.show()
+    if output_dir:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        plt.savefig(output_dir / f"{metadata['label']}_timeseries.png")
+    if show_plot:
+        plt.show()