diff --git a/README.md b/README.md index d3f34cd..463eed5 100644 --- a/README.md +++ b/README.md @@ -1 +1,23 @@ -# compressor +# ClimateBenchPress + +This repository contains the main functionality for the ClimateBenchPress compression benchmark. + +## Getting Started + +This project uses the uv package manager to handle dependencies. If you don't already have it installed follow the instructions at . + +Next, clone this repository and within the project directory install all the necessary dependencies with: +```bash +uv sync +uv pip install -e "." +``` + +### Downloading the Data + +Make sure you have all the necessary data downloaded by following the instructions at . + +## Funding + +ClimateBenchPress has been developed as part of [Embed2Scale](https://embed2scale.eu/) and [ESiWACE3](https://www.esiwace.eu/). + +Funded by the European Union. This work has received funding from the European High Performance Computing Joint Undertaking (JU) under grant agreement No 101093054 and EU’s Horizon Europe program under grant agreement number 101131841. This work also received funding from [UK Research and Innovation (UKRI)](https://www.ukri.org/). diff --git a/docs/index.md b/docs/index.md index d3f34cd..406f59b 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1 +1,27 @@ -# compressor +# ClimateBenchPress + +This repository contains the main functionality for the ClimateBenchPress compression benchmark. + +## Getting Started + +This project uses the uv package manager to handle dependencies. If you don't already have it installed follow the instructions at . + +Next, clone this repository and within the project directory install all the necessary dependencies with: +```bash +uv sync +uv pip install -e "." +``` + +### Downloading the Data + +Make sure you have all the necessary data downloaded by following the instructions at . + +## Using the Benchmark + +Further details on how to run the benchmark evaluation code. + +## Funding + +ClimateBenchPress has been developed as part of [Embed2Scale](https://embed2scale.eu/) and [ESiWACE3](https://www.esiwace.eu/). + +Funded by the European Union. This work has received funding from the European High Performance Computing Joint Undertaking (JU) under grant agreement No 101093054 and EU’s Horizon Europe program under grant agreement number 101131841. This work also received funding from [UK Research and Innovation (UKRI)](https://www.ukri.org/). diff --git a/docs/run_benchmark.md b/docs/run_benchmark.md new file mode 100644 index 0000000..bc51edb --- /dev/null +++ b/docs/run_benchmark.md @@ -0,0 +1,85 @@ +# Evaluating the Benchmark Results + +To evaluate the benchmark results, ensure you have data at `path/to/data-loader/datasets`, which should be downloaded using the [data loader](https://github.com/ClimateBenchPress/data-loader). + +On a high-level the benchmark evaluation pipeline progresses in the following steps: + +1. Compute three error bound levels for each input dataset. +2. Compress each input dataset with all the benchmark compressors for all three error bounds. +3. Compute compressor performance metrics for evaluation purposes. +4. Optional: Create summary plots of the benchmark results. + +We will now go through each of these steps in more detail. As you work through these steps, the pipeline will progressively populate the directories `datasets-error-bounds`, `compressed-datasets`, `metrics`, and `plots`. + +## Create Error Bounds + +Begin by creating the error bounds for each dataset using the following command: +```bash +uv run python -m climatebenchpress.compressor.scripts.create_error_bounds \ + --data-loader-basepath=path/to/data-loader +``` +This step creates three error bounds for each variable in the datasets and stores the information in the `datasets-error-bounds` directory. + +## Compress Input Datasets + +Next, compress all the input datasets by running: +```bash +uv run python -m climatebenchpress.compressor.scripts.compress \ + --data-loader-basepath=path/to/data-loader +``` +This command will populate the `compressed-datasets` directory with the following structure: +``` +compressed-datasets/ + dataset1/ + {var_name}-{err_bound_type}={low_err_bound_val}_{var_name2}-{err_bound_type2}={low_err_bound_val2} + compressor1/ + decompressed.zarr + measurements.json + compressor2/ + ... + {var_name}-{err_bound_type}={mid_err_bound_val}_{var_name2}-{err_bound_type2}={mid_err_bound_val2}/ + ... + {var_name}-{err_bound_type}={high_err_bound_val}_{var_name2}-{err_bound_type2}={high_err_bound_val2}/ + ... + dataset2/ + ... + ... +``` +For each dataset, the results for the three different error bounds are stored in different directories. The `var_name` indicates the variable(s) in the dataset that are being compressed, while `err_bound_type` will be either `abs_error` or `rel_error`. + +You can use additional arguments to control which compressors and datasets are processed: `--exclude-compressor` and `--exclude-dataset` to avoid using certain compressors and datasets, or `--include-compressor` and `--include-dataset` to only use selected compressors on selected datasets. +For example, the command +```bash +uv run python -m climatebenchpress.compressor.scripts.compress \ + --data-loader-basepath=path/to/data-loader \ + --include-compressor sz3 jpeg2000 \ + --include-dataset era5 +``` +compresses the era5 data with the compressors SZ3 and JPEG2000. +These arguments are particularly useful if you wish to parallelize the benchmark evaluation using tools such as `xargs`. + +## Compute Metrics + +After compression, evaluate compression metrics on the compressed datasets using: +```bash +uv run python -m climatebenchpress.compressor.scripts.compute_metrics \ + --data-loader-basepath=path/to/data-loader +``` +You can apply the same filtering options with `--exclude-compressor`, `--exclude-dataset`, `--include-compressor` and `--include-dataset` arguments as used in the compression step. + +Once the metrics are computed, combine all the metrics into a single CSV file by running: +```bash +uv run python -m climatebenchpress.compressor.scripts.concatenate_metrics +``` +This will create the `metrics/all_results.csv` file which contains all the results. + +## Optional: Create Plots + +Finally, generate visualization plots with the following command: +```bash +uv run python -m climatebenchpress.compressor.plotting.plot_metrics \ + --data-loader-basepath=path/to/data-loader +``` +This will create plots in the `plots` directory. By default, this assumes access to a LaTeX compiler. If you do not have one on your system, you can add the `--avoid-latex` flag to this command. + +Note that the full plotting process can take quite a lot of time because it generates individual plots for each error bound-compressor-dataset combination. If you want to avoid generating individual plots for certain datasets, you can do so with the `--exclude-dataset` command line option. diff --git a/mkdocs.yml b/mkdocs.yml index 30876fc..c909e60 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -26,6 +26,8 @@ theme: nav: - Home: index.md + - Tutorials: + - Run the benchmark: run_benchmark.md - Links: - GitHub: https://github.com/ClimateBenchPress/compressor/ - PyPI: https://pypi.org/project/climatebenchpress-compressor/ @@ -40,7 +42,7 @@ plugins: source_dirs: - nav_heading: [Documentation] base: src - ignore: ["cf.py"] + ignore: ["cf.py", "monitor.py", "variable_plotters.py", "error_dist_plotter.py", "__init__.py"] - mkdocstrings: enable_inventory: true handlers: @@ -49,7 +51,7 @@ plugins: docstring_section_style: list docstring_style: numpy show_if_no_docstring: true - filters: ["!^_$", "!^_[^_]", "!^__", "__init__", "__call__", "!^cf$"] + filters: ["!^_$", "!^_[^_]", "!^__", "__init__", "__call__", "!^cf$", "!^parser$", "!^args$"] members_order: source group_by_category: false show_source: false @@ -59,11 +61,7 @@ plugins: show_root_toc_entry: false merge_init_into_class: true annotations_path: source - summary: - attributes: false - classes: true - functions: true - modules: true + summary: false inventories: - https://docs.python.org/3.12/objects.inv - https://numpy.org/doc/2.2/objects.inv diff --git a/src/climatebenchpress/compressor/compressors/abc.py b/src/climatebenchpress/compressor/compressors/abc.py index 1dd2575..ffbcd95 100644 --- a/src/climatebenchpress/compressor/compressors/abc.py +++ b/src/climatebenchpress/compressor/compressors/abc.py @@ -3,8 +3,8 @@ from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Mapping -from functools import partial from dataclasses import dataclass +from functools import partial from types import MappingProxyType from typing import Callable, Optional @@ -19,12 +19,35 @@ @dataclass class NamedPerVariableCodec: + """Dataclass representing a codec for one dataset and compressor. + + Attributes + ---------- + name : str + Name of the error bound used to create the codecs, a combination of variable + names and error bounds. + codecs : dict[VariableName, Callable[[], Codec]] + Dictionary mapping variable names to codec constructors. + """ + name: ErrorBoundName codecs: dict[VariableName, Callable[[], Codec]] @dataclass class ErrorBound: + """Dataclass representing an error bound for a variable. + + Can only have one of `abs_error` or `rel_error` set, not both. + + Attributes + ---------- + abs_error : Optional[float] + Absolute error bound for the variable. + rel_error : Optional[float] + Relative error bound for the variable. + """ + abs_error: Optional[float] = None rel_error: Optional[float] = None @@ -52,7 +75,8 @@ class VariantErrorBoundPerVariable: class Compressor(ABC): - # Abstract interface, must be implemented by subclasses + """Abstract base class for compressors.""" + name: str description: str @@ -65,6 +89,7 @@ def abs_bound_codec( data_min: Optional[float] = None, data_max: Optional[float] = None, ) -> Codec: + """Create a codec with an absolute error bound.""" pass @staticmethod @@ -76,6 +101,7 @@ def rel_bound_codec( data_min: Optional[float] = None, data_max: Optional[float] = None, ) -> Codec: + """Create a codec with a relative error bound.""" pass @classmethod diff --git a/src/climatebenchpress/compressor/compressors/bitround.py b/src/climatebenchpress/compressor/compressors/bitround.py index 501d8f8..feb83ce 100644 --- a/src/climatebenchpress/compressor/compressors/bitround.py +++ b/src/climatebenchpress/compressor/compressors/bitround.py @@ -9,6 +9,13 @@ class BitRound(Compressor): + """Bit Rounding compressor. + + This compressor applies bit rounding to the data, which reduces the precision of the data + while preserving its overall structure. It then applies the Zstandard lossless codec + for further compression. + """ + name = "bitround" description = "Bit Rounding" diff --git a/src/climatebenchpress/compressor/compressors/bitround_pco.py b/src/climatebenchpress/compressor/compressors/bitround_pco.py index e1f34aa..a20fe03 100644 --- a/src/climatebenchpress/compressor/compressors/bitround_pco.py +++ b/src/climatebenchpress/compressor/compressors/bitround_pco.py @@ -10,6 +10,12 @@ class BitRoundPco(Compressor): + """Bit Rounding + PCodec compressor. + + This compressor first applies bit rounding to the data, which reduces the precision of the data + while preserving its overall structure. After that, it uses PCodec for further compression. + """ + name = "bitround-pco" description = "Bit Rounding + PCodec" diff --git a/src/climatebenchpress/compressor/compressors/jpeg2000.py b/src/climatebenchpress/compressor/compressors/jpeg2000.py index 1d73bba..a076ad2 100644 --- a/src/climatebenchpress/compressor/compressors/jpeg2000.py +++ b/src/climatebenchpress/compressor/compressors/jpeg2000.py @@ -12,6 +12,21 @@ class Jpeg2000(Compressor): + """JPEG2000 compressor. + + Note that JPEG2000 does not guarantee pointwise error bounds, but only average error bounds + through specifying a target Peak Signal to Noise Ratio (PSNR). We convert + the absolute error bound to a PSNR value using the formula: + ``` + PSNR = 20 * (log10(data_range) - log10(error_bound)) + ``` + where `data_range = max(data) - min(data)`. + + Additionally, JPEG2000 expects integer data, not floating point, so we linearly quantize the + data into integers ranging between 0 and 2**25 - 1, with 2**25-1 the maximum integer + value accepted by JPEG2000. + """ + name = "jpeg2000" description = "JPEG 2000" diff --git a/src/climatebenchpress/compressor/compressors/sperr.py b/src/climatebenchpress/compressor/compressors/sperr.py index d97625e..05d1523 100644 --- a/src/climatebenchpress/compressor/compressors/sperr.py +++ b/src/climatebenchpress/compressor/compressors/sperr.py @@ -6,6 +6,8 @@ class Sperr(Compressor): + """SPERR compressor.""" + name = "sperr" description = "SPERR" diff --git a/src/climatebenchpress/compressor/compressors/stochround.py b/src/climatebenchpress/compressor/compressors/stochround.py index dad28e7..576fef9 100644 --- a/src/climatebenchpress/compressor/compressors/stochround.py +++ b/src/climatebenchpress/compressor/compressors/stochround.py @@ -9,6 +9,12 @@ class StochRound(Compressor): + """Stochastic Rounding + PCodec compressor. + + This compressor first applies stochastic rounding to the data, which adds noise to the data + while rounding it. After that, it uses Zstandard for further compression. + """ + name = "stochround" description = "Stochastic Rounding" diff --git a/src/climatebenchpress/compressor/compressors/stochround_pco.py b/src/climatebenchpress/compressor/compressors/stochround_pco.py index 6f25180..fab75fa 100644 --- a/src/climatebenchpress/compressor/compressors/stochround_pco.py +++ b/src/climatebenchpress/compressor/compressors/stochround_pco.py @@ -9,6 +9,12 @@ class StochRoundPco(Compressor): + """Stochastic Rounding + PCodec compressor. + + This compressor first applies stochastic rounding to the data, which adds noise to the data + while rounding it. After that, it uses PCodec for further compression. + """ + name = "stochround-pco" description = "Stochastic Rounding + PCodec" diff --git a/src/climatebenchpress/compressor/compressors/sz3.py b/src/climatebenchpress/compressor/compressors/sz3.py index b4f6c08..6d00300 100644 --- a/src/climatebenchpress/compressor/compressors/sz3.py +++ b/src/climatebenchpress/compressor/compressors/sz3.py @@ -6,6 +6,8 @@ class Sz3(Compressor): + """SZ3 compressor.""" + name = "sz3" description = "SZ3" diff --git a/src/climatebenchpress/compressor/compressors/tthresh.py b/src/climatebenchpress/compressor/compressors/tthresh.py index 401186f..f983d4a 100644 --- a/src/climatebenchpress/compressor/compressors/tthresh.py +++ b/src/climatebenchpress/compressor/compressors/tthresh.py @@ -6,6 +6,8 @@ class Tthresh(Compressor): + """Tthresh compressor.""" + name = "tthresh" description = "tthresh" diff --git a/src/climatebenchpress/compressor/compressors/zfp.py b/src/climatebenchpress/compressor/compressors/zfp.py index 9e28cb8..18077c5 100644 --- a/src/climatebenchpress/compressor/compressors/zfp.py +++ b/src/climatebenchpress/compressor/compressors/zfp.py @@ -6,6 +6,8 @@ class Zfp(Compressor): + """ZFP compressor.""" + name = "zfp" description = "ZFP" diff --git a/src/climatebenchpress/compressor/compressors/zfp_round.py b/src/climatebenchpress/compressor/compressors/zfp_round.py index bc55e81..89c7ee0 100644 --- a/src/climatebenchpress/compressor/compressors/zfp_round.py +++ b/src/climatebenchpress/compressor/compressors/zfp_round.py @@ -6,6 +6,12 @@ class ZfpRound(Compressor): + """ZFP-ROUND compressor. + + This is an adjusted version of the ZFP compressor with an improved rounding mechanism + for the transform coefficients. + """ + name = "zfp-round" description = "ZFP-ROUND" diff --git a/src/climatebenchpress/compressor/metrics/abc.py b/src/climatebenchpress/compressor/metrics/abc.py index 4d19674..ab45cbf 100644 --- a/src/climatebenchpress/compressor/metrics/abc.py +++ b/src/climatebenchpress/compressor/metrics/abc.py @@ -6,6 +6,8 @@ class Metric(ABC): + """Base class for metrics.""" + @abstractmethod def __call__(self, x: xr.DataArray, y: xr.DataArray) -> float: """ diff --git a/src/climatebenchpress/compressor/metrics/spectral_error.py b/src/climatebenchpress/compressor/metrics/spectral_error.py index 1162aca..dd1a510 100644 --- a/src/climatebenchpress/compressor/metrics/spectral_error.py +++ b/src/climatebenchpress/compressor/metrics/spectral_error.py @@ -1,35 +1,33 @@ -""" -This module contains code derived from the PySTEPS library. -BSD 3-Clause License - -Copyright (c) 2019, PySteps developers -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -* Redistributions of source code must retain the above copyright notice, this - list of conditions and the following disclaimer. - -* Redistributions in binary form must reproduce the above copyright notice, - this list of conditions and the following disclaimer in the documentation - and/or other materials provided with the distribution. - -* Neither the name of the copyright holder nor the names of its - contributors may be used to endorse or promote products derived from - this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -""" +# This module contains code derived from the PySTEPS library. +# BSD 3-Clause License + +# Copyright (c) 2019, PySteps developers +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. __all__ = ["SpectralError"] diff --git a/src/climatebenchpress/compressor/plotting/plot_metrics.py b/src/climatebenchpress/compressor/plotting/plot_metrics.py index 0b0f468..379b601 100644 --- a/src/climatebenchpress/compressor/plotting/plot_metrics.py +++ b/src/climatebenchpress/compressor/plotting/plot_metrics.py @@ -11,7 +11,7 @@ from .error_dist_plotter import ErrorDistPlotter from .variable_plotters import PLOTTERS -COMPRESSOR2LINEINFO = [ +_COMPRESSOR2LINEINFO = [ ("jpeg2000", ("#EE7733", "-")), ("sperr", ("#117733", ":")), ("zfp-round", ("#DDAA33", "--")), @@ -25,15 +25,15 @@ ] -def get_lineinfo(compressor: str) -> tuple[str, str]: +def _get_lineinfo(compressor: str) -> tuple[str, str]: """Get the line color and style for a given compressor.""" - for comp, (color, linestyle) in COMPRESSOR2LINEINFO: + for comp, (color, linestyle) in _COMPRESSOR2LINEINFO: if compressor.startswith(comp): return color, linestyle raise ValueError(f"Unknown compressor: {compressor}") -COMPRESSOR2LEGEND_NAME = [ +_COMPRESSOR2LEGEND_NAME = [ ("jpeg2000", "JPEG2000"), ("sperr", "SPERR"), ("zfp-round", "ZFP-ROUND"), @@ -47,9 +47,9 @@ def get_lineinfo(compressor: str) -> tuple[str, str]: ] -def get_legend_name(compressor: str) -> str: +def _get_legend_name(compressor: str) -> str: """Get the legend name for a given compressor.""" - for comp, name in COMPRESSOR2LEGEND_NAME: + for comp, name in _COMPRESSOR2LEGEND_NAME: if compressor.startswith(comp): return name @@ -66,6 +66,29 @@ def plot_metrics( tiny_datasets: bool = False, use_latex: bool = True, ): + """Create diagnostic plots for the metrics computed by the compressors. + + Parameters + ---------- + basepath: Path + Assumes compressed datasets are stored in `basepath / compressed-datasets` + and metrics in `basepath / metrics`. + data_loader_basepath: Path | None + Assumes datasets are stored in `data_loader_basepath / datasets`. + If None, uses `basepath / datasets`. + bound_names: list[str] + Names of the error bounds to use for plotting. + normalizer: str + Compressor to use for normalization of the metrics. + exclude_dataset: list[str] + List of dataset names to exclude from the plotting. + exclude_compressor: list[str] + List of compressor names to exclude from the plotting. + 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. + """ metrics_path = basepath / "metrics" plots_path = basepath / "plots" datasets = (data_loader_basepath or basepath) / "datasets" @@ -80,24 +103,24 @@ def plot_metrics( filter_tiny = is_tiny if tiny_datasets else ~is_tiny df = df[filter_tiny] - plot_per_variable_metrics( + _plot_per_variable_metrics( datasets=datasets, compressed_datasets=compressed_datasets, plots_path=plots_path, all_results=df, ) - df = rename_compressors(df) - normalized_df = normalize(df, bound_normalize="mid", normalizer=normalizer) - plot_bound_violations( + df = _rename_compressors(df) + normalized_df = _normalize(df, bound_normalize="mid", normalizer=normalizer) + _plot_bound_violations( normalized_df, bound_names, plots_path / "bound_violations.pdf" ) - plot_throughput(df, plots_path / "throughput.pdf") - plot_instruction_count(df, plots_path / "instruction_count.pdf") + _plot_throughput(df, plots_path / "throughput.pdf") + _plot_instruction_count(df, plots_path / "instruction_count.pdf") for metric in ["Relative MAE", "Relative DSSIM", "Relative MaxAbsError"]: with plt.rc_context(rc={"text.usetex": use_latex}): - plot_aggregated_rd_curve( + _plot_aggregated_rd_curve( normalized_df, normalizer=normalizer, compression_metric="Relative CR", @@ -108,7 +131,7 @@ def plot_metrics( ) -def rename_compressors(df): +def _rename_compressors(df): """Give compressors consistent names. They sometimes have suffixes if they are applied on a converted error bound. The three patterns are: - {compressor_name} @@ -122,7 +145,7 @@ def rename_compressors(df): return df -def sort_error_bounds(error_bounds: list[str]) -> list[str]: +def _sort_error_bounds(error_bounds: list[str]) -> list[str]: """Each error bound has the format {variable_name}-{bound_type}={bound_value}_{variable_name2}-{bound_type2}={bound_value2} @@ -136,7 +159,7 @@ def sort_error_bounds(error_bounds: list[str]) -> list[str]: ) -def normalize(data, bound_normalize="mid", normalizer=None): +def _normalize(data, bound_normalize="mid", normalizer=None): """Generate normalized metrics for each compressor and variable. The normalization is done either with respect to either a user provided compressor or the compressor with the highest average rank over all variables (ranked by @@ -187,7 +210,7 @@ def get_normalizer(row): return normalized -def plot_per_variable_metrics( +def _plot_per_variable_metrics( datasets: Path, compressed_datasets: Path, plots_path: Path, @@ -206,7 +229,7 @@ def plot_per_variable_metrics( metric_name = dist_metric.lower().replace(" ", "_") if df[df["Variable"] == var][dist_metric].isnull().all(): continue - plot_variable_rd_curve( + _plot_variable_rd_curve( df[df["Variable"] == var], distortion_metric=dist_metric, outfile=dataset_plots_path @@ -214,7 +237,7 @@ def plot_per_variable_metrics( ) error_bounds = df[df["Dataset"] == dataset]["Error Bound"].unique() - error_bounds = sort_error_bounds(error_bounds) + error_bounds = _sort_error_bounds(error_bounds) error_dist_plotter = ErrorDistPlotter( variables=variables, @@ -250,7 +273,7 @@ def plot_per_variable_metrics( error_bound_vals[var][0], ) - plot_variable_error( + _plot_variable_error( ds[var], ds_new[var], dataset, @@ -264,16 +287,16 @@ def plot_per_variable_metrics( variables, compressors, error_bound_vals, - get_legend_name, - get_lineinfo, + _get_legend_name, + _get_lineinfo, ) fig, _ = error_dist_plotter.get_final_figure() - savefig(dataset_plots_path / f"error_histograms_{dataset}.pdf") + _savefig(dataset_plots_path / f"error_histograms_{dataset}.pdf") plt.close(fig) -def plot_variable_error( +def _plot_variable_error( uncompressed_data: xr.DataArray, compressed_data: xr.DataArray, dataset_name: str, @@ -294,7 +317,7 @@ def plot_variable_error( print(f"No plotter found for dataset {dataset_name}") -def plot_variable_rd_curve(df, distortion_metric, outfile: None | Path = None): +def _plot_variable_rd_curve(df, distortion_metric, outfile: None | Path = None): plt.figure(figsize=(8, 6)) compressors = df["Compressor"].unique() for comp in compressors: @@ -305,11 +328,11 @@ def plot_variable_rd_curve(df, distortion_metric, outfile: None | Path = None): for i in sorting_ixs ] distortion = [compressor_data[distortion_metric].iloc[i] for i in sorting_ixs] - color, linestyle = get_lineinfo(comp) + color, linestyle = _get_lineinfo(comp) plt.plot( compr_ratio, distortion, - label=get_legend_name(comp), + label=_get_legend_name(comp), marker="s", color=color, linestyle=linestyle, @@ -349,11 +372,11 @@ def plot_variable_rd_curve(df, distortion_metric, outfile: None | Path = None): plt.tight_layout() if outfile is not None: - savefig(outfile) + _savefig(outfile) plt.close() -def plot_aggregated_rd_curve( +def _plot_aggregated_rd_curve( normalized_df, normalizer, compression_metric, @@ -381,11 +404,11 @@ def plot_aggregated_rd_curve( agg_distortion.loc[(bound, comp), distortion_metric] for bound in bound_names ] - color, linestyle = get_lineinfo(comp) + color, linestyle = _get_lineinfo(comp) plt.plot( compr_ratio, distortion, - label=get_legend_name(comp), + label=_get_legend_name(comp), marker="s", color=color, linestyle=linestyle, @@ -423,7 +446,7 @@ def plot_aggregated_rd_curve( right=True, ) - normalizer_label = get_legend_name(normalizer) + normalizer_label = _get_legend_name(normalizer) if "MAE" in distortion_metric: plt.legend( title="Compressor", @@ -503,11 +526,11 @@ def plot_aggregated_rd_curve( plt.tight_layout() if outfile is not None: - savefig(outfile) + _savefig(outfile) plt.close() -def plot_throughput(df, outfile: None | Path = None): +def _plot_throughput(df, outfile: None | Path = None): # Transform throughput measurements from raw B/s to s/MB for better comparison # with instruction count measurements. encode_col = "Encode Throughput [raw B / s]" @@ -519,8 +542,8 @@ def plot_throughput(df, outfile: None | Path = None): new_df[transformed_decode_col] = 1e6 / new_df[decode_col] encode_col, decode_col = transformed_encode_col, transformed_decode_col - grouped_df = get_median_and_quantiles(new_df, encode_col, decode_col) - plot_grouped_df( + grouped_df = _get_median_and_quantiles(new_df, encode_col, decode_col) + _plot_grouped_df( grouped_df, title="", ylabel="Throughput [s / MB]", @@ -528,11 +551,11 @@ def plot_throughput(df, outfile: None | Path = None): ) -def plot_instruction_count(df, outfile: None | Path = None): +def _plot_instruction_count(df, outfile: None | Path = None): encode_col = "Encode Instructions [# / raw B]" decode_col = "Decode Instructions [# / raw B]" - grouped_df = get_median_and_quantiles(df, encode_col, decode_col) - plot_grouped_df( + grouped_df = _get_median_and_quantiles(df, encode_col, decode_col) + _plot_grouped_df( grouped_df, title="", ylabel="Instructions [# / raw B]", @@ -540,7 +563,7 @@ def plot_instruction_count(df, outfile: None | Path = None): ) -def get_median_and_quantiles(df, encode_column, decode_column): +def _get_median_and_quantiles(df, encode_column, decode_column): return df.groupby(["Compressor", "Error Bound Name"])[ [encode_column, decode_column] ].agg( @@ -565,13 +588,13 @@ def get_median_and_quantiles(df, encode_column, decode_column): ) -def plot_grouped_df(grouped_df, title, ylabel, outfile: None | Path = None): +def _plot_grouped_df(grouped_df, title, ylabel, outfile: None | Path = None): fig, axes = plt.subplots(1, 3, figsize=(18, 6), sharex=True, sharey=True) # 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_legend_name(c) for c in compressors] x_positions = range(len(x_labels)) error_bounds = ["low", "mid", "high"] @@ -590,7 +613,7 @@ def plot_grouped_df(grouped_df, title, ylabel, outfile: None | Path = None): bound_data["encode_upper_quantile"], ], label="Encoding", - color=[get_lineinfo(comp)[0] for comp in compressors], + color=[_get_lineinfo(comp)[0] for comp in compressors], ) # Plot decode throughput @@ -603,7 +626,7 @@ def plot_grouped_df(grouped_df, title, ylabel, outfile: None | Path = None): bound_data["decode_upper_quantile"], ], label="Decoding", - edgecolor=[get_lineinfo(comp)[0] for comp in compressors], + edgecolor=[_get_lineinfo(comp)[0] for comp in compressors], fill=False, linewidth=4, ) @@ -632,16 +655,16 @@ def plot_grouped_df(grouped_df, title, ylabel, outfile: None | Path = None): fig.tight_layout() if outfile is not None: - savefig(outfile) + _savefig(outfile) plt.close() -def plot_bound_violations(df, bound_names, outfile: None | Path = None): +def _plot_bound_violations(df, bound_names, outfile: None | Path = None): fig, axs = plt.subplots(1, 3, figsize=(len(bound_names) * 6, 6), sharey=True) 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_legend_name) pass_fail = df_bound.pivot( index="Compressor", columns="Variable", values="Satisfies Bound (Passed)" ) @@ -674,11 +697,11 @@ def plot_bound_violations(df, bound_names, outfile: None | Path = None): fig.tight_layout() if outfile is not None: - savefig(outfile) + _savefig(outfile) plt.close() -def savefig(outfile: Path): +def _savefig(outfile: Path): ispdf = outfile.suffix == ".pdf" if ispdf: # Saving a PDF with the alternative code below leads to a corrupted file. diff --git a/src/climatebenchpress/compressor/scripts/compress.py b/src/climatebenchpress/compressor/scripts/compress.py index 24059cb..a771010 100644 --- a/src/climatebenchpress/compressor/scripts/compress.py +++ b/src/climatebenchpress/compressor/scripts/compress.py @@ -29,6 +29,28 @@ def compress( data_loader_basepath: None | Path = None, progress: bool = True, ): + """Compress datasets with compressors. + + Parameters + ---------- + basepath : Path + Compressed dataset will be stored in `basepath / compressed-datasets`. + exclude_dataset : Container[str] + Datasets to exclude from compression. + include_dataset : None | Container[str] + Datasets to include in compression. If `None`, all datasets are included. + If specified, only datasets in `include_dataset` will be compressed. + exclude_compressor : Container[str] + Compressors to exclude from compression. + include_compressor : None | Container[str] + Compressors to include in compression. If `None`, all compressors are included. + If specified, only compressors in `include_compressor` will be used. + data_loader_basepath : None | Path + Base path for the data loader datasets. If `None`, defaults to `basepath / .. / data-loader`. + Input datasets will be loaded from `data_loader_basepath / datasets`. + progress : bool + Whether to show a progress bar during compression. + """ datasets = (data_loader_basepath or basepath) / "datasets" compressed_datasets = basepath / "compressed-datasets" datasets_error_bounds = basepath / "datasets-error-bounds" diff --git a/src/climatebenchpress/compressor/scripts/compute_metrics.py b/src/climatebenchpress/compressor/scripts/compute_metrics.py index 8db5bc3..9a93411 100644 --- a/src/climatebenchpress/compressor/scripts/compute_metrics.py +++ b/src/climatebenchpress/compressor/scripts/compute_metrics.py @@ -34,6 +34,27 @@ def compute_metrics( exclude_compressor: Iterable[str] = tuple(), include_compressor: None | Iterable[str] = None, ): + """Compute evaluation metrics for compressors. + + Parameters + ---------- + basepath : Path + Assumes the compressed datasets are stored in `basepath / compressed-datasets`. + Computed metrics will be stored in `basepath / metrics`. + data_loader_basepath : None | Path + Base path for the data loader datasets. If `None`, defaults to `basepath / .. / data-loader`. + Input datasets will be loaded from `data_loader_basepath / datasets`. + exclude_dataset : Iterable[str] + Datasets to exclude from evaluation. + include_dataset : None | Iterable[str] + Datasets to include in evaluation. If `None`, all datasets are included. + If specified, only datasets in `include_dataset` will be evaluated. + exclude_compressor : Iterable[str] + Compressors to exclude from evaluation. + include_compressor : None | Iterable[str] + Compressors to include in evaluation. If `None`, all compressors are included. + If specified, only compressors in `include_compressor` will be evaluated. + """ exclude_compressor = add_compressor_suffixes(exclude_compressor) include_compressor = add_compressor_suffixes(include_compressor) diff --git a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py index 02459d0..5816ae5 100644 --- a/src/climatebenchpress/compressor/scripts/concatenate_metrics.py +++ b/src/climatebenchpress/compressor/scripts/concatenate_metrics.py @@ -11,6 +11,14 @@ def concatenate_metrics(basepath: Path = Path()): + """Concatenate metrics from all datasets and compressors into a single CSV file. + + Parameters + ---------- + 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. + """ compressed_datasets = basepath / "compressed-datasets" error_bounds_dir = basepath / "datasets-error-bounds" metrics_dir = basepath / "metrics" diff --git a/src/climatebenchpress/compressor/scripts/create_error_bounds.py b/src/climatebenchpress/compressor/scripts/create_error_bounds.py index 057dfba..e624ec0 100644 --- a/src/climatebenchpress/compressor/scripts/create_error_bounds.py +++ b/src/climatebenchpress/compressor/scripts/create_error_bounds.py @@ -76,6 +76,17 @@ def create_error_bounds( basepath: Path = Path(), data_loader_basepath: None | Path = None, ): + """Create three error bounds for all datasets and the variables in them. + + Parameters + ---------- + basepath : Path + The base path where the error bounds will be stored. + The error bounds will be stored in `basepath / datasets-error-bounds`. + data_loader_basepath : Path, optional + The base path where the datasets are stored. If not provided, it defaults to `basepath / .. / data-loader`. + The datasets will be loaded from `data_loader_basepath / datasets`. + """ datasets = (data_loader_basepath or basepath) / "datasets" datasets_error_bounds = basepath / "datasets-error-bounds" diff --git a/src/climatebenchpress/compressor/tests/abc.py b/src/climatebenchpress/compressor/tests/abc.py index 3aba8da..71fb25d 100644 --- a/src/climatebenchpress/compressor/tests/abc.py +++ b/src/climatebenchpress/compressor/tests/abc.py @@ -4,6 +4,8 @@ class Test(ABC): + """Base class for pass/fail tests for compressors.""" + @abstractmethod def __call__(self, x: xr.DataArray, y: xr.DataArray) -> tuple[bool, float]: """ diff --git a/src/climatebenchpress/compressor/tests/error_bound.py b/src/climatebenchpress/compressor/tests/error_bound.py index 895b7a5..c6ad78d 100644 --- a/src/climatebenchpress/compressor/tests/error_bound.py +++ b/src/climatebenchpress/compressor/tests/error_bound.py @@ -5,7 +5,24 @@ class ErrorBound(Test): - def __init__(self, error_type: str, threshold: float = 0.05): + """Tests whether the absolute or relative error between two arrays is below a threshold. + + For relative error, we use the formula `|x - y| <= |x| *threshold` to ensure + that the error is also defined when the input `x` is zero. Consequently, the + relative error also checks that the input and output have matching 0s. + + Additionally, this test checks that locations of `nan` and `inf` values in + `x` and `y` match. + + Parameters + ---------- + error_type : str + Either "abs_error" for absolute error or "rel_error" for relative error. + threshold : float + The threshold value for the error. + """ + + def __init__(self, error_type: str, threshold: float): self.threshold = threshold assert error_type in [ "abs_error", diff --git a/src/climatebenchpress/compressor/tests/r2_correlation.py b/src/climatebenchpress/compressor/tests/r2_correlation.py index 1af9586..39942ca 100644 --- a/src/climatebenchpress/compressor/tests/r2_correlation.py +++ b/src/climatebenchpress/compressor/tests/r2_correlation.py @@ -6,6 +6,13 @@ class R2(Test): + """Test whether the R^2 correlation coefficient between two arrays exceeds a threshold. + + Default threshold is taken from [1]. + + > [1] Baker, Allison H., Haiying Xu, Dorit M. Hammerling, Shaomeng Li, and John P. Clyne. "Toward a multi-method approach: Lossy data compression for climate simulation data." In International conference on high Performance computing, pp. 30-42. Cham: Springer International Publishing, 201 + """ + def __init__(self, threshold: float = 0.99999): self.threshold = threshold diff --git a/src/climatebenchpress/compressor/tests/spatial_relative_error.py b/src/climatebenchpress/compressor/tests/spatial_relative_error.py index 24221dc..2c8cd5b 100644 --- a/src/climatebenchpress/compressor/tests/spatial_relative_error.py +++ b/src/climatebenchpress/compressor/tests/spatial_relative_error.py @@ -5,6 +5,23 @@ class SRE(Test): + """Test whether the spatial relative error between two arrays exceeds a threshold. + + The spatial relative error is defined as the percentage of grid points where the relative error + exceeds a given delta, i.e. $\\frac{1}{N} \\sum_i \\mathbb{I}[|x_i - y_i| / |x_i| > \\delta]$. + + Default threshold and delta is taken from [1]. + + > [1] Baker, Allison H., Haiying Xu, Dorit M. Hammerling, Shaomeng Li, and John P. Clyne. "Toward a multi-method approach: Lossy data compression for climate simulation data." In International conference on high Performance computing, pp. 30-42. Cham: Springer International Publishing, 201 + + Parameters + ---------- + delta : float + The relative error threshold for each grid point. + threshold : float + The maximum proportion of grid points that can exceed the relative error threshold. + """ + def __init__(self, delta: float = 1e-4, threshold: float = 0.05): self.delta = delta self.threshold = threshold