From 6d9a8fa07a083b6296ec69351dde46bc7e0c32b9 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 28 Dec 2023 00:47:37 +0530 Subject: [PATCH] feat: initial array object implementation (#1) * feat: initial array object implementation * do something with no errors * bump minimal support version to 3.9 for a typing feature * Implemented a few more features. * Move 'import cupy' to an _import module. * Added property and method docstrings and stubs from the specification. * Added in-place and reflected operator implementations. * Move code out of __init__.py. * Final directory rearrangement. * Make files distinct by adding docstrings. --- .github/workflows/ci.yml | 2 +- .pre-commit-config.yaml | 8 - pyproject.toml | 27 +- src/ragged/__init__.py | 13 +- src/ragged/common/__init__.py | 14 + src/ragged/common/_import.py | 22 + src/ragged/common/_obj.py | 879 ++++++++++++++++++++++++ src/ragged/common/_typing.py | 50 ++ src/ragged/v202212/__init__.py | 16 + src/ragged/v202212/_obj.py | 14 + tests/test_0001_initial_array_object.py | 10 + tests/test_package.py | 9 - 12 files changed, 1035 insertions(+), 29 deletions(-) create mode 100644 src/ragged/common/__init__.py create mode 100644 src/ragged/common/_import.py create mode 100644 src/ragged/common/_obj.py create mode 100644 src/ragged/common/_typing.py create mode 100644 src/ragged/v202212/__init__.py create mode 100644 src/ragged/v202212/_obj.py create mode 100644 tests/test_0001_initial_array_object.py delete mode 100644 tests/test_package.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1b0252c..1fc948f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -40,7 +40,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.8", "3.12"] + python-version: ["3.9", "3.12"] runs-on: [ubuntu-latest, macos-latest, windows-latest] include: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b81dad2..204c757 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -65,14 +65,6 @@ repos: hooks: - id: shellcheck - - repo: local - hooks: - - id: disallow-caps - name: Disallow improper capitalization - language: pygrep - entry: PyBind|Numpy|Cmake|CCache|Github|PyTest - exclude: .pre-commit-config.yaml - - repo: https://github.com/abravalheri/validate-pyproject rev: v0.15 hooks: diff --git a/pyproject.toml b/pyproject.toml index dadf3b7..6b01751 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ authors = [ description = "Ragged array library, complying with Python API specification." readme = "README.md" license.file = "LICENSE" -requires-python = ">=3.8" +requires-python = ">=3.9" classifiers = [ "Development Status :: 1 - Planning", "Intended Audience :: Science/Research", @@ -21,7 +21,6 @@ classifiers = [ "Programming Language :: Python", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", @@ -30,7 +29,9 @@ classifiers = [ "Typing :: Typed", ] dynamic = ["version"] -dependencies = [] +dependencies = [ + "awkward", +] [project.optional-dependencies] test = [ @@ -87,7 +88,7 @@ report.exclude_also = [ [tool.mypy] files = ["src", "tests"] -python_version = "3.8" +python_version = "3.9" warn_unused_configs = true strict = true enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"] @@ -100,6 +101,17 @@ module = "ragged.*" disallow_untyped_defs = true disallow_incomplete_defs = true +[[tool.mypy.overrides]] +module = "numpy.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "cupy.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "awkward.*" +ignore_missing_imports = true [tool.ruff] src = ["src"] @@ -132,6 +144,8 @@ ignore = [ "PLR09", # Too many <...> "PLR2004", # Magic value used in comparison "ISC001", # Conflicts with formatter + "RET505", # I like my if (return) elif (return) else (return) pattern + "PLR5501", # I like my if (return) elif (return) else (return) pattern ] isort.required-imports = ["from __future__ import annotations"] # Uncomment if using a _compat.typing backport @@ -143,7 +157,7 @@ isort.required-imports = ["from __future__ import annotations"] [tool.pylint] -py-version = "3.8" +py-version = "3.9" ignore-paths = [".*/_version.py"] reports.output-format = "colorized" similarities.ignore-imports = "yes" @@ -153,4 +167,7 @@ messages_control.disable = [ "line-too-long", "missing-module-docstring", "wrong-import-position", + "missing-class-docstring", + "missing-function-docstring", + "R1705", # I like my if (return) elif (return) else (return) pattern ] diff --git a/src/ragged/__init__.py b/src/ragged/__init__.py index dbc9dfa..d8d04be 100644 --- a/src/ragged/__init__.py +++ b/src/ragged/__init__.py @@ -1,12 +1,13 @@ -""" -Copyright (c) 2023 Jim Pivarski. All rights reserved. +# BSD 3-Clause License; see https://github.com/scikit-hep/ragged/blob/main/LICENSE -ragged: Ragged array library, complying with Python API specification. """ +Ragged array module. +FIXME: needs more documentation! -from __future__ import annotations +Version 2022.12 is current, so `ragged.v202212.*` is identical to `ragged.*`. +""" -from ._version import version as __version__ +from __future__ import annotations -__all__ = ["__version__"] +from .v202212 import * # noqa: F403 diff --git a/src/ragged/common/__init__.py b/src/ragged/common/__init__.py new file mode 100644 index 0000000..dcc9188 --- /dev/null +++ b/src/ragged/common/__init__.py @@ -0,0 +1,14 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/ragged/blob/main/LICENSE + +""" +Generic definitions used by the version-specific modules, such as +`ragged.v202212`. + +https://data-apis.org/array-api/latest/API_specification/ +""" + +from __future__ import annotations + +from ._obj import array + +__all__ = ["array"] diff --git a/src/ragged/common/_import.py b/src/ragged/common/_import.py new file mode 100644 index 0000000..05c73ef --- /dev/null +++ b/src/ragged/common/_import.py @@ -0,0 +1,22 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/ragged/blob/main/LICENSE + +from __future__ import annotations + +from typing import Any + + +def cupy() -> Any: + try: + import cupy as cp # pylint: disable=C0415 + + return cp + except ModuleNotFoundError as err: + error_message = """to use the "cuda" backend, you must install cupy: + + pip install cupy + +or + + conda install -c conda-forge cupy +""" + raise ModuleNotFoundError(error_message) from err diff --git a/src/ragged/common/_obj.py b/src/ragged/common/_obj.py new file mode 100644 index 0000000..a6aafc6 --- /dev/null +++ b/src/ragged/common/_obj.py @@ -0,0 +1,879 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/ragged/blob/main/LICENSE + +from __future__ import annotations + +import enum +from numbers import Real +from typing import TYPE_CHECKING, Any, Union + +import awkward as ak +import numpy as np +from awkward.contents import ( + Content, + ListArray, + ListOffsetArray, + NumpyArray, + RegularArray, +) + +from . import _import +from ._typing import ( + Device, + Dtype, + NestedSequence, + PyCapsule, + Shape, + SupportsDLPack, +) + + +def _shape_dtype(layout: Content) -> tuple[Shape, Dtype]: + node = layout + shape: Shape = (len(layout),) + while isinstance(node, (ListArray, ListOffsetArray, RegularArray)): + if isinstance(node, RegularArray): + shape = (*shape, node.size) + else: + shape = (*shape, None) + node = node.content + + if isinstance(node, NumpyArray): + shape = shape + node.data.shape[1:] + return shape, node.data.dtype + + msg = f"Awkward Array type must have regular and irregular lists only, not {layout.form.type!s}" + raise TypeError(msg) + + +# https://github.com/python/typing/issues/684#issuecomment-548203158 +if TYPE_CHECKING: + from enum import Enum + + class ellipsis(Enum): # pylint: disable=C0103 + Ellipsis = "..." # pylint: disable=C0103 + + Ellipsis = ellipsis.Ellipsis # pylint: disable=W0622 + +else: + ellipsis = type(...) # pylint: disable=C0103 + +GetSliceKey = Union[ + int, + slice, + ellipsis, + None, + tuple[Union[int, slice, ellipsis, None], ...], + "array", +] + +SetSliceKey = Union[ + int, slice, ellipsis, tuple[Union[int, slice, ellipsis], ...], "array" +] + + +class array: # pylint: disable=C0103 + """ + Ragged array class and constructor. + + https://data-apis.org/array-api/latest/API_specification/array_object.html + """ + + # Constructors, internal functions, and other methods that are unbound by + # the Array API specification. + + _impl: ak.Array | SupportsDLPack # ndim > 0 ak.Array or ndim == 0 NumPy or CuPy + _shape: Shape + _dtype: Dtype + _device: Device + + @classmethod + def _new(cls, impl: ak.Array, shape: Shape, dtype: Dtype, device: Device) -> array: + """ + Simple/fast array constructor for internal code. + """ + + out = cls.__new__(cls) + out._impl = impl + out._shape = shape + out._dtype = dtype + out._device = device + return out + + def __init__( + self, + array_like: ( + array + | ak.Array + | SupportsDLPack + | bool + | int + | float + | NestedSequence[bool | int | float] + ), + dtype: None | Dtype | type | str = None, + device: None | Device = None, + ): + """ + Primary array constructor. + + Args: + array_like: Data to use as or convert into a ragged array. + dtype: NumPy dtype describing the data (subclass of `np.number`, + without `shape` or `fields`). + device: If `"cpu"`, the array is backed by NumPy and resides in + main memory; if `"cuda"`, the array is backed by CuPy and + resides in CUDA global memory. + """ + + if isinstance(array_like, array): + self._impl = array_like._impl + self._shape, self._dtype = array_like._shape, array_like._dtype + + elif isinstance(array_like, ak.Array): + self._impl = array_like + self._shape, self._dtype = _shape_dtype(self._impl.layout) + + elif isinstance(array_like, (bool, Real)): + self._impl = np.array(array_like) + self._shape, self._dtype = (), self._impl.dtype + + else: + self._impl = ak.Array(array_like) + self._shape, self._dtype = _shape_dtype(self._impl.layout) + + if not isinstance(dtype, np.dtype): + dtype = np.dtype(dtype) + + if dtype is not None and dtype != self._dtype: + if isinstance(self._impl, ak.Array): + self._impl = ak.values_astype(self._impl, dtype) + self._shape, self._dtype = _shape_dtype(self._impl.layout) + else: + self._impl = np.array(array_like, dtype=dtype) + self._dtype = dtype + + if self._dtype.fields is not None: + msg = f"dtype must not have fields: dtype.fields = {self._dtype.fields}" + raise TypeError(msg) + + if self._dtype.shape != (): + msg = f"dtype must not have a shape: dtype.shape = {self._dtype.shape}" + raise TypeError(msg) + + if not issubclass(self._dtype.type, np.number): + msg = f"dtype must be numeric: dtype.type = {self._dtype.type}" + raise TypeError(msg) + + if device is not None: + if isinstance(self._impl, ak.Array) and device != ak.backend(self._impl): + self._impl = ak.to_backend(self._impl, device) + elif isinstance(self._impl, np.ndarray) and device == "cuda": + cp = _import.cupy() + self._impl = cp.array(self._impl.item()) + + def __str__(self) -> str: + """ + String representation of the array. + """ + + if len(self._shape) == 0: + return f"{self._impl.item()}" + elif len(self._shape) == 1: + return f"{ak._prettyprint.valuestr(self._impl, 1, 80)}" + else: + prep = ak._prettyprint.valuestr(self._impl, 20, 80 - 4)[1:-1].replace( + "\n ", "\n " + ) + return f"[\n {prep}\n]" + + def __repr__(self) -> str: + """ + REPL-string representation of the array. + """ + + if len(self._shape) == 0: + return f"ragged.array({self._impl.item()})" + elif len(self._shape) == 1: + return f"ragged.array({ak._prettyprint.valuestr(self._impl, 1, 80 - 14)})" + else: + prep = ak._prettyprint.valuestr(self._impl, 20, 80 - 4)[1:-1].replace( + "\n ", "\n " + ) + return f"ragged.array([\n {prep}\n])" + + # Attributes: https://data-apis.org/array-api/latest/API_specification/array_object.html#attributes + + @property + def dtype(self) -> Dtype: + """ + Data type of the array elements. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.dtype.html + """ + + return self._dtype + + @property + def device(self) -> Device: + """ + Hardware device the array data resides on. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.device.html + """ + + return self._device + + @property + def mT(self) -> array: + """ + Transpose of a matrix (or a stack of matrices). + + Raises: + ValueError: If any ragged dimension's lists are not sorted from longest + to shortest, which is the only way that left-aligned ragged + transposition is possible. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.mT.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + @property + def ndim(self) -> int: + """ + Number of array dimensions (axes). + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.ndim.html + """ + + return len(self._shape) + + @property + def shape(self) -> Shape: + """ + Array dimensions. + + Regular dimensions are represented by `int` values in the `shape` and + irregular (ragged) dimensions are represented by `None`. + + According to the specification, "An array dimension must be `None` if + and only if a dimension is unknown," which is a different + interpretation than we are making here. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.shape.html + """ + + return self._shape + + @property + def size(self) -> None | int: + """ + Number of elements in an array. + + This property never returns `None` because we do not consider + dimensions to be unknown, and numerical values within ragged + lists can be counted. + + Example: + An array like `ragged.array([[1.1, 2.2, 3.3], [], [4.4, 5.5]])` has + a size of 5 because it contains 5 numerical values. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.size.html + """ + + if len(self._shape) == 0: + return 1 + else: + return int(ak.count(self._impl)) + + @property + def T(self) -> array: + """ + Transpose of the array. + + Raises: + ValueError: If any ragged dimension's lists are not sorted from longest + to shortest, which is the only way that left-aligned ragged + transposition is possible. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.T.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + # methods: https://data-apis.org/array-api/latest/API_specification/array_object.html#methods + + def __abs__(self) -> array: + """ + Calculates the absolute value for each element of an array instance. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__abs__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __add__(self, other: int | float | array, /) -> array: + """ + Calculates the sum for each element of an array instance with the + respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__add__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __and__(self, other: int | bool | array, /) -> array: + """ + Evaluates `self_i & other_i` for each element of an array instance with + the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__and__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __array_namespace__(self, *, api_version: None | str = None) -> Any: + """ + Returns an object that has all the array API functions on it. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__array_namespace__.html + """ + + assert api_version is None, "FIXME" + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __bool__(self) -> bool: # FIXME pylint: disable=E0304 + """ + Converts a zero-dimensional array to a Python `bool` object. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__bool__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __complex__(self) -> complex: + """ + Converts a zero-dimensional array to a Python `complex` object. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__complex__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __dlpack__(self, *, stream: None | int | Any = None) -> PyCapsule: + """ + Exports the array for consumption by `from_dlpack()` as a DLPack + capsule. + + Args: + stream: CuPy Stream object (https://docs.cupy.dev/en/stable/reference/generated/cupy.cuda.Stream.html) + if not `None`. + + Raises: + ValueError: If any dimensions are ragged. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html + """ + + assert stream is None, "FIXME" + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __dlpack_device__(self) -> tuple[enum.Enum, int]: + """ + Returns device type and device ID in DLPack format. + + Raises: + ValueError: If any dimensions are ragged. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack_device__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __eq__(self, other: int | float | bool | array, /) -> array: # type: ignore[override] + """ + Computes the truth value of `self_i == other_i` for each element of an + array instance with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__eq__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __float__(self) -> float: + """ + Converts a zero-dimensional array to a Python `float` object. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__float__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __floordiv__(self, other: int | float | array, /) -> array: + """ + Evaluates `self_i // other_i` for each element of an array instance + with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__floordiv__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __ge__(self, other: int | float | array, /) -> array: + """ + Computes the truth value of `self_i >= other_i` for each element of an + array instance with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__ge__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __getitem__(self, key: GetSliceKey, /) -> array: + """ + Returns self[key]. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__getitem__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __gt__(self, other: int | float | array, /) -> array: + """ + Computes the truth value of `self_i > other_i` for each element of an + array instance with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__gt__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __index__(self) -> int: # FIXME pylint: disable=E0305 + """ + Converts a zero-dimensional integer array to a Python `int` object. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__index__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __int__(self) -> int: + """ + Converts a zero-dimensional array to a Python `int` object. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__int__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __invert__(self) -> array: + """ + Evaluates `~self_i` for each element of an array instance. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__invert__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __le__(self, other: int | float | array, /) -> array: + """ + Computes the truth value of `self_i <= other_i` for each element of an + array instance with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__le__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __lshift__(self, other: int | array, /) -> array: + """ + Evaluates `self_i << other_i` for each element of an array instance + with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__lshift__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __lt__(self, other: int | float | array, /) -> array: + """ + Computes the truth value of `self_i < other_i` for each element of an + array instance with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__lt__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __matmul__(self, other: array, /) -> array: + """ + Computes the matrix product. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__matmul__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __mod__(self, other: int | float | array, /) -> array: + """ + Evaluates `self_i % other_i` for each element of an array instance with + the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__mod__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __mul__(self, other: int | float | array, /) -> array: + """ + Calculates the product for each element of an array instance with the + respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__mul__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __ne__(self, other: int | float | bool | array, /) -> array: # type: ignore[override] + """ + Computes the truth value of `self_i != other_i` for each element of an + array instance with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__ne__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __neg__(self) -> array: + """ + Evaluates `-self_i` for each element of an array instance. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__neg__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __or__(self, other: int | bool | array, /) -> array: + """ + Evaluates `self_i | other_i` for each element of an array instance with + the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__or__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __pos__(self) -> array: + """ + Evaluates `+self_i` for each element of an array instance. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__pos__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __pow__(self, other: int | float | array, /) -> array: + """ + Calculates an implementation-dependent approximation of exponentiation + by raising each element (the base) of an array instance to the power of + `other_i` (the exponent), where `other_i` is the corresponding element + of the array `other`. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__pow__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __rshift__(self, other: int | array, /) -> array: + """ + Evaluates `self_i >> other_i` for each element of an array instance + with the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__rshift__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __setitem__( + self, key: SetSliceKey, value: int | float | bool | array, / + ) -> None: + """ + Sets `self[key]` to value. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__setitem__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __sub__(self, other: int | float | array, /) -> array: + """ + Calculates the difference for each element of an array instance with + the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__sub__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __truediv__(self, other: int | float | array, /) -> array: + """ + Evaluates `self_i / other_i` for each element of an array instance with + the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__truediv__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def __xor__(self, other: int | bool | array, /) -> array: + """ + Evaluates `self_i ^ other_i` for each element of an array instance with + the respective element of the array other. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__xor__.html + """ + + msg = "not implemented yet, but will be" + raise RuntimeError(msg) + + def to_device(self, device: Device, /, *, stream: None | int | Any = None) -> array: + """ + Copy the array from the device on which it currently resides to the + specified device. + + Args: + device: If `"cpu"`, the array is backed by NumPy and resides in + main memory; if `"cuda"`, the array is backed by CuPy and + resides in CUDA global memory. + stream: CuPy Stream object (https://docs.cupy.dev/en/stable/reference/generated/cupy.cuda.Stream.html) + for `device="cuda"`. + + https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.to_device.html + """ + + if isinstance(self._impl, ak.Array) and device != ak.backend(self._impl): + assert stream is None, "FIXME: use CuPy stream" + impl = ak.to_backend(self._impl, device) + + elif isinstance(self._impl, np.ndarray): + if device == "cuda": + assert stream is None, "FIXME: use CuPy stream" + cp = _import.cupy() + impl = cp.array(self._impl.item()) + else: + impl = self._impl + + else: + impl = np.array(self._impl.item()) if device == "cpu" else self._impl + + return self._new(impl, self._shape, self._dtype, device) + + # in-place operators: https://data-apis.org/array-api/2022.12/API_specification/array_object.html#in-place-operators + + def __iadd__(self, other: int | float | array, /) -> array: + """ + Calculates `self = self + other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self + other + self._impl, self._device = out._impl, out._device + return self + + def __isub__(self, other: int | float | array, /) -> array: + """ + Calculates `self = self - other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self - other + self._impl, self._device = out._impl, out._device + return self + + def __imul__(self, other: int | float | array, /) -> array: + """ + Calculates `self = self * other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self * other + self._impl, self._device = out._impl, out._device + return self + + def __itruediv__(self, other: int | float | array, /) -> array: + """ + Calculates `self = self / other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self / other + self._impl, self._device = out._impl, out._device + return self + + def __ifloordiv__(self, other: int | float | array, /) -> array: + """ + Calculates `self = self // other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self // other + self._impl, self._device = out._impl, out._device + return self + + def __ipow__(self, other: int | float | array, /) -> array: + """ + Calculates `self = self ** other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self**other + self._impl, self._device = out._impl, out._device + return self + + def __imod__(self, other: int | float | array, /) -> array: + """ + Calculates `self = self % other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self % other + self._impl, self._device = out._impl, out._device + return self + + def __imatmul__(self, other: array, /) -> array: + """ + Calculates `self = self @ other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self @ other + self._impl, self._device = out._impl, out._device + return self + + def __iand__(self, other: int | bool | array, /) -> array: + """ + Calculates `self = self & other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self & other + self._impl, self._device = out._impl, out._device + return self + + def __ior__(self, other: int | bool | array, /) -> array: + """ + Calculates `self = self | other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self | other + self._impl, self._device = out._impl, out._device + return self + + def __ixor__(self, other: int | bool | array, /) -> array: + """ + Calculates `self = self ^ other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self ^ other + self._impl, self._device = out._impl, out._device + return self + + def __ilshift__(self, other: int | array, /) -> array: + """ + Calculates `self = self << other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self << other + self._impl, self._device = out._impl, out._device + return self + + def __irshift__(self, other: int | array, /) -> array: + """ + Calculates `self = self >> other` in-place. + + (Internal arrays are immutable; this only replaces the array that the + Python object points to.) + """ + + out = self >> other + self._impl, self._device = out._impl, out._device + return self + + # reflected operators: https://data-apis.org/array-api/2022.12/API_specification/array_object.html#reflected-operators + + __radd__ = __add__ + __rsub__ = __sub__ + __rmul__ = __mul__ + __rtruediv__ = __truediv__ + __rfloordiv__ = __floordiv__ + __rpow__ = __pow__ + __rmod__ = __mod__ + __rmatmul__ = __matmul__ + __rand__ = __and__ + __ror__ = __or__ + __rxor__ = __xor__ + __rlshift__ = __lshift__ + __rrshift__ = __rshift__ diff --git a/src/ragged/common/_typing.py b/src/ragged/common/_typing.py new file mode 100644 index 0000000..72bbca8 --- /dev/null +++ b/src/ragged/common/_typing.py @@ -0,0 +1,50 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/ragged/blob/main/LICENSE + +from __future__ import annotations + +import numbers +from typing import Any, Literal, Optional, Protocol, TypeVar, Union + +import numpy as np + +T_co = TypeVar("T_co", covariant=True) + + +# not actually checked because of https://github.com/python/typing/discussions/1145 +class NestedSequence(Protocol[T_co]): + def __getitem__(self, key: int, /) -> T_co | NestedSequence[T_co]: + ... + + def __len__(self, /) -> int: + ... + + +PyCapsule = Any + + +class SupportsDLPack(Protocol): + def __dlpack__(self, /, *, stream: None = ...) -> PyCapsule: + ... + + def item(self) -> numbers.Number: + ... + + +Shape = tuple[Optional[int], ...] + +Dtype = np.dtype[ + Union[ + np.int8, + np.int16, + np.int32, + np.int64, + np.uint8, + np.uint16, + np.uint32, + np.uint64, + np.float32, + np.float64, + ] +] + +Device = Union[Literal["cpu"], Literal["cuda"]] diff --git a/src/ragged/v202212/__init__.py b/src/ragged/v202212/__init__.py new file mode 100644 index 0000000..b49396c --- /dev/null +++ b/src/ragged/v202212/__init__.py @@ -0,0 +1,16 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/ragged/blob/main/LICENSE + +""" +Defines a ragged array module that is compliant with version 2022.12 of the +Array API. + +This is the current default: `ragged.v202212.*` is imported into `ragged.*`. + +https://data-apis.org/array-api/2022.12/API_specification/ +""" + +from __future__ import annotations + +from ._obj import array + +__all__ = ["array"] diff --git a/src/ragged/v202212/_obj.py b/src/ragged/v202212/_obj.py new file mode 100644 index 0000000..d9b59ac --- /dev/null +++ b/src/ragged/v202212/_obj.py @@ -0,0 +1,14 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/ragged/blob/main/LICENSE + +from __future__ import annotations + +from ..common._obj import array as common_array + + +class array(common_array): # pylint: disable=C0103 + """ + Ragged array class and constructor for data-apis.org/array-api/2022.12. + """ + + +__all__ = ["array"] diff --git a/tests/test_0001_initial_array_object.py b/tests/test_0001_initial_array_object.py new file mode 100644 index 0000000..d2138dd --- /dev/null +++ b/tests/test_0001_initial_array_object.py @@ -0,0 +1,10 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/ragged/blob/main/LICENSE + +from __future__ import annotations + +import ragged + + +def test(): + a = ragged.array([[1, 2], [3]]) + assert a is not None diff --git a/tests/test_package.py b/tests/test_package.py deleted file mode 100644 index 481ac2b..0000000 --- a/tests/test_package.py +++ /dev/null @@ -1,9 +0,0 @@ -from __future__ import annotations - -import importlib.metadata - -import ragged as m - - -def test_version(): - assert importlib.metadata.version("ragged") == m.__version__