diff --git a/HARK/interpolation.py b/HARK/interpolation.py index a4b0b5e3d..4ccada536 100644 --- a/HARK/interpolation.py +++ b/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): diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index e7c65152a..d0d4095f5 100644 --- a/docs/CHANGELOG.md +++ b/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) - Fixes `HARK.dual_measure` indexing the Q income process one period ahead of P whenever `cycles != 1`: `_draw_Q_shocks_indshock` chose `IncShkDstn_Q[t]` where `get_shocks` uses `IncShkDstn[t - 1]`, so an infinite-horizon agent with `T_cycle > 1` drew the Q sample from the wrong period's distribution and scaled it by the wrong `PermGroFac`. Invisible until now because every fixture used `T_cycle == 1`, where indices 0 and -1 name the same element. - Fixes `MarkovProcess.draw(shuffle=True)` returning uninitialized memory for an agent whose source state has no row in the transition matrix. The output buffer is now sentinel-filled and verified, so those agents raise `IndexError` (as the unshuffled path already did) instead of silently inheriting the previous period's `Mrkv` values. - Fixes a division by zero at `LivPrb == 1` in `compute_mean_pLvl`, and a wrong limit in the corresponding guard in `compute_pLvl_factor`. Both compute the newborn share of a stationary population; it is now one shared helper returning `1 / T_age` at the no-mortality limit rather than `nan` or `0`. diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 33ac743f2..1f57b86cd 100644 --- a/tests/test_interpolation.py +++ b/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)