diff --git a/benchmarks/test_bench_aggregate.py b/benchmarks/test_bench_aggregate.py index 4c5c169..8d8e27c 100644 --- a/benchmarks/test_bench_aggregate.py +++ b/benchmarks/test_bench_aggregate.py @@ -232,6 +232,36 @@ def test_user_disaggregate(benchmark, wide_result): ) +@pytest.fixture(scope="module") +def sliced_result(): + da = make_data(365, 32, 12) + return aggregate(da, time_dim="time", cluster_dim="variable", n_clusters=12) + + +def test_user_disaggregate_sliced(benchmark, sliced_result): + benchmark.pedantic( + sliced_result.clustering.disaggregate, + args=(sliced_result.cluster_representatives,), + **CONFIG_OPTS, + ) + + +def test_user_disaggregate_segmented(benchmark): + da = make_data(365, 32, 1) + result = aggregate( + da, + time_dim="time", + cluster_dim="variable", + n_clusters=12, + segments=tsam.SegmentConfig(n_segments=6), + ) + benchmark.pedantic( + result.clustering.disaggregate, + args=(result.cluster_representatives,), + **CONFIG_OPTS, + ) + + # --- large tier (production-sized, opt-in via --large) ------------------------- LARGE_OPTS = dict(rounds=2, iterations=1, warmup_rounds=0) diff --git a/src/tsam_xarray/_clustering.py b/src/tsam_xarray/_clustering.py index c391992..415c66f 100644 --- a/src/tsam_xarray/_clustering.py +++ b/src/tsam_xarray/_clustering.py @@ -429,11 +429,17 @@ def disaggregate(self, data: xr.DataArray) -> xr.DataArray: slice_coords = {d: data.coords[d].values for d in slice_dims} keys = list(itertools.product(*(slice_coords[d] for d in slice_dims))) + crs = [_lookup_clustering(self.clusterings, key) for key in keys] + + if _is_gatherable(crs, data, slice_dims): + return _disaggregate_gather( + crs, data, self.dim_names, slice_dims, slice_coords + ) + results = [] - for key in keys: + for key, cr in zip(keys, crs, strict=True): sel = dict(zip(slice_dims, key, strict=True)) data_slice = data.sel(sel) - cr = _lookup_clustering(self.clusterings, key) results.append(_disaggregate_single(cr, data_slice, self.dim_names)) return _concat_along_dims(results, slice_dims, slice_coords) @@ -704,6 +710,156 @@ def _apply_single( ) +def _is_gatherable( + crs: list[tsam.ClusteringResult], + data: xr.DataArray, + slice_dims: list[str], +) -> bool: + """Whether all slices can be disaggregated by one vectorized gather. + + The gather produces a single rectangular array with one shared time + axis, so it needs unsegmented clusterings of identical shape whose time + indices agree. + + Args: + crs: Per-slice clusterings, in the order the slices appear + along ``slice_dims``. + data: The payload to disaggregate. + slice_dims: Dimension(s) aggregated independently. + + Returns: + ``True`` if one gather covers every slice, ``False`` to fall + back to the per-slice tsam path. + """ + if any(d not in data.dims for d in slice_dims): + return False + first = crs[0] + if first.segment_durations is not None: + return False + return all( + cr.segment_durations is None + and cr.n_timesteps_per_period == first.n_timesteps_per_period + and len(cr.cluster_assignments) == len(first.cluster_assignments) + and _same_time_index(cr.time_index, first.time_index) + for cr in crs[1:] + ) + + +def _same_time_index(a: pd.Index | None, b: pd.Index | None) -> bool: + """Whether two stored time indices would produce the same time axis. + + Args: + a: A clustering's stored time index, or ``None``. + b: The index to compare against, or ``None``. + + Returns: + ``True`` if both are absent or hold equal values. + """ + if a is None or b is None: + return a is None and b is None + return len(a) == len(b) and bool(np.array_equal(np.asarray(a), np.asarray(b))) + + +def _validate_gather_input( + cr: tsam.ClusteringResult, + clusters: np.ndarray, + n_timesteps: int, +) -> None: + """Mirror tsam's disaggregate input checks for the vectorized path.""" + expected = set(cr.cluster_assignments) + got = set(np.asarray(clusters).tolist()) + if got != expected: + parts = [] + if expected - got: + parts.append(f"missing clusters {sorted(expected - got)}") + if got - expected: + parts.append(f"unexpected clusters {sorted(got - expected)}") + msg = ( + f"Cluster IDs in data do not match this clustering: " + f"{', '.join(parts)}. " + f"Expected {sorted(expected)}, got {sorted(got)}." + ) + raise ValueError(msg) + + if n_timesteps != cr.n_timesteps_per_period: + msg = ( + f"data has {n_timesteps} timesteps per cluster, " + f"expected {cr.n_timesteps_per_period}" + ) + raise ValueError(msg) + + +def _disaggregate_gather( + crs: list[tsam.ClusteringResult], + data: xr.DataArray, + dim_names: DimNames, + slice_dims: list[str], + slice_coords: dict[str, Any], +) -> xr.DataArray: + """Disaggregate unsegmented clusterings with a single vectorized gather. + + Expanding an unsegmented clustering is a gather along the cluster axis: + every original period takes the values of its assigned representative. + tsam's ``disaggregate()`` expresses this as an + ``unstack``/``.loc``/``stack`` round-trip through pandas, which dominates + the runtime and is repeated once per slice. Doing it in numpy covers all + slices at once. + + Output dim order matches the per-slice path: ``(*slice_dims, time, + *other_dims)``. + """ + cluster_dim = dim_names.cluster + timestep_dim = dim_names.timestep + other_dims = [ + str(d) + for d in data.dims + if d not in (cluster_dim, timestep_dim) and d not in slice_dims + ] + ordered = data.transpose(cluster_dim, timestep_dim, *slice_dims, *other_dims) + + clusters = ordered.coords[cluster_dim].values + n_clusters = len(clusters) + n_timesteps = ordered.sizes[timestep_dim] + slice_sizes = tuple(ordered.sizes[d] for d in slice_dims) + other_sizes = ordered.shape[2 + len(slice_dims) :] + + positions = pd.Index(clusters) + index = np.empty((len(crs), len(crs[0].cluster_assignments)), dtype=np.intp) + for i, cr in enumerate(crs): + _validate_gather_input(cr, clusters, n_timesteps) + index[i] = positions.get_indexer(np.asarray(cr.cluster_assignments)) + + n_slices = int(np.prod(slice_sizes, dtype=int)) + n_periods = index.shape[1] + + # (cluster, timestep, slices, others) -> (cluster, slices, timestep, others) + values = ordered.values.reshape(n_clusters, n_timesteps, n_slices, -1) + values = values.transpose(0, 2, 1, 3) + gathered = values[index.T, np.arange(n_slices)[None, :]] + gathered = gathered.transpose(1, 0, 2, 3) + gathered = gathered.reshape(*slice_sizes, n_periods * n_timesteps, *other_sizes) + + n_time = n_periods * n_timesteps + stored = crs[0].time_index + time_index: pd.Index = ( + stored + if stored is not None and len(stored) == n_time + else pd.RangeIndex(n_time) + ) + + result = xr.DataArray( + gathered, + dims=[*slice_dims, "time", *other_dims], + coords={"time": time_index}, + ) + for d in (*slice_dims, *other_dims): + if d in slice_coords: + result = result.assign_coords({d: slice_coords[d]}) + elif d in data.coords: + result = result.assign_coords({d: data.coords[d]}) + return result + + def _disaggregate_single( cr: tsam.ClusteringResult, data: xr.DataArray, @@ -711,9 +867,15 @@ def _disaggregate_single( ) -> xr.DataArray: """Disaggregate a single (non-sliced) DataArray using a ClusteringResult. - Relies on tsam's ``cr.disaggregate()`` to return a DataFrame indexed - by the original ``DatetimeIndex`` stored on the clustering. + Segmented clusterings go through tsam's ``cr.disaggregate()``, which + returns a DataFrame indexed by the original ``DatetimeIndex`` stored on + the clustering. Unsegmented ones take the vectorized gather instead. """ + if cr.segment_durations is None: + return _disaggregate_gather( + [cr], data, dim_names, slice_dims=[], slice_coords={} + ) + cluster_dim = dim_names.cluster timestep_dim = dim_names.timestep other_dims = [str(d) for d in data.dims if d not in (cluster_dim, timestep_dim)] @@ -726,19 +888,12 @@ def _disaggregate_single( flat = ordered.values.reshape(n_clusters * n_timesteps, -1) - if cr.segment_durations is not None: - idx_tuples = [ - (int(c), seg, int(dur)) - for c in clusters - for seg, dur in enumerate(cr.segment_durations[int(c)]) - ] - mi = pd.MultiIndex.from_tuples( - idx_tuples, names=["cluster", "segment", "duration"] - ) - else: - mi = pd.MultiIndex.from_product( - [clusters, range(n_timesteps)], names=["cluster", "timestep"] - ) + idx_tuples = [ + (int(c), seg, int(dur)) + for c in clusters + for seg, dur in enumerate(cr.segment_durations[int(c)]) + ] + mi = pd.MultiIndex.from_tuples(idx_tuples, names=["cluster", "segment", "duration"]) df = pd.DataFrame(flat, index=mi, columns=range(flat.shape[1])) expanded = cr.disaggregate(df) diff --git a/src/tsam_xarray/_result.py b/src/tsam_xarray/_result.py index fc2523a..1e9451f 100644 --- a/src/tsam_xarray/_result.py +++ b/src/tsam_xarray/_result.py @@ -231,70 +231,4 @@ def disaggregate(self, data: xr.DataArray) -> xr.DataArray: Data with ``cluster`` and ``timestep`` replaced by the original ``time`` dimension. """ - # Use stored slice_dims for canonical ordering - slice_dims = self.clustering.slice_dims - if not slice_dims: - return self._disaggregate_single(data) - - import itertools - - from tsam_xarray._core import _concat_along_dims - - slice_coords = {d: data.coords[d].values for d in slice_dims} - keys = list(itertools.product(*(slice_coords[d] for d in slice_dims))) - results = [] - for key in keys: - sel = dict(zip(slice_dims, key, strict=True)) - data_slice = data.sel(sel) - result_slice = self._make_slice_view(sel) - results.append(result_slice._disaggregate_single(data_slice)) - - return _concat_along_dims(results, slice_dims, slice_coords) - - def _make_slice_view(self, sel: dict[str, object]) -> AggregationResult: - """Create a view of this result for a single slice.""" - from tsam_xarray._clustering import ( - ClusteringResult as CR, - ) - from tsam_xarray._clustering import ( - _lookup_clustering, - ) - - # Build key in stored slice_dims order - key = tuple(sel[d] for d in self.clustering.slice_dims) - cr = _lookup_clustering(self.clustering.clusterings, key) - - return AggregationResult( - cluster_representatives=self.cluster_representatives.sel(sel), - cluster_assignments=self.cluster_assignments.sel(sel), - cluster_counts=self.cluster_counts.sel(sel), - segment_durations=( - self.segment_durations.sel(sel) - if self.segment_durations is not None - else None - ), - accuracy=AccuracyMetrics( - rmse=self.accuracy.rmse.sel(sel), - mae=self.accuracy.mae.sel(sel), - rmse_duration=self.accuracy.rmse_duration.sel(sel), - weighted_rmse=self.accuracy.weighted_rmse.sel(sel), - weighted_mae=self.accuracy.weighted_mae.sel(sel), - weighted_rmse_duration=self.accuracy.weighted_rmse_duration.sel(sel), - ), - reconstructed=self.reconstructed.sel(sel), - original=self.original.sel(sel), - clustering=CR( - time_dim=self.clustering.time_dim, - cluster_dim=self.clustering.cluster_dim, - slice_dims=[], - clusterings={(): cr}, - dim_names=self.clustering.dim_names, - ), - ) - - def _disaggregate_single(self, data: xr.DataArray) -> xr.DataArray: - """Disaggregate without slice dims.""" - from tsam_xarray._clustering import _disaggregate_single - - cr = self.clustering.clusterings[()] - return _disaggregate_single(cr, data, self.clustering.dim_names) + return self.clustering.disaggregate(data) diff --git a/test/test_aggregate.py b/test/test_aggregate.py index bf912b4..d478843 100644 --- a/test/test_aggregate.py +++ b/test/test_aggregate.py @@ -1668,6 +1668,101 @@ def test_1d_disaggregate(self, tmp_path): expected = result.disaggregate(result.cluster_representatives) xr.testing.assert_allclose(dis, expected) + def test_sliced_dim_order_is_slice_dims_first(self): + """Sliced disaggregate keeps slice dims leading, then time.""" + da = _make_da(scenarios=["low", "high"]) + result = tsam_xarray.aggregate( + da, + time_dim="time", + cluster_dim=["variable", "region"], + n_clusters=4, + ) + dis = result.clustering.disaggregate(result.cluster_representatives) + + assert dis.dims[:2] == ("scenario", "time") + assert dis.shape[:2] == (2, da.sizes["time"]) + + def test_cluster_coords_are_labels_not_positions(self): + """A reordered cluster axis is gathered by label, not by position.""" + da = _make_da() + da_flat = da.isel(region=0).drop_vars("region") + result = tsam_xarray.aggregate( + da_flat, time_dim="time", cluster_dim="variable", n_clusters=4 + ) + reps = result.cluster_representatives + shuffled = reps.isel(cluster=[3, 1, 0, 2]) + + dis = result.clustering.disaggregate(shuffled) + expected = result.clustering.disaggregate(reps) + xr.testing.assert_identical(dis, expected) + + def test_preserves_integer_dtype(self): + """Disaggregate does not upcast integer payloads.""" + da = _make_da() + da_flat = da.isel(region=0).drop_vars("region") + result = tsam_xarray.aggregate( + da_flat, time_dim="time", cluster_dim="variable", n_clusters=4 + ) + payload = (result.cluster_representatives * 100).astype(np.int64) + + assert result.clustering.disaggregate(payload).dtype == np.int64 + + def test_rejects_mismatched_clusters(self): + """Missing cluster IDs raise, as tsam's own disaggregate does.""" + da = _make_da() + da_flat = da.isel(region=0).drop_vars("region") + result = tsam_xarray.aggregate( + da_flat, time_dim="time", cluster_dim="variable", n_clusters=4 + ) + reps = result.cluster_representatives + + with pytest.raises(ValueError, match="missing clusters"): + result.clustering.disaggregate(reps.isel(cluster=[0, 1, 2])) + with pytest.raises(ValueError, match="timesteps"): + result.clustering.disaggregate(reps.isel(timestep=slice(0, 12))) + + def test_differing_time_index_falls_back(self): + """Slices whose stored time indices disagree decline the gather.""" + import dataclasses + + from tsam_xarray import ClusteringResult + from tsam_xarray._clustering import _disaggregate_single, _is_gatherable + + da = _make_da(scenarios=["low", "high"]) + result = tsam_xarray.aggregate( + da, + time_dim="time", + cluster_dim=["variable", "region"], + n_clusters=4, + ) + reps = result.cluster_representatives + + stored = dict(result.clustering.clusterings) + key = ("high",) + stored[key] = dataclasses.replace( + stored[key], + time_index=stored[key].time_index + pd.Timedelta(days=365), + ) + clustering = ClusteringResult( + time_dim=result.clustering.time_dim, + cluster_dim=result.clustering.cluster_dim, + slice_dims=result.clustering.slice_dims, + clusterings=stored, + dim_names=result.clustering.dim_names, + ) + + assert not _is_gatherable(list(stored.values()), reps, clustering.slice_dims) + + dis = clustering.disaggregate(reps) + for (scenario,), cr in stored.items(): + expected = _disaggregate_single( + cr, reps.sel(scenario=scenario), clustering.dim_names + ) + actual = dis.sel(scenario=scenario, time=expected.indexes["time"]) + xr.testing.assert_allclose( + actual.drop_vars("scenario"), expected.drop_vars("scenario") + ) + class TestSliceEdgeCases: def test_cluster_count_mismatch_raises(self):