Compute Katz distances without integer wraparound

src/antropy/fractal.py

@@ -11,6 +11,15 @@
 __all__ = ["petrosian_fd", "katz_fd", "higuchi_fd", "detrended_fluctuation"]
 
 
+def _abs_integer_difference(x, y):
+    """Compute integer magnitudes before converting them to floating point."""
+    high = np.maximum(x, y).astype(np.uint64)
+    low = np.minimum(x, y).astype(np.uint64)
+    # uint64 subtraction is modulo 2**64. Even when signed values straddle
+    # zero, the ordered difference fits in uint64 and is recovered exactly.
+    return (high - low).astype(np.float64)
+
+
 def petrosian_fd(x, axis=-1):
     """Petrosian fractal dimension.
 
@@ -134,6 +143,9 @@ def katz_fd(x, axis=-1):
     Original code from the `mne-features <https://mne.tools/mne-features/>`_
     package by Jean-Baptiste Schiratti and Alexandre Gramfort.
 
+    Integer and boolean distances are computed before conversion to floating
+    point, avoiding integer overflow and loss of nearby large integer values.
+
     References
     ----------
     * https://ieeexplore.ieee.org/abstract/document/904882
@@ -186,15 +198,21 @@ def katz_fd(x, axis=-1):
     """
     x = np.asarray(x)
 
+    if x.dtype.kind in "biu":
+        samples = np.moveaxis(x, axis, -1)
+        dists = _abs_integer_difference(samples[..., 1:], samples[..., :-1])
+        d = np.max(_abs_integer_difference(samples, samples[..., :1]), axis=-1)
+        axis = -1
+    else:
+        dists = np.abs(np.diff(x, axis=axis))
+        d = np.max(np.abs(x - x.take([0], axis=axis)), axis=axis)
+
     # Define total length of curve
-    dists = np.abs(np.diff(x, axis=axis))
     L = np.sum(dists, axis=axis)
 
     # Average distance between successive points
     a = np.mean(dists, axis=axis)
 
-    # Compute the farthest distance between starting point and any other point
-    d = np.max(np.abs(x - x.take([0], axis=axis)), axis=axis)
     kfd_val = np.log10(L / a) / (np.log10(d / a))
     return kfd_val
 

tests/test_katz_integer.py

@@ -0,0 +1,62 @@
+"""Katz distances must be computed before narrowing integer differences."""
+
+import math
+
+import numpy as np
+import pytest
+
+from antropy import katz_fd
+
+CASES = [np.array([1, 5, 2, 8, 3], dtype=dtype) for dtype in (np.uint8, np.uint16, np.uint32)] + [
+    np.array([2**64 - 8 + i for i in (1, 5, 2, 7, 6)], dtype=np.uint64),
+    np.array([0, 2**64 - 1, 2**63, 0, 2**63], dtype=np.uint64),
+    np.array([-128, 127, 0, -127, 126], dtype=np.int8),
+    np.array([2**63 - 8 + i for i in (1, 5, 2, 7, 6)], dtype=np.int64),
+    np.array([-(2**63), 2**63 - 1, 0, -(2**63) + 1, 2**63 - 2], dtype=np.int64),
+    np.array([True, False, False, False, True]),
+]
+
+
+def integer_reference(sequence):
+    """Use Python integers for distances and only then evaluate the existing formula."""
+    values = [int(value) for value in sequence]
+    length = sum(abs(b - a) for a, b in zip(values, values[1:]))
+    average = float(length) / (len(values) - 1)
+    diameter = max(abs(value - values[0]) for value in values)
+    return math.log10(len(values) - 1) / math.log10(float(diameter) / average)
+
+
+@pytest.mark.parametrize("sequence", CASES)
+@pytest.mark.parametrize("layout", ["vector", "rows", "columns"])
+def test_katz_integer_distances(sequence, layout):
+    """Support ADC values, signed extrema and nearby large 64-bit integers."""
+    rows = np.stack([sequence, sequence[::-1]])
+    expected = np.array([integer_reference(row) for row in rows])
+    if layout == "vector":
+        data, axis, expected = sequence, -1, expected[0]
+    elif layout == "rows":
+        data, axis = rows, -1
+    else:
+        data, axis = rows.T, 0
+    original = data.copy()
+    data.flags.writeable = False
+
+    result = katz_fd(data, axis=axis)
+
+    np.testing.assert_allclose(result, expected, rtol=1e-12)
+    np.testing.assert_array_equal(data, original)
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("axis", [-1, 0])
+def test_katz_floating_path_unchanged(dtype, axis):
+    """Fractional input follows the existing floating-point calculation."""
+    data = np.array([[0.25, 1.5, -2.0, 3.25, -1.0]], dtype=dtype)
+    if axis == 0:
+        data = data.T
+    distances = np.abs(np.diff(data, axis=axis))
+    average = np.mean(distances, axis=axis)
+    diameter = np.max(np.abs(data - data.take([0], axis=axis)), axis=axis)
+    expected = np.log10(np.sum(distances, axis=axis) / average) / np.log10(diameter / average)
+
+    np.testing.assert_array_equal(katz_fd(data, axis=axis), expected)