Added computation of mean frontal melt rates.

analysis/compute_basins_stats.py

@@ -19,7 +19,7 @@
 Compute basins.
 """
 
-# pylint: disable=redefined-outer-name
+# pylint: disable=redefined-outer-name,unused-import
 
 import re
 import time
@@ -31,6 +31,7 @@
 import dask
 import geopandas as gp
 import numpy as np
+import pint_xarray
 import xarray as xr
 from dask.distributed import Client, progress
 
@@ -141,20 +142,6 @@
             ds = ds.sel(
                 time=slice(options.temporal_range[0], options.temporal_range[1])
             )
-    bmb_var = "tendency_of_ice_mass_due_to_basal_mass_flux"
-    if bmb_var in ds:
-        bmb_grounded_da = ds[bmb_var].where(ds["mask"] == 2)
-        bmb_grounded_da.name = "tendency_of_ice_mass_due_to_basal_mass_flux_grounded"
-        bmb_floating_da = ds[bmb_var].where(ds["mask"] == 3)
-        bmb_floating_da.name = "tendency_of_ice_mass_due_to_basal_mass_flux_floating"
-        ds = xr.merge([ds, bmb_grounded_da, bmb_floating_da])
-
-    del ds["time"].attrs["bounds"]
-    ds = ds.drop_vars(
-        ["time_bounds", "timestamp"], errors="ignore"
-    ).rio.set_spatial_dims(x_dim="x", y_dim="y")
-    ds.rio.write_crs(crs, inplace=True)
-
     p_config = ds["pism_config"]
     p_run_stats = ds["run_stats"]
 
@@ -195,6 +182,45 @@
         coords={"run_stats_axis": rs_keys, "exp_id": m_id},
         name="run_stats",
     )
+
+    bmb_var = "tendency_of_ice_mass_due_to_basal_mass_flux"
+    if bmb_var in ds:
+        bmb_grounded_da = ds[bmb_var].where(ds["mask"] == 2)
+        bmb_grounded_da.name = "tendency_of_ice_mass_due_to_basal_mass_flux_grounded"
+        bmb_floating_da = ds[bmb_var].where(ds["mask"] == 3)
+        bmb_floating_da.name = "tendency_of_ice_mass_due_to_basal_mass_flux_floating"
+        ds = xr.merge([ds, bmb_grounded_da, bmb_floating_da])
+
+    fm_var = "tendency_of_ice_mass_due_to_frontal_melt"
+    if fm_var in ds:
+        fm_da = ds[fm_var].pint.quantify()
+        fm_da = fm_da.where(fm_da != 0, np.nan)
+        config_ice_density = float(
+            pism_config.sel({"pism_config_axis": "constants.ice.density"}).values
+        )
+        dx = float(pism_config.sel({"pism_config_axis": "grid.dx"}).values)
+        dy = float(pism_config.sel({"pism_config_axis": "grid.dy"}).values)
+
+        area = xr.DataArray(dx * dy).pint.quantify("m^2")
+        ice_density = (
+            xr.DataArray(config_ice_density).pint.quantify("kg m^-3").pint.to("Gt m^-3")
+        )
+        fm_rate_da = (fm_da / area / ice_density).pint.to("m day^-1").pint.dequantify()
+        fm_rate_da.name = "frontal_melt_rate"
+        fm_rate_ds = (
+            fm_rate_da.to_dataset()
+            .drop_vars(["time_bounds", "timestamp"], errors="ignore")
+            .rio.set_spatial_dims(x_dim="x", y_dim="y")
+        )
+        fm_rate_ds.rio.write_crs(crs, inplace=True)
+        del fm_rate_ds["time"].attrs["bounds"]
+
+    del ds["time"].attrs["bounds"]
+    ds = ds.drop_vars(
+        ["time_bounds", "timestamp"], errors="ignore"
+    ).rio.set_spatial_dims(x_dim="x", y_dim="y")
+    ds.rio.write_crs(crs, inplace=True)
+
     ds = xr.merge(
         [
             ds[mb_vars].drop_vars(["pism_config", "run_stats"], errors="ignore"),
@@ -216,13 +242,37 @@
     )
     basin_names = ["GRACE"] + [basin["SUBREGION1"] for _, basin in basins.iterrows()]
     n_basins = len(basin_names)
-    futures = client.map(compute_basin, basins_ds_scattered, basin_names)
+    futures = client.map(compute_basin, basins_ds_scattered, basin_names, reduce="sum")
     progress(futures)
+
     basin_sums = (
         xr.concat(client.gather(futures), dim="basin")
         .drop_vars(["mapping", "spatial_ref"])
         .sortby(["basin", "pism_config_axis"])
     )
+
+    if fm_var in ds:
+        basins_ds_scattered = client.scatter(
+            [fm_rate_ds]
+            + [fm_rate_ds.rio.clip([basin.geometry]) for _, basin in basins.iterrows()]
+        )
+        basin_names = ["GRACE"] + [
+            basin["SUBREGION1"] for _, basin in basins.iterrows()
+        ]
+        n_basins = len(basin_names)
+        futures = client.map(
+            compute_basin, basins_ds_scattered, basin_names, reduce="mean"
+        )
+        progress(futures)
+
+        fm_basin_sums = (
+            xr.concat(client.gather(futures), dim="basin")
+            .drop_vars(["mapping", "spatial_ref"])
+            .sortby(["basin"])
+        )
+
+        basin_sums = xr.merge([basin_sums, fm_basin_sums])
+
     if cf:
         basin_sums["basin"] = basin_sums["basin"].astype(f"S{n_basins}")
         basin_sums.attrs["Conventions"] = "CF-1.8"

environment-dev.yml

@@ -10,6 +10,7 @@ dependencies:
   - ncview
   - pre-commit
   - mypy
+  - black
   - pylint
   - pytest
   - pip

pism_ragis/plotting.py

@@ -24,7 +24,6 @@
 import warnings
 from importlib.resources import files
 from pathlib import Path
-from typing import Dict, List, Union
 
 import matplotlib as mpl
 import matplotlib.pylab as plt
@@ -50,8 +49,9 @@
 @timeit
 def plot_posteriors(
     df: pd.DataFrame,
-    order: List[str] | None = None,
-    fig_dir: Union[str, Path] = "figures",
+    order: list[str] | None = None,
+    figsize: tuple[float, float] | None = (6.4, 5.2),
+    fig_dir: str | Path = "figures",
     fontsize: float = 4,
 ):
     """
