Draft: Extend idxmin and idxmax to accept multiple dimensions
#10125
Conversation
This file contains 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
fhis was almost totally done by Claude Code (though our files are super-inefficient, such that it used 17M tokens! which still only costs $8.44 but I guess it could be a couple of orders of magnitude less than that)
my main contribution was trying to find a better way to express the:
```
if (dim is ... or (isinstance(dim, Iterable) and not isinstance(dim, str))) and not isinstance(
dim, tuple
):
```
we seem to have a lot of functions which touch on this: `infix_dims`, `parse_ordered_dims` (which only seems to be called in our tests?!), `parse_dims_as_tuple`, `parse_dims_as_set`. Let me know if there's a recommended way
for more information, see https://pre-commit.ci
max-sixty
changed the title
idxmin and idxmax to accept a multiple dimensionsidxmin and idxmax to accept multiple dimensions
shoyer
reviewed
Mar 14, 2025
Comment on lines
+5986
to
+5990
| >>> array.idxmin(dim=["x", "y"]) | ||
| {'x': <xarray.DataArray 'x' ()> Size: 8B | ||
| array(0.), | ||
| 'y': <xarray.DataArray 'y' ()> Size: 8B | ||
| array(0)} |
There was a problem hiding this comment.
This example looks right to me, but does the code actually implement this behavior?
The unit tests don't include anything like this. 2d idxmin looks like it reduces over each dimension independently, and would return a dict of 1D results.
There was a problem hiding this comment.
Sorry, I think my doctest run didn't work, you're right.
Moving this to draft
shoyer
reviewed
Mar 14, 2025
Comment on lines
+6004
to
+6009
| result[k] = self.idxmin( | ||
| dim=k, # type: ignore[arg-type] # k is Hashable from self.dims | ||
| skipna=skipna, | ||
| fill_value=fill_value, | ||
| keep_attrs=keep_attrs, | ||
| ) |
There was a problem hiding this comment.
Shouldn't this be a recursive call to idxmin of some kind?
max-sixty
changed the title
idxmin and idxmax to accept multiple dimensionsidxmin and idxmax to accept multiple dimensions
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.
fhis was almost totally done by Claude Code (though our files are super-inefficient, such that it used 17M tokens! which still only costs $8.44 but I guess it could be a couple of orders of magnitude less than that)
my main contribution was trying to find a better way to express the:
we seem to have a lot of functions which touch on this:
infix_dims,parse_ordered_dims(which only seems to be called in our tests?!),parse_dims_as_tuple,parse_dims_as_set. Let me know if there's a recommended way