diff --git a/analysis/compute_basins_stats.py b/analysis/compute_basins_stats.py index 6a73891..e15cecb 100755 --- a/analysis/compute_basins_stats.py +++ b/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" diff --git a/environment-dev.yml b/environment-dev.yml index 7cd5971..ad70253 100644 --- a/environment-dev.yml +++ b/environment-dev.yml @@ -10,6 +10,7 @@ dependencies: - ncview - pre-commit - mypy + - black - pylint - pytest - pip diff --git a/pism_ragis/plotting.py b/pism_ragis/plotting.py index 2cd401e..89a1f40 100644 --- a/pism_ragis/plotting.py +++ b/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. diff --git a/pism_ragis/processing.py b/pism_ragis/processing.py index 2ffdc23..a2f6c6f 100644 --- a/pism_ragis/processing.py +++ b/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()