Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

implement the PintIndex #163

Merged
merged 61 commits into from
Jul 9, 2024
Merged

implement the PintIndex #163

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
61 commits
Select commit Hold shift + click to select a range
b9baa9c
add a `PintMetaIndex` that for now can only `sel`
keewis Mar 25, 2022
89c5e2a
add a function to compare indexers
keewis Mar 27, 2022
17e9aec
expect indexer dicts for strip_indexer_units
keewis Mar 27, 2022
30a1d80
move the indexer comparison function to the utils
keewis Mar 27, 2022
34caf09
change extract_indexer_units to expect a dict
keewis Mar 27, 2022
05aa5a6
fix a few calls to extract_indexer_units
keewis Mar 27, 2022
2e0e5bd
one more call
keewis Mar 27, 2022
818db6c
Merge branch 'main' into pint-meta-index
keewis Sep 11, 2023
a049d03
implement `create_variables` and `from_variables`
keewis Sep 13, 2023
0706bc6
use the new index to attach units to dimension coordinates
keewis Sep 13, 2023
a603860
pass the dictionary of indexers instead iterating manually
keewis Sep 13, 2023
dea881c
use `Coordinates._construct_direct`
keewis Sep 13, 2023
9427b20
Merge branch 'main' into pint-meta-index
keewis Sep 14, 2023
a54b94a
delegate `isel` to the wrapped index and wrap the result
keewis Sep 14, 2023
23d1f76
Merge branch 'main' into pint-meta-index
keewis Sep 16, 2023
f0d0890
add a inline `repr` for the index
keewis Dec 6, 2023
fa9f1b3
stubs for the remaining methods
keewis Dec 9, 2023
fb01e32
rename the index class to `PintIndex`
keewis Dec 9, 2023
9278a2d
add a utility method to wrap the output of the wrapped index's methods
keewis Dec 9, 2023
3200bc8
implement `equals`
keewis Dec 9, 2023
281f03c
implement `roll`, `rename`, and `__getitem__` by forwarding
keewis Dec 9, 2023
c5e9022
start adding tests
keewis Dec 10, 2023
2c2c814
add tests for `create_variables`
keewis Dec 10, 2023
e5d8369
add tests for `sel`
keewis Dec 10, 2023
50a7287
add tests for `isel`
keewis Dec 10, 2023
2b3c5bb
improve the tests for `sel`
keewis Dec 10, 2023
57ea8e5
add tests for `equals`
keewis Dec 10, 2023
3eed8c9
add tests for `roll`
keewis Dec 10, 2023
aebaf37
add tests for `rename`
keewis Dec 10, 2023
58c540f
add tests for `__getitem__`
keewis Dec 10, 2023
9cb7e91
add tests for `_repr_inline_`
keewis Dec 10, 2023
9822520
configure coverage, just in case
keewis Dec 11, 2023
55ccb00
use `_replace` instead of manually constructing the new index
keewis Dec 11, 2023
6bd6726
explicitly check that the pint index gets created
keewis Dec 22, 2023
c7d523b
also verify that non-quantity variables don't become `PintIndex`ed
keewis Dec 22, 2023
bae3c3e
Merge branch 'main' into pint-meta-index
keewis Jun 23, 2024
9dde67d
don't use `.pint.sel`
keewis Jun 23, 2024
235ca0e
Merge branch 'main' into pint-meta-index
keewis Jun 23, 2024
b927436
fix `PintIndex.from_variables` and properly test it
keewis Jun 23, 2024
c31e6b0
quantify the test data
keewis Jun 23, 2024
415d059
explicity quantify the input of the `interp_like` tests
keewis Jun 25, 2024
1939b2d
also strip the units of `other`
keewis Jul 6, 2024
2538104
change expectations in the conversion tests
keewis Jul 6, 2024
eb2c405
refactor `attach_units_dataset`
keewis Jul 6, 2024
0d46b66
get `convert_units` to accept indexes
keewis Jul 6, 2024
caf4668
strip indexes as well
keewis Jul 6, 2024
7303960
change the `.pint.to` tests to not include indexes
keewis Jul 6, 2024
e88738d
extract the units of `other` in `.pint.interp_like`
keewis Jul 6, 2024
0b400d0
quantify the input and expected data in the `reindex` tests
keewis Jul 6, 2024
5bd3ec7
remove the left-over explicit quantification in the `interp` tests
keewis Jul 6, 2024
77eef6d
get `.pint.reindex` to work by explicitly converting, stripping, and …
keewis Jul 6, 2024
c38eb5a
quantify the input and expected objects in the `reindex_like` tests
keewis Jul 6, 2024
7277eb5
get `reindex_like` to work with indexes
keewis Jul 6, 2024
c7cf340
quantify expected only if we expect to make use of it
keewis Jul 6, 2024
948d20f
quantify input and expected objects in the `sel` and `loc` tests
keewis Jul 6, 2024
8c76cbc
get `.pint.sel` and `.pint.loc` to work with the indexes
keewis Jul 6, 2024
f9cb15c
remove the warning about indexed coordinates
keewis Jul 6, 2024
49942bf
preserve the order of the variables
keewis Jul 6, 2024
20dd15c
remove the remaining uses of `Coordinates._construct_direct`
keewis Jul 8, 2024
5efb318
whats-new entry
keewis Jul 9, 2024
f53539a
expose the index
keewis Jul 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 27 additions & 68 deletions pint_xarray/accessors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,10 +160,9 @@ def __getitem__(self, indexers):
raise NotImplementedError("pandas-style indexing is not supported, yet")

