xarray-contrib / cupy-xarray Goto Github PK

View Code? Open in Web Editor NEW

57.0 8.0 11.0 271 KB

Interface for using cupy in xarray, providing convenience accessors.

Home Page: https://cupy-xarray.readthedocs.io

License: Apache License 2.0

Python 100.00%

cupy xarray python

cupy-xarray's Introduction

cupy-xarray

Important

⚠️ This project is looking for maintainers and contributors. Come help out!

Interface for using cupy in xarray, providing convenience accessors.

Installation

From anaconda:

conda install cupy-xarray -c conda-forge

From PyPI:

pip install cupy-xarray

The latest version from Github:

pip install git+https://github.com/xarray-contrib/cupy-xarray.git

Usage

import xarray as xr
import cupy_xarray  # This registers the `DataSet.cupy` and `DataArray.cupy` namespaces but is not used directly

ds = xr.tutorial.load_dataset("air_temperature.nc")
type(ds.air.data)  # numpy.ndarray

%timeit ds.air.mean()  # 8.56 ms ± 15.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

ds = ds.cupy.as_cupy()  # Also available via convenience method ds.as_cupy()
type(ds.air.data)  # cupy.core.core.ndarray

%timeit ds.air.mean()  # 2.14 ms ± 21.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

ds = ds.as_numpy()
type(ds.air.data)  # numpy.ndarray

cupy-xarray's People

Contributors

Stargazers

Watchers

Forkers

quasiben raybellwaves mmccarty cdgaete dcherian aazuspan kadykov negin513 tsanona saishashank85

cupy-xarray's Issues

add duck array tests

Once pydata/xarray#4972 is merged we should create a test_cupy.py

BUG: import cupy_xarray does not work with xarray v2022.11+

It seems that xarray commit broke dask_array_type import from xarray.core.pycompat

import cupy as cp
import cupy_xarray  # Adds .cupy to Xarray objects
import numpy as np
import xarray as xr

gives the following error:

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[1], line 5
      3 import xarray as xr
      4 import numpy as np
----> 5 import cupy_xarray  # Adds .cupy to Xarray objects

File c:\ProgramData\Miniconda3\envs\cupy-xarray-original\lib\site-packages\cupy_xarray\__init__.py:2
      [1](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/__init__.py?line=0) from . import _version
----> [2](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/__init__.py?line=1) from .accessors import CupyDataArrayAccessor, CupyDatasetAccessor  # noqa
      [4](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/__init__.py?line=3) __version__ = _version.get_versions()["version"]

File c:\ProgramData\Miniconda3\envs\cupy-xarray-original\lib\site-packages\cupy_xarray\accessors.py:8
      [1](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=0) import cupy as cp
      [2](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=1) from xarray import (
      [3](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=2)     DataArray,
      [4](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=3)     Dataset,
      [5](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=4)     register_dataarray_accessor,
      [6](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=5)     register_dataset_accessor,
      [7](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=6) )
----> [8](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=7) from xarray.core.pycompat import dask_array_type
     [11](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=10) @register_dataarray_accessor("cupy")
     [12](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=11) class CupyDataArrayAccessor:
     [13](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=12)     """
     [14](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=13)     Access methods for DataArrays using Cupy.
     [15](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=14)     Methods and attributes can be accessed through the `.cupy` attribute.
     [16](file:///c%3A/ProgramData/Miniconda3/envs/cupy-xarray-original/lib/site-packages/cupy_xarray/accessors.py?line=15)     """

ImportError: cannot import name 'dask_array_type' from 'xarray.core.pycompat' (c:\ProgramData\Miniconda3\envs\cupy-xarray-original\lib\site-packages\xarray\core\pycompat.py)

python version: 3.9.15
cupy version: 11.4.0
xarray version: 2202.12.0
dask version: 2022.12.1
cupy-xarray version: 0.1.2

Problem importing DuckArrayModule

ModuleNotFoundError Traceback (most recent call last)
Cell In [13], line 4
2 import pandas as pd
3 import xarray as xr
----> 4 import cupy_xarray
5 xr.set_options(display_expand_data=False)
7 rng = np.random.default_rng(seed=0) # we'll use this later

File /usr/local/lib/python3.9/dist-packages/cupy_xarray/init.py:2
1 from . import _version
----> 2 from .accessors import CupyDataArrayAccessor, CupyDatasetAccessor # noqa
4 version = _version.get_versions()["version"]

