From 91bc620a1a99b00a84f3161e287a50958383a6e3 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 23 Jul 2026 23:07:33 +0200 Subject: [PATCH] perf: compute accuracy and reconstructed on first access tsam v4 defers metric computation until result.accuracy is read, but the wrapper materialized both fields eagerly during result conversion, paying the cost even for callers that never look at them. Store a factory per field and let cached_property populate it on first access; repr reports the metrics only once computed instead of forcing them. On tsam v3 the metrics are computed inside tsam.aggregate() either way, so this defers only the xarray conversion there. 730d x 256 cols, Distribution(scope=global) + extremes, metrics unread: tsam develop-v4 2.59s -> 1.83s (-29%); tsam 3.4.2 unchanged. Co-Authored-By: Claude Fable 5 --- src/tsam_xarray/_clustering.py | 33 ++++++++++-------- src/tsam_xarray/_core.py | 39 +++++++++++---------- src/tsam_xarray/_result.py | 39 +++++++++++++++++---- test/test_lazy_result.py | 64 ++++++++++++++++++++++++++++++++++ 4 files changed, 136 insertions(+), 39 deletions(-) create mode 100644 test/test_lazy_result.py diff --git a/src/tsam_xarray/_clustering.py b/src/tsam_xarray/_clustering.py index c391992..dd4a259 100644 --- a/src/tsam_xarray/_clustering.py +++ b/src/tsam_xarray/_clustering.py @@ -650,7 +650,9 @@ def _apply_single( typical = _representatives_to_da( tsam_result.cluster_representatives, col_dims, dim_names ) - reconstructed = _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims) + + def _make_reconstructed() -> xr.DataArray: + return _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims) cw = _cluster_counts(tsam_result) cluster_ids = np.array(sorted(cw.keys())) @@ -668,18 +670,19 @@ def _apply_single( if isinstance(df.columns, pd.MultiIndex): col_names = [str(n) for n in df.columns.names] - accuracy = AccuracyMetrics( - rmse=_metric_to_da(tsam_result.accuracy.rmse, col_dims, col_names), - mae=_metric_to_da(tsam_result.accuracy.mae, col_dims, col_names), - rmse_duration=_metric_to_da( - tsam_result.accuracy.rmse_duration, col_dims, col_names - ), - weighted_rmse=xr.DataArray(tsam_result.accuracy.weighted_rmse), - weighted_mae=xr.DataArray(tsam_result.accuracy.weighted_mae), - weighted_rmse_duration=xr.DataArray( - tsam_result.accuracy.weighted_rmse_duration - ), - ) + def _make_accuracy() -> AccuracyMetrics: + return AccuracyMetrics( + rmse=_metric_to_da(tsam_result.accuracy.rmse, col_dims, col_names), + mae=_metric_to_da(tsam_result.accuracy.mae, col_dims, col_names), + rmse_duration=_metric_to_da( + tsam_result.accuracy.rmse_duration, col_dims, col_names + ), + weighted_rmse=xr.DataArray(tsam_result.accuracy.weighted_rmse), + weighted_mae=xr.DataArray(tsam_result.accuracy.weighted_mae), + weighted_rmse_duration=xr.DataArray( + tsam_result.accuracy.weighted_rmse_duration + ), + ) seg_durations = _segment_durations_to_da(tsam_result.segment_durations, dim_names) @@ -696,8 +699,8 @@ def _apply_single( cluster_assignments=assignments_da, cluster_counts=cluster_counts_da, segment_durations=seg_durations, - accuracy=accuracy, - reconstructed=reconstructed, + _accuracy_factory=_make_accuracy, + _reconstructed_factory=_make_reconstructed, original=da, clustering=clustering_info, is_transferred=True, diff --git a/src/tsam_xarray/_core.py b/src/tsam_xarray/_core.py index f6f25aa..e791dcc 100644 --- a/src/tsam_xarray/_core.py +++ b/src/tsam_xarray/_core.py @@ -724,8 +724,10 @@ def _result_from_tsam( typical = _representatives_to_da( tsam_result.cluster_representatives, col_dims, dim_names ) - reconstructed = _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims) - reconstructed = reconstructed.transpose(*da.dims).reindex_like(da) + + def _make_reconstructed() -> xr.DataArray: + rec = _reconstructed_to_da(tsam_result.reconstructed, time_dim, col_dims) + return rec.transpose(*da.dims).reindex_like(da) cw = _cluster_counts(tsam_result) cluster_ids = np.array(sorted(cw.keys())) @@ -743,18 +745,19 @@ def _result_from_tsam( if isinstance(df.columns, pd.MultiIndex): col_names = [str(n) for n in df.columns.names] - accuracy = AccuracyMetrics( - rmse=_metric_to_da(tsam_result.accuracy.rmse, col_dims, col_names), - mae=_metric_to_da(tsam_result.accuracy.mae, col_dims, col_names), - rmse_duration=_metric_to_da( - tsam_result.accuracy.rmse_duration, col_dims, col_names - ), - weighted_rmse=xr.DataArray(tsam_result.accuracy.weighted_rmse), - weighted_mae=xr.DataArray(tsam_result.accuracy.weighted_mae), - weighted_rmse_duration=xr.DataArray( - tsam_result.accuracy.weighted_rmse_duration - ), - ) + def _make_accuracy() -> AccuracyMetrics: + return AccuracyMetrics( + rmse=_metric_to_da(tsam_result.accuracy.rmse, col_dims, col_names), + mae=_metric_to_da(tsam_result.accuracy.mae, col_dims, col_names), + rmse_duration=_metric_to_da( + tsam_result.accuracy.rmse_duration, col_dims, col_names + ), + weighted_rmse=xr.DataArray(tsam_result.accuracy.weighted_rmse), + weighted_mae=xr.DataArray(tsam_result.accuracy.weighted_mae), + weighted_rmse_duration=xr.DataArray( + tsam_result.accuracy.weighted_rmse_duration + ), + ) seg_durations = _segment_durations_to_da(tsam_result.segment_durations, dim_names) @@ -773,8 +776,8 @@ def _result_from_tsam( cluster_assignments=assignments_da, cluster_counts=cluster_counts_da, segment_durations=seg_durations, - accuracy=accuracy, - reconstructed=reconstructed, + _accuracy_factory=_make_accuracy, + _reconstructed_factory=_make_reconstructed, original=da, clustering=clustering_info, ) @@ -860,7 +863,7 @@ def _acc_field(field_name: str) -> xr.DataArray: cluster_assignments=_field("cluster_assignments"), cluster_counts=_field("cluster_counts"), segment_durations=_optional_field("segment_durations"), - accuracy=AccuracyMetrics( + _accuracy_factory=lambda: AccuracyMetrics( rmse=_acc_field("rmse"), mae=_acc_field("mae"), rmse_duration=_acc_field("rmse_duration"), @@ -880,7 +883,7 @@ def _acc_field(field_name: str) -> xr.DataArray: slice_coords, ), ), - reconstructed=_field("reconstructed"), + _reconstructed_factory=lambda: _field("reconstructed"), original=_field("original"), clustering=merged_clustering, is_transferred=first.is_transferred, diff --git a/src/tsam_xarray/_result.py b/src/tsam_xarray/_result.py index fc2523a..6e6fcf7 100644 --- a/src/tsam_xarray/_result.py +++ b/src/tsam_xarray/_result.py @@ -3,13 +3,16 @@ from __future__ import annotations import warnings -from dataclasses import dataclass +from dataclasses import dataclass, field +from functools import cached_property from typing import TYPE_CHECKING import pandas as pd import xarray as xr if TYPE_CHECKING: + from collections.abc import Callable + from tsam_xarray._clustering import ClusteringResult from tsam_xarray._dim_names import DimNames @@ -75,8 +78,12 @@ class AggregationResult: ``None``. Dims: ``(cluster, timestep, *slice_dims)``. accuracy: Per-column and weighted accuracy metrics. + Computed on first access; on a tsam that defers + metric computation (v4), never reading it skips + the computation entirely. reconstructed: Reconstructed time series (same shape and dim order as ``original``). + Computed on first access, like ``accuracy``. original: The input data. clustering: Reusable clustering metadata. See `ClusteringResult`. @@ -88,23 +95,43 @@ class AggregationResult: cluster_assignments: xr.DataArray cluster_counts: xr.DataArray segment_durations: xr.DataArray | None - accuracy: AccuracyMetrics - reconstructed: xr.DataArray original: xr.DataArray clustering: ClusteringResult is_transferred: bool = False + _accuracy_factory: Callable[[], AccuracyMetrics] = field( + kw_only=True, repr=False, compare=False + ) + _reconstructed_factory: Callable[[], xr.DataArray] = field( + kw_only=True, repr=False, compare=False + ) + + @cached_property + def accuracy(self) -> AccuracyMetrics: + """Per-column and weighted accuracy metrics, computed on first access.""" + return self._accuracy_factory() + + @cached_property + def reconstructed(self) -> xr.DataArray: + """Reconstructed series on the original time axis, computed on first access.""" + return self._reconstructed_factory() def __repr__(self) -> str: c = self.clustering slices = f", slice_dims={c.slice_dims}" if c.slice_dims else "" seg = f", n_segments={self.n_segments}" if self.n_segments else "" + # cached_property stores into __dict__; checking it avoids forcing + # the computation just to print the result + if "accuracy" in self.__dict__: + acc = f"weighted_rmse={float(self.accuracy.weighted_rmse.mean()):.4f}" + else: + acc = "accuracy=" return ( f"AggregationResult(" f"n_clusters={self.n_clusters}, " f"n_periods={c.n_original_periods}, " f"cluster_dim={c.cluster_dim}" f"{slices}{seg}, " - f"weighted_rmse={float(self.accuracy.weighted_rmse.mean()):.4f})" + f"{acc})" ) @property @@ -273,7 +300,7 @@ def _make_slice_view(self, sel: dict[str, object]) -> AggregationResult: if self.segment_durations is not None else None ), - accuracy=AccuracyMetrics( + _accuracy_factory=lambda: AccuracyMetrics( rmse=self.accuracy.rmse.sel(sel), mae=self.accuracy.mae.sel(sel), rmse_duration=self.accuracy.rmse_duration.sel(sel), @@ -281,7 +308,7 @@ def _make_slice_view(self, sel: dict[str, object]) -> AggregationResult: weighted_mae=self.accuracy.weighted_mae.sel(sel), weighted_rmse_duration=self.accuracy.weighted_rmse_duration.sel(sel), ), - reconstructed=self.reconstructed.sel(sel), + _reconstructed_factory=lambda: self.reconstructed.sel(sel), original=self.original.sel(sel), clustering=CR( time_dim=self.clustering.time_dim, diff --git a/test/test_lazy_result.py b/test/test_lazy_result.py new file mode 100644 index 0000000..8491aa4 --- /dev/null +++ b/test/test_lazy_result.py @@ -0,0 +1,64 @@ +"""Laziness of AggregationResult.accuracy and .reconstructed.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from tsam_xarray import aggregate + + +def _data(n_slices: int = 1) -> xr.DataArray: + rng = np.random.default_rng(0) + n_t = 14 * 24 + dims = ["variable", "time"] + shape: tuple[int, ...] = (3, n_t) + coords: dict[str, object] = { + "variable": ["a", "b", "c"], + "time": pd.date_range("2023-01-01", periods=n_t, freq="h"), + } + if n_slices > 1: + dims = ["scenario", *dims] + shape = (n_slices, *shape) + coords["scenario"] = [f"s{i}" for i in range(n_slices)] + return xr.DataArray(rng.random(shape), dims=dims, coords=coords, name="load") + + +@pytest.mark.parametrize("n_slices", [1, 3]) +def test_accuracy_and_reconstructed_are_deferred(n_slices): + result = aggregate( + _data(n_slices), time_dim="time", cluster_dim="variable", n_clusters=4 + ) + + assert "accuracy" not in result.__dict__ + assert "reconstructed" not in result.__dict__ + + repr(result) + assert "accuracy" not in result.__dict__ + + accuracy = result.accuracy + reconstructed = result.reconstructed + assert "accuracy" in result.__dict__ + assert "reconstructed" in result.__dict__ + + assert result.accuracy is accuracy + assert result.reconstructed is reconstructed + + +def test_deferred_values_match_shapes(): + da = _data(3) + result = aggregate(da, time_dim="time", cluster_dim="variable", n_clusters=4) + + assert result.reconstructed.dims == da.dims + assert result.reconstructed.shape == da.shape + assert set(result.accuracy.rmse.dims) == {"variable", "scenario"} + assert not result.residuals.isnull().all() + + +def test_repr_shows_metrics_once_computed(): + result = aggregate(_data(), time_dim="time", cluster_dim="variable", n_clusters=4) + assert "not computed" in repr(result) + _ = result.accuracy + assert "weighted_rmse=" in repr(result)