dims = self.ds.dims
indexer_units = conversion.extract_indexer_units(indexers)
indexer_units = {
name: conversion.extract_indexer_units(indexer)
for name, indexer in indexers.items()
if name in dims
name: indexer for name, indexer in indexer_units.items() if name in dims
}

# convert the indexes to the indexer's units
Expand All @@ -173,10 +172,7 @@ def __getitem__(self, indexers):
raise KeyError(*e.args) from e

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(indexers)
return converted.loc[stripped_indexers]


Expand All @@ -191,10 +187,9 @@ def __getitem__(self, indexers):
raise NotImplementedError("pandas-style indexing is not supported, yet")

dims = self.da.dims
indexer_units = conversion.extract_indexer_units(indexers)
indexer_units = {
name: conversion.extract_indexer_units(indexer)
for name, indexer in indexers.items()
if name in dims
name: indexer for name, indexer in indexer_units.items() if name in dims
}

# convert the indexes to the indexer's units
Expand All @@ -204,10 +199,7 @@ def __getitem__(self, indexers):
raise KeyError(*e.args) from e

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(indexers)
return converted.loc[stripped_indexers]

def __setitem__(self, indexers, values):
Expand All @@ -227,10 +219,7 @@ def __setitem__(self, indexers, values):
raise KeyError(*e.args) from e

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in converted.items()
}
stripped_indexers = conversion.strip_indexer_units(converted)
self.da.loc[stripped_indexers] = values


