Add pint arrays support

cupy_xarray/accessors.py

@@ -7,8 +7,8 @@
 )
 from xarray.core.pycompat import DuckArrayModule
 
-dsk = DuckArrayModule("dask")
-dask_array_type = dsk.type
+dask_array_type = DuckArrayModule("dask").type
+pint_array_type = DuckArrayModule("pint").type
 
 
 @register_dataarray_accessor("cupy")
@@ -26,6 +26,8 @@ def is_cupy(self):
         """bool: The underlying data is a cupy array."""
         if isinstance(self.da.data, dask_array_type):
             return isinstance(self.da.data._meta, cp.ndarray)
+        if isinstance(self.da.data, pint_array_type):
+            return isinstance(self.da.data.magnitude, cp.ndarray)
         return isinstance(self.da.data, cp.ndarray)
 
     def as_cupy(self):
@@ -52,20 +54,14 @@ def as_cupy(self):
 
         """
         if isinstance(self.da.data, dask_array_type):
-            return DataArray(
+            return self._as_dataarray(
                 data=self.da.data.map_blocks(cp.asarray),
-                coords=self.da.coords,
-                dims=self.da.dims,
-                name=self.da.name,
-                attrs=self.da.attrs,
             )
-        return DataArray(
-            data=cp.asarray(self.da.data),
-            coords=self.da.coords,
-            dims=self.da.dims,
-            name=self.da.name,
-            attrs=self.da.attrs,
-        )
+        if isinstance(self.da.data, pint_array_type):
+            return self._as_dataarray(
+                data=(self.da.data.units * cp.asarray(self.da.data.magnitude)),
+            )
+        return self._as_dataarray(data=cp.asarray(self.da.data))
 
     def as_numpy(self):
         """
@@ -79,27 +75,26 @@ def as_numpy(self):
         """
         if self.is_cupy:
             if isinstance(self.da.data, dask_array_type):
-                return DataArray(
+                return self._as_dataarray(
                     data=self.da.data.map_blocks(
                         lambda block: block.get(), dtype=self.da.data._meta.dtype
                     ),
-                    coords=self.da.coords,
-                    dims=self.da.dims,
-                    name=self.da.name,
-                    attrs=self.da.attrs,
                 )
-            return DataArray(
-                data=self.da.data.get(),
-                coords=self.da.coords,
-                dims=self.da.dims,
-                name=self.da.name,
-                attrs=self.da.attrs,
-            )
+            return self._as_dataarray(data=self.da.data.get())
         return self.da.as_numpy()
 
     def get(self):
         return self.da.data.get()
 
+    def _as_dataarray(self, data):
+        return DataArray(
+            data=data,
+            coords=self.da.coords,
+            dims=self.da.dims,
+            name=self.da.name,
+            attrs=self.da.attrs,
+        )
+
 
 @register_dataset_accessor("cupy")
 class CupyDatasetAccessor:

cupy_xarray/tests/test_accessors.py

@@ -1,10 +1,18 @@
 import numpy as np
 import pytest
 import xarray as xr
-from xarray.core.pycompat import dask_array_type
+from xarray.core.pycompat import DuckArrayModule
+from xarray.tests import requires_pint
+
+dask_array_type = DuckArrayModule("dask").type
+pint_array_type = DuckArrayModule("pint").type
+
+from pint import UnitRegistry
 
 import cupy_xarray  # noqa: F401
 
+ureg = UnitRegistry()
+
 
 @pytest.fixture
 def tutorial_ds_air():
@@ -26,6 +34,16 @@ def tutorial_da_air_dask(tutorial_ds_air_dask):
     return tutorial_ds_air_dask.air
 
 
+@pytest.fixture
+def tutorial_ds_air_pint():
+    return xr.tutorial.load_dataset("air_temperature") * ureg.Quantity("degree_Kelvin")
+
+
+@pytest.fixture
+def tutorial_da_air_pint(tutorial_ds_air_pint):
+    return tutorial_ds_air_pint.air
+
+
 def test_data_set_accessor(tutorial_ds_air):
     ds = tutorial_ds_air
     assert hasattr(ds, "cupy")
@@ -64,3 +82,17 @@ def test_data_array_accessor_dask(tutorial_da_air_dask):
 
     da = da.cupy.as_numpy()
     assert not da.cupy.is_cupy
+
+
+@requires_pint
+def test_data_array_accessor_pint(tutorial_da_air_pint):
+    da = tutorial_da_air_pint
+    assert hasattr(da, "cupy")
+    assert not da.cupy.is_cupy
+
+    da = da.as_cupy()
+    assert da.cupy.is_cupy
+    assert isinstance(da.data, pint_array_type)
+
+    da = da.cupy.as_numpy()
+    assert not da.cupy.is_cupy