From 30ca5bb5e1821c43bd2c8cf703f89edc0b47021d Mon Sep 17 00:00:00 2001 From: chrisholder Date: Sun, 21 Jun 2026 14:17:15 +0100 Subject: [PATCH 1/2] feat: add gradient-based soft barycentre averaging (method="soft") Add `soft_barycenter_average` (in `_ba_soft.py`): a gradient-based barycentre that minimises a soft elastic objective with `scipy.optimize.minimize` (L-BFGS-B) using the private soft-DTW / soft-MSM gradients. Wired into `elastic_barycenter_average` as `method="soft"` and into the averaging callable registry, exported alongside `VALID_SOFT_BA_METHODS`. Clean separation (no soft distances in discrete BA): soft distances are only averaged via `method="soft"`. `soft_dtw` is removed from `VALID_BA_DISTANCE_METHODS` and from `_get_alignment_path`, and early validation rejects `method="soft"` with a hard distance and discrete methods with a soft distance. The discrete-BA test list drops `soft_dtw` and the temporary `test_kasba` skip is reverted. Fixes over the prototype: - `return_distances_to_center` now recomputes distances at the optimised barycentre instead of returning the initial values. - Multivariate soft-MSM averaging accumulates the gradient per channel (independent), matching the soft-MSM distance; the prototype assigned a 1D gradient into a 2D slot and never worked for soft-MSM. - Threading uses `numba_thread_handler`; docstring covers both soft distances. - `_ba_setup` gains `compute_previous_cost` so the soft path skips the upfront pairwise-cost computation. Tests parametrise over soft-DTW and soft-MSM, and validate the soft-DTW barycentre against tslearn's `softdtw_barycenter`. --- aeon/clustering/averaging/__init__.py | 4 + aeon/clustering/averaging/_averaging.py | 2 + aeon/clustering/averaging/_ba_soft.py | 246 ++++++++++++++++++ aeon/clustering/averaging/_ba_utils.py | 25 +- .../averaging/_barycenter_averaging.py | 38 ++- aeon/clustering/averaging/tests/test_kasba.py | 4 +- .../averaging/tests/test_soft_ba.py | 113 ++++++++ aeon/testing/utils/_distance_parameters.py | 11 +- 8 files changed, 430 insertions(+), 13 deletions(-) create mode 100644 aeon/clustering/averaging/_ba_soft.py create mode 100644 aeon/clustering/averaging/tests/test_soft_ba.py diff --git a/aeon/clustering/averaging/__init__.py b/aeon/clustering/averaging/__init__.py index 59bcca1b4a..6e5d1bd5b4 100644 --- a/aeon/clustering/averaging/__init__.py +++ b/aeon/clustering/averaging/__init__.py @@ -5,18 +5,22 @@ "mean_average", "petitjean_barycenter_average", "subgradient_barycenter_average", + "soft_barycenter_average", "VALID_BA_METHODS", "VALID_BA_DISTANCE_METHODS", + "VALID_SOFT_BA_METHODS", "shift_invariant_average", "kasba_average", ] from aeon.clustering.averaging._averaging import mean_average from aeon.clustering.averaging._ba_petitjean import petitjean_barycenter_average +from aeon.clustering.averaging._ba_soft import soft_barycenter_average from aeon.clustering.averaging._ba_subgradient import subgradient_barycenter_average from aeon.clustering.averaging._ba_utils import ( VALID_BA_DISTANCE_METHODS, VALID_BA_METHODS, + VALID_SOFT_BA_METHODS, ) from aeon.clustering.averaging._barycenter_averaging import elastic_barycenter_average from aeon.clustering.averaging._kasba_average import kasba_average diff --git a/aeon/clustering/averaging/_averaging.py b/aeon/clustering/averaging/_averaging.py index 97b17eb122..5a0206d7df 100644 --- a/aeon/clustering/averaging/_averaging.py +++ b/aeon/clustering/averaging/_averaging.py @@ -7,6 +7,7 @@ import numpy as np from aeon.clustering.averaging._ba_petitjean import petitjean_barycenter_average +from aeon.clustering.averaging._ba_soft import soft_barycenter_average from aeon.clustering.averaging._ba_subgradient import subgradient_barycenter_average from aeon.clustering.averaging._kasba_average import kasba_average from aeon.clustering.averaging._shift_scale_invariant_averaging import ( @@ -38,6 +39,7 @@ def mean_average(X: np.ndarray) -> np.ndarray: "subgradient": subgradient_barycenter_average, "kasba": kasba_average, "petitjean": petitjean_barycenter_average, + "soft": soft_barycenter_average, } diff --git a/aeon/clustering/averaging/_ba_soft.py b/aeon/clustering/averaging/_ba_soft.py new file mode 100644 index 0000000000..adb6c9cd89 --- /dev/null +++ b/aeon/clustering/averaging/_ba_soft.py @@ -0,0 +1,246 @@ +"""Gradient-based soft barycentre averaging (soft-DTW and soft-MSM).""" + +import warnings + +import numpy as np +from numba import njit, prange +from scipy.optimize import minimize + +from aeon.clustering.averaging._ba_utils import _ba_setup +from aeon.distances.elastic.soft._soft_dtw import _soft_dtw_grad_x +from aeon.distances.elastic.soft._soft_msm import _soft_msm_grad_x +from aeon.utils.decorators.numba_threading import numba_thread_handler + + +@numba_thread_handler +def soft_barycenter_average( + X, + distance="soft_dtw", + max_iters=30, + tol=1e-5, + init_barycenter="mean", + weights=None, + precomputed_medoids_pairwise_distance: np.ndarray | None = None, + verbose=False, + minimise_method="L-BFGS-B", + random_state: int | None = None, + n_jobs: int = 1, + return_distances_to_center: bool = False, + return_cost: bool = False, + **kwargs, +): + """Compute a soft barycentre of a collection of time series. + + Computes a barycentre by minimising a differentiable soft elastic objective + using gradient-based optimisation (Cuturi & Blondel, 2017 [1]_). Unlike DBA, + which performs discrete realignment updates, the barycentre is the minimiser + of the smooth objective, with the gradient with respect to the barycentre + obtained from the soft-minimum dynamic programming recursion. + + Both ``"soft_dtw"`` and ``"soft_msm"`` are supported. soft-DTW uses a + dependent multivariate cost; soft-MSM uses an independent (per-channel) + cost, so its gradient is accumulated channel-by-channel. + + Parameters + ---------- + X : np.ndarray of shape (n_cases, n_channels, n_timepoints) or (n_cases, + n_timepoints) + Collection of time series to average. If a 2D array is provided, it is + internally reshaped to ``(n_cases, 1, n_timepoints)``. + distance : {"soft_dtw", "soft_msm"}, default="soft_dtw" + Soft distance function to minimise. + max_iters : int, default=30 + Maximum number of optimiser iterations. + tol : float, default=1e-5 + Optimiser tolerance on the change in objective value. + init_barycenter : {"mean", "medoids", "random"} or np.ndarray of shape \ + (n_channels, n_timepoints), default="mean" + Initial barycentre, or the strategy used to construct it. + weights : np.ndarray of shape (n_cases,), default=None + Optional non-negative weights for each time series. If None, all series + receive weight 1. + precomputed_medoids_pairwise_distance : np.ndarray of shape (n_cases, n_cases), \ + default=None + Optional pairwise distance matrix used when ``init_barycenter="medoids"``. + verbose : bool, default=False + If True, prints progress information during optimisation. + minimise_method : str, default="L-BFGS-B" + The optimisation method passed to :func:`scipy.optimize.minimize`. + random_state : int or None, default=None + Random seed used for stochastic initialisations (e.g., ``"random"``). + n_jobs : int, default=1 + Number of parallel jobs for the gradient evaluations. + return_distances_to_center : bool, default=False + If True, also return the distance from each series in ``X`` to the final + barycentre. + return_cost : bool, default=False + If True, also return the final objective value. + **kwargs + Additional keyword arguments forwarded to the soft distance gradient + (e.g. ``gamma``, ``c``, ``window``). + + Returns + ------- + barycenter : np.ndarray of shape (n_channels, n_timepoints) + The soft barycentre minimising the smooth objective. + distances_to_center : np.ndarray of shape (n_cases,), optional + Returned if ``return_distances_to_center=True``. Distances from each + series to the final barycentre. + cost : float, optional + Returned if ``return_cost=True``. The final objective value. + + References + ---------- + .. [1] Cuturi, M. & Blondel, M. "Soft-DTW: a Differentiable Loss Function + for Time-Series." ICML 2017. + """ + if len(X) <= 1: + center = X[0] if X.ndim == 3 else X + if return_distances_to_center and return_cost: + return center, np.zeros(X.shape[0]), 0.0 + elif return_distances_to_center: + return center, np.zeros(X.shape[0]) + elif return_cost: + return center, 0.0 + return center + + ( + _X, + barycenter, + prev_barycenter, + cost, + _, + distances_to_center, + _, + random_state, + n_jobs, + weights, + ) = _ba_setup( + X, + distance=distance, + weights=weights, + init_barycenter=init_barycenter, + previous_cost=None, + previous_distance_to_center=None, + precomputed_medoids_pairwise_distance=precomputed_medoids_pairwise_distance, + n_jobs=n_jobs, + random_state=random_state, + compute_previous_cost=False, + ) + + latest = {"f": None, "g_inf": None} + it = {"k": 0} + + def _func(Z): + f, g, _ = _soft_barycenter_one_iter( + barycenter=Z.reshape(*barycenter.shape), + X=_X, + weights=weights, + distance=distance, + **kwargs, + ) + latest["f"] = float(f) + latest["g_inf"] = float(np.max(np.abs(g))) + return f, g.ravel() + + def _cb(xk): + it["k"] += 1 + print( # noqa: T201 + f"[Soft-BA] iter={it['k']} cost={latest['f']:.6f} " + f"||g||={latest['g_inf']:.3e}" + ) + + res = minimize( + _func, + barycenter.ravel(), + method=minimise_method, + jac=True, + tol=tol, + options=dict(maxiter=max_iters), + callback=_cb if verbose else None, + ) + + if res.success is False: + warnings.warn( + f"Optimisation failed to converge. Reason given by method: " + f"{res.message}. For more detail set verbose=True.", + RuntimeWarning, + stacklevel=2, + ) + + barycenter = res.x.reshape(*barycenter.shape) + + # Recompute distances at the optimised barycentre rather than returning the + # initial values from ``_ba_setup``. + if return_distances_to_center: + _, _, distances_to_center = _soft_barycenter_one_iter( + barycenter=barycenter, + X=_X, + weights=weights, + distance=distance, + **kwargs, + ) + + if return_distances_to_center and return_cost: + return barycenter, distances_to_center, res.fun + elif return_distances_to_center: + return barycenter, distances_to_center + elif return_cost: + return barycenter, res.fun + return barycenter + + +@njit(cache=True, fastmath=True) +def _soft_msm_grad_x_nd(barycenter, ts, gamma, c): + """Per-channel (independent) soft-MSM gradient for a multivariate series.""" + n_channels = barycenter.shape[0] + grad = np.zeros_like(barycenter) + total = 0.0 + for ch in range(n_channels): + g, d = _soft_msm_grad_x(barycenter[ch : ch + 1], ts[ch : ch + 1], c, gamma) + grad[ch] = g + total += d + return grad, total + + +@njit(cache=True, fastmath=True, parallel=True) +def _soft_barycenter_one_iter( + barycenter: np.ndarray, + X: np.ndarray, + weights: np.ndarray, + distance: str, + window: float | None = None, + gamma: float = 1.0, + c: float = 1.0, +): + X_size = len(X) + local_jacobian_products = np.zeros( + (X_size, barycenter.shape[0], barycenter.shape[1]) + ) + local_distances = np.zeros(X_size) + distances_to_center = np.zeros(X_size) + + if distance == "soft_dtw": + for i in prange(X_size): + local_jacobian_products[i], curr_dist = _soft_dtw_grad_x( + barycenter, X[i], gamma=gamma, window=window + ) + local_distances[i] = curr_dist + distances_to_center[i] = curr_dist + elif distance == "soft_msm": + for i in prange(X_size): + local_jacobian_products[i], curr_dist = _soft_msm_grad_x_nd( + barycenter, X[i], gamma, c + ) + local_distances[i] = curr_dist + distances_to_center[i] = curr_dist + else: + raise ValueError(f"Distance '{distance}' not supported for soft barycenter.") + + jacobian_product = np.zeros_like(barycenter) + total_distance = 0.0 + for i in range(X_size): + jacobian_product += local_jacobian_products[i] * weights[i] + total_distance += local_distances[i] * weights[i] + + return total_distance, jacobian_product, distances_to_center diff --git a/aeon/clustering/averaging/_ba_utils.py b/aeon/clustering/averaging/_ba_utils.py index a5d03495ab..691df00cbc 100644 --- a/aeon/clustering/averaging/_ba_utils.py +++ b/aeon/clustering/averaging/_ba_utils.py @@ -12,7 +12,6 @@ msm_alignment_path, pairwise_distance, shape_dtw_alignment_path, - soft_dtw_alignment_path, twe_alignment_path, wddtw_alignment_path, wdtw_alignment_path, @@ -140,8 +139,6 @@ def _get_alignment_path( ) elif distance == "adtw": return adtw_alignment_path(ts, center, window=window, warp_penalty=warp_penalty) - elif distance == "soft_dtw": - return soft_dtw_alignment_path(ts, center, gamma=gamma, window=window) else: # When numba version > 0.57 add more informative error with what method # was passed. @@ -158,6 +155,7 @@ def _ba_setup( n_jobs: int = 1, random_state: int | None = None, weights: np.ndarray | None = None, + compute_previous_cost: bool = True, **kwargs, ): if X.ndim == 3: @@ -187,11 +185,12 @@ def _ba_setup( **kwargs, ) - pw_dist = pairwise_distance( - _X, init_barycenter, method=distance, n_jobs=n_jobs, **kwargs - ) - previous_cost = np.sum(pw_dist) - previous_distance_to_center = pw_dist.flatten() + if compute_previous_cost: + pw_dist = pairwise_distance( + _X, init_barycenter, method=distance, n_jobs=n_jobs, **kwargs + ) + previous_cost = np.sum(pw_dist) + previous_distance_to_center = pw_dist.flatten() barycenter = np.copy(init_barycenter) prev_barycenter = np.copy(init_barycenter) @@ -223,8 +222,13 @@ def _ba_setup( "subgradient", "kasba", "petitjean", + "soft", ] +# Distances usable with the discrete BA methods (petitjean/subgradient/kasba). +# Soft distances are deliberately excluded: they have an alignment distribution +# rather than a single hard path, so discrete realignment is ill-defined. Use +# ``method="soft"`` (gradient-based) for soft distances instead. VALID_BA_DISTANCE_METHODS = [ "adtw", "dtw", @@ -235,5 +239,10 @@ def _ba_setup( "twe", "msm", "shape_dtw", +] + +# Soft distances usable with ``method="soft"`` (gradient-based barycentre). +VALID_SOFT_BA_METHODS = [ "soft_dtw", + "soft_msm", ] diff --git a/aeon/clustering/averaging/_barycenter_averaging.py b/aeon/clustering/averaging/_barycenter_averaging.py index 4b49e3961a..eef7089750 100644 --- a/aeon/clustering/averaging/_barycenter_averaging.py +++ b/aeon/clustering/averaging/_barycenter_averaging.py @@ -5,8 +5,12 @@ import numpy as np from aeon.clustering.averaging._ba_petitjean import petitjean_barycenter_average +from aeon.clustering.averaging._ba_soft import soft_barycenter_average from aeon.clustering.averaging._ba_subgradient import subgradient_barycenter_average -from aeon.clustering.averaging._ba_utils import VALID_BA_METHODS +from aeon.clustering.averaging._ba_utils import ( + VALID_BA_METHODS, + VALID_SOFT_BA_METHODS, +) from aeon.clustering.averaging._kasba_average import kasba_average @@ -30,6 +34,7 @@ def elastic_barycenter_average( return_cost: bool = False, return_distances_to_center: bool = False, n_jobs: int = 1, + minimise_method: str = "L-BFGS-B", **kwargs, ): """ @@ -113,6 +118,20 @@ def elastic_barycenter_average( "Rock the KASBA: Blazingly Fast and Accurate Time Series Clustering." arXiv:2411.17838, 2024. """ + # Soft distances are only valid with the gradient-based ``method="soft"``, + # and ``method="soft"`` only accepts soft distances. + if method == "soft": + if distance not in VALID_SOFT_BA_METHODS: + raise ValueError( + f"method='soft' requires a soft distance, one of " + f"{VALID_SOFT_BA_METHODS}, but got distance={distance!r}." + ) + elif distance in VALID_SOFT_BA_METHODS: + raise ValueError( + f"Soft distance {distance!r} can only be averaged with " + f"method='soft' (gradient-based), not method={method!r}." + ) + if method == "petitjean": return petitjean_barycenter_average( X, @@ -171,6 +190,23 @@ def elastic_barycenter_average( return_distances_to_center=return_distances_to_center, **kwargs, ) + elif method == "soft": + return soft_barycenter_average( + X, + distance=distance, + max_iters=max_iters, + tol=tol, + init_barycenter=init_barycenter, + weights=weights, + precomputed_medoids_pairwise_distance=precomputed_medoids_pairwise_distance, + verbose=verbose, + minimise_method=minimise_method, + random_state=random_state, + n_jobs=n_jobs, + return_cost=return_cost, + return_distances_to_center=return_distances_to_center, + **kwargs, + ) else: raise ValueError( f"Invalid method: {method}. Please use one of the following: " diff --git a/aeon/clustering/averaging/tests/test_kasba.py b/aeon/clustering/averaging/tests/test_kasba.py index 5b5b517dae..8770dfe4ec 100644 --- a/aeon/clustering/averaging/tests/test_kasba.py +++ b/aeon/clustering/averaging/tests/test_kasba.py @@ -92,7 +92,7 @@ def test_kasba_ba_uni(distance, init_barycenter): assert average_ts_uni.shape == X_train_uni[0].shape assert np.allclose(average_ts_uni, call_directly_average_ts_uni) - if distance not in ["shape_dtw", "edr", "soft_dtw"]: + if distance not in ["shape_dtw", "edr"]: # Test not just returning the init barycenter assert not np.array_equal(average_ts_uni, init_barycenter) @@ -133,7 +133,7 @@ def test_kasba_ba_multi(distance, init_barycenter): assert average_ts_multi.shape == X_train_multi[0].shape assert np.allclose(average_ts_multi, call_directly_average_ts_multi) # EDR and shape_dtw with random values don't update the barycenter so skipping - if distance not in ["shape_dtw", "edr", "soft_dtw"]: + if distance not in ["shape_dtw", "edr"]: # Test not just returning the init barycenter assert not np.array_equal(average_ts_multi, init_barycenter) diff --git a/aeon/clustering/averaging/tests/test_soft_ba.py b/aeon/clustering/averaging/tests/test_soft_ba.py new file mode 100644 index 0000000000..f16cb6e3f7 --- /dev/null +++ b/aeon/clustering/averaging/tests/test_soft_ba.py @@ -0,0 +1,113 @@ +"""Tests for gradient-based soft barycentre averaging.""" + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +from aeon.clustering.averaging import ( + elastic_barycenter_average, + soft_barycenter_average, +) +from aeon.testing.data_generation import make_example_3d_numpy +from aeon.testing.utils._distance_parameters import TEST_SOFT_DISTANCES_WITH_PARAMS + + +@pytest.mark.parametrize("distance,params", TEST_SOFT_DISTANCES_WITH_PARAMS) +@pytest.mark.parametrize("n_channels", [1, 3]) +def test_soft_ba_shape_and_finite(distance, params, n_channels): + """Soft barycentre has the series shape and finite values, uni/multivariate.""" + X = make_example_3d_numpy( + n_cases=6, + n_channels=n_channels, + n_timepoints=12, + return_y=False, + random_state=1, + ) + bary = soft_barycenter_average(X, distance=distance, **params) + assert bary.shape == (n_channels, 12) + assert np.all(np.isfinite(bary)) + + +@pytest.mark.parametrize("distance,params", TEST_SOFT_DISTANCES_WITH_PARAMS) +def test_soft_ba_return_cost_and_distances(distance, params): + """``return_cost`` / ``return_distances_to_center`` give finite values.""" + X = make_example_3d_numpy( + n_cases=5, n_channels=2, n_timepoints=10, return_y=False, random_state=2 + ) + bary, dists, cost = soft_barycenter_average( + X, + distance=distance, + return_distances_to_center=True, + return_cost=True, + **params, + ) + assert dists.shape == (5,) + # Distances are recomputed at the optimised barycentre (not the inf init). + assert np.all(np.isfinite(dists)) + assert np.isfinite(cost) + + +@pytest.mark.parametrize("distance,params", TEST_SOFT_DISTANCES_WITH_PARAMS) +def test_soft_ba_weights_change_result(distance, params): + """Non-uniform weights should change the barycentre.""" + X = make_example_3d_numpy( + n_cases=6, n_channels=1, n_timepoints=10, return_y=False, random_state=3 + ) + unweighted = soft_barycenter_average(X, distance=distance, **params) + weights = np.array([5.0, 1.0, 1.0, 1.0, 1.0, 1.0]) + weighted = soft_barycenter_average(X, distance=distance, weights=weights, **params) + assert not np.allclose(unweighted, weighted) + + +def test_soft_ba_single_series_returned_unchanged(): + """A single-series collection returns that series.""" + X = make_example_3d_numpy( + n_cases=1, n_channels=1, n_timepoints=8, return_y=False, random_state=4 + ) + bary = soft_barycenter_average(X, distance="soft_dtw") + assert_allclose(bary, X[0]) + + +def test_soft_ba_only_soft_distances_via_method_soft(): + """``method="soft"`` requires a soft distance; discrete methods reject soft.""" + X = make_example_3d_numpy( + n_cases=5, n_channels=1, n_timepoints=10, return_y=False, random_state=5 + ) + with pytest.raises(ValueError, match="requires a soft distance"): + elastic_barycenter_average(X, method="soft", distance="dtw") + with pytest.raises(ValueError, match="only be averaged with"): + elastic_barycenter_average(X, method="petitjean", distance="soft_dtw") + + +@pytest.mark.parametrize("distance,params", TEST_SOFT_DISTANCES_WITH_PARAMS) +def test_soft_ba_via_elastic_dispatch(distance, params): + """``elastic_barycenter_average(method='soft')`` matches the direct call.""" + X = make_example_3d_numpy( + n_cases=5, n_channels=1, n_timepoints=10, return_y=False, random_state=6 + ) + direct = soft_barycenter_average(X, distance=distance, **params) + dispatched = elastic_barycenter_average( + X, method="soft", distance=distance, **params + ) + assert_allclose(direct, dispatched) + + +def test_soft_dtw_barycenter_matches_tslearn(): + """soft-DTW barycentre should match tslearn's ``softdtw_barycenter``. + + Both minimise the same smooth soft-DTW objective with L-BFGS from the mean + initialisation, so the optima should agree closely. + """ + ts_bary = pytest.importorskip("tslearn.barycenters") + X = make_example_3d_numpy( + n_cases=5, n_channels=1, n_timepoints=12, return_y=False, random_state=7 + ) + gamma = 1.0 + aeon_bary = soft_barycenter_average( + X, distance="soft_dtw", gamma=gamma, max_iters=100, tol=1e-6 + ) + # tslearn expects (n_cases, n_timepoints, n_channels); aeon uses + # (n_cases, n_channels, n_timepoints). + X_tslearn = np.transpose(X, (0, 2, 1)) + tslearn_bary = ts_bary.softdtw_barycenter(X_tslearn, gamma=gamma, max_iter=100) + assert_allclose(aeon_bary, tslearn_bary.T, rtol=1e-2, atol=1e-2) diff --git a/aeon/testing/utils/_distance_parameters.py b/aeon/testing/utils/_distance_parameters.py index c82dd14314..b197b84ae0 100644 --- a/aeon/testing/utils/_distance_parameters.py +++ b/aeon/testing/utils/_distance_parameters.py @@ -22,11 +22,18 @@ ("shift_scale", {"max_shift": 2}), ] -# All the distances that return a full alignment path. +# Distances usable with the discrete barycentre-averaging methods. Soft +# distances are excluded: they are only averaged via ``method="soft"``. TEST_DISTANCES_WITH_FULL_ALIGNMENT_PATH = [ (name, params) for name, params in TEST_DISTANCE_WITH_PARAMS - if name in ["dtw", "wdtw", "edr", "twe", "msm", "shape_dtw", "adtw", "soft_dtw"] + if name in ["dtw", "wdtw", "edr", "twe", "msm", "shape_dtw", "adtw"] +] + +# Soft distances (with params) for the gradient-based soft barycentre tests. +TEST_SOFT_DISTANCES_WITH_PARAMS = [ + ("soft_dtw", {"gamma": 0.1}), + ("soft_msm", {"gamma": 0.1}), ] From 4a175f45b0ff18988296182a71a0ad176ff19162 Mon Sep 17 00:00:00 2001 From: chrisholder Date: Tue, 23 Jun 2026 19:19:14 +0100 Subject: [PATCH 2/2] fix: forward window/itakura to soft-MSM BA, document method="soft", add tests - `_soft_barycenter_one_iter` now threads `window` and `itakura_max_slope` through to both gradient branches. Previously `window` was silently ignored for `distance="soft_msm"` (and `itakura_max_slope` for both), despite being listed in the docstring kwargs. - document `method="soft"` and the `minimise_method` parameter in `elastic_barycenter_average` (both were undocumented). - replace the unused `_ba_setup` unpack targets with `_`. - add tests: the optimiser moves the barycentre away from its init, and a MULTITHREAD_TESTING-gated determinism test (n_jobs must not change the result), mirroring the discrete-BA threaded tests. --- aeon/clustering/averaging/_ba_soft.py | 30 +++++++++++------ .../averaging/_barycenter_averaging.py | 14 ++++++-- .../averaging/tests/test_soft_ba.py | 32 +++++++++++++++++++ 3 files changed, 65 insertions(+), 11 deletions(-) diff --git a/aeon/clustering/averaging/_ba_soft.py b/aeon/clustering/averaging/_ba_soft.py index adb6c9cd89..34ffdfdad4 100644 --- a/aeon/clustering/averaging/_ba_soft.py +++ b/aeon/clustering/averaging/_ba_soft.py @@ -107,13 +107,13 @@ def soft_barycenter_average( ( _X, barycenter, - prev_barycenter, - cost, + _, + _, _, distances_to_center, _, - random_state, - n_jobs, + _, + _, weights, ) = _ba_setup( X, @@ -143,7 +143,7 @@ def _func(Z): latest["g_inf"] = float(np.max(np.abs(g))) return f, g.ravel() - def _cb(xk): + def _cb(_xk): it["k"] += 1 print( # noqa: T201 f"[Soft-BA] iter={it['k']} cost={latest['f']:.6f} " @@ -191,13 +191,20 @@ def _cb(xk): @njit(cache=True, fastmath=True) -def _soft_msm_grad_x_nd(barycenter, ts, gamma, c): +def _soft_msm_grad_x_nd(barycenter, ts, gamma, c, window=None, itakura_max_slope=None): """Per-channel (independent) soft-MSM gradient for a multivariate series.""" n_channels = barycenter.shape[0] grad = np.zeros_like(barycenter) total = 0.0 for ch in range(n_channels): - g, d = _soft_msm_grad_x(barycenter[ch : ch + 1], ts[ch : ch + 1], c, gamma) + g, d = _soft_msm_grad_x( + barycenter[ch : ch + 1], + ts[ch : ch + 1], + c, + gamma, + window, + itakura_max_slope, + ) grad[ch] = g total += d return grad, total @@ -210,6 +217,7 @@ def _soft_barycenter_one_iter( weights: np.ndarray, distance: str, window: float | None = None, + itakura_max_slope: float | None = None, gamma: float = 1.0, c: float = 1.0, ): @@ -223,14 +231,18 @@ def _soft_barycenter_one_iter( if distance == "soft_dtw": for i in prange(X_size): local_jacobian_products[i], curr_dist = _soft_dtw_grad_x( - barycenter, X[i], gamma=gamma, window=window + barycenter, + X[i], + gamma=gamma, + window=window, + itakura_max_slope=itakura_max_slope, ) local_distances[i] = curr_dist distances_to_center[i] = curr_dist elif distance == "soft_msm": for i in prange(X_size): local_jacobian_products[i], curr_dist = _soft_msm_grad_x_nd( - barycenter, X[i], gamma, c + barycenter, X[i], gamma, c, window, itakura_max_slope ) local_distances[i] = curr_dist distances_to_center[i] = curr_dist diff --git a/aeon/clustering/averaging/_barycenter_averaging.py b/aeon/clustering/averaging/_barycenter_averaging.py index eef7089750..c8023fea35 100644 --- a/aeon/clustering/averaging/_barycenter_averaging.py +++ b/aeon/clustering/averaging/_barycenter_averaging.py @@ -48,10 +48,15 @@ def elastic_barycenter_average( - "subgradient": Stochastic subgradient DBA algorithm [2]_. - "kasba": KASBA algorithm [3]_, a fast stochastic variant that samples subsets of time series during each iteration. + - "soft": Gradient-based soft barycentre averaging that minimises a smooth soft + elastic objective (soft-DTW or soft-MSM) with ``scipy.optimize.minimize``. Petitjean is slower but more reliable at converging to the optimal solution. Subgradient is faster but not guaranteed to converge optimally. KASBA is designed for large datasets, trading off some accuracy for a much faster runtime. + The "soft" method is required for soft distances (e.g. ``"soft_dtw"``, + ``"soft_msm"``), which are not valid with the discrete methods; conversely the + discrete methods are not valid with soft distances. Parameters ---------- @@ -74,8 +79,10 @@ def elastic_barycenter_average( - "mean": Uses the mean of the time series. - "medoids": Uses the medoid of the time series. - "random": Uses a randomly selected time series instance. - method : {"petitjean", "subgradient", "kasba"}, default="petitjean" - The algorithm to use for barycenter averaging. + method : {"petitjean", "subgradient", "kasba", "soft"}, default="petitjean" + The algorithm to use for barycenter averaging. ``"soft"`` is required for + (and only valid with) soft distances; the other methods are only valid with + ordinary (hard) distances. initial_step_size : float, default=0.05 Initial step size for gradient-based methods ("subgradient" and "kasba"). final_step_size : float, default=0.005 @@ -98,6 +105,9 @@ def elastic_barycenter_average( If True, prints progress information. random_state : int or None, default=None Random seed for reproducibility. + minimise_method : str, default="L-BFGS-B" + The ``scipy.optimize.minimize`` solver used by ``method="soft"``. Ignored by + the other methods. **kwargs Additional keyword arguments passed to the chosen distance function. diff --git a/aeon/clustering/averaging/tests/test_soft_ba.py b/aeon/clustering/averaging/tests/test_soft_ba.py index f16cb6e3f7..cc496ef3cd 100644 --- a/aeon/clustering/averaging/tests/test_soft_ba.py +++ b/aeon/clustering/averaging/tests/test_soft_ba.py @@ -9,6 +9,7 @@ soft_barycenter_average, ) from aeon.testing.data_generation import make_example_3d_numpy +from aeon.testing.testing_config import MULTITHREAD_TESTING from aeon.testing.utils._distance_parameters import TEST_SOFT_DISTANCES_WITH_PARAMS @@ -92,6 +93,37 @@ def test_soft_ba_via_elastic_dispatch(distance, params): assert_allclose(direct, dispatched) +@pytest.mark.parametrize("distance,params", TEST_SOFT_DISTANCES_WITH_PARAMS) +def test_soft_ba_moves_away_from_init(distance, params): + """The optimiser should move the barycentre away from its initialisation.""" + X = make_example_3d_numpy( + n_cases=6, n_channels=1, n_timepoints=12, return_y=False, random_state=8 + ) + mean_init = X.mean(axis=0) + bary = soft_barycenter_average( + X, distance=distance, init_barycenter="mean", **params + ) + assert not np.allclose(bary, mean_init) + + +@pytest.mark.skipif(not MULTITHREAD_TESTING, reason="Only run on multithread testing") +@pytest.mark.parametrize("distance,params", TEST_SOFT_DISTANCES_WITH_PARAMS) +@pytest.mark.parametrize("n_jobs", [2, -1]) +def test_soft_ba_threaded(distance, params, n_jobs): + """Threaded soft BA is deterministic: n_jobs must not change the result.""" + X = make_example_3d_numpy( + n_cases=10, n_channels=3, n_timepoints=10, return_y=False, random_state=2 + ) + serial = soft_barycenter_average(X, distance=distance, n_jobs=1, **params) + parallel = soft_barycenter_average(X, distance=distance, n_jobs=n_jobs, **params) + dispatched = elastic_barycenter_average( + X, method="soft", distance=distance, n_jobs=n_jobs, **params + ) + assert serial.shape == parallel.shape + assert_allclose(serial, parallel) + assert_allclose(serial, dispatched) + + def test_soft_dtw_barycenter_matches_tslearn(): """soft-DTW barycentre should match tslearn's ``softdtw_barycenter``.