File /usr/local/lib/python3.9/dist-packages/cupy_xarray/accessors.py:8
1 import cupy as cp
2 from xarray import (
3 DataArray,
4 Dataset,
5 register_dataarray_accessor,
6 register_dataset_accessor,
7 )
----> 8 from xarray.core.pycompat import DuckArrayModule
10 dsk = DuckArrayModule("dask")
11 dask_array_type = dsk.type

ModuleNotFoundError: No module named 'xarray.core.pycompat'

As i check from the xarray module.
The DuckArrayModule is present in xarray/namedarray/pycompat.py and not in core.

Also im not installing cupy from requirements since i've already installed cupy-cuda11x as per my cuda version.

Updating CI badges in the README

I have noticed the links for some the CI badges are not set up correctly in the README which results in a fail status on the build status badge.
The docs badge also not redirecting the correct readthedocs link.

Dont use tutorial datasets in tests

Avoid the download and just create a fake dataset.

Update Documentations

@dcherian and I brainstormed on the topics that should be included in the docs as following:

Basic Concepts :

How to convert numpy and cupy.
.mean , .std and basic operations.
Creating Arrays
Checking for cupy arrays
Accessing the underlying array
Moving data between CPU and GPU (or host and device)

Fundamental Xarray + Cupy:

High-level Xarray functions

Make the docs as notebooks.

ImportError: cannot import name 'dask_array_type' from 'xarray.core.pycompat'

Hi cupy-xarray developers,

Thank you for this great package! I am new to the cupy GPU computing world. On the server, cupy and xarray are working as expected for creating array. However, cupy_xarray cannot be imported. The error message is below

Python 3.9.16 (main, Jan 11 2023, 16:05:54) 
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import cupy as cp
>>> cp.array([1,2,3])
array([1, 2, 3])
>>> A = cp.array([1,2,3])
>>> import cupy_xarray
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/6embdqs6/.conda/envs/vint/lib/python3.9/site-packages/cupy_xarray/__init__.py", line 2, in <module>
    from .accessors import CupyDataArrayAccessor, CupyDatasetAccessor  # noqa
  File "/home/6embdqs6/.conda/envs/vint/lib/python3.9/site-packages/cupy_xarray/accessors.py", line 8, in <module>
    from xarray.core.pycompat import dask_array_type
ImportError: cannot import name 'dask_array_type' from 'xarray.core.pycompat' (~/.conda/envs/vint/lib/python3.9/site-packages/xarray/core/pycompat.py)

I went to xarray.core.pycompat on the current main branch of Xarray repo and cannot seem to find dask_array_type. I installed the package through condo-forge channel. My conda env is shown below.

