diff --git a/src/tsam_xarray/_core.py b/src/tsam_xarray/_core.py index f6f25aa..97be5a7 100644 --- a/src/tsam_xarray/_core.py +++ b/src/tsam_xarray/_core.py @@ -3,6 +3,7 @@ from __future__ import annotations import itertools +import os import warnings from collections.abc import Hashable, Sequence from typing import Any, cast @@ -42,6 +43,7 @@ def aggregate( weights: Weights = None, cluster_on: ClusterOn = None, dim_names: DimNames | None = None, + n_jobs: int | None = None, **tsam_kwargs: Any, ) -> AggregationResult: """Aggregate an xarray DataArray using tsam. @@ -111,6 +113,15 @@ def aggregate( avoid collisions with the caller's own dimension names. See `DimNames`. + n_jobs: Number of threads for aggregating slices in + parallel. ``None`` or ``1`` (default) runs + sequentially; ``-1`` uses all CPUs, ``-2`` all but + one (joblib convention); a positive N uses exactly + N threads. Only used when slice dims are present. + Slices share no state, and the heavy numpy/scipy + work releases the GIL, so threads suffice — no + processes, no pickling. + **tsam_kwargs: Additional keyword arguments passed to ``tsam.aggregate()``. """ @@ -141,13 +152,10 @@ def aggregate( slice_coords = {d: da.coords[d].values for d in slice_dims} slice_keys = list(itertools.product(*(slice_coords[d] for d in slice_dims))) - results: list[AggregationResult] = [] - - for key in slice_keys: + def _run_slice(key: tuple[Any, ...]) -> AggregationResult: sel = dict(zip(slice_dims, key, strict=True)) - da_slice = da.sel(sel) - r = _aggregate_single( - da_slice, + return _aggregate_single( + da.sel(sel), n_clusters, time_dim, col_dims, @@ -156,7 +164,15 @@ def aggregate( tsam_kwargs, resolved_dim_names, ) - results.append(r) + + n_workers = _resolve_n_workers(n_jobs, len(slice_keys)) + if n_workers <= 1: + results = [_run_slice(key) for key in slice_keys] + else: + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=n_workers) as executor: + results = list(executor.map(_run_slice, slice_keys)) # Validate consistent cluster counts (can differ with extremes="append") _validate_consistent_cluster_counts(results, slice_keys) @@ -164,6 +180,22 @@ def aggregate( return _concat_results(results, slice_dims, slice_coords, slice_keys) +def _resolve_n_workers(n_jobs: int | None, n_slices: int) -> int: + """Resolve n_jobs to a worker count, following the joblib convention. + + ``None``/``1`` mean sequential, ``-1`` all CPUs, ``-2`` all but one, and + a positive N exactly N workers, capped at the number of slices. + """ + if n_jobs is None or n_jobs == 1: + return 1 + if n_jobs < 0: + cpus = os.cpu_count() or 1 + workers = cpus + 1 + n_jobs + else: + workers = n_jobs + return max(1, min(workers, n_slices)) + + def _resolve_cluster_dim( cluster_dim: Sequence[str] | str, ) -> list[str]: diff --git a/test/test_parallel.py b/test/test_parallel.py new file mode 100644 index 0000000..225347d --- /dev/null +++ b/test/test_parallel.py @@ -0,0 +1,61 @@ +"""Parallel slice aggregation via n_jobs.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from tsam_xarray import aggregate +from tsam_xarray._core import _resolve_n_workers + + +def _data(n_slices: int = 4) -> xr.DataArray: + rng = np.random.default_rng(7) + n_t = 21 * 24 + return xr.DataArray( + rng.random((n_slices, 3, n_t)), + dims=["scenario", "variable", "time"], + coords={ + "scenario": [f"s{i}" for i in range(n_slices)], + "variable": ["a", "b", "c"], + "time": pd.date_range("2023-01-01", periods=n_t, freq="h"), + }, + name="load", + ) + + +@pytest.mark.parametrize("n_jobs", [2, -1]) +def test_parallel_matches_sequential(n_jobs): + da = _data() + sequential = aggregate(da, time_dim="time", cluster_dim="variable", n_clusters=4) + parallel = aggregate( + da, time_dim="time", cluster_dim="variable", n_clusters=4, n_jobs=n_jobs + ) + + xr.testing.assert_identical( + sequential.cluster_representatives, parallel.cluster_representatives + ) + xr.testing.assert_identical( + sequential.cluster_assignments, parallel.cluster_assignments + ) + xr.testing.assert_identical(sequential.reconstructed, parallel.reconstructed) + xr.testing.assert_identical(sequential.accuracy.rmse, parallel.accuracy.rmse) + + +def test_n_jobs_without_slice_dims_is_ignored(): + da = _data(1).isel(scenario=0, drop=True) + result = aggregate( + da, time_dim="time", cluster_dim="variable", n_clusters=4, n_jobs=-1 + ) + assert result.n_clusters == 4 + + +def test_resolve_n_workers(): + assert _resolve_n_workers(None, 8) == 1 + assert _resolve_n_workers(1, 8) == 1 + assert _resolve_n_workers(3, 8) == 3 + assert _resolve_n_workers(16, 8) == 8 + assert _resolve_n_workers(-1, 100) >= 1 + assert _resolve_n_workers(-1, 100) == _resolve_n_workers(-2, 100) + 1