[ENH] Gradient-based soft barycentre averaging

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

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,
 }
 
 

aeon/clustering/averaging/_ba_soft.py

@@ -0,0 +1,258 @@
+"""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,
+        _,
+        _,
+        _,
+        distances_to_center,
+        _,
+        _,
+        _,
+        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, 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,
+            window,
+            itakura_max_slope,
+        )
+        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,
+    itakura_max_slope: 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,
+                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, window, itakura_max_slope
+            )
+            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

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",
 ]