feat: n_jobs for threaded slice aggregation

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,14 +164,38 @@ 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)
 
     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]:

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