Speed up interpolation: pack grid indices into one key instead of np.unique(axis=0)

HARK/interpolation.py

@@ -103,10 +103,49 @@ def _iter_unique_pairs(*positions):
     """
     if not positions or positions[0].size == 0:
         return
-    stacked = np.column_stack(positions)
-    combos, inverse = np.unique(stacked, axis=0, return_inverse=True)
-    for k, combo in enumerate(combos):
-        yield (*(int(v) for v in combo), inverse == k)
+
+    # np.unique(stacked, axis=0) builds a void-dtype view of the rows and
+    # lexsorts it, and that sort dominated the aggregate-shock solver: in a
+    # profile of test_small_open_economy, argsort was 71.3s of 143.7s and this
+    # function 76.2s cumulative. These positions are bounded grid indices, so
+    # the rows pack losslessly into one integer key, and np.unique on a 1-D
+    # integer array takes a much cheaper path.
+    #
+    # The packing is mixed-radix with each axis's own stride, so key order is
+    # the same lexicographic order np.unique(axis=0) returned. Callers index
+    # by mask and do not depend on that order, but matching it keeps the
+    # change bit-identical rather than merely equivalent.
+    strides = [int(p.max()) + 1 if p.size else 1 for p in positions]
+
+    # Range check BEFORE any int64 arithmetic. The largest key the packing can
+    # produce is prod(strides) - 1, and strides are Python ints, so this
+    # product is computed at arbitrary precision and cannot itself wrap.
+    # Unreachable for grid indices; a wrapped key would merge distinct cells
+    # rather than fail, which is not a failure mode worth leaving open.
+    total = 1
+    for stride in strides:
+        total *= stride
+    if total - 1 > np.iinfo(np.int64).max:
+        combos, inverse = np.unique(
+            np.column_stack(positions), axis=0, return_inverse=True
+        )
+        for k, combo in enumerate(combos):
+            yield (*(int(v) for v in combo), inverse == k)
+        return
+
+    # copy=False: these come from np.searchsorted, so they are already intp
+    # (int64 on 64-bit) and the copy would be pure overhead in a hot path.
+    # Safe because key is never written in place -- `key = key * stride + pos`
+    # rebinds to a fresh array, and the single-axis case (the common one)
+    # skips the loop entirely and only reads key.
+    key = positions[0].astype(np.int64, copy=False)
+    for pos, stride in zip(positions[1:], strides[1:]):
+        key = key * stride + pos
+
+    _, first, inverse = np.unique(key, return_index=True, return_inverse=True)
+    inverse = inverse.reshape(-1)
+    for k, idx in enumerate(first):
+        yield (*(int(p[idx]) for p in positions), inverse == k)
 
 
 def _envelope_partial(envelope, args, deriv_attr):

docs/CHANGELOG.md

@@ -27,6 +27,7 @@ Release Date: TBD
 
 #### Minor Changes
 
+- Speeds up multi-dimensional interpolation by replacing `np.unique(..., axis=0)` in `_iter_unique_pairs` with a packed integer key. The old form built a void-dtype view of the rows and lexsorted it, which dominated solver time: profiling the slowest test in the suite, `argsort` accounted for 71.3s of 143.7s. Most callers pass a single position array, where the row-lexsort machinery was being paid to sort one column. Measured on the same machine: that test 147.62s to 70.11s (2.11x), and the full test suite 413.02s to 242.46s (1.70x). Results are unchanged; the helper is shared by the 2D/3D/4D interpolators, so any model doing multi-dimensional interpolation benefits. [#1817](https://github.com/econ-ark/HARK/pull/1817)
 - Declares the public API of the recently added modules: `__all__` for `HARK.simulation.normalization` and `HARK.ConsumptionSaving.ConsAggIndMarkovModel`, and an API-reference section for `HARK.simulation.normalization` on the Simulation tools page. [#1811](https://github.com/econ-ark/HARK/pull/1811)
 - Excludes scipy 1.18.0, whose `PPoly`-family objects (e.g. `CubicHermiteSpline`) cannot be `deepcopy`-ed (`TypeError: cannot pickle 'module' object`), breaking `ValueFuncCRRA` construction and the existing test suite wherever that scipy version is resolved. [#1788](https://github.com/econ-ark/HARK/pull/1788)
 - Adds opt-in `markov_shuffle` and `balanced_transitions` parameters to `MarkovConsumerType.get_markov_states`: quota-exact Markov transitions via `MarkovProcess.draw(shuffle=True)`, optionally with systematic sampling by pLvl. Default False; the default call is unchanged. [#1793](https://github.com/econ-ark/HARK/pull/1793)

tests/test_interpolation.py

@@ -3,6 +3,7 @@
 """
 
 from HARK.interpolation import (
+    _iter_unique_pairs,
     IdentityFunction,
     LinearInterp,
     CubicInterp,
@@ -1107,3 +1108,57 @@ def test_scalar_array_reversed(self):
         result_array = lioi(np.array([2.0, 2.0, 2.0]), y)
 
         np.testing.assert_array_almost_equal(result_scalar, result_array)
+
+
+class TestIterUniquePairs(unittest.TestCase):
+    """The integer packing in _iter_unique_pairs and its overflow fallback."""
+
+    @staticmethod
+    def _collect(*positions):
+        return sorted(
+            (tuple(idx), tuple(np.flatnonzero(mask)))
+            for *idx, mask in _iter_unique_pairs(*positions)
+        )
+
+    @staticmethod
+    def _reference(*positions):
+        """What np.unique(axis=0) produced before the packing was introduced."""
+        combos, inverse = np.unique(
+            np.column_stack(positions), axis=0, return_inverse=True
+        )
+        inverse = inverse.reshape(-1)
+        return sorted(
+            (tuple(int(v) for v in combo), tuple(np.flatnonzero(inverse == k)))
+            for k, combo in enumerate(combos)
+        )
+
+    def test_matches_unique_axis0_on_one_two_and_three_axes(self):
+        rng = np.random.default_rng(0)
+        for n_axes in (1, 2, 3):
+            pos = [rng.integers(0, 7, size=500) for _ in range(n_axes)]
+            self.assertEqual(self._collect(*pos), self._reference(*pos), n_axes)
+
+    def test_empty_and_single_cell(self):
+        self.assertEqual(self._collect(np.array([], dtype=int)), [])
+        self.assertEqual(self._collect(np.zeros(4, dtype=int)), [((0,), (0, 1, 2, 3))])
+
+    def test_overflow_falls_back_and_stays_correct(self):
+        # Constructed so the packed key genuinely COLLIDES without the guard,
+        # not merely so the guard's condition is true. With three axes of
+        # stride 2**22, key = a * 2**44 + b * 2**22 + c, so the row
+        # (2**20, 0, 0) packs to 2**64, which wraps to exactly 0 and merges
+        # with (0, 0, 0). An earlier version of this test used rows whose
+        # wrapped keys stayed distinct, so it passed with the guard removed.
+        top = 2**22 - 1  # present in each axis, so every stride is 2**22
+        pos = [
+            np.array([0, 2**20, top]),
+            np.array([0, 0, top]),
+            np.array([0, 0, top]),
+        ]
+        self.assertGreater(2**66 - 1, np.iinfo(np.int64).max)
+
+        got = self._collect(*pos)
+        self.assertEqual(got, self._reference(*pos))
+        # Three distinct rows must remain three distinct cells; under the
+        # wrapped key the first two merge into one.
+        self.assertEqual(len(got), 3)