@@ -61,9 +61,11 @@ def plot_posteriors(
     ----------
     df : pd.DataFrame
         DataFrame containing the data to plot.
-    order : List[str] or None, optional
+    order : list[str] or None, optional
         Order of the basins for the y-axis, by default None.
-    fig_dir : Union[str, Path], optional
+    figsize : tuple[float, float] or None, optional
+        Size of the figure, by default (6.4, 5.2).
+    fig_dir : str or Path, optional
         Directory to save the figures, by default "figures".
     fontsize : float, optional
         Font size for the plot, by default 4.
@@ -87,7 +89,7 @@ def plot_posteriors(
             3,
             6,
             sharey=True,
-            figsize=[6.4, 5.2],
+            figsize=figsize,
         )
         fig.subplots_adjust(hspace=0.1, wspace=0.1)
         for k, v in enumerate(
@@ -135,7 +137,8 @@ def plot_posteriors(
 @timeit
 def plot_prior_posteriors(
     df: pd.DataFrame,
-    fig_dir: Union[str, Path] = "figures",
+    figsize: tuple[float, float] | None = (6.4, 3.2),
+    fig_dir: str | Path = "figures",
     fontsize: float = 4,
     bins_dict: dict = {},
 ):
@@ -146,7 +149,9 @@ def plot_prior_posteriors(
     ----------
     df : pd.DataFrame
         DataFrame containing the data to plot.
-    fig_dir : Union[str, Path], optional
+    figsize : tuple[float, float] or None, optional
+        Size of the figure, by default (6.4, 3.2).
+    fig_dir : str or Path, optional
         Directory to save the figures, by default "figures".
     fontsize : float, optional
         Font size for the plot, by default 4.
@@ -179,7 +184,7 @@ def plot_prior_posteriors(
                     3,
                     6,
                     sharey=True,
-                    figsize=[6.4, 3.2],
+                    figsize=figsize,
                 )
                 fig.subplots_adjust(hspace=0.5, wspace=0.1)
                 for k, v in enumerate(m_df.drop(columns=["ensemble"]).columns):
@@ -235,11 +240,12 @@ def plot_basins(
     prior: xr.Dataset,
     posterior: xr.Dataset,
     filter_var: str,
-    filter_range: List[int] = [1990, 2019],
-    fig_dir: Union[str, Path] = "figures",
+    figsize: tuple[float, float] | None = None,
+    filter_range: list[int] = [1990, 2019],
+    fig_dir: str | Path = "figures",
     fudge_factor: float = 3.0,
-    plot_range: List[int] = [1980, 2020],
-    config: Dict = {},
+    plot_range: list[int] = [1980, 2020],
+    config: dict = {},
 ):
     """
     Plot basins using observed, prior, and posterior datasets.
@@ -258,16 +264,18 @@ def plot_basins(
         The posterior dataset.
     filter_var : str
         The variable used for filtering.
-    filter_range : List[int], optional
+    figsize : tuple[float, float] or None, optional
+        Size of the figure, by default None.
+    filter_range : list[int], optional
         A list containing the start and end years for filtering, by default [1990, 2019].
-    fig_dir : Union[str, Path], optional
+    fig_dir : str or Path, optional
         The directory where figures will be saved, by default "figures".
     fudge_factor : float, optional
         A multiplicative factor applied to the observed standard deviation to widen the likelihood function,
         allowing for greater tolerance in the matching process, by default 3.0.
-    plot_range : List[int], optional
+    plot_range : list[int], optional
         A list containing the start and end years for plotting, by default [1980, 2020].
-    config : Dict, optional
+    config : dict, optional
         Configuration dictionary, by default {}.
 
     Examples
@@ -296,6 +304,7 @@ def plot_basins(
                 config=config,
                 filter_var=filter_var,
                 filter_range=filter_range,
+                figsize=figsize,
                 fig_dir=fig_dir,
                 fudge_factor=fudge_factor,
                 obs_alpha=obs_alpha,
@@ -311,8 +320,9 @@ def plot_sensitivity_indices(
     indices_var: str = "S1",
     indices_conf_var: str = "S1_conf",
     basin: str = "",
+    figsize: tuple[float, float] | None = (3.2, 1.8),
     filter_var: str = "",
-    fig_dir: Union[str, Path] = "figures",
+    fig_dir: str | Path = "figures",
     fontsize: float = 6,
 ):
     """
@@ -330,9 +340,11 @@ def plot_sensitivity_indices(
         The variable name for confidence intervals of sensitivity indices in the dataset, by default "S1_conf".
     basin : str, optional
         The basin parameter to be used in the plot, by default "".
+    figsize : tuple[float, float] or None, optional
+        Size of the figure, by default (3.2, 1.8).
     filter_var : str, optional
         The variable used for filtering, by default "".
-    fig_dir : Union[str, Path], optional
+    fig_dir : str or Path, optional
         The directory where the figures will be saved, by default "figures".
     fontsize : float, optional
         The font size for the plot, by default 6.
@@ -346,7 +358,7 @@ def plot_sensitivity_indices(
 
     with mpl.rc_context({"font.size": fontsize}):
 
-        fig, ax = plt.subplots(1, 1, figsize=(3.2, 1.8))
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
         for g in ds[dim]:
             indices_da = ds[indices_var].sel({dim: g})
             conf_da = ds[indices_conf_var].sel({dim: g})
@@ -375,14 +387,15 @@ def plot_obs_sims(
     sim_posterior: xr.Dataset,
     config: dict,
     filter_var: str,
-    filter_range: List[int] = [1990, 2019],
-    fig_dir: Union[str, Path] = "figures",
+    filter_range: list[int] = [1990, 2019],
+    figsize: tuple[float, float] | None = None,
+    fig_dir: str | Path = "figures",
     fudge_factor: float = 3.0,
     reference_year: float = 1986.0,
     sim_alpha: float = 0.4,
     obs_alpha: float = 1.0,
     sigma: float = 2,
-    percentiles: List[float] = [0.025, 0.975],
+    percentiles: list[float] = [0.025, 0.975],
     fontsize: float = 6,
 ) -> None:
     """
@@ -400,9 +413,11 @@ def plot_obs_sims(
         Configuration dictionary containing variable names.
     filter_var : str
         Variable used for filtering.
-    filter_range : List[int], optional
+    filter_range : list[int], optional
         Range of years for filtering, by default [1990, 2019].
-    fig_dir : Union[str, Path], optional
+    figsize : tuple[float, float] or None, optional
+        Size of the figure, by default None.
+    fig_dir : str or Path, optional
         Directory to save the figures, by default "figures".
     fudge_factor : float, optional
         A multiplicative factor applied to the observed standard deviation to widen the likelihood function,
@@ -415,7 +430,7 @@ def plot_obs_sims(
         Alpha value for observation plots, by default 1.0.
     sigma : float, optional
         Sigma value for uncertainty, by default 2.
-    percentiles : List[float], optional
+    percentiles : list[float], optional
         Percentiles for credibility interval, by default [0.025, 0.975].
     fontsize : float, optional
         Font size for the plot, by default 6.
@@ -458,7 +473,7 @@ def plot_obs_sims(
             3,
             1,
             sharex=True,
-            figsize=(6.4, 3.6),
+            figsize=figsize,
             height_ratios=[2, 1, 1],
         )
         fig.subplots_adjust(hspace=0.05, wspace=0.05)
@@ -649,7 +664,7 @@ def plot_obs_sims(
 def plot_outliers(
     filtered_da: xr.DataArray,
     outliers_da: xr.DataArray,
-    filename: Union[Path, str],
+    filename: Path | str,
     fontsize: int = 6,
 ):
     """
@@ -665,7 +680,7 @@ def plot_outliers(
         The DataArray containing the filtered data.
     outliers_da : xr.DataArray
         The DataArray containing the outliers.
-    filename : Union[Path, str]
+    filename : Path or str
         The path or filename where the plot will be saved.
     fontsize : int, optional
         The font size for the plot, by default 6.

pism_ragis/processing.py

@@ -330,7 +330,10 @@ def preprocess_scalar_nc(
 
 
 def compute_basin(
-    ds: xr.Dataset, name: str = "basin", dim: List = ["x", "y"]
+    ds: xr.Dataset,
+    name: str = "basin",
+    dim: List = ["x", "y"],
+    reduce: Literal["sum", "mean"] = "sum",
 ) -> xr.Dataset:
     """
     Compute the sum of the dataset over the 'x' and 'y' dimensions and add a new dimension 'basin'.
@@ -343,6 +346,8 @@ def compute_basin(
         The name to assign to the new 'basin' dimension.
     dim : List
         The dimensions to sum over.
+    reduce : {"sum", "mean"}, optional
+        The reduction method to apply, by default "sum".
 
     Returns
     -------
@@ -354,7 +359,12 @@ def compute_basin(
     >>> ds = xr.Dataset({'var': (('x', 'y'), np.random.rand(5, 5))})
     >>> compute_basin(ds, 'new_basin')
     """
-    ds = ds.sum(dim=dim).expand_dims("basin", axis=-1)
+    if reduce == "sum":
+        ds = ds.sum(dim=dim).expand_dims("basin")
+    elif reduce == "mean":
+        ds = ds.mean(dim=dim).expand_dims("basin")
+    else:
+        raise ValueError("Invalid reduce method. Use 'sum' or 'mean'.")
     ds["basin"] = [name]
     return ds.compute()