diff --git a/src/climatebenchpress/compressor/plotting/constants.py b/src/climatebenchpress/compressor/plotting/constants.py new file mode 100644 index 0000000..8bd4cce --- /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_compressor_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/error_dist_plotter.py b/src/climatebenchpress/compressor/plotting/error_dist_plotter.py index dfbb736..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,10 +66,12 @@ 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) + 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 408b098..55ec1e8 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -1,66 +1,40 @@ import argparse from pathlib import Path +import matplotlib.colors as mcolors import matplotlib.pyplot as plt 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 .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 -_COMPRESSOR2LINEINFO = [ - ("jpeg2000", ("#EE7733", "-")), - ("sperr", ("#117733", ":")), - ("zfp-round", ("#DDAA33", "--")), - ("zfp", ("#EE3377", "--")), - ("sz3", ("#CC3311", "-.")), - ("bitround-pco", ("#0077BB", ":")), - ("bitround", ("#33BBEE", "-")), - ("stochround-pco", ("#BBBBBB", "--")), - ("stochround", ("#009988", "--")), - ("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: - if compressor.startswith(comp): - return color, linestyle - raise ValueError(f"Unknown compressor: {compressor}") - - -_COMPRESSOR2LEGEND_NAME = [ - ("jpeg2000", "JPEG2000"), - ("sperr", "SPERR"), - ("zfp-round", "ZFP-ROUND"), - ("zfp", "ZFP"), - ("sz3", "SZ3"), - ("bitround-pco", "BitRound + PCO"), - ("bitround", "BitRound + Zstd"), - ("stochround-pco", "StochRound + PCO"), - ("stochround", "StochRound + Zstd"), - ("tthresh", "TTHRESH"), -] - -DISTORTION2LEGEND_NAME = { - "Relative MAE": "Mean Absolute Error", - "Relative DSSIM": "DSSIM", - "Relative MaxAbsError": "Max Absolute Error", - "Spectral Error": "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. + +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_compressor_legend_name(compressor), + ) def plot_metrics( @@ -69,9 +43,11 @@ 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, + per_variable_plots: bool = True, ): """Create diagnostic plots for the metrics computed by the compressors. @@ -89,10 +65,16 @@ 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 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" @@ -103,6 +85,8 @@ def plot_metrics( # Filter out excluded datasets and compressors df = df[~df["Compressor"].isin(exclude_compressor)] + 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 @@ -111,16 +95,24 @@ 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. + 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" ) @@ -290,7 +282,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( @@ -298,7 +290,7 @@ def _plot_per_variable_metrics( variables, compressors, error_bound_vals, - _get_legend_name, + _get_compressor_legend_name, _get_lineinfo, ) @@ -344,6 +336,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 +349,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 +379,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 +428,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 +438,28 @@ 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 + marker_alpha = 0.8 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, + alpha=marker_alpha, + ) + legend_handles.append( + _make_legend_handle(comp, color, linestyle, marker, line_alpha) ) if remove_outliers: @@ -502,9 +519,10 @@ def _plot_aggregated_rd_curve( fontsize=16, ) 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, ) @@ -646,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"] @@ -717,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)" ) @@ -771,8 +789,23 @@ 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( + "--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", @@ -785,7 +818,9 @@ 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), + per_variable_plots=(not args.skip_per_variable_plots), ) diff --git a/src/climatebenchpress/compressor/plotting/scorecards.py b/src/climatebenchpress/compressor/plotting/scorecards.py new file mode 100644 index 0000000..a0e9230 --- /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_compressor_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_compressor_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_compressor_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", + ) diff --git a/src/climatebenchpress/compressor/plotting/variable_plotters.py b/src/climatebenchpress/compressor/plotting/variable_plotters.py index aac7a94..9ceec03 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_compressor_legend_name + class Plotter(ABC): datasets: list[str] @@ -46,9 +48,15 @@ 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_compressor_legend_name(compressor)}", + fontsize=self.title_fontsize + 4, + y=0.88, + ) if outfile is not None: + format = outfile.suffix[1:] # Remove the leading dot with outfile.open("wb") as f: - fig.savefig(f, dpi=300) + fig.savefig(f, dpi=300, format=format) plt.close() @@ -163,6 +171,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: @@ -209,6 +218,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"] @@ -360,7 +414,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() @@ -443,6 +497,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: 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, ) diff --git a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py index 2c8d272..44b48da 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 import pandas as pd @@ -9,7 +10,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 @@ -17,6 +18,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" @@ -42,16 +46,39 @@ 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.stem + / 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 @@ -71,7 +98,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" @@ -131,8 +158,10 @@ def merge_metrics( .merge( test_per_variable.reset_index(), on="Variable", + how="outer", ), on="Variable", + how="left", ) @@ -189,6 +218,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)