# packages in environment at .conda/envs/vint:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                        main  
_openmp_mutex             5.1                       1_gnu  
anyio                     3.5.0            py39h06a4308_0  
argon2-cffi               21.3.0             pyhd3eb1b0_0  
argon2-cffi-bindings      21.2.0           py39h7f8727e_0  
asttokens                 2.0.5              pyhd3eb1b0_0  
attrs                     22.1.0           py39h06a4308_0  
babel                     2.11.0           py39h06a4308_0  
backcall                  0.2.0              pyhd3eb1b0_0  
beautifulsoup4            4.11.1           py39h06a4308_0  
blas                      1.0                    openblas  
bleach                    4.1.0              pyhd3eb1b0_0  
bokeh                     2.4.3            py39h06a4308_0  
bottleneck                1.3.5            py39h7deecbd_0  
brotli                    1.0.9                h5eee18b_7  
brotli-bin                1.0.9                h5eee18b_7  
brotlipy                  0.7.0           py39h27cfd23_1003  
bzip2                     1.0.8                h7b6447c_0  
c-ares                    1.18.1               h7f8727e_0  
ca-certificates           2022.12.7            ha878542_0    conda-forge
certifi                   2022.12.7          pyhd8ed1ab_0    conda-forge
cffi                      1.15.1           py39h5eee18b_3  
cftime                    1.5.1.1          py39hce1f21e_0  
charset-normalizer        2.0.4              pyhd3eb1b0_0  
click                     8.0.4            py39h06a4308_0  
cloudpickle               2.0.0              pyhd3eb1b0_0  
comm                      0.1.2            py39h06a4308_0  
contourpy                 1.0.5            py39hdb19cb5_0  
cryptography              38.0.4           py39h9ce1e76_0  
cudatoolkit               10.2.89             h713d32c_10    conda-forge
cupy                      11.5.0           py39h505621e_0    conda-forge
cupy-xarray               0.1.2              pyhd8ed1ab_0    conda-forge
curl                      7.87.0               h5eee18b_0  
cycler                    0.11.0             pyhd3eb1b0_0  
cytoolz                   0.12.0           py39h5eee18b_0  
dask                      2023.1.1           pyhd8ed1ab_0    conda-forge
dask-core                 2023.1.1           pyhd8ed1ab_0    conda-forge
dask-cuda                 23.2.0             pyhd8ed1ab_0    conda-forge
dask-jobqueue             0.8.1              pyhd8ed1ab_0    conda-forge
dbus                      1.13.18              hb2f20db_0  
debugpy                   1.5.1            py39h295c915_0  
decorator                 5.1.1              pyhd3eb1b0_0  
defusedxml                0.7.1              pyhd3eb1b0_0  
distributed               2023.1.1           pyhd8ed1ab_0    conda-forge
entrypoints               0.4              py39h06a4308_0  
executing                 0.8.3              pyhd3eb1b0_0  
expat                     2.4.9                h6a678d5_0  
fastrlock                 0.8              py39h5a03fae_2    conda-forge
fftw                      3.3.9                h27cfd23_1  
flit-core                 3.6.0              pyhd3eb1b0_0  
fontconfig                2.14.1               h52c9d5c_1  
fonttools                 4.25.0             pyhd3eb1b0_0  
freetype                  2.12.1               h4a9f257_0  
fsspec                    2022.11.0        py39h06a4308_0  
giflib                    5.2.1                h5eee18b_1  
glib                      2.69.1               he621ea3_2  
gst-plugins-base          1.14.0               h8213a91_2  
gstreamer                 1.14.0               h28cd5cc_2  
hdf4                      4.2.13               h3ca952b_2  
hdf5                      1.10.6               h3ffc7dd_1  
heapdict                  1.0.1              pyhd3eb1b0_0  
icu                       58.2                 he6710b0_3  
idna                      3.4              py39h06a4308_0  
importlib-metadata        4.11.3           py39h06a4308_0  
ipykernel                 6.19.2           py39hb070fc8_0  
ipython                   8.10.0           py39h06a4308_0  
ipython_genutils          0.2.0              pyhd3eb1b0_1  
jedi                      0.18.1           py39h06a4308_1  
jinja2                    3.1.2            py39h06a4308_0  
jpeg                      9e                   h7f8727e_0  
json5                     0.9.6              pyhd3eb1b0_0  
jsonschema                4.17.3           py39h06a4308_0  
jupyter_client            7.3.4            py39h06a4308_0  
jupyter_core              5.2.0            py39h06a4308_0  
jupyter_server            1.23.4           py39h06a4308_0  
jupyterlab                3.5.3            py39h06a4308_0  
jupyterlab_pygments       0.1.2                      py_0  
jupyterlab_server         2.16.5           py39h06a4308_0  
kiwisolver                1.4.4            py39h6a678d5_0  
krb5                      1.19.4               h568e23c_0  
lcms2                     2.12                 h3be6417_0  
ld_impl_linux-64          2.38                 h1181459_1  
lerc                      3.0                  h295c915_0  
libblas                   3.9.0           15_linux64_openblas    conda-forge
libbrotlicommon           1.0.9                h5eee18b_7  
libbrotlidec              1.0.9                h5eee18b_7  
libbrotlienc              1.0.9                h5eee18b_7  
libcblas                  3.9.0           15_linux64_openblas    conda-forge
libclang                  10.0.1          default_hb85057a_2  
libcurl                   7.87.0               h91b91d3_0  
libdeflate                1.8                  h7f8727e_5  
libedit                   3.1.20221030         h5eee18b_0  
libev                     4.33                 h7f8727e_1  
libevent                  2.1.12               h8f2d780_0  
libffi                    3.4.2                h6a678d5_6  
libgcc-ng                 11.2.0               h1234567_1  
libgfortran-ng            12.2.0              h69a702a_19    conda-forge
libgfortran5              12.2.0              h337968e_19    conda-forge
libgomp                   11.2.0               h1234567_1  
liblapack                 3.9.0           15_linux64_openblas    conda-forge
libllvm10                 10.0.1               hbcb73fb_5  
libllvm11                 11.1.0               h9e868ea_6  
libnetcdf                 4.8.1                h8322cc2_2  
libnghttp2                1.46.0               hce63b2e_0  
libopenblas               0.3.20          pthreads_h78a6416_0    conda-forge
libpng                    1.6.37               hbc83047_0  
libpq                     12.9                 h16c4e8d_3  
libsodium                 1.0.18               h7b6447c_0  
libssh2                   1.10.0               h8f2d780_0  
libstdcxx-ng              11.2.0               h1234567_1  
libtiff                   4.5.0                h6a678d5_1  
libuuid                   1.41.5               h5eee18b_0  
libwebp                   1.2.4                h11a3e52_0  
libwebp-base              1.2.4                h5eee18b_0  
libxcb                    1.15                 h7f8727e_0  
libxkbcommon              1.0.1                hfa300c1_0  
libxml2                   2.9.14               h74e7548_0  
libxslt                   1.1.35               h4e12654_0  
libzip                    1.8.0                h5cef20c_0  
llvmlite                  0.39.1           py39he621ea3_0  
locket                    1.0.0            py39h06a4308_0  
lxml                      4.9.1            py39h1edc446_0  
lz4                       3.1.3            py39h27cfd23_0  
lz4-c                     1.9.4                h6a678d5_0  
markupsafe                2.1.1            py39h7f8727e_0  
matplotlib                3.6.2            py39h06a4308_0  
matplotlib-base           3.6.2            py39h945d387_0  
matplotlib-inline         0.1.6            py39h06a4308_0  
mistune                   0.8.4           py39h27cfd23_1000  
msgpack-python            1.0.3            py39hd09550d_0  
munkres                   1.1.4                      py_0  
nbclassic                 0.4.8            py39h06a4308_0  
nbclient                  0.5.13           py39h06a4308_0  
nbconvert                 6.5.4            py39h06a4308_0  
nbformat                  5.7.0            py39h06a4308_0  
ncurses                   6.4                  h6a678d5_0  
nest-asyncio              1.5.6            py39h06a4308_0  
netcdf4                   1.6.2            py39h89d13dc_0  
notebook                  6.5.2            py39h06a4308_0  
notebook-shim             0.2.2            py39h06a4308_0  
nspr                      4.33                 h295c915_0  
nss                       3.74                 h0370c37_0  
numba                     0.56.4           py39h417a72b_0  
numexpr                   2.8.4            py39hd2a5715_0  
numpy                     1.22.3           py39hc58783e_2    conda-forge
openssl                   1.1.1t               h7f8727e_0  
packaging                 22.0             py39h06a4308_0  
pandas                    1.5.2            py39h417a72b_0  
pandocfilters             1.5.0              pyhd3eb1b0_0  
parso                     0.8.3              pyhd3eb1b0_0  
partd                     1.2.0              pyhd3eb1b0_1  
pcre                      8.45                 h295c915_0  
pexpect                   4.8.0              pyhd3eb1b0_3  
pickleshare               0.7.5           pyhd3eb1b0_1003  
pillow                    9.3.0            py39h6a678d5_2  
pip                       22.3.1           py39h06a4308_0  
platformdirs              2.5.2            py39h06a4308_0  
ply                       3.11             py39h06a4308_0  
prometheus_client         0.14.1           py39h06a4308_0  
prompt-toolkit            3.0.36           py39h06a4308_0  
psutil                    5.9.0            py39h5eee18b_0  
ptyprocess                0.7.0              pyhd3eb1b0_2  
pure_eval                 0.2.2              pyhd3eb1b0_0  
pycparser                 2.21               pyhd3eb1b0_0  
pygments                  2.11.2             pyhd3eb1b0_0  
pynvml                    11.5.0             pyhd8ed1ab_0    conda-forge
pyopenssl                 22.0.0             pyhd3eb1b0_0  
pyparsing                 3.0.9            py39h06a4308_0  
pyqt                      5.15.7           py39h6a678d5_1  
pyqt5-sip                 12.11.0          py39h6a678d5_1  
pyrsistent                0.18.0           py39heee7806_0  
pysocks                   1.7.1            py39h06a4308_0  
python                    3.9.16               h7a1cb2a_0  
python-dateutil           2.8.2              pyhd3eb1b0_0  
python-fastjsonschema     2.16.2           py39h06a4308_0  
python_abi                3.9                      2_cp39    conda-forge
pytz                      2022.7           py39h06a4308_0  
pyyaml                    6.0              py39h5eee18b_1  
pyzmq                     23.2.0           py39h6a678d5_0  
qt-main                   5.15.2               h327a75a_7  
qt-webengine              5.15.9               hd2b0992_4  
qtwebkit                  5.212                h4eab89a_4  
readline                  8.2                  h5eee18b_0  
requests                  2.28.1           py39h06a4308_0  
scipy                     1.7.3            py39hf838250_2  
send2trash                1.8.0              pyhd3eb1b0_1  
setuptools                65.6.3           py39h06a4308_0  
sip                       6.6.2            py39h6a678d5_0  
six                       1.16.0             pyhd3eb1b0_1  
sniffio                   1.2.0            py39h06a4308_1  
sortedcontainers          2.4.0              pyhd3eb1b0_0  
soupsieve                 2.3.2.post1      py39h06a4308_0  
sqlite                    3.40.1               h5082296_0  
stack_data                0.2.0              pyhd3eb1b0_0  
tbb                       2021.6.0             hdb19cb5_1  
tblib                     1.7.0              pyhd3eb1b0_0  
terminado                 0.17.1           py39h06a4308_0  
tinycss2                  1.2.1            py39h06a4308_0  
tk                        8.6.12               h1ccaba5_0  
toml                      0.10.2             pyhd3eb1b0_0  
tomli                     2.0.1            py39h06a4308_0  
toolz                     0.12.0           py39h06a4308_0  
tornado                   6.1              py39h27cfd23_0  
traitlets                 5.7.1            py39h06a4308_0  
typing-extensions         4.4.0            py39h06a4308_0  
typing_extensions         4.4.0            py39h06a4308_0  
tzdata                    2022g                h04d1e81_0  
urllib3                   1.26.14          py39h06a4308_0  
wcwidth                   0.2.5              pyhd3eb1b0_0  
webencodings              0.5.1            py39h06a4308_1  
websocket-client          0.58.0           py39h06a4308_4  
wheel                     0.38.4           py39h06a4308_0  
xarray                    2023.2.0           pyhd8ed1ab_0    conda-forge
xz                        5.2.10               h5eee18b_1  
yaml                      0.2.5                h7b6447c_0  
zeromq                    4.3.4                h2531618_0  
zict                      2.1.0            py39h06a4308_0  
zipp                      3.11.0           py39h06a4308_0  
zlib                      1.2.13               h5eee18b_0  
zstd                      1.5.2                ha4553b6_0

