Preserve exact comparisons in fast permutation entropy

src/antropy/entropy.py

@@ -31,12 +31,11 @@
 # the ordinal index from a small pre-computed table.  This avoids the O(m
 # log m) sort and replaces it with O(m) bitwise operations.
 #
-# Note: equal values (ties) are encoded as False by '<', which may assign
-# them to a different ordinal pattern than argsort would.  For the typical
-# case of continuous data this never occurs in practice.
+# Comparisons use '<=' because the left column always has the earlier
+# position. This breaks ties by position, matching stable argsort.
 # ---------------------------------------------------------------------------
 
-# order=3 — key bits: bit2=(a<b), bit1=(a<c), bit0=(b<c)
+# order=3 — key bits: bit2=(a<=b), bit1=(a<=c), bit0=(b<=c)
 # The 8-entry _PE3_LOOKUP maps the 3-bit key to the hashmult hash value
 # (sum of argsort * [1,3,9]) used by the general path; _PE3_REMAP then
 # compresses those 6 sparse hash values to the dense range 0..5.
@@ -52,8 +51,8 @@
     _PE3_REMAP[_v] = _i
 del _i, _v
 
-# order=4 — key bits: bit5=(a<b), bit4=(a<c), bit3=(a<d), bit2=(b<c),
-#                     bit1=(b<d), bit0=(c<d)
+# order=4 — key bits: bit5=(a<=b), bit4=(a<=c), bit3=(a<=d), bit2=(b<=c),
+#                     bit1=(b<=d), bit0=(c<=d)
 # The 64-entry table directly maps each valid 6-bit key to ordinal 0..23.
 # Invalid keys (impossible comparison patterns) keep the sentinel value -1.
 _PE4_LOOKUP = np.full(64, -1, dtype=np.int64)
@@ -87,24 +86,26 @@ def _perm_entropy_fast(x, order, delay, normalize):
     n, m = x.shape
     n_embed = m - (order - 1) * delay
 
-    # Apply a positional epsilon jitter — adding i*eps to column i — so that
-    # any tied values are broken by column index, exactly matching argsort's
-    # behaviour.  This handles integer signals, quantized data, zero-padded
-    # epochs, and any other input with exact duplicate values.
-    # eps is scaled to the signal magnitude so the jitter never affects the
-    # ordering of distinct values.
-    eps = np.finfo(np.float64).eps * (float(np.abs(x).max()) + 1)
-    cols = [
-        x[:, i * delay : i * delay + n_embed].astype(np.float64) + i * eps for i in range(order)
-    ]
+    # NaN comparisons do not define an ordering and can produce invalid keys.
+    # Process valid epochs separately so a missing sample cannot affect them.
+    # Self-comparison detects NaN without requiring a float dtype.
+    valid_epochs = (x == x).all(axis=1)
+    if not valid_epochs.all():
+        result = np.full(n, np.nan)
+        result[valid_epochs] = _perm_entropy_fast(x[valid_epochs], order, delay, normalize)
+        return float(result[0]) if is_1d else result
+
+    # Keep the original values: a numeric jitter can reorder close values,
+    # and conversion to float64 can lose differences between large integers.
+    cols = [x[:, i * delay : i * delay + n_embed] for i in range(order)]
 
     if order == 3:
         col0, col1, col2 = cols
         # Encode the 3 pairwise comparisons as bits of an integer key
         bit_key = (
-            ((col0 < col1).astype(np.uint8) << 2)
-            | ((col0 < col2).astype(np.uint8) << 1)
-            | (col1 < col2).astype(np.uint8)
+            ((col0 <= col1).astype(np.uint8) << 2)
+            | ((col0 <= col2).astype(np.uint8) << 1)
+            | (col1 <= col2).astype(np.uint8)
         )
         keys = _PE3_REMAP[_PE3_LOOKUP[bit_key]]  # ordinal indices 0..5
         n_perms = 6
@@ -113,12 +114,12 @@ def _perm_entropy_fast(x, order, delay, normalize):
         col0, col1, col2, col3 = cols
         # Encode the 6 pairwise comparisons as bits of an integer key
         bit_key = (
-            ((col0 < col1).astype(np.uint8) << 5)
-            | ((col0 < col2).astype(np.uint8) << 4)
-            | ((col0 < col3).astype(np.uint8) << 3)
-            | ((col1 < col2).astype(np.uint8) << 2)
-            | ((col1 < col3).astype(np.uint8) << 1)
-            | (col2 < col3).astype(np.uint8)
+            ((col0 <= col1).astype(np.uint8) << 5)
+            | ((col0 <= col2).astype(np.uint8) << 4)
+            | ((col0 <= col3).astype(np.uint8) << 3)
+            | ((col1 <= col2).astype(np.uint8) << 2)
+            | ((col1 <= col3).astype(np.uint8) << 1)
+            | (col2 <= col3).astype(np.uint8)
         )
         keys = _PE4_LOOKUP[bit_key]  # ordinal indices 0..23
         n_perms = 24
@@ -162,7 +163,8 @@ def perm_entropy(x, order=3, delay=1, normalize=False):
     -------
     pe : float or np.array
         Permutation entropy. Returns a scalar for 1D input, or an array of
