Add enforcement for np.sort and np.argsort

strax/__init__.py

@@ -4,6 +4,7 @@
 # Glue the package together
 # See https://www.youtube.com/watch?v=0oTh1CXRaQ0 if this confuses you
 # The order of subpackes is not invariant, since we use strax.xxx inside strax
+from .sort_enforcement import *
 from .utils import *
 from .chunk import *
 from .dtypes import *

strax/sort_enforcement.py

@@ -0,0 +1,79 @@
+import numpy as np
+import functools
+import warnings
+
+# Store original functions
+original_sort = np.sort
+original_argsort = np.argsort
+
+
+class SortingError(Exception):
+    """Custom exception for sorting violations."""
+
+    pass
+
+
+def enforce_mergesort(func_name):
+    """Raises warning if not using mergesort."""
+    warnings.warn(
+        f"Direct use of {func_name} detected.\n"
+        "Please use the mergesort wrapper to ensure deterministic behavior.",
+        UserWarning,
+        stacklevel=3,
+    )
+
+
+def create_sort_wrapper(original_func, func_name):
+    """Creates a wrapper that enforces mergesort usage."""
+
+    @functools.wraps(original_func)
+    def wrapper(arr, *args, **kwargs):
+        # Check if explicit quicksort is requested
+        if kwargs.get("kind") == "quicksort" or kwargs.get("kind") == "heapsort":
+            raise SortingError(
+                "quicksort and heapsort are not allowed due to non-deterministic behavior.\n"
+                "Please use mergesort explicitely, "
+                "or remove the 'kind' parameter to use mergesort by default."
+            )
+
+        # Always use mergesort
+        kwargs["kind"] = "mergesort"
+
+        # Log usage of direct sorting functions
+        if original_func in {np.sort, np.argsort}:
+            enforce_mergesort(func_name)
+
+        return original_func(arr, *args, **kwargs)
+
+    return wrapper
+
+
+# Create wrappers for all sorting functions
+sort_wrapper = create_sort_wrapper(original_sort, "np.sort")
+argsort_wrapper = create_sort_wrapper(original_argsort, "np.argsort")
+
+# Export wrapped versions for explicit usage
+mergesort = sort_wrapper
+mergesort_argsort = argsort_wrapper
+
+
+def enable_safe_sorting():
+    """Patches all NumPy sorting methods to enforce mergesort.
+
+    Returns a function to restore original behavior if needed.
+
+    """
+    # Replace all sorting methods
+    np.sort = sort_wrapper
+    np.argsort = argsort_wrapper
+
+    def restore_original_sorts():
+        """Restores original NumPy sorting behavior."""
+        np.sort = original_sort
+        np.argsort = original_argsort
+
+    return restore_original_sorts
+
+
+# Enable safe sorting immediately
+restore_sorts = enable_safe_sorting()

tests/test_sort.py

@@ -0,0 +1,73 @@
+import unittest
+import numpy as np
+import warnings
+from strax.sort_enforcement import SortingError, mergesort, mergesort_argsort, restore_sorts
+
+
+class TestSortEnforcement(unittest.TestCase):
+    def setUp(self):
+        """Set up test fixtures before each test method."""
+        self.arr = np.array([3, 1, 4, 1, 5, 9, 2, 6])
+        # Store expected sorted array and indices for comparison
+        self.expected_sorted = np.array([1, 1, 2, 3, 4, 5, 6, 9])
+        self.expected_argsort = np.array([1, 3, 6, 0, 2, 4, 7, 5])
+
+    def test_explicit_mergesort(self):
+        """Test explicit mergesort function (should not warn)"""
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")  # Turn warnings into errors
+            sorted_arr = mergesort(self.arr)
+            np.testing.assert_array_equal(sorted_arr, self.expected_sorted)
+
+    def test_explicit_mergesort_argsort(self):
+        """Test explicit mergesort_argsort function (should not warn)"""
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")  # Turn warnings into errors
+            sorted_indices = mergesort_argsort(self.arr)
+            np.testing.assert_array_equal(sorted_indices, self.expected_argsort)
+
+    def test_quicksort_rejection(self):
+        """Test that quicksort and heapsort raise errors for both sort and argsort."""
+        # Test np.sort
+        with self.assertRaises(SortingError):
+            np.sort(self.arr, kind="quicksort")
+        with self.assertRaises(SortingError):
+            np.sort(self.arr, kind="heapsort")
+
+        # Test np.argsort
+        with self.assertRaises(SortingError):
+            np.argsort(self.arr, kind="quicksort")
+        with self.assertRaises(SortingError):
+            np.argsort(self.arr, kind="heapsort")
+
+    def test_sort_stability(self):
+        """Test that sorting is stable (mergesort property)"""
+        # Create array with duplicate values
+        arr = np.array(
+            [(1, "a"), (2, "b"), (1, "c"), (2, "d")], dtype=[("num", int), ("letter", "U1")]
+        )
+        sorted_arr = np.sort(arr, order="num")
+        # Check that relative order of equal elements is preserved
+        self.assertEqual(sorted_arr[0]["letter"], "a")
+        self.assertEqual(sorted_arr[1]["letter"], "c")
+        self.assertEqual(sorted_arr[2]["letter"], "b")
+        self.assertEqual(sorted_arr[3]["letter"], "d")
+
+    def test_restore_functionality(self):
+        """Test that restore_sorts function works correctly."""
+        # First verify mergesort is enforced
+        with self.assertRaises(SortingError):
+            np.sort(self.arr, kind="quicksort")
+
+        # Restore original behavior
+        restore_sorts()
+
+        # Now quicksort should work without error
+        try:
+            np.sort(self.arr, kind="quicksort")
+        except SortingError:
+            self.fail("quicksort raised SortingError after restore")
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)