Fix ChunkCachedArray indexer support - #2921
VeckoTheGecko wants to merge 4 commits into
Conversation
Also bounds-check indices (raising IndexError, matching numpy) and raise NotImplementedError for slices in vectorized indexers.
| # Step 0: Broadcast the index arrays and flatten them into a list of points, | ||
| # restoring the broadcast shape at the end. | ||
| broadcast = np.broadcast_arrays(*indices) | ||
| out_shape = broadcast[0].shape | ||
| indices = tuple(idx.ravel() for idx in broadcast) | ||
| n_points = int(np.prod(out_shape)) | ||
| if n_points == 0: | ||
| return np.empty(out_shape, dtype=self.array.dtype) |
There was a problem hiding this comment.
Previously we were assuming for vectorized indexing only 1D arrays for vectorized indexing, which doesn't match 1-to-1 with numpy
Numpy allows vectorized indexing as long as the shapes broadcast against each other
>>> import numpy as np
>>> np.array([[1,2,3,4], [7,8,9,10]])[[0,0],[1,1]]
array([2, 2])
>>> a = np.array([0,0])
>>> b = np.array([[1,1]])
>>> np.array([[1,2,3,4], [7,8,9,10]])[a, b]
array([[2, 2]])
>>> np.broadcast([1,2,3],[[1,2,3]])
<numpy.broadcast object at 0x1066040c0>
>>> np.broadcast([1,2,3],[[1,2,3,4]])
Traceback (most recent call last):
File "<python-input-13>", line 1, in <module>
np.broadcast([1,2,3],[[1,2,3,4]])
~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^
ValueError: shape mismatch: objects cannot be broadcast to a single shape. Mismatch is between arg 0 with shape (3,) and arg 1 with shape (1, 4).
>>> Us supporting this helps with the reliability of our implementation, and (AFAICT) is at little to no cost
VeckoTheGecko
changed the title
erikvansebille
left a comment
erikvansebille
left a comment
There was a problem hiding this comment.
looks good, two comments below
| out_of_bounds = (idx < -size) | (idx >= size) | ||
| if out_of_bounds.any(): | ||
| raise IndexError(f"index {idx[out_of_bounds][0]} is out of bounds for axis {d} with size {size}") | ||
| normalized.append(np.where(idx < 0, idx + size, idx)) |
There was a problem hiding this comment.
Should we not pre-allocate the normalised list? We do know how long it will be, don't we?
There was a problem hiding this comment.
I don't think this will have much of a measurable benefit. len(indices) is only 3 or 4
There was a problem hiding this comment.
Is len(indices) not the number of particles, which can be order millions? What is indices then?
|
you're able to quickly check if this works for you by setting in your pixi.toml |
I can't test this now because my laptop died. Will have to wait until at least next week |
Yes, this seems to be working now! |
This branch has not been deployed
3 participants
Description
An LLM generated summary of the changes:
Note here that slice support isn't offered. I'm not sure what the usecase for this even was (cc @erikvansebille ?). Either way - now it clearly errors out so that authors of interpolators have a better idea of what's happening
I'm not sure the performance cost of the bounds checks above are small
Checklist
ChunkCachedArray._raw_vindexraisesIndexErrorwhen given empty index arrays #2906, and fixes TypeError in ChunkCachedArray: '<' not supported between instances of 'slice' and 'int' #2897mainfor normal development,v3-supportfor v3 support)AI Disclosure