Add pint arrays support

[kadykov]

This commit fixes the following issue of the wrong behavior of is_cupy property with pint xarrays:

import xarray as xr
import cupy as cp
import pint_xarray
import cupy_xarray
from pint_xarray import unit_registry as ureg

cp_da = xr.DataArray(
    cp.linspace(0, 1, 11),
    coords=dict(
        x=np.linspace(0, 10, 11),
    ),
)

pint_da = cp_da.pint.quantify("meter")

print(f"is_cupy response: {pint_da.cupy.is_cupy}")
print(f"pint_da type {type(pint_da.pint.magnitude)}")

Currently, the output is:

is_cupy response: False
pint_da type <class 'cupy.ndarray'>

[dcherian]

Thanks!

Can you add a test please? You'll need to have pint as an optional dependency. SO we'll need to copy some of this code from Xarray: https://github.com/pydata/xarray/blob/6c5840e1198707cdcf7dc459f27ea9510eb76388/xarray/tests/__init__.py#L83

The test should look like:

@requires_pint
def test_is_cupy_pint():
	pass

[dcherian]

Let's copy this over to tests/__init__.py from xarray and similarly from DuckArrayModule

[negin513]

This test checks this case:

pint + xr.DataArray --> cupy --> numpy

[negin513]

This test checks the following:

 pint + Dask.DataArray --> cupy --> numpy

[negin513]

This pull request addresses an issue where the is_cupy property returns an incorrect value (False) for pint xarrays that are backed by CuPy arrays. The proposed modifications effectively resolve this issue and the changes are clear. Great job. 👍

Recommendations
I have two recommendations regarding the tests:

• To ensure all the edge cases are adequately resolved and to prevent potential future problems, please ensure that the test cases cover all relevant edge cases and scenarios, as this will help to guarantee the robustness of the implemented fix. I added more comments to go over this.

• I suggest improving the naming of test functions in the accessors module to better reflect their purpose and the functionality being tested. I also want to recommend adding docstrings to each test function, explaining the specific case or scenario that each test covers. Currently the name of the tests (and tests) are very confusing and this caused me some confusions on which cases are covered.

The followings are the lists of tests of their purpose:

Test Function Name	Description
test_data_set_accessor	xr.dataset --> cupy ---> np
test_data_array_accessor	xr.DataArray --> cupy --> np
test_data_array_accessor_dask	dask.DataArray --> cupy --> np
test_data_array_accessor_pint	pint + xr.DataArray --> cupy --> np
test_data_array_accessor_pint_dask	pint + Dask.DataArray --> cupy --> np

For example, the name test_xr_dataset_accessor might be a better name than test_data_set_accessor. This way we can have test_dask_dataarray_accessor , etc.

I also suggest covering the following cases that are not covered:
pint + xr.dataset --> cupy --> np

We have tutorial_ds_air_pint already defined above that can be used in this test.

[kadykov]

Thank you for your comments and for the good review of the changes.
I re-implemented the tests with the clearer naming and wider coverage of input data.

[kadykov]

Also closes this issue #31

[negin513]

Hey @dcherian , Do you have any additional comments for this PR? I think we can merge this PR.

[dcherian]

Sorry for the massive delay here. @kadykov

I cleaned up the combinatorial explosion of variables ;) but I don't have a machine to actually run the tests on. Can either you or @negin513 confirm that the tests work please?

[dcherian]

Actually I have access now. fixing things...

[dcherian]

@keewis is this expectation right for DataArray(pint(np.ndarray)).chunk()

[keewis]

be aware that DataArray.chunk() simply passes the wrapped duck array to dask.array.from_array if it was not already a dask collection, which means that it would produce a dask(pint(np.ndarray)). This is not what we want, so there's DataArray.pint.chunk to work around this.