Is there anything that I missed when installing cupy-xarray which caused this error? Thank you!

Transfer repo to the xarray-contrib

Thank you for setting this up, @jacobtomlinson! I'd like to help out with ironing out some of the details/development of cupy-xarray. In light of @tomchor's pydata/xarray#4212 (comment), I'm wondering if you would be interested in transferring this repo over to the @xarray-contrib org to increase visibility?

Add conda-forge build

Naive question

Could the new dask backend using cupy implements by itself the cupy-xarray machinery ?

%%time
import dask.array as da
import cupy as cp
import dask
cp.cuda.Device(0).use()
with dask.config.set({"array.backend": "cupy"}):
    x = da.random.normal(10,1, (100_000, 100_000), chunks=(10_000, 2000))
    y = x.mean(axis=0)[::100]
    z = y.compute()
    print(z.shape, type(z))

Or is it too naive and xarray is using numpy directly at many places in the code ?

Gaël,

Add optional-deps test env

With pint installed

Add CI workflows

Add pre-commit configuration
~~Remove versioneer and use setuptools-scm~~
Add Github workflow for testing infrastructure
Rename branch from master -> main

add API to docs build

autodoc/autosummary import the package to automatically add the docstrings.

This means cupy-xarray and cupy are imported. importing cupy fails on RTD saying it can't find libcuda.

It'd be nice to find a workaround.

Worst case, I guess we can run autodoc locally and then push the generated html?

(note that installing cuda brings in cudatoolkit which is a 800MB download; it'd be nice to just not do that)

Requirements for better Xarray integration

This is a tracking issue for stuff that needs to happen for better Xarray integration. For some background, see here and here

sliding_window_view for rolling:
- cupy/cupy#6956 and cupy/cupy#7575
- cupy/cupy#4825
flox support for groupby/groupby_bins/resample:
- xarray-contrib/flox#128 and cupy/cupy#7561
- also requires manual dispatching in Xarray: https://github.com/pydata/xarray/blob/ebfa77311280a15d8c455ceb5c8ddc25a2efef39/xarray/util/generate_reductions.py#L354
Better reprs:
- cupy/cupy#6926;
- An HTML repr like dask would also be nice.
- Also a _repr_inline?
test out & document apply_ufunc with vectorize=True (cupy already supports vectorize)
Experiment with cuDF indexes instead of Pandas indexes