Expand Down Expand Up @@ -619,10 +608,9 @@ def reindex(
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex")

dims = self.da.dims
indexer_units = conversion.extract_indexer_units(indexers)
indexer_units = {
name: conversion.extract_indexer_units(indexer)
for name, indexer in indexers.items()
if name in dims
name: indexer for name, indexer in indexer_units.items() if name in dims
}

# TODO: handle tolerance
Expand All @@ -632,10 +620,7 @@ def reindex(
converted = conversion.convert_units(self.da, indexer_units)

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(indexers)
indexed = converted.reindex(
stripped_indexers,
method=method,
Expand Down Expand Up @@ -702,10 +687,9 @@ def interp(
indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp")

dims = self.da.dims
indexer_units = conversion.extract_indexer_units(indexers)
indexer_units = {
name: conversion.extract_indexer_units(indexer)
for name, indexer in indexers.items()
if name in dims
name: indexer for name, indexer in indexer_units.items() if name in dims
}

# convert the indexes to the indexer's units
Expand All @@ -714,10 +698,7 @@ def interp(
stripped = conversion.strip_units(converted)

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(indexers)
interpolated = stripped.interp(
stripped_indexers,
method=method,
Expand Down Expand Up @@ -775,10 +756,9 @@ def sel(
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel")

dims = self.da.dims
indexer_units = conversion.extract_indexer_units(indexers)
indexer_units = {
name: conversion.extract_indexer_units(indexer)
for name, indexer in indexers.items()
if name in dims
name: indexer for name, indexer in indexer_units.items() if name in dims
}

# TODO: handle tolerance
Expand All @@ -790,10 +770,7 @@ def sel(
raise KeyError(*e.args) from e

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(indexers)
indexed = converted.sel(
stripped_indexers,
method=method,
Expand Down Expand Up @@ -843,10 +820,7 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
raise KeyError(*e.args) from e

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in converted_indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(converted_indexers)
indexed = self.da.drop_sel(
stripped_indexers,
errors=errors,
Expand Down Expand Up @@ -1342,10 +1316,9 @@ def reindex(
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex")

dims = self.ds.dims
indexer_units = conversion.extract_indexer_units(indexers)
indexer_units = {
name: conversion.extract_indexer_units(indexer)
for name, indexer in indexers.items()
if name in dims
name: indexer for name, indexer in indexer_units.items() if name in dims
}

# TODO: handle tolerance
Expand All @@ -1355,10 +1328,7 @@ def reindex(
converted = conversion.convert_units(self.ds, indexer_units)

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(indexers)
indexed = converted.reindex(
stripped_indexers,
method=method,
Expand Down Expand Up @@ -1425,10 +1395,9 @@ def interp(
indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp")

dims = self.ds.dims
indexer_units = conversion.extract_indexer_units(indexers)
indexer_units = {
name: conversion.extract_indexer_units(indexer)
for name, indexer in indexers.items()
if name in dims
name: indexer for name, indexer in indexer_units.items() if name in dims
}

# convert the indexes to the indexer's units
Expand All @@ -1437,10 +1406,7 @@ def interp(
stripped = conversion.strip_units(converted)

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(indexers)
interpolated = stripped.interp(
stripped_indexers,
method=method,
Expand Down Expand Up @@ -1498,10 +1464,9 @@ def sel(
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel")

dims = self.ds.dims
indexer_units = conversion.extract_indexer_units(indexers)
indexer_units = {
name: conversion.extract_indexer_units(indexer)
for name, indexer in indexers.items()
if name in dims
name: indexer for name, indexer in indexer_units.items() if name in dims
}

# TODO: handle tolerance
Expand All @@ -1513,10 +1478,7 @@ def sel(
raise KeyError(*e.args) from e

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(indexers)
indexed = converted.sel(
stripped_indexers,
method=method,
Expand Down Expand Up @@ -1568,10 +1530,7 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
raise KeyError(*e.args) from e

# index
stripped_indexers = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in converted_indexers.items()
}
stripped_indexers = conversion.strip_indexer_units(converted_indexers)
indexed = self.ds.drop_sel(
stripped_indexers,
errors=errors,
Expand Down
46 changes: 26 additions & 20 deletions pint_xarray/conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -381,25 +381,31 @@ def convert(indexer, units):
return converted


def extract_indexer_units(indexer):
if isinstance(indexer, slice):
return slice_extract_units(indexer)
elif isinstance(indexer, (DataArray, Variable)):
return array_extract_units(indexer.data)
else:
return array_extract_units(indexer)
def extract_indexer_units(indexers):
def extract(indexer):
if isinstance(indexer, slice):
return slice_extract_units(indexer)
elif isinstance(indexer, (DataArray, Variable)):
return array_extract_units(indexer.data)
else:
return array_extract_units(indexer)

return {name: extract(indexer) for name, indexer in indexers.items()}

def strip_indexer_units(indexer):
if isinstance(indexer, slice):
return slice(
array_strip_units(indexer.start),
array_strip_units(indexer.stop),
array_strip_units(indexer.step),
)
elif isinstance(indexer, DataArray):
return strip_units(indexer)
elif isinstance(indexer, Variable):
return strip_units_variable(indexer)
else:
return array_strip_units(indexer)

def strip_indexer_units(indexers):
def strip(indexer):
if isinstance(indexer, slice):
return slice(
array_strip_units(indexer.start),
array_strip_units(indexer.stop),
array_strip_units(indexer.step),
)
elif isinstance(indexer, DataArray):
return strip_units(indexer)
elif isinstance(indexer, Variable):
return strip_units_variable(indexer)
else:
return array_strip_units(indexer)

return {name: strip(indexer) for name, indexer in indexers.items()}
30 changes: 30 additions & 0 deletions pint_xarray/index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
from xarray.core.indexes import Index

from . import conversion


class PintMetaIndex(Index):
# TODO: inherit from MetaIndex once that exists
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm actually I'm not sure how a MetaIndex class would look like. So far we used the generic term "meta-index" to refer to indexes that would wrap one or several indexes, but I don't know if there will be a need to provide a generic class for that.

Copy link
Collaborator Author

@keewis keewis Mar 28, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree, it doesn't really look like we actually need a base class for that, but I noticed that a few methods don't make sense for meta-indexes, from_variables for example. It's probably fine to use the default for those, though.

def __init__(self, *, index, units):
"""create a unit-aware MetaIndex

Parameters
----------
index : xarray.Index
The wrapped index object.
units : mapping of hashable to unit-like
The units of the indexed coordinates
"""
self.index = index
self.units = units

# don't need `from_variables`: we're always *wrapping* an existing index

def sel(self, labels):
converted_labels = conversion.convert_indexer_units(labels, self.units)
stripped_labels = {
name: conversion.strip_indexer_units(indexer)
for name, indexer in converted_labels.items()
}

return self.index.sel(stripped_labels)
Loading