From fd0ad115265cd344bc40bfac65ece42b88414373 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Tue, 26 May 2026 11:51:49 +0100 Subject: [PATCH 01/16] Use markers to distinguish compressors --- .../compressor/plotting/plot_metrics.py | 104 +++++++++++++----- 1 file changed, 75 insertions(+), 29 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 408b098..75f5e80 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -1,7 +1,9 @@ import argparse from pathlib import Path +import matplotlib.colors as mcolors import matplotlib.pyplot as plt +from matplotlib.lines import Line2D import numpy as np import pandas as pd import seaborn as sns @@ -12,24 +14,24 @@ from .variable_plotters import PLOTTERS _COMPRESSOR2LINEINFO = [ - ("jpeg2000", ("#EE7733", "-")), - ("sperr", ("#117733", ":")), - ("zfp-round", ("#DDAA33", "--")), - ("zfp", ("#EE3377", "--")), - ("sz3", ("#CC3311", "-.")), - ("bitround-pco", ("#0077BB", ":")), - ("bitround", ("#33BBEE", "-")), - ("stochround-pco", ("#BBBBBB", "--")), - ("stochround", ("#009988", "--")), - ("tthresh", ("#882255", "-.")), + ("jpeg2000", ("#EE7733", "-", "o")), + ("sperr", ("#117733", ":", "s")), + ("zfp-round", ("#DDAA33", "--", "D")), + ("zfp", ("#EE3377", "--", "^")), + ("sz3", ("#CC3311", "-.", "v")), + ("bitround-pco", ("#0077BB", ":", "P")), + ("bitround", ("#33BBEE", "-", "X")), + ("stochround-pco", ("#BBBBBB", "--", "d")), + ("stochround", ("#009988", "--", "h")), + ("tthresh", ("#882255", "-.", "<")), ] -def _get_lineinfo(compressor: str) -> tuple[str, str]: - """Get the line color and style for a given compressor.""" - for comp, (color, linestyle) in _COMPRESSOR2LINEINFO: +def _get_lineinfo(compressor: str) -> tuple[str, str, str]: + """Get the line color, style, and marker for a given compressor.""" + for comp, (color, linestyle, marker) in _COMPRESSOR2LINEINFO: if compressor.startswith(comp): - return color, linestyle + return color, linestyle, marker raise ValueError(f"Unknown compressor: {compressor}") @@ -54,6 +56,23 @@ def _get_lineinfo(compressor: str) -> tuple[str, str]: } +def _make_legend_handle(compressor, color, linestyle, marker, line_alpha): + """Proxy artist combining the alpha-faded line and the opaque marker.""" + faded = mcolors.to_rgba(color, alpha=line_alpha) + return Line2D( + [0], + [0], + color=faded, + linestyle=linestyle, + linewidth=4, + marker=marker, + markersize=12, + markerfacecolor=color, + markeredgecolor=color, + label=_get_legend_name(compressor), + ) + + def _get_legend_name(compressor: str) -> str: """Get the legend name for a given compressor.""" for comp, name in _COMPRESSOR2LEGEND_NAME: @@ -103,6 +122,8 @@ def plot_metrics( # Filter out excluded datasets and compressors df = df[~df["Compressor"].isin(exclude_compressor)] + df = df[~df["Compressor"].str.startswith("safeguarded-")] + df = df[~df["Compressor"].str.startswith("rp")] df = df[~df["Dataset"].isin(exclude_dataset)] is_tiny = df["Dataset"].str.endswith("-tiny") filter_tiny = is_tiny if tiny_datasets else ~is_tiny @@ -111,13 +132,13 @@ def plot_metrics( filter_chunked = is_chunked if chunked_datasets else ~is_chunked df = df[filter_chunked] - _plot_per_variable_metrics( - datasets=datasets, - compressed_datasets=compressed_datasets, - plots_path=plots_path, - all_results=df, - rd_curves_metrics=["Max Absolute Error", "MAE", "DSSIM", "Spectral Error"], - ) + # _plot_per_variable_metrics( + # datasets=datasets, + # compressed_datasets=compressed_datasets, + # plots_path=plots_path, + # all_results=df, + # rd_curves_metrics=["Max Absolute Error", "MAE", "DSSIM", "Spectral Error"], + # ) df = _rename_compressors(df) normalized_df = _normalize(df) @@ -344,6 +365,7 @@ def _plot_variable_rd_curve( ): plt.figure(figsize=(8, 6)) compressors = df["Compressor"].unique() + legend_handles = [] for comp in compressors: compressor_data = df[df["Compressor"] == comp] assert len(compressor_data) == len(bounds) @@ -356,16 +378,26 @@ def _plot_variable_rd_curve( for i in bound_ixs ] distortion = [compressor_data[distortion_metric].loc[i] for i in bound_ixs] - color, linestyle = _get_lineinfo(comp) + color, linestyle, marker = _get_lineinfo(comp) + line_alpha = 0.5 plt.plot( compr_ratio, distortion, - label=_get_legend_name(comp), - marker="s", color=color, linestyle=linestyle, linewidth=4, - markersize=8, + alpha=line_alpha, + ) + plt.plot( + compr_ratio, + distortion, + marker=marker, + color=color, + linestyle="None", + markersize=12, + ) + legend_handles.append( + _make_legend_handle(comp, color, linestyle, marker, line_alpha) ) plt.xlabel("Compression Ratio [raw B / enc B]", fontsize=14) @@ -376,6 +408,7 @@ def _plot_variable_rd_curve( plt.ylabel(distortion_metric, fontsize=14) plt.legend( + handles=legend_handles, title="Compressor", fontsize=10, title_fontsize=12, @@ -424,6 +457,7 @@ def _plot_aggregated_rd_curve( [compression_metric, distortion_metric] ].agg(agg) + legend_handles = [] for comp in compressors: compr_ratio = [ agg_distortion.loc[(bound, comp), compression_metric] @@ -433,16 +467,27 @@ def _plot_aggregated_rd_curve( agg_distortion.loc[(bound, comp), distortion_metric] for bound in bound_names ] - color, linestyle = _get_lineinfo(comp) + color, linestyle, marker = _get_lineinfo(comp) + line_alpha = 0.6 plt.plot( compr_ratio, distortion, - label=_get_legend_name(comp), - marker="s", color=color, linestyle=linestyle, + # linestyle="-", linewidth=4, - markersize=8, + alpha=line_alpha, + ) + plt.plot( + compr_ratio, + distortion, + marker=marker, + color=color, + linestyle="None", + markersize=12, + ) + legend_handles.append( + _make_legend_handle(comp, color, linestyle, marker, line_alpha) ) if remove_outliers: @@ -502,6 +547,7 @@ def _plot_aggregated_rd_curve( fontsize=16, ) plt.legend( + handles=legend_handles, title="Compressor", loc="upper right", bbox_to_anchor=(0.8, 0.99), From c5b3c68df8623885357ca7af277c45f3a6095ce5 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Tue, 26 May 2026 16:38:21 +0100 Subject: [PATCH 02/16] Allow for compressor names with '.' in their name --- .../compressor/scripts/concatenate_metrics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py index 56836ac..d6a2399 100644 --- a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py +++ b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py @@ -50,7 +50,7 @@ def concatenate_metrics(basepath: Path = Path()): compressed_datasets / dataset.name / error_bound.name - / compressor.stem + / compressor.name ) measurements = load_measurements(compressed_dataset, compressor) @@ -72,7 +72,7 @@ def load_measurements(compressed_dataset: Path, compressor: Path) -> pd.DataFram for var, variable_measurements in measurements.items(): rows.append( { - "Compressor": compressor.stem, + "Compressor": compressor.name, "Variable": var, "Compression Ratio [raw B / enc B]": variable_measurements[ "decoded_bytes" From cef1251cddd48996da118f55664e1825dde1e266 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Tue, 26 May 2026 20:55:16 +0100 Subject: [PATCH 03/16] Add EBCC configs --- src/climatebenchpress/compressor/plotting/plot_metrics.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 75f5e80..8adce44 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -3,11 +3,11 @@ import matplotlib.colors as mcolors import matplotlib.pyplot as plt -from matplotlib.lines import Line2D import numpy as np import pandas as pd import seaborn as sns import xarray as xr +from matplotlib.lines import Line2D from ..scripts.compute_metrics import parse_error_bounds from .error_dist_plotter import ErrorDistPlotter @@ -24,6 +24,7 @@ ("stochround-pco", ("#BBBBBB", "--", "d")), ("stochround", ("#009988", "--", "h")), ("tthresh", ("#882255", "-.", "<")), + ("ebcc", ("#AA4444", "-.", "8")), ] @@ -46,6 +47,7 @@ def _get_lineinfo(compressor: str) -> tuple[str, str, str]: ("stochround-pco", "StochRound + PCO"), ("stochround", "StochRound + Zstd"), ("tthresh", "TTHRESH"), + ("ebcc", "EBCC"), ] DISTORTION2LEGEND_NAME = { From b31018764c85534e87d67c7f5408a45cb929fc32 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Wed, 27 May 2026 19:52:51 +0100 Subject: [PATCH 04/16] Handle missing files in concatenation --- .../compressor/scripts/concatenate_metrics.py | 45 +++++++++++++++++-- 1 file changed, 41 insertions(+), 4 deletions(-) diff --git a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py index d6a2399..d0590d3 100644 --- a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py +++ b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py @@ -2,6 +2,7 @@ import argparse import json +import warnings from pathlib import Path from typing import Optional @@ -10,7 +11,7 @@ from .compute_metrics import parse_error_bounds -def concatenate_metrics(basepath: Path = Path()): +def concatenate_metrics(basepath: Path = Path(), skip_missing: bool = False): """Concatenate metrics from all datasets and compressors into a single CSV file. Parameters @@ -18,6 +19,9 @@ def concatenate_metrics(basepath: Path = Path()): basepath : Path Assumes that the metrics are stored in `basepath / metrics`. The script will create a `basepath / metrics / all_results.csv` file containing the concatenated results. + skip_missing : bool + If True, skip missing `metrics.csv`, `tests.csv`, or `measurements.json` files + and emit a warning instead of raising. Missing fields will be filled with NaN. """ compressed_datasets = basepath / "compressed-datasets" error_bounds_dir = basepath / "datasets-error-bounds" @@ -43,16 +47,41 @@ def concatenate_metrics(basepath: Path = Path()): for compressor in error_bound.iterdir(): metrics_csv = compressor / "metrics.csv" - metrics = pd.read_csv(metrics_csv) tests_csv = compressor / "tests.csv" - tests = pd.read_csv(tests_csv) compressed_dataset = ( compressed_datasets / dataset.name / error_bound.name / compressor.name ) + measurements_json = compressed_dataset / "measurements.json" + + if skip_missing: + missing = [ + p + for p in (metrics_csv, tests_csv, measurements_json) + if not p.exists() + ] + for p in missing: + warnings.warn(f"Skipping missing file: {p}") + if measurements_json in missing: + # Without measurements.json we have no variable list, so we + # cannot construct any rows for this compressor. + continue + measurements = load_measurements(compressed_dataset, compressor) + metrics = ( + pd.read_csv(metrics_csv) + if metrics_csv.exists() + else pd.DataFrame(columns=["Variable", "Metric", "Error"]) + ) + tests = ( + pd.read_csv(tests_csv) + if tests_csv.exists() + else pd.DataFrame( + columns=["Variable", "Test", "Passed", "Value"] + ) + ) df = merge_metrics(measurements, metrics, tests) df["Dataset"] = dataset.name @@ -132,8 +161,10 @@ def merge_metrics( .merge( test_per_variable.reset_index(), on="Variable", + how="outer", ), on="Variable", + how="left", ) @@ -190,6 +221,12 @@ def get_error_bound_name( if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--basepath", type=Path, default=Path()) + parser.add_argument( + "--skip-missing", + action="store_true", + help="Skip missing metrics/tests/measurements files (with a warning) " + "and fill the corresponding fields with NaN.", + ) args = parser.parse_args() - concatenate_metrics(basepath=args.basepath) + concatenate_metrics(basepath=args.basepath, skip_missing=args.skip_missing) From ab00eead0771a19406bf3bd87cf98c93fe2488a4 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Wed, 27 May 2026 19:53:05 +0100 Subject: [PATCH 05/16] Plotting adjustments --- .../compressor/plotting/error_dist_plotter.py | 2 +- .../compressor/plotting/plot_metrics.py | 21 ++++++++++++------- .../compressor/plotting/variable_plotters.py | 2 +- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/error_dist_plotter.py b/src/climatebenchpress/compressor/plotting/error_dist_plotter.py index dfbb736..ae86332 100644 --- a/src/climatebenchpress/compressor/plotting/error_dist_plotter.py +++ b/src/climatebenchpress/compressor/plotting/error_dist_plotter.py @@ -63,7 +63,7 @@ def plot_error_bound_histograms( compressors = [comp for comp in compressors if "-pco" not in comp] for var in variables: for comp in compressors: - color, linestyle = get_line_info(comp) + color, linestyle, _ = get_line_info(comp) label = get_legend_name(comp) # Don't state the lossless compressor in the legend. if label.startswith("BitRound"): diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 8adce44..57d8314 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -19,12 +19,14 @@ ("zfp-round", ("#DDAA33", "--", "D")), ("zfp", ("#EE3377", "--", "^")), ("sz3", ("#CC3311", "-.", "v")), + ("sz3-abs", ("#CC3311", "-.", "v")), ("bitround-pco", ("#0077BB", ":", "P")), ("bitround", ("#33BBEE", "-", "X")), ("stochround-pco", ("#BBBBBB", "--", "d")), ("stochround", ("#009988", "--", "h")), ("tthresh", ("#882255", "-.", "<")), ("ebcc", ("#AA4444", "-.", "8")), + ("ebcc-abs", ("#AA7744", "-.", "8")), ] @@ -41,12 +43,14 @@ def _get_lineinfo(compressor: str) -> tuple[str, str, str]: ("sperr", "SPERR"), ("zfp-round", "ZFP-ROUND"), ("zfp", "ZFP"), + ("sz3-abs", "SZ3-Abs"), ("sz3", "SZ3"), ("bitround-pco", "BitRound + PCO"), ("bitround", "BitRound + Zstd"), ("stochround-pco", "StochRound + PCO"), ("stochround", "StochRound + Zstd"), ("tthresh", "TTHRESH"), + ("ebcc-abs", "EBCC-Abs"), ("ebcc", "EBCC"), ] @@ -134,13 +138,13 @@ def plot_metrics( filter_chunked = is_chunked if chunked_datasets else ~is_chunked df = df[filter_chunked] - # _plot_per_variable_metrics( - # datasets=datasets, - # compressed_datasets=compressed_datasets, - # plots_path=plots_path, - # all_results=df, - # rd_curves_metrics=["Max Absolute Error", "MAE", "DSSIM", "Spectral Error"], - # ) + _plot_per_variable_metrics( + datasets=datasets, + compressed_datasets=compressed_datasets, + plots_path=plots_path, + all_results=df, + rd_curves_metrics=["Max Absolute Error", "MAE", "DSSIM", "Spectral Error"], + ) df = _rename_compressors(df) normalized_df = _normalize(df) @@ -248,6 +252,9 @@ def _plot_per_variable_metrics( ): """Creates all the plots which only depend on a single variable.""" for dataset in all_results["Dataset"].unique(): + if dataset != "ifs-uncompressed": + continue + df = all_results[all_results["Dataset"] == dataset] dataset_plots_path = plots_path / dataset dataset_plots_path.mkdir(parents=True, exist_ok=True) diff --git a/src/climatebenchpress/compressor/plotting/variable_plotters.py b/src/climatebenchpress/compressor/plotting/variable_plotters.py index aac7a94..f25c66f 100644 --- a/src/climatebenchpress/compressor/plotting/variable_plotters.py +++ b/src/climatebenchpress/compressor/plotting/variable_plotters.py @@ -360,7 +360,7 @@ class CamsPlotter(Plotter): datasets = ["cams-nitrogen-dioxide-tiny", "cams-nitrogen-dioxide"] def plot_fields(self, fig, ax, ds, ds_new, dataset_name, var, err_bound): - selector = dict(valid_time=0, hybrid=3) + selector = dict(valid_time=0, pressure_level=3) in_min = ds.isel(**selector).min().values.item() in_max = ds.isel(**selector).max().values.item() out_min = ds_new.isel(**selector).min().values.item() From 81ae67182943e98c7ad8d159c92158f03c4d7ba4 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Fri, 29 May 2026 14:59:39 +0100 Subject: [PATCH 06/16] Add cloud ice plotter --- .../compressor/plotting/plot_metrics.py | 10 ++-- .../compressor/plotting/variable_plotters.py | 47 +++++++++++++++++++ 2 files changed, 52 insertions(+), 5 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 57d8314..15f66ff 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -18,15 +18,15 @@ ("sperr", ("#117733", ":", "s")), ("zfp-round", ("#DDAA33", "--", "D")), ("zfp", ("#EE3377", "--", "^")), + ("sz3-abs", ("#CC3311", "-.", "p")), ("sz3", ("#CC3311", "-.", "v")), - ("sz3-abs", ("#CC3311", "-.", "v")), ("bitround-pco", ("#0077BB", ":", "P")), ("bitround", ("#33BBEE", "-", "X")), ("stochround-pco", ("#BBBBBB", "--", "d")), ("stochround", ("#009988", "--", "h")), ("tthresh", ("#882255", "-.", "<")), + ("ebcc-abs", ("#AA4444", "-.", "X")), ("ebcc", ("#AA4444", "-.", "8")), - ("ebcc-abs", ("#AA7744", "-.", "8")), ] @@ -252,7 +252,7 @@ def _plot_per_variable_metrics( ): """Creates all the plots which only depend on a single variable.""" for dataset in all_results["Dataset"].unique(): - if dataset != "ifs-uncompressed": + if dataset != "nextgems-icon": continue df = all_results[all_results["Dataset"] == dataset] @@ -558,8 +558,8 @@ def _plot_aggregated_rd_curve( plt.legend( handles=legend_handles, title="Compressor", - loc="upper right", - bbox_to_anchor=(0.8, 0.99), + loc="upper left", + ncol=2, fontsize=12, title_fontsize=14, ) diff --git a/src/climatebenchpress/compressor/plotting/variable_plotters.py b/src/climatebenchpress/compressor/plotting/variable_plotters.py index f25c66f..3af3bbe 100644 --- a/src/climatebenchpress/compressor/plotting/variable_plotters.py +++ b/src/climatebenchpress/compressor/plotting/variable_plotters.py @@ -209,6 +209,51 @@ def plot_fields(self, fig, ax, ds, ds_new, dataset_name, var, err_bound): self.error_title = "Absolute Error" +class IFSCIWCPlotter(Plotter): + datasets = ["ifs-cloud-ice-water-content"] + + def plot_fields(self, fig, ax, ds, ds_new, dataset_name, var, err_bound): + selector = dict(time=0, level=80) + # Calculate shared vmin and vmax for consistent color ranges + data_orig = ds.isel(**selector) + data_new = ds_new.isel(**selector) + vmax = float(np.nanpercentile(data_orig.values, 98)) + vmin = float(np.nanmin(data_orig.values)) + norm = mcolors.PowerNorm(gamma=0.4, vmin=vmin, vmax=max(vmax, 1e-6)) + cmap = "Blues" + + data_orig.plot(ax=ax[0], transform=ccrs.PlateCarree(), norm=norm, cmap=cmap) + data_new.plot( + ax=ax[1], + transform=ccrs.PlateCarree(), + norm=norm, + cmap=cmap, + rasterized=True, + ) + error = data_orig - data_new + non_zero_mask = np.abs(data_orig) > 0.0 + # Check where both original and new data are zero + both_zero_mask = (np.abs(data_orig) == 0.0) & (np.abs(data_new) == 0.0) + rel_error = xr.where( + both_zero_mask, + 0.0, + xr.where(non_zero_mask, error / np.abs(data_orig), 1e12), + ) + + _, bound_value = err_bound + vmin_error, vmax_error = -bound_value, bound_value + rel_error.plot( + ax=ax[2], + transform=ccrs.PlateCarree(), + rasterized=True, + vmin=vmin_error, + vmax=vmax_error, + cbar_kwargs={"ticks": [-bound_value, 0, bound_value]}, + cmap="seismic", + ) + self.error_title = "Relative Error" + + class Era5Plotter(Plotter): datasets = ["era5-tiny", "era5", "ifs-uncompressed"] @@ -443,6 +488,8 @@ def plot_fields(self, fig, ax, ds, ds_new, dataset_name, var, err_bound): Era5Plotter, EsaBiomassPlotter, NextGEMSPlotter, + IFSHumidityPlotter, + IFSCIWCPlotter, ] PLOTTERS: dict[str, type[Plotter]] = dict() for plotter_cls in plotter_clss: From f9f8daff2daba49b9cdbb9f87e3ccb9c403e777c Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Wed, 3 Jun 2026 12:31:58 +0100 Subject: [PATCH 07/16] Adjust plotting for -abs variants --- .../compressor/plotting/error_dist_plotter.py | 10 +++++++++- .../compressor/plotting/plot_metrics.py | 6 ++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/error_dist_plotter.py b/src/climatebenchpress/compressor/plotting/error_dist_plotter.py index ae86332..3d60dc4 100644 --- a/src/climatebenchpress/compressor/plotting/error_dist_plotter.py +++ b/src/climatebenchpress/compressor/plotting/error_dist_plotter.py @@ -26,6 +26,12 @@ def compute_errors(self, compressor, ds, ds_new, var, err_bound_type): if "-pco" in compressor: return + if "-abs" in compressor and err_bound_type == "abs_error": + # The compressors with a "-abs" suffix are versions of compressors + # that have their relative error bound option removed. For absolute + # error bounds, the errors are the same as their non "-abs" counterparts. + return + error = robust_error(ds[var], ds_new[var]) if err_bound_type == "abs_error": error = error.compute().values @@ -60,7 +66,9 @@ def plot_error_bound_histograms( # We only plot bitround and stochround once because the lossless compressor # does not change the error plot distribution. Hence, we ignore the PCO # compressors here. - compressors = [comp for comp in compressors if "-pco" not in comp] + compressors = [ + comp for comp in compressors if "-pco" not in comp and "-abs" not in comp + ] for var in variables: for comp in compressors: color, linestyle, _ = get_line_info(comp) diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 15f66ff..47db53a 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -25,7 +25,7 @@ ("stochround-pco", ("#BBBBBB", "--", "d")), ("stochround", ("#009988", "--", "h")), ("tthresh", ("#882255", "-.", "<")), - ("ebcc-abs", ("#AA4444", "-.", "X")), + ("ebcc-abs", ("#AA4444", "-.", ">")), ("ebcc", ("#AA4444", "-.", "8")), ] @@ -252,7 +252,7 @@ def _plot_per_variable_metrics( ): """Creates all the plots which only depend on a single variable.""" for dataset in all_results["Dataset"].unique(): - if dataset != "nextgems-icon": + if dataset != "ifs-uncompressed": continue df = all_results[all_results["Dataset"] == dataset] @@ -478,6 +478,7 @@ def _plot_aggregated_rd_curve( ] color, linestyle, marker = _get_lineinfo(comp) line_alpha = 0.6 + marker_alpha = 0.8 plt.plot( compr_ratio, distortion, @@ -494,6 +495,7 @@ def _plot_aggregated_rd_curve( color=color, linestyle="None", markersize=12, + alpha=marker_alpha, ) legend_handles.append( _make_legend_handle(comp, color, linestyle, marker, line_alpha) From b9f1630b62c175c52a01f28e773944b046069d95 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Tue, 14 Jul 2026 08:59:40 +0100 Subject: [PATCH 08/16] Use proper legend name in figure title --- .../compressor/plotting/constants.py | 56 +++++++++++++++++ .../compressor/plotting/plot_metrics.py | 62 +------------------ .../compressor/plotting/variable_plotters.py | 18 +++++- 3 files changed, 75 insertions(+), 61 deletions(-) create mode 100644 src/climatebenchpress/compressor/plotting/constants.py diff --git a/src/climatebenchpress/compressor/plotting/constants.py b/src/climatebenchpress/compressor/plotting/constants.py new file mode 100644 index 0000000..801366c --- /dev/null +++ b/src/climatebenchpress/compressor/plotting/constants.py @@ -0,0 +1,56 @@ +_COMPRESSOR2LINEINFO = [ + ("jpeg2000", ("#EE7733", "-", "o")), + ("sperr", ("#117733", ":", "s")), + ("zfp-round", ("#DDAA33", "--", "D")), + ("zfp", ("#EE3377", "--", "^")), + ("sz3-abs", ("#CC3311", "-.", "p")), + ("sz3", ("#CC3311", "-.", "v")), + ("bitround-pco", ("#0077BB", ":", "P")), + ("bitround", ("#33BBEE", "-", "X")), + ("stochround-pco", ("#BBBBBB", "--", "d")), + ("stochround", ("#009988", "--", "h")), + ("tthresh", ("#882255", "-.", "<")), + ("ebcc-abs", ("#AA4444", "-.", ">")), + ("ebcc", ("#AA4444", "-.", "8")), +] + + +def _get_lineinfo(compressor: str) -> tuple[str, str, str]: + """Get the line color, style, and marker for a given compressor.""" + for comp, (color, linestyle, marker) in _COMPRESSOR2LINEINFO: + if compressor.startswith(comp): + return color, linestyle, marker + raise ValueError(f"Unknown compressor: {compressor}") + + +_COMPRESSOR2LEGEND_NAME = [ + ("jpeg2000", "JPEG2000"), + ("sperr", "SPERR"), + ("zfp-round", "ZFP-ROUND"), + ("zfp", "ZFP"), + ("sz3-abs", "SZ3-Abs"), + ("sz3", "SZ3"), + ("bitround-pco", "BitRound + PCO"), + ("bitround", "BitRound + Zstd"), + ("stochround-pco", "StochRound + PCO"), + ("stochround", "StochRound + Zstd"), + ("tthresh", "TTHRESH"), + ("ebcc-abs", "EBCC-Abs"), + ("ebcc", "EBCC"), +] + +DISTORTION2LEGEND_NAME = { + "Relative MAE": "Mean Absolute Error", + "Relative DSSIM": "DSSIM", + "Relative MaxAbsError": "Max Absolute Error", + "Relative SpectralError": "Spectral Error", +} + + +def _get_legend_name(compressor: str) -> str: + """Get the legend name for a given compressor.""" + for comp, name in _COMPRESSOR2LEGEND_NAME: + if compressor.startswith(comp): + return name + + return compressor # Fallback to the compressor name if not found in the mapping. diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 47db53a..4927275 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -10,57 +10,10 @@ from matplotlib.lines import Line2D from ..scripts.compute_metrics import parse_error_bounds +from .constants import DISTORTION2LEGEND_NAME, _get_legend_name, _get_lineinfo from .error_dist_plotter import ErrorDistPlotter from .variable_plotters import PLOTTERS -_COMPRESSOR2LINEINFO = [ - ("jpeg2000", ("#EE7733", "-", "o")), - ("sperr", ("#117733", ":", "s")), - ("zfp-round", ("#DDAA33", "--", "D")), - ("zfp", ("#EE3377", "--", "^")), - ("sz3-abs", ("#CC3311", "-.", "p")), - ("sz3", ("#CC3311", "-.", "v")), - ("bitround-pco", ("#0077BB", ":", "P")), - ("bitround", ("#33BBEE", "-", "X")), - ("stochround-pco", ("#BBBBBB", "--", "d")), - ("stochround", ("#009988", "--", "h")), - ("tthresh", ("#882255", "-.", "<")), - ("ebcc-abs", ("#AA4444", "-.", ">")), - ("ebcc", ("#AA4444", "-.", "8")), -] - - -def _get_lineinfo(compressor: str) -> tuple[str, str, str]: - """Get the line color, style, and marker for a given compressor.""" - for comp, (color, linestyle, marker) in _COMPRESSOR2LINEINFO: - if compressor.startswith(comp): - return color, linestyle, marker - raise ValueError(f"Unknown compressor: {compressor}") - - -_COMPRESSOR2LEGEND_NAME = [ - ("jpeg2000", "JPEG2000"), - ("sperr", "SPERR"), - ("zfp-round", "ZFP-ROUND"), - ("zfp", "ZFP"), - ("sz3-abs", "SZ3-Abs"), - ("sz3", "SZ3"), - ("bitround-pco", "BitRound + PCO"), - ("bitround", "BitRound + Zstd"), - ("stochround-pco", "StochRound + PCO"), - ("stochround", "StochRound + Zstd"), - ("tthresh", "TTHRESH"), - ("ebcc-abs", "EBCC-Abs"), - ("ebcc", "EBCC"), -] - -DISTORTION2LEGEND_NAME = { - "Relative MAE": "Mean Absolute Error", - "Relative DSSIM": "DSSIM", - "Relative MaxAbsError": "Max Absolute Error", - "Spectral Error": "Spectral Error", -} - def _make_legend_handle(compressor, color, linestyle, marker, line_alpha): """Proxy artist combining the alpha-faded line and the opaque marker.""" @@ -79,15 +32,6 @@ def _make_legend_handle(compressor, color, linestyle, marker, line_alpha): ) -def _get_legend_name(compressor: str) -> str: - """Get the legend name for a given compressor.""" - for comp, name in _COMPRESSOR2LEGEND_NAME: - if compressor.startswith(comp): - return name - - return compressor # Fallback to the compressor name if not found in the mapping. - - def plot_metrics( basepath: Path = Path(), data_loader_basepath: None | Path = None, @@ -252,7 +196,7 @@ def _plot_per_variable_metrics( ): """Creates all the plots which only depend on a single variable.""" for dataset in all_results["Dataset"].unique(): - if dataset != "ifs-uncompressed": + if dataset != "cmip6-access-tos": continue df = all_results[all_results["Dataset"] == dataset] @@ -320,7 +264,7 @@ def _plot_per_variable_metrics( comp, var, error_bound_vals[var], - outfile=err_bound_path / f"{var}_{comp}.png", + outfile=err_bound_path / f"{var}_{comp}.pdf", ) error_dist_plotter.plot_error_bound_histograms( diff --git a/src/climatebenchpress/compressor/plotting/variable_plotters.py b/src/climatebenchpress/compressor/plotting/variable_plotters.py index 3af3bbe..1238485 100644 --- a/src/climatebenchpress/compressor/plotting/variable_plotters.py +++ b/src/climatebenchpress/compressor/plotting/variable_plotters.py @@ -8,6 +8,8 @@ import xarray as xr import xarray.plot.utils as xplot_utils +from .constants import _get_legend_name + class Plotter(ABC): datasets: list[str] @@ -46,9 +48,20 @@ def plot( ax[2].set_title(self.error_title, fontsize=self.title_fontsize) # fig.suptitle(f"{var} Error for {dataset_name} ({compressor})") fig.tight_layout() + fig.suptitle( + f"{_get_legend_name(compressor)}", + fontsize=self.title_fontsize + 4, + y=0.88, + ) if outfile is not None: - with outfile.open("wb") as f: - fig.savefig(f, dpi=300) + if outfile.suffix == ".pdf": + # Passing a file handle hides the suffix from matplotlib, so it + # falls back to the default PNG format and writes PNG bytes into + # a .pdf file. Pass the Path directly so the format is inferred. + fig.savefig(outfile, dpi=300, bbox_inches="tight") + else: + with outfile.open("wb") as f: + fig.savefig(f, dpi=300) plt.close() @@ -163,6 +176,7 @@ def plot_fields(self, fig, ax, ds, ds_new, dataset_name, var, err_bound): transform=ccrs.PlateCarree(), add_colorbar=False, cmap=plt.cm.colors.ListedColormap(["yellow"]), + rasterized=True, ) for a in ax: From 37aee4337d73fbc29665db31614db10ebadb1558 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Tue, 1 Sep 2026 14:10:24 +0100 Subject: [PATCH 09/16] Add the scorecard plots to the plot_metrics script --- .../compressor/plotting/plot_metrics.py | 8 + .../compressor/plotting/scorecards.py | 380 ++++++++++++++++++ 2 files changed, 388 insertions(+) create mode 100644 src/climatebenchpress/compressor/plotting/scorecards.py diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 4927275..e422cc9 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -12,6 +12,7 @@ from ..scripts.compute_metrics import parse_error_bounds from .constants import DISTORTION2LEGEND_NAME, _get_legend_name, _get_lineinfo from .error_dist_plotter import ErrorDistPlotter +from .scorecards import converted_bound_cells, plot_scorecards from .variable_plotters import PLOTTERS @@ -90,8 +91,15 @@ def plot_metrics( rd_curves_metrics=["Max Absolute Error", "MAE", "DSSIM", "Spectral Error"], ) + # The conversion markers are encoded in the compressor name suffixes, so they have + # to be collected before the names are normalized. + converted_cells = converted_bound_cells(df) + df = _rename_compressors(df) normalized_df = _normalize(df) + plot_scorecards( + df, plots_path / "scorecards", converted_cells, bound_names=bound_names + ) _plot_bound_violations( normalized_df, bound_names, plots_path / "bound_violations.pdf" ) diff --git a/src/climatebenchpress/compressor/plotting/scorecards.py b/src/climatebenchpress/compressor/plotting/scorecards.py new file mode 100644 index 0000000..2bcfeee --- /dev/null +++ b/src/climatebenchpress/compressor/plotting/scorecards.py @@ -0,0 +1,380 @@ +"""Scorecard figures summarizing the benchmark results. + +For each error bound level one scorecard is emitted, split into two rows of +metrics. Each cell shows the raw metric value, coloured by its relative +difference to a reference compressor. +""" + +from pathlib import Path + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns + +from .constants import _get_legend_name + +METRICS2NAME = { + "MAE": "Mean Absolute Error", + "Spatial Relative Error (Value)": "SRE", + "Compression Ratio [raw B / enc B]": "Compression Ratio", + "Satisfies Bound (Value)": r"% of Pixels Exceeding Error Bound", +} + +VARIABLE2NAME = { + "10m_u_component_of_wind": "10u", + "10m_v_component_of_wind": "10v", + "mean_sea_level_pressure": "msl", +} + +HIGHER_BETTER_METRICS = ["DSSIM", "Compression Ratio [raw B / enc B]"] + +# DSSIM and Spectral Error are unreliable for variables with large NaN regions. +UNRELIABLE_NAN_METRICS = {"DSSIM", "Spectral Error"} +UNRELIABLE_NAN_VARIABLES = {"ta", "tos"} + +_CONSERVATIVE_SUFFIXES = ("-conservative-abs", "-conservative-rel") + + +def converted_bound_cells(df: pd.DataFrame) -> set[tuple[str, str]]: + """Find the (compressor, variable) pairs whose error bound had to be converted. + + Must be called on the raw results, i.e. before `_rename_compressors` strips the + `-conservative-abs` / `-conservative-rel` suffixes which encode the conversion. + The returned compressor names are the stripped ones, so they match the renamed + frame the scorecards are built from. + """ + converted = set() + for compressor, variable in zip(df["Compressor"], df["Variable"]): + for suffix in _CONSERVATIVE_SUFFIXES: + if compressor.endswith(suffix): + converted.add((compressor.removesuffix(suffix), variable)) + return converted + + +def _create_data_matrix( + df: pd.DataFrame, + error_bound: str, + metrics: list[str], + ref_compressor: str, +) -> tuple[np.ndarray, list[str], list[str]]: + df_filtered = df[df["Error Bound Name"] == error_bound].copy() + # Convert to percentage. + df_filtered["Satisfies Bound (Value)"] = ( + df_filtered["Satisfies Bound (Value)"] * 100 + ) + + variables = sorted(df_filtered["Variable"].unique()) + compressors = sorted(df_filtered["Compressor"].unique()) + compressors = [ref_compressor] + [c for c in compressors if c != ref_compressor] + + column_labels = [f"{v}\n{m}" for m in metrics for v in variables] + data_matrix = np.full((len(compressors), len(column_labels)), np.nan) + + for i, compressor in enumerate(compressors): + for j, metric in enumerate(metrics): + for k, variable in enumerate(variables): + subset = df_filtered[ + (df_filtered["Compressor"] == compressor) + & (df_filtered["Variable"] == variable) + ] + if subset.empty: + print(f"No data for Compressor: {compressor}, Variable: {variable}") + continue + if ( + metric in UNRELIABLE_NAN_METRICS + and variable in UNRELIABLE_NAN_VARIABLES + ): + continue + + col_idx = j * len(variables) + k + if metric in subset.columns: + values = subset[metric] + if len(values) == 1: + data_matrix[i, col_idx] = values.iloc[0] + + return data_matrix, compressors, variables + + +def _create_compression_scorecard( + data_matrix: np.ndarray, + compressors: list[str], + variables: list[str], + metrics: list[str], + converted_cells: set[tuple[str, str]], + cbar: bool = True, + ref_compressor: str = "bitround", + higher_better_metrics: list[str] = HIGHER_BETTER_METRICS, + save_fn: str | Path | None = None, +): + """Create a scorecard plot of relative metric differences vs a reference.""" + ref_idx = compressors.index(ref_compressor) + ref_values = data_matrix[ref_idx, :] + + relative_matrix = np.full_like(data_matrix, np.nan) + for i in range(len(compressors)): + for j in range(data_matrix.shape[1]): + if np.isnan(data_matrix[i, j]) or np.isnan(ref_values[j]): + continue + ref_val = np.abs(ref_values[j]) + if ref_val == 0.0: + ref_val = 1e-10 + metric = metrics[j // len(variables)] + if metric in higher_better_metrics: + relative_matrix[i, j] = ( + (ref_values[j] - data_matrix[i, j]) / ref_val * 100 + ) + elif metric == "Satisfies Bound (Value)": + relative_matrix[i, j] = 100 if data_matrix[i, j] != 0 else 0 + else: + relative_matrix[i, j] = ( + (data_matrix[i, j] - ref_values[j]) / ref_val * 100 + ) + + reds = sns.color_palette("Reds", 6) + blues = sns.color_palette("Blues_r", 6) + cmap = mpl.colors.ListedColormap(blues + [(0.95, 0.95, 0.95)] + reds) + cb_levels = [-100, -75, -50, -25, -10, -1, 1, 10, 25, 50, 75, 100] + norm = mpl.colors.BoundaryNorm(cb_levels, cmap.N, extend="both") + + ncompressors = len(compressors) + nvariables = len(variables) + nmetrics = len(metrics) + + panel_width = (2.5 / 5) * nvariables + label_width = 1.5 * panel_width + padding_right = 0.1 + panel_height = panel_width / nvariables + + title_height = panel_height * 1.25 + cbar_height = panel_height * 2 + spacing_height = panel_height * 0.1 + spacing_width = panel_height * 0.2 + + total_width = ( + label_width + + nmetrics * panel_width + + (nmetrics - 1) * spacing_width + + padding_right + ) + total_height = ( + title_height + + cbar_height + + ncompressors * panel_height + + (ncompressors - 1) * spacing_height + ) + + fig = plt.figure(figsize=(total_width, total_height)) + gs = mpl.gridspec.GridSpec( + ncompressors, + nmetrics, + figure=fig, + left=label_width / total_width, + right=1 - padding_right / total_width, + top=1 - (title_height / total_height), + bottom=cbar_height / total_height, + hspace=spacing_height / panel_height, + wspace=spacing_width / panel_width, + ) + + img = None + border_targets: list[tuple[mpl.axes.Axes, int]] = [] + for row, compressor in enumerate(compressors): + for col, metric in enumerate(metrics): + ax = fig.add_subplot(gs[row, col]) + + start_col = col * nvariables + end_col = start_col + nvariables + rel_values = relative_matrix[row, start_col:end_col].reshape(1, -1) + abs_values = data_matrix[row, start_col:end_col] + + img = ax.imshow(rel_values, aspect="auto", cmap=cmap, norm=norm) + + ax.set_xticks([]) + ax.set_xticklabels([]) + ax.set_yticks([]) + ax.set_yticklabels([]) + + for i in range(nvariables): + rect = mpl.patches.Rectangle( + (i - 0.5, -0.5), + 1, + 1, + linewidth=1, + edgecolor="white", + facecolor="none", + ) + ax.add_patch(rect) + + if (compressor, variables[i]) in converted_cells: + border_targets.append((ax, i)) + + for i, val in enumerate(abs_values): + color = "black" if abs(rel_values[0, i]) < 75 else "white" + fontsize = 10 + if ( + metric in UNRELIABLE_NAN_METRICS + and variables[i] in UNRELIABLE_NAN_VARIABLES + ): + text = "N/A" + color = "black" + elif np.isnan(val): + text = "Fail" + color = "black" + elif abs(val) > 10_000: + text = f"{val:.1e}" + fontsize = 8 + elif abs(val) > 10: + text = f"{val:.0f}" + elif abs(val) > 1: + text = f"{val:.1f}" + elif val == 0: + text = "0.0" + elif abs(val) < 0.01: + text = f"{val:.1e}" + fontsize = 8 + else: + text = f"{val:.2f}" + ax.text( + i, + 0, + text, + ha="center", + va="center", + fontsize=fontsize, + color=color, + ) + + if col == 0: + ax.set_ylabel( + _get_legend_name(compressor), + rotation=0, + ha="right", + va="center", + labelpad=10, + fontsize=14, + ) + + if row == 0: + ax.set_title(METRICS2NAME.get(metric, metric), fontsize=16, pad=10) + ax.tick_params(top=True, labeltop=True, bottom=False, labelbottom=False) + ax.set_xticks(range(nvariables)) + ax.set_xticklabels( + [VARIABLE2NAME.get(v, v) for v in variables], + rotation=45, + ha="left", + fontsize=12, + ) + + for spine in ax.spines.values(): + spine.set_color("0.7") + + # Mark converted cells with a small black triangle in the upper right corner. + # Drawn last so they sit on top of the white grid rectangles and the axes spines. + triangle_size = 0.3 + for ax, i in border_targets: + x_right = i + 0.5 + y_top = -0.5 + triangle = mpl.patches.Polygon( + [ + (x_right - triangle_size, y_top), + (x_right, y_top), + (x_right, y_top + triangle_size), + ], + closed=True, + facecolor="black", + edgecolor="none", + zorder=10, + clip_on=False, + ) + ax.add_patch(triangle) + + if cbar and img is not None: + rel_cbar_height = cbar_height / total_height + cax = fig.add_axes((0.4, rel_cbar_height * 0.3, 0.5, rel_cbar_height * 0.2)) + cb = fig.colorbar(img, cax=cax, orientation="horizontal") + cb.ax.set_xticks(cb_levels) + cb.ax.set_xlabel( + f"Better ← % difference vs {_get_legend_name(ref_compressor)} → Worse", + fontsize=16, + ) + + plt.tight_layout() + + if save_fn: + # bbox_inches="tight" is needed here because the row labels and the rotated + # column labels stick out of the figure box. + plt.savefig(save_fn, dpi=300, bbox_inches="tight") + plt.close() + else: + plt.show() + + +def plot_scorecards( + df: pd.DataFrame, + plots_path: Path, + converted_cells: set[tuple[str, str]], + bound_names: list[str] = ["low", "mid", "high"], + ref_compressor: str = "bitround", + metrics: list[str] = [ + "DSSIM", + "MAE", + "Max Absolute Error", + "Spectral Error", + "Compression Ratio [raw B / enc B]", + "Satisfies Bound (Value)", + ], +): + """Create one scorecard per error bound, each split into two rows of metrics. + + Parameters + ---------- + df: pd.DataFrame + Results with the compressor names already normalized by + `plot_metrics._rename_compressors`. + plots_path: Path + Directory the scorecards are written to. Created if it does not exist. + converted_cells: set[tuple[str, str]] + (compressor, variable) pairs to flag as having a converted error bound, as + returned by `converted_bound_cells` on the raw results. + """ + if ref_compressor not in df["Compressor"].values: + print( + f"Reference compressor {ref_compressor} is missing from the results, " + "skipping the scorecards." + ) + return + + plots_path.mkdir(parents=True, exist_ok=True) + + nrow1 = len(metrics) // 2 + for bound in bound_names: + if df[df["Error Bound Name"] == bound].empty: + print(f"No results for the {bound} error bound, skipping its scorecard.") + continue + + print(f"Creating scorecard for {bound} bound...") + data_matrix, compressors, variables = _create_data_matrix( + df, bound, metrics, ref_compressor + ) + split = nrow1 * len(variables) + _create_compression_scorecard( + data_matrix[:, :split], + compressors, + variables, + metrics[:nrow1], + converted_cells, + ref_compressor=ref_compressor, + cbar=False, + save_fn=plots_path / f"{bound}_scorecard_row1.pdf", + ) + _create_compression_scorecard( + data_matrix[:, split:], + compressors, + variables, + metrics[nrow1:], + converted_cells, + ref_compressor=ref_compressor, + save_fn=plots_path / f"{bound}_scorecard_row2.pdf", + ) From 16b6e7e5e0bfff1dfa845b780b3412833ef98fce Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Tue, 1 Sep 2026 14:15:36 +0100 Subject: [PATCH 10/16] Add flag to disable per variable plots --- .../compressor/plotting/plot_metrics.py | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index e422cc9..bac22a7 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -42,6 +42,7 @@ def plot_metrics( tiny_datasets: bool = False, chunked_datasets: bool = False, use_latex: bool = True, + per_variable_plots: bool = True, ): """Create diagnostic plots for the metrics computed by the compressors. @@ -63,6 +64,9 @@ def plot_metrics( If True, only plot the tiny datasets. Defaults to False. use_latex: bool If True, use LaTeX for rendering text in the plots. Defaults to True. + per_variable_plots: bool + If True, create the per-variable plots, which require reading the compressed + datasets and are hence by far the most expensive ones. Defaults to True. """ metrics_path = basepath / "metrics" plots_path = basepath / "plots" @@ -83,13 +87,14 @@ def plot_metrics( filter_chunked = is_chunked if chunked_datasets else ~is_chunked df = df[filter_chunked] - _plot_per_variable_metrics( - datasets=datasets, - compressed_datasets=compressed_datasets, - plots_path=plots_path, - all_results=df, - rd_curves_metrics=["Max Absolute Error", "MAE", "DSSIM", "Spectral Error"], - ) + if per_variable_plots: + _plot_per_variable_metrics( + datasets=datasets, + compressed_datasets=compressed_datasets, + plots_path=plots_path, + all_results=df, + rd_curves_metrics=["Max Absolute Error", "MAE", "DSSIM", "Spectral Error"], + ) # The conversion markers are encoded in the compressor name suffixes, so they have # to be collected before the names are normalized. @@ -782,6 +787,13 @@ def _savefig(outfile: Path, fig=None): parser.add_argument("--exclude-compressor", type=str, nargs="+", default=[]) parser.add_argument("--tiny-datasets", action="store_true", default=False) parser.add_argument("--avoid-latex", action="store_true", default=False) + parser.add_argument( + "--skip-per-variable-plots", + action="store_true", + default=False, + help="Skip the per-variable plots, which require reading the compressed " + "datasets and are hence by far the most expensive ones.", + ) parser.add_argument("--basepath", type=Path, default=Path()) parser.add_argument( "--data-loader-basepath", @@ -797,4 +809,5 @@ def _savefig(outfile: Path, fig=None): exclude_dataset=args.exclude_dataset, tiny_datasets=args.tiny_datasets, use_latex=(not args.avoid_latex), + per_variable_plots=(not args.skip_per_variable_plots), ) From c8633d7fa3acefd4f5915ec71ec1724efc18e8ac Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Tue, 1 Sep 2026 14:26:10 +0100 Subject: [PATCH 11/16] Allow compress to overwrite existing results --- .../compressor/scripts/compress.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/climatebenchpress/compressor/scripts/compress.py b/src/climatebenchpress/compressor/scripts/compress.py index 8130874..df0f5f2 100644 --- a/src/climatebenchpress/compressor/scripts/compress.py +++ b/src/climatebenchpress/compressor/scripts/compress.py @@ -3,6 +3,7 @@ import argparse import json import math +import shutil import traceback from collections.abc import Callable, Container, Mapping from pathlib import Path @@ -34,6 +35,7 @@ def compress( include_variable: None | Container[str] = None, data_loader_basepath: None | Path = None, chunked: bool = False, + overwrite: bool = False, progress: bool = True, ): """Compress datasets with compressors. @@ -62,6 +64,10 @@ def compress( Input datasets will be loaded from `data_loader_basepath / datasets`. chunked : bool Whether to chunk the input data. + overwrite : bool + Whether to overwrite existing decompressed datasets. If `False`, any + compressor-dataset combination with an existing `decompressed.zarr` is + skipped. progress : bool Whether to show a progress bar during compression. """ @@ -154,7 +160,9 @@ def compress( compressed_dataset_path = compressed_dataset / "decompressed.zarr" if compressed_dataset_path.exists(): - continue + if not overwrite: + continue + _cleanup(compressed_dataset) print( f"Compressing {dataset.parent.name} with {compressor.description} ..." @@ -185,6 +193,13 @@ def compress( ).compute() +def _cleanup(compressed_dataset: Path): + """Delete existing output files for a compressor-dataset combination.""" + shutil.rmtree(compressed_dataset / "decompressed.zarr", ignore_errors=True) + (compressed_dataset / "measurements.json").unlink(missing_ok=True) + (compressed_dataset / "error.out").unlink(missing_ok=True) + + def compress_decompress( codecs: dict[str, Callable[[], Codec]], ds: xr.Dataset, @@ -458,6 +473,7 @@ def revise(dimension, guess): "--data-loader-basepath", type=Path, default=Path() / ".." / "data-loader" ) parser.add_argument("--chunked", action="store_true", default=False) + parser.add_argument("--overwrite", action="store_true", default=False) args = parser.parse_args() compress( @@ -470,5 +486,6 @@ def revise(dimension, guess): include_variable=args.include_variable, data_loader_basepath=args.data_loader_basepath, chunked=args.chunked, + overwrite=args.overwrite, progress=True, ) From 75f2e7463c8fd0eb547aea94a15244dfa86a314d Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Tue, 1 Sep 2026 15:16:33 +0100 Subject: [PATCH 12/16] Fix formatting --- .../compressor/scripts/concatenate_metrics.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py index c0d2b59..44b48da 100644 --- a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py +++ b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py @@ -77,9 +77,7 @@ def concatenate_metrics(basepath: Path = Path(), skip_missing: bool = False): tests = ( pd.read_csv(tests_csv) if tests_csv.exists() - else pd.DataFrame( - columns=["Variable", "Test", "Passed", "Value"] - ) + else pd.DataFrame(columns=["Variable", "Test", "Passed", "Value"]) ) df = merge_metrics(measurements, metrics, tests) From 62a3573bd224bb3916df9e1b8411b2c32afefd98 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Wed, 2 Sep 2026 10:12:10 +0100 Subject: [PATCH 13/16] _get_legend_name -> _get_compressor_legend_name --- .../compressor/plotting/constants.py | 2 +- .../compressor/plotting/plot_metrics.py | 14 +++++++++----- .../compressor/plotting/scorecards.py | 6 +++--- .../compressor/plotting/variable_plotters.py | 4 ++-- 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/constants.py b/src/climatebenchpress/compressor/plotting/constants.py index 801366c..8bd4cce 100644 --- a/src/climatebenchpress/compressor/plotting/constants.py +++ b/src/climatebenchpress/compressor/plotting/constants.py @@ -47,7 +47,7 @@ def _get_lineinfo(compressor: str) -> tuple[str, str, str]: } -def _get_legend_name(compressor: str) -> str: +def _get_compressor_legend_name(compressor: str) -> str: """Get the legend name for a given compressor.""" for comp, name in _COMPRESSOR2LEGEND_NAME: if compressor.startswith(comp): diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index bac22a7..0799780 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -10,7 +10,11 @@ from matplotlib.lines import Line2D from ..scripts.compute_metrics import parse_error_bounds -from .constants import DISTORTION2LEGEND_NAME, _get_legend_name, _get_lineinfo +from .constants import ( + DISTORTION2LEGEND_NAME, + _get_compressor_legend_name, + _get_lineinfo, +) from .error_dist_plotter import ErrorDistPlotter from .scorecards import converted_bound_cells, plot_scorecards from .variable_plotters import PLOTTERS @@ -29,7 +33,7 @@ def _make_legend_handle(compressor, color, linestyle, marker, line_alpha): markersize=12, markerfacecolor=color, markeredgecolor=color, - label=_get_legend_name(compressor), + label=_get_compressor_legend_name(compressor), ) @@ -285,7 +289,7 @@ def _plot_per_variable_metrics( variables, compressors, error_bound_vals, - _get_legend_name, + _get_compressor_legend_name, _get_lineinfo, ) @@ -660,7 +664,7 @@ def _plot_grouped_df( # Bar width bar_width = 0.35 compressors = grouped_df.index.levels[0].tolist() - x_labels = [_get_legend_name(c) for c in compressors] + x_labels = [_get_compressor_legend_name(c) for c in compressors] x_positions = range(len(x_labels)) error_bounds = ["low", "mid", "high"] @@ -731,7 +735,7 @@ def _plot_bound_violations(df, bound_names, outfile: None | Path = None): for i, bound_name in enumerate(bound_names): df_bound = df[df["Error Bound Name"] == bound_name].copy() - df_bound["Compressor"] = df_bound["Compressor"].map(_get_legend_name) + df_bound["Compressor"] = df_bound["Compressor"].map(_get_compressor_legend_name) pass_fail = df_bound.pivot( index="Compressor", columns="Variable", values="Satisfies Bound (Passed)" ) diff --git a/src/climatebenchpress/compressor/plotting/scorecards.py b/src/climatebenchpress/compressor/plotting/scorecards.py index 2bcfeee..a0e9230 100644 --- a/src/climatebenchpress/compressor/plotting/scorecards.py +++ b/src/climatebenchpress/compressor/plotting/scorecards.py @@ -13,7 +13,7 @@ import pandas as pd import seaborn as sns -from .constants import _get_legend_name +from .constants import _get_compressor_legend_name METRICS2NAME = { "MAE": "Mean Absolute Error", @@ -248,7 +248,7 @@ def _create_compression_scorecard( if col == 0: ax.set_ylabel( - _get_legend_name(compressor), + _get_compressor_legend_name(compressor), rotation=0, ha="right", va="center", @@ -296,7 +296,7 @@ def _create_compression_scorecard( cb = fig.colorbar(img, cax=cax, orientation="horizontal") cb.ax.set_xticks(cb_levels) cb.ax.set_xlabel( - f"Better ← % difference vs {_get_legend_name(ref_compressor)} → Worse", + f"Better ← % difference vs {_get_compressor_legend_name(ref_compressor)} → Worse", fontsize=16, ) diff --git a/src/climatebenchpress/compressor/plotting/variable_plotters.py b/src/climatebenchpress/compressor/plotting/variable_plotters.py index 1238485..1902314 100644 --- a/src/climatebenchpress/compressor/plotting/variable_plotters.py +++ b/src/climatebenchpress/compressor/plotting/variable_plotters.py @@ -8,7 +8,7 @@ import xarray as xr import xarray.plot.utils as xplot_utils -from .constants import _get_legend_name +from .constants import _get_compressor_legend_name class Plotter(ABC): @@ -49,7 +49,7 @@ def plot( # fig.suptitle(f"{var} Error for {dataset_name} ({compressor})") fig.tight_layout() fig.suptitle( - f"{_get_legend_name(compressor)}", + f"{_get_compressor_legend_name(compressor)}", fontsize=self.title_fontsize + 4, y=0.88, ) From 56229e1c206d376b66ec06a9c2fdb0544bede7d4 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Wed, 2 Sep 2026 10:20:19 +0100 Subject: [PATCH 14/16] Make prefix filtering CLI argument --- .../compressor/plotting/plot_metrics.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 0799780..1e8a0f8 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -43,6 +43,7 @@ def plot_metrics( bound_names: list[str] = ["low", "mid", "high"], exclude_dataset: list[str] = [], exclude_compressor: list[str] = [], + exclude_compressor_prefix: list[str] = ["safeguarded-", "rp"], tiny_datasets: bool = False, chunked_datasets: bool = False, use_latex: bool = True, @@ -64,6 +65,9 @@ def plot_metrics( List of dataset names to exclude from the plotting. exclude_compressor: list[str] List of compressor names to exclude from the plotting. + exclude_compressor_prefix: list[str] + List of prefixes of compressor names to exclude from the plotting. Defaults + to the safeguarded and random projection variants. tiny_datasets: bool If True, only plot the tiny datasets. Defaults to False. use_latex: bool @@ -81,8 +85,8 @@ def plot_metrics( # Filter out excluded datasets and compressors df = df[~df["Compressor"].isin(exclude_compressor)] - df = df[~df["Compressor"].str.startswith("safeguarded-")] - df = df[~df["Compressor"].str.startswith("rp")] + if exclude_compressor_prefix: + df = df[~df["Compressor"].str.startswith(tuple(exclude_compressor_prefix))] df = df[~df["Dataset"].isin(exclude_dataset)] is_tiny = df["Dataset"].str.endswith("-tiny") filter_tiny = is_tiny if tiny_datasets else ~is_tiny @@ -445,7 +449,6 @@ def _plot_aggregated_rd_curve( distortion, color=color, linestyle=linestyle, - # linestyle="-", linewidth=4, alpha=line_alpha, ) @@ -789,6 +792,14 @@ def _savefig(outfile: Path, fig=None): parser = argparse.ArgumentParser() parser.add_argument("--exclude-dataset", type=str, nargs="+", default=[]) parser.add_argument("--exclude-compressor", type=str, nargs="+", default=[]) + parser.add_argument( + "--exclude-compressor-prefix", + type=str, + nargs="*", + default=[], + help="Exclude all compressors whose name starts with one of these prefixes. " + "Pass with no values to keep all compressors.", + ) parser.add_argument("--tiny-datasets", action="store_true", default=False) parser.add_argument("--avoid-latex", action="store_true", default=False) parser.add_argument( @@ -810,6 +821,7 @@ def _savefig(outfile: Path, fig=None): basepath=args.basepath, data_loader_basepath=args.data_loader_basepath, exclude_compressor=args.exclude_compressor, + exclude_compressor_prefix=args.exclude_compressor_prefix, exclude_dataset=args.exclude_dataset, tiny_datasets=args.tiny_datasets, use_latex=(not args.avoid_latex), From 61bdaec23fabd5474df840c747b342e0b8d5390d Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Wed, 2 Sep 2026 10:25:59 +0100 Subject: [PATCH 15/16] Support virtual file system for PDF plots --- .../compressor/plotting/variable_plotters.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/variable_plotters.py b/src/climatebenchpress/compressor/plotting/variable_plotters.py index 1902314..9ceec03 100644 --- a/src/climatebenchpress/compressor/plotting/variable_plotters.py +++ b/src/climatebenchpress/compressor/plotting/variable_plotters.py @@ -54,14 +54,9 @@ def plot( y=0.88, ) if outfile is not None: - if outfile.suffix == ".pdf": - # Passing a file handle hides the suffix from matplotlib, so it - # falls back to the default PNG format and writes PNG bytes into - # a .pdf file. Pass the Path directly so the format is inferred. - fig.savefig(outfile, dpi=300, bbox_inches="tight") - else: - with outfile.open("wb") as f: - fig.savefig(f, dpi=300) + format = outfile.suffix[1:] # Remove the leading dot + with outfile.open("wb") as f: + fig.savefig(f, dpi=300, format=format) plt.close() From ef48faf0e6e0054498a8042f85a944ca157c2975 Mon Sep 17 00:00:00 2001 From: Tim Reichelt Date: Wed, 2 Sep 2026 10:27:29 +0100 Subject: [PATCH 16/16] Remove hardcoded dataset filtering --- src/climatebenchpress/compressor/plotting/plot_metrics.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 1e8a0f8..55ec1e8 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -217,9 +217,6 @@ def _plot_per_variable_metrics( ): """Creates all the plots which only depend on a single variable.""" for dataset in all_results["Dataset"].unique(): - if dataset != "cmip6-access-tos": - continue - df = all_results[all_results["Dataset"] == dataset] dataset_plots_path = plots_path / dataset dataset_plots_path.mkdir(parents=True, exist_ok=True)