-        shape ``(n_epochs,)`` for 2D input.
+        shape ``(n_epochs,)`` for 2D input. For orders 3 and 4, an epoch
+        containing NaN returns NaN without affecting other epochs.
 
     Notes
     -----

tests/test_entropy.py

@@ -27,8 +27,8 @@ def _perm_entropy_orig(x, order=3, delay=1, normalize=False):
     """Original argsort-based implementation with no fast path, used as reference.
 
     Uses kind='stable' so that ties are broken by position (earlier index ranks
-    lower), matching the behaviour of the positional epsilon jitter in the fast
-    path.  For tie-free data stable and quicksort give identical results.
+    lower), matching the position-based comparisons in the fast path.
+    For tie-free data stable and quicksort give identical results.
     """
     from math import factorial
 
@@ -82,8 +82,7 @@ def test_perm_entropy(self):
     def test_perm_entropy_ties(self):
         """Fast path must agree with the original argsort implementation on signals with ties.
 
-        A positional epsilon jitter is applied to all inputs so ties are broken
-        by column index, exactly matching argsort's position-based tiebreaking.
+        Ties are broken by column index, matching stable argsort.
         This applies to integer arrays, float arrays with duplicate values,
         zero-padded signals, and quantized data stored as float.
         """
@@ -94,8 +93,7 @@ def test_perm_entropy_ties(self):
             perm_entropy(x_const, order=3), _perm_entropy_orig(x_const, order=3)
         )
 
-        # Ties at different window positions: without jitter the fast path would
-        # collapse both tied windows to the same pattern; with jitter it matches argsort.
+        # Ties at different window positions must remain distinct patterns.
         # x = [3, 2, 2, 1, 3]:
         #   window [3,2,2] → argsort [1,2,0]  (b<c<a, tie broken by position)
         #   window [2,2,1] → argsort [2,0,1]  (c<a<b, tie broken by position)
@@ -106,7 +104,7 @@ def test_perm_entropy_ties(self):
             atol=1e-12,
         )
 
-        # Same signal stored as float — jitter must apply regardless of dtype
+        # Same signal stored as float: tie handling is independent of dtype.
         x_ties_float = x_ties.astype(float)
         np.testing.assert_allclose(
             perm_entropy(x_ties_float, order=3),
@@ -145,6 +143,66 @@ def test_perm_entropy_ties(self):
                 err_msg=f"Mismatch on quantized float signal for order={order}",
             )
 
+    def test_perm_entropy_preserves_ordering(self):
+        x = np.array([2, 0, 3, 1, 4, 1, 5, 2, 3, 0, 1, 4], dtype=np.int64)
+        examples = [
+            x,
+            x + 2**60,
+            x.astype(np.uint64) + 2**63,
+            np.array([int(value) + 2**100 for value in x], dtype=object),
+            x * 1e-18,
+            1.0 + x * np.finfo(float).eps,
+        ]
+        for order in [3, 4]:
+            for delay in [1, 2, 3]:
+                for normalize in [False, True]:
+                    # All examples preserve x's ordinal patterns. The original
+                    # embedding helper casts integers to float64, so use x as
+                    # the reference before adding large offsets.
+                    reference = _perm_entropy_orig(x, order, delay, normalize)
+                    for values in examples:
+                        np.testing.assert_allclose(
+                            perm_entropy(values, order, delay, normalize),
+                            reference,
+                            rtol=0,
+                            atol=1e-14,
+                        )
+                    batch = np.vstack([x, x + 2**60])
+                    np.testing.assert_allclose(
+                        perm_entropy(batch, order, delay, normalize),
+                        [reference, reference],
+                        rtol=0,
+                        atol=1e-14,
+                    )
+
+    def test_perm_entropy_nan_epoch_isolation(self):
+        values = np.array([1, 0, 2, 3, 5, 4, 1, 2.0])
+        missing = values.copy()
+        missing[1] = np.nan
+        for order in [3, 4]:
+            for normalize in [False, True]:
+                self.assertTrue(np.isnan(perm_entropy(missing, order, normalize=normalize)))
+                self.assertTrue(
+                    np.isnan(perm_entropy(missing.astype(object), order, normalize=normalize))
+                )
+                self.assertTrue(np.isnan(perm_entropy(np.vstack([missing, missing]), order)).all())
+                expected = _perm_entropy_orig(values, order, normalize=normalize)
+                for position in [0, 1, 2]:
+                    batch = np.vstack([values, values, values])
+                    batch[position] = missing
+                    before = batch.copy()
+                    actual = perm_entropy(batch, order, normalize=normalize)
+                    self.assertTrue(np.isnan(actual[position]))
+                    np.testing.assert_allclose(np.delete(actual, position), expected)
+                    np.testing.assert_array_equal(batch, before)
+
+                infinite = values.copy()
+                infinite[:2] = [-np.inf, np.inf]
+                np.testing.assert_allclose(
+                    perm_entropy(infinite, order, normalize=normalize),
+                    _perm_entropy_orig(infinite, order, normalize=normalize),
+                )
+
     def test_perm_entropy_fast_path(self):
         """Fast paths (order=3/4) must match the original argsort implementation for 1D input."""
         rng = np.random.default_rng(0)