[MRG] Fix GAK returning NaN for time series longer than ~405 samples (#450)

CHANGELOG.md

@@ -27,6 +27,7 @@ Changelogs for this project are recorded in this file since v0.2.0.
 ### Fixed
 
 * `lcss` and `lcss_path_from_metric` now respect the `global_constraint` / Sakoe-Chiba / Itakura band, which was previously ignored. ([#526](https://github.com/tslearn-team/tslearn/issues/526))
+* `gak` and `cdist_gak` no longer return `NaN` for time series longer than about 405 samples. They are now normalized in log space, and the Global Alignment Kernel recursion falls back to a log-space accumulation when its value leaves the range of a 64-bit float. This also fixes `TimeSeriesSVC` / `TimeSeriesSVR` and `KernelKMeans` with `kernel="gak"` on long time series. ([#450](https://github.com/tslearn-team/tslearn/issues/450))
 
 ## [v0.9.0]
 

docs/user_guide/kernel.rst

@@ -87,8 +87,12 @@ This estimate is made available in ``tslearn`` through
     sigma = sigma_gak(X)
     k_01 = gak(X[0], X[1], sigma=sigma)
 
-Note however that, on long time series, this estimate can lead to numerical
-overflows, which smaller values can avoid.
+``gak`` and ``cdist_gak`` are normalized in log space, so they stay accurate
+whatever the length of the time series. The unnormalized kernel
+:math:`k(\mathbf{x}, \mathbf{y})` returned by ``unnormalized_gak``, however,
+sums over every possible alignment and therefore leaves the range of a 64-bit
+float for time series longer than about 405 samples, in which case ``inf`` is
+returned.
 
 Finally, the unnormalized GAK is related to :ref:`softDTW <dtw-softdtw>` [3]_ through the
 following formula:

tests/test_metrics.py

@@ -546,6 +546,61 @@ def test_masks():
             )
 
 
+def test_gak_log_and_linear_recursions_agree():
+    # The log-space recursion is only used when the linear one overflows, so
+    # check that both give the same answer on inputs where both are valid.
+    from tslearn.metrics._gak import (
+        _njit_gak_from_gram_matrix,
+        _njit_log_gak_from_gram_matrix,
+    )
+
+    rng = np.random.RandomState(0)
+    for sz1, sz2 in [(1, 1), (1, 7), (7, 1), (5, 5), (30, 20)]:
+        log_gram = -np.abs(rng.randn(sz1, sz2))
+        np.testing.assert_allclose(
+            _njit_log_gak_from_gram_matrix(log_gram),
+            np.log(_njit_gak_from_gram_matrix(np.exp(log_gram))),
+            rtol=1e-10,
+        )
+
+
+def test_gak_long_time_series():
+    # Non-regression test for
+    # https://github.com/tslearn-team/tslearn/issues/450
+    # The unnormalized GAK recursion sums over every alignment path, so it
+    # grows like the central Delannoy numbers and leaves the range of a
+    # 64-bit float for time series longer than about 405 samples. The
+    # normalization used to be computed from those overflowed values, which
+    # turned every GAK value into NaN past that length.
+    sz = 500
+    rng = np.random.RandomState(0)
+    s1 = np.zeros((sz, 1))
+    s2 = rng.randn(sz, 1) * 0.01
+
+    # The unnormalized kernel really is beyond float64 range here, so this
+    # test does exercise the overflow.
+    assert np.isinf(tslearn.metrics.unnormalized_gak(s1, s1, sigma=2.0))
+
+    # The normalized kernel, however, is perfectly well defined.
+    np.testing.assert_allclose(tslearn.metrics.gak(s1, s1, sigma=2.0), 1.0)
+    np.testing.assert_allclose(tslearn.metrics.gak(s2, s2, sigma=2.0), 1.0)
+
+    value = tslearn.metrics.gak(s1, s2, sigma=2.0)
+    assert np.isfinite(value)
+    assert 0.0 < value <= 1.0
+
+    dataset = np.stack([s1, s2])
+    matrix = tslearn.metrics.cdist_gak(dataset, sigma=2.0)
+    assert np.isfinite(matrix).all()
+    np.testing.assert_allclose(np.diag(matrix), 1.0)
+    np.testing.assert_allclose(matrix, matrix.T)
+
+    # The one-dataset and two-dataset code paths must agree.
+    np.testing.assert_allclose(
+        matrix, tslearn.metrics.cdist_gak(dataset, dataset, sigma=2.0)
+    )
+
+
 def test_gak():
 
     with pytest.raises(ZeroDivisionError):

tests/test_svm.py

@@ -43,6 +43,26 @@ def test_gamma_value_svm():
             estimator.fit(X, y)
 
 
+def test_gak_svm_long_time_series():
+    # Non-regression test for
+    # https://github.com/tslearn-team/tslearn/issues/450
+    # The GAK kernel matrix used to be filled with NaN for time series longer
+    # than about 405 samples, which made both estimators fail at fit time with
+    # "ValueError: Input X contains NaN".
+    n, sz, d = 6, 500, 1
+    rng = np.random.RandomState(0)
+    time_series = rng.randn(n, sz, d)
+    labels = np.array([0, 0, 0, 1, 1, 1])
+
+    cdist_mat = cdist_gak(time_series, sigma=1.0)
+    assert np.isfinite(cdist_mat).all()
+
+    for ModelClass in [TimeSeriesSVC, TimeSeriesSVR]:
+        estimator = ModelClass(kernel="gak")
+        estimator.fit(time_series, labels)
+        assert np.isfinite(estimator.predict(time_series)).all()
+
+
 def test_attributes():
     n, sz, d = 5, 10, 3
     rng = np.random.RandomState(0)

tslearn/metrics/_gak.py

@@ -166,13 +166,16 @@ def gak(s1, s2, sigma=1.0, be=None):
     s1 = to_time_series(s1, remove_nans=True, be=be)
     s2 = to_time_series(s2, remove_nans=True, be=be)
 
-    denom = be.sqrt(
-        _unnormalized_gak(s1, s1, sigma=sigma, backend=be)
-    ) * be.sqrt(
-        _unnormalized_gak(s2, s2, sigma=sigma, backend=be)
+    # Normalizing in log space keeps the ratio representable even when the
+    # unnormalized values themselves overflow (see issue #450).
+    log_denom = 0.5 * (
+        _log_unnormalized_gak(s1, s1, sigma=sigma, backend=be)
+        + _log_unnormalized_gak(s2, s2, sigma=sigma, backend=be)
     )
 
-    return _unnormalized_gak(s1, s2, sigma=sigma, backend=be) / denom
+    return be.exp(
+        _log_unnormalized_gak(s1, s2, sigma=sigma, backend=be) - log_denom
+    )
 
 
 def unnormalized_gak(s1, s2, sigma=1.0, be=None):
@@ -213,6 +216,14 @@ def unnormalized_gak(s1, s2, sigma=1.0, be=None):
     float
         Kernel value
 
+    Notes
+    -----
+    The unnormalized kernel sums over every possible alignment between the two
+    series, so it grows extremely fast with their length and leaves the range
+    of a 64-bit float for time series longer than about 405 samples, in which
+    case `inf` is returned. Use :func:`gak`, whose normalization is computed in
+    log space, when a value is needed for long time series.
+
     Examples
     --------
     >>> unnormalized_gak([1, 2, 3],
@@ -244,16 +255,46 @@ def unnormalized_gak(s1, s2, sigma=1.0, be=None):
     return _unnormalized_gak(s1, s2, sigma, be)
 
 
+def _log_gram_matrix(s1, s2, sigma, backend):
+    log_gram = -backend.cdist(s1, s2, "sqeuclidean") / (2 * sigma * sigma)
+    log_gram -= backend.log(2 - backend.exp(log_gram))
+    return log_gram
+
+
 def _unnormalized_gak(s1, s2, sigma, backend):
-    gram = -backend.cdist(s1, s2, "sqeuclidean") / (2 * sigma * sigma)
-    gram -= backend.log(2 - backend.exp(gram))
-    gram = backend.exp(gram)
+    gram = backend.exp(_log_gram_matrix(s1, s2, sigma, backend))
     if backend.is_numpy:
         return _njit_gak_from_gram_matrix(gram)
     else:
         return _gak_from_gram_matrix(gram)
 
 
+def _log_unnormalized_gak(s1, s2, sigma, backend):
+    """Compute the natural logarithm of the unnormalized GAK value.
+
+    The unnormalized kernel sums over every alignment path, so it grows like
+    the central Delannoy numbers (:math:`\\approx (3 + 2\\sqrt{2})^{sz}`) and
+    leaves the range of a 64-bit float for time series longer than about 405
+    samples. The recursion is therefore evaluated in linear space first, which
+    is cheaper, and only re-run in log space when that result is not usable.
+    """
+    log_gram = _log_gram_matrix(s1, s2, sigma, backend)
+    if backend.is_numpy:
+        gak_from_gram = _njit_gak_from_gram_matrix
+        log_gak_from_gram = _njit_log_gak_from_gram_matrix
+    else:
+        gak_from_gram = _gak_from_gram_matrix
+        log_gak_from_gram = _log_gak_from_gram_matrix
+
+    value = gak_from_gram(backend.exp(log_gram))
+    if 0.0 < value < backend.inf:
+        return backend.log(value)
+
+    # The value has overflowed to `inf` (or underflowed to 0): redo the
+    # recursion in log space, where it stays representable.
+    return log_gak_from_gram(log_gram)
+
+
 def __make_gak_from_gram_matrix(backend):
 
     def _gak_from_gram_matrix_generic(
@@ -286,6 +327,65 @@ def _gak_from_gram_matrix_generic(
     _gak_from_gram_matrix = _njit_gak_from_gram_matrix
 
 
+def __make_log_gak_from_gram_matrix(backend):
+
+    def _log_gak_from_gram_matrix_generic(
+        log_gram
+    ):
+
+        sz1, sz2 = log_gram.shape
+        neg_inf = -backend.inf
+
+        cum_sum = backend.full(
+            (sz1 + 1, sz2 + 1), neg_inf, dtype=log_gram.dtype
+        )
+        cum_sum[0, 0] = 0.0
+
+        for i in range(sz1):
+            for j in range(sz2):
+                # log-sum-exp of the three predecessors, factoring out their
+                # maximum so that the exponentials stay in [0, 1].
+                top = cum_sum[i, j + 1]
+                left = cum_sum[i + 1, j]
+                diag = cum_sum[i, j]
+
+                max_pred = top
+                if left > max_pred:
+                    max_pred = left
+                if diag > max_pred:
+                    max_pred = diag
+
+                if max_pred == neg_inf:
+                    # Every path reaching this cell has zero weight.
+                    continue
+
+                cum_sum[i + 1, j + 1] = (
+                    max_pred
+                    + backend.log(
+                        backend.exp(top - max_pred)
+                        + backend.exp(left - max_pred)
+                        + backend.exp(diag - max_pred)
+                    )
+                    + log_gram[i, j]
+                )
+
+        return cum_sum[-1, -1]
+
+    if backend is numpy:
+        return njit(nogil=True)(_log_gak_from_gram_matrix_generic)
+    else:
+        return _log_gak_from_gram_matrix_generic
+
+
+_njit_log_gak_from_gram_matrix = __make_log_gak_from_gram_matrix(numpy)
+if HAS_TORCH:
+    _log_gak_from_gram_matrix = __make_log_gak_from_gram_matrix(
+        instantiate_backend("torch")
+    )
+else:
+    _log_gak_from_gram_matrix = _njit_log_gak_from_gram_matrix
+
+
 def cdist_gak(
     dataset1,
     dataset2=None,
@@ -387,8 +487,8 @@ def _cdist_gak(
     if be is None:
        be = instantiate_backend(dataset1, dataset2)
 
-    unnormalized_matrix = _cdist_generic(
-       dist_fun=_unnormalized_gak,
+    log_unnormalized_matrix = _cdist_generic(
+       dist_fun=_log_unnormalized_gak,
        dataset1=dataset1,
        dataset2=dataset2,
        n_jobs=n_jobs,
@@ -400,28 +500,38 @@ def _cdist_gak(
        sigma = sigma
     )
     if dataset2 is None:
-       diagonal = be.diag(be.sqrt(1.0 / be.diag(unnormalized_matrix)))
-       diagonal_left = diagonal_right = diagonal
+       log_diagonal_left = be.diag(log_unnormalized_matrix)
+       log_diagonal_right = log_diagonal_left
     else:
-       diagonal_left = Parallel(n_jobs=n_jobs, prefer="threads", verbose=verbose)(
-           delayed(_unnormalized_gak)(
+       log_diagonal_left = Parallel(
+           n_jobs=n_jobs, prefer="threads", verbose=verbose
+       )(
+           delayed(_log_unnormalized_gak)(
                _to_time_series(dataset1[i], remove_nans=True, backend=be),
                _to_time_series(dataset1[i], remove_nans=True, backend=be),
                sigma=sigma,
                backend=be
            )
            for i in range(len(dataset1))
        )
-       diagonal_right = Parallel(n_jobs=n_jobs, prefer="threads", verbose=verbose)(
-           delayed(_unnormalized_gak)(
+       log_diagonal_right = Parallel(
+           n_jobs=n_jobs, prefer="threads", verbose=verbose
+       )(
+           delayed(_log_unnormalized_gak)(
                _to_time_series(dataset2[j], remove_nans=True, backend=be),
                _to_time_series(dataset2[j], remove_nans=True, backend=be),
                sigma=sigma,
                backend=be
            )
            for j in range(len(dataset2))
        )
-       diagonal_left = be.diag(1.0 / be.sqrt(diagonal_left))
-       diagonal_right = be.diag(1.0 / be.sqrt(diagonal_right))
-
-    return diagonal_left @ unnormalized_matrix @ diagonal_right
+       log_diagonal_left = be.array(log_diagonal_left)
+       log_diagonal_right = be.array(log_diagonal_right)
+
+    # Normalize in log space so that the result stays finite even when the
+    # unnormalized kernel values overflow (see issue #450).
+    return be.exp(
+       log_unnormalized_matrix
+       - 0.5 * log_diagonal_left[:, None]
+       - 0.5 * log_diagonal_right[None, :]
+    )