Speed up interpolation: pack grid indices into one key instead of np.unique(axis=0) - #1817
Open
alanlujan91 wants to merge 5 commits into
Open
Speed up interpolation: pack grid indices into one key instead of np.unique(axis=0)#1817alanlujan91 wants to merge 5 commits into
alanlujan91 wants to merge 5 commits into
Conversation
_iter_unique_pairs took np.unique(stacked, axis=0), which builds a
void-dtype view of the rows and lexsorts it. That sort was the single
largest cost in HARK's test suite.
Profiling tests/ConsumptionSaving/test_ConsAggShockModel.py::
testAggShockMarkovConsumerType::test_small_open_economy, the slowest
test in the suite: argsort was 71.3s of 143.7s wall clock, and
_iter_unique_pairs 76.2s cumulative. Over half the run was inside one
sort.
The positions are bounded grid indices, so the rows pack losslessly into
a single 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 reproduces the lexicographic order np.unique(axis=0) returned.
Callers index by mask and do not depend on that order, but matching it
makes the change bit-identical rather than merely equivalent.
Most callers pass a SINGLE position array, which is where the waste was
worst: np.unique(..., axis=0) on an (N, 1) array pays the whole
row-lexsort machinery to sort one column.
Measured, same machine, same run:
test_small_open_economy 147.62s -> 70.11s 2.11x
full test suite 413.02s -> 242.46s 1.70x
839 passed either way. The gain is not confined to the aggregate-shock
models: _iter_unique_pairs is shared by the 2D/3D/4D interpolators, so
anything doing multi-dimensional interpolation benefits.
An int64 overflow guard falls back to the original path. It is
unreachable for grid indices, but a silently wrapped key would merge
distinct cells rather than fail, and that is not a failure mode worth
leaving open.
Also measured and NOT included: replacing .flatten() with np.ravel() in
_broadcast_eval, despite 112442 calls. It moved the same test 70.11s ->
69.19s, inside noise, and ravel returns views where flatten copies, so
it changes aliasing for no measurable gain.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR optimizes a hot path in HARK/interpolation.py by replacing np.unique(..., axis=0) over stacked index rows with a mixed-radix packing of per-axis grid indices into a 1-D int64 key, enabling a faster np.unique path while preserving lexicographic ordering and providing an overflow fallback.
Changes:
- Pack multi-axis grid positions into a single integer key and run
np.uniqueon the 1-D key instead of row-wisenp.unique(axis=0). - Preserve the previous lexicographic ordering of unique combinations via mixed-radix strides.
- Add an
int64overflow guard that falls back to the originalnp.unique(axis=0)approach.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Review feedback: astype(np.int64) copies unconditionally, in a path called tens of thousands of times per solve. Applied, and it is the common case rather than an edge one. These positions come from _locate_clipped, i.e. np.searchsorted, which returns intp -- int64 on 64-bit platforms. Verified directly: astype(np.int64, copy=False) returns the same object, so the copy was pure overhead on every call. Safe because key is never written in place: `key = key * stride + pos` rebinds to a fresh array, so the aliased input is only ever read. The single-axis case, which is the common one, skips the loop entirely and passes key straight to np.unique. Not separately measurable at this scale (one O(N) allocation against a 70s test), so no timing claim is made for it; the 2.11x / 1.70x figures in the PR description are unchanged. It is strictly less work either way. Full suite 839 passed.
Review feedback: the overflow guard ran after the loop that builds key, so any wrap had already happened by the time it was checked. The reported collision was not reachable. strides are Python ints, so total is arbitrary-precision and cannot wrap, and the fallback recomputes from positions via np.unique(axis=0) while discarding the wrapped key -- results were correct either way. But the ordering was still wrong on its own terms: it computed a key it was about to throw away, and put the check somewhere that made the intent hard to read. Moved ahead of all int64 arithmetic. Bound corrected too. The largest key the packing can produce is prod(strides) - 1, so the test is against iinfo(int64).max rather than max // 2, which was conservative by a factor of two for no reason. The guard was marked "pragma: no cover", which is its own problem: an untested guard is worth about as much as no guard. It is now covered, along with the packing itself against the np.unique(axis=0) behavior it replaced, for one, two and three axes plus the empty and single-cell cases. Worth recording how the overflow test was arrived at, because the first version was useless. It used three axes whose strides multiply past int64 -- which makes the guard's CONDITION true, but whose particular rows still packed to distinct keys once wrapped, so it passed with the guard removed. The version here is constructed so the keys genuinely collide: with stride 2**22 the row (2**20, 0, 0) packs to 2**64, which wraps to exactly 0 and merges with (0, 0, 0). Rejection-tested, and it now fails without the guard by merging two cells into one. Full suite 842 passed.
0.17.3 is close, and a 1.70x suite-wide speedup is the kind of thing that belongs in release notes. States the measured figures and that results are unchanged, and notes the helper is shared by the 2D/3D/4D interpolators so the gain is not confined to the aggregate-shock models.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.
Suppressed comments (1)
HARK/interpolation.py:118
- _iter_unique_pairs always computes per-axis maxima and the mixed-radix
totalproduct even when called with a singlepositionsarray (the most common case per the PR description). That adds an extra full scan (p.max()) plus Python-loop overhead on the hot path, despite 1Dnp.uniquealready being sufficient and preserving the oldnp.unique(axis=0)ordering for a single column.
strides = [int(p.max()) + 1 if p.size else 1 for p in positions]
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Standalone and independent of #1814 / #1815 -- it touches only
HARK/interpolation.py, which neither of those branches modifies.What was slow
_iter_unique_pairsusednp.unique(stacked, axis=0). That builds a void-dtype view of the rows and lexsorts it, and the sort turned out to be the single largest cost in the test suite.Profiling the slowest test in the suite (
test_ConsAggShockModel.py::testAggShockMarkovConsumerType::test_small_open_economy):numpy.ndarray.argsortinterpolation.py _evaluate_segment_valuesTotal wall clock was 143.7s, so over half the run was inside one sort, reached through
_iter_unique_pairs(76.2s cumulative).The waste is worst in the most common case: most callers pass a single position array, so
np.unique(..., axis=0)was paying the entire row-lexsort machinery to sort one column.The change
These positions are bounded grid indices, so the rows pack losslessly into a single integer key and
np.uniqueon a 1-D integer array takes a much cheaper path. The packing is mixed-radix with each axis's own stride, so key order reproduces the lexicographic ordernp.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.An
int64overflow guard falls back to the original path. Unreachable for grid indices, but a silently wrapped key would merge distinct cells rather than fail.Measured
Same machine, same run,
-p no:randomly:test_small_open_economy-n auto)839 passed either way.
argsortno longer appears in the profile's top entries at all.The gain is not confined to the aggregate-shock models.
_iter_unique_pairsis shared by the 2D/3D/4D interpolators, so anything doing multi-dimensional interpolation benefits -- which is why the whole suite moves, not just the four slow tests.Measured and deliberately not included
Replacing
.flatten()withnp.ravel()in_broadcast_eval, despite 112,442 calls: it moved the same test 70.11s -> 69.19s, inside noise, andravelreturns views whereflattencopies, so it changes aliasing semantics for no measurable gain.What is still hot, for anyone continuing
After this change the top costs are
BilinearInterp._evaluate(12.5s, already fully vectorized -- four fancy-index gathers),_segment_values(8.5s), andLinearInterpOnInterp1D._linear_y_blend. That last one loops per y-cell because each cell needs a different pair of 1-D interpolators; vectorizing it would mean restructuring the class to hold its knots as arrays, not a local tweak.