mparray 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mparray-0.1.0/LICENSE +21 -0
- mparray-0.1.0/PKG-INFO +103 -0
- mparray-0.1.0/README.md +62 -0
- mparray-0.1.0/pyproject.toml +75 -0
- mparray-0.1.0/src/mparray/__init__.py +9 -0
- mparray-0.1.0/src/mparray/_mparray.py +730 -0
- mparray-0.1.0/src/mparray/special/__init__.py +1 -0
- mparray-0.1.0/src/mparray/special/_special.py +230 -0
- mparray-0.1.0/src/mparray/testing/__init__.py +1 -0
- mparray-0.1.0/src/mparray/testing/_testing.py +14 -0
mparray-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Matt Haberland
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
mparray-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,103 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: mparray
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Array API-compatible, arbitrary-precision arrays
|
|
5
|
+
Author: Matt Haberland
|
|
6
|
+
Author-email: Matt Haberland <matt.haberland@gmail.com>
|
|
7
|
+
License: MIT License
|
|
8
|
+
|
|
9
|
+
Copyright (c) 2026 Matt Haberland
|
|
10
|
+
|
|
11
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
12
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
13
|
+
in the Software without restriction, including without limitation the rights
|
|
14
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
15
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
16
|
+
furnished to do so, subject to the following conditions:
|
|
17
|
+
|
|
18
|
+
The above copyright notice and this permission notice shall be included in all
|
|
19
|
+
copies or substantial portions of the Software.
|
|
20
|
+
|
|
21
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
22
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
23
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
24
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
25
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
26
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
27
|
+
SOFTWARE.
|
|
28
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
29
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
30
|
+
Requires-Dist: numpy
|
|
31
|
+
Requires-Dist: mpmath
|
|
32
|
+
Requires-Dist: scipy
|
|
33
|
+
Requires-Dist: jupyter-book ; extra == 'docs'
|
|
34
|
+
Requires-Dist: ghp-import ; extra == 'docs'
|
|
35
|
+
Requires-Dist: pytest ; extra == 'test'
|
|
36
|
+
Requires-Python: >=3.13
|
|
37
|
+
Project-URL: Home, https://github.com/mdhaber/mparray
|
|
38
|
+
Provides-Extra: docs
|
|
39
|
+
Provides-Extra: test
|
|
40
|
+
Description-Content-Type: text/markdown
|
|
41
|
+
|
|
42
|
+
MPArray is an [Array API Standard](https://data-apis.org/array-api/latest/) compatible array library that features arbitrary precision arithmetic.
|
|
43
|
+
|
|
44
|
+
Install with `pip`:
|
|
45
|
+
|
|
46
|
+
```shell
|
|
47
|
+
pip install mparray
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
Import to access the array namespace:
|
|
51
|
+
|
|
52
|
+
```shell
|
|
53
|
+
import mparray as xp
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
Use any feature defined by the 2025.12 version of the standard:
|
|
57
|
+
|
|
58
|
+
```shell
|
|
59
|
+
x = xp.arange(5)
|
|
60
|
+
# MPArray([0, 1, 2, 3, 4], dtype=int64)
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
Integer arrays are filled with Python `int`s; consequently, elements never overflow.
|
|
64
|
+
|
|
65
|
+
```shell
|
|
66
|
+
x = x[:2] + 10**50
|
|
67
|
+
# MPArray([100000000000000000000000000000000000000000000000000,
|
|
68
|
+
# 100000000000000000000000000000000000000000000000001],
|
|
69
|
+
# dtype=int64)
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
Dtypes follow array API promotion rules for compatibility, but do not limit the underlying representation. For instance, `int32` arrays also use Python `int`s.
|
|
73
|
+
|
|
74
|
+
```shell
|
|
75
|
+
xp.astype(x, xp.int32)
|
|
76
|
+
# MPArray([100000000000000000000000000000000000000000000000000,
|
|
77
|
+
# 100000000000000000000000000000000000000000000000001],
|
|
78
|
+
# dtype=int32)
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
Real and complex floating point arrays are backed by [`mpmath`](https://github.com/mpmath/mpmath).
|
|
82
|
+
|
|
83
|
+
```shell
|
|
84
|
+
from mpmath import mp
|
|
85
|
+
mp.dps = 55 # set the desired precision
|
|
86
|
+
y = xp.astype(x, xp.float64)
|
|
87
|
+
# MPArray([mpf('100000000000000000000000000000000000000000000000000.0'),
|
|
88
|
+
# mpf('100000000000000000000000000000000000000000000000001.0')],
|
|
89
|
+
# dtype=float64)
|
|
90
|
+
xp.exp(y)
|
|
91
|
+
# MPArray([mpf('4.535356657536074105363661926352788592352627616652483210606e+43429448190325182765112891891660508229439700580366'),
|
|
92
|
+
# mpf('1.232837758776106335937289338644047193491770865081040599056e+43429448190325182765112891891660508229439700580367')],
|
|
93
|
+
# dtype=float64)
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
Arbitrary-precision equivalents of some SciPy special functions are available:
|
|
97
|
+
|
|
98
|
+
```shell
|
|
99
|
+
xp.special.ndtr(-y)
|
|
100
|
+
# MPArray([mpf('8.998129248551223738202088675072137970394711995443364311445e-2171472409516259138255644594583025411471985029018332830572230087781144236110599295949020894164292471'),
|
|
101
|
+
# mpf('6.461085845597172042837958170074610008600868697345677551181e-2171472409516259138255644594583025411471985029018376260020420755455996799425617236333573042128793971')],
|
|
102
|
+
# dtype=float64)
|
|
103
|
+
```
|
mparray-0.1.0/README.md
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
MPArray is an [Array API Standard](https://data-apis.org/array-api/latest/) compatible array library that features arbitrary precision arithmetic.
|
|
2
|
+
|
|
3
|
+
Install with `pip`:
|
|
4
|
+
|
|
5
|
+
```shell
|
|
6
|
+
pip install mparray
|
|
7
|
+
```
|
|
8
|
+
|
|
9
|
+
Import to access the array namespace:
|
|
10
|
+
|
|
11
|
+
```shell
|
|
12
|
+
import mparray as xp
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
Use any feature defined by the 2025.12 version of the standard:
|
|
16
|
+
|
|
17
|
+
```shell
|
|
18
|
+
x = xp.arange(5)
|
|
19
|
+
# MPArray([0, 1, 2, 3, 4], dtype=int64)
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
Integer arrays are filled with Python `int`s; consequently, elements never overflow.
|
|
23
|
+
|
|
24
|
+
```shell
|
|
25
|
+
x = x[:2] + 10**50
|
|
26
|
+
# MPArray([100000000000000000000000000000000000000000000000000,
|
|
27
|
+
# 100000000000000000000000000000000000000000000000001],
|
|
28
|
+
# dtype=int64)
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
Dtypes follow array API promotion rules for compatibility, but do not limit the underlying representation. For instance, `int32` arrays also use Python `int`s.
|
|
32
|
+
|
|
33
|
+
```shell
|
|
34
|
+
xp.astype(x, xp.int32)
|
|
35
|
+
# MPArray([100000000000000000000000000000000000000000000000000,
|
|
36
|
+
# 100000000000000000000000000000000000000000000000001],
|
|
37
|
+
# dtype=int32)
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
Real and complex floating point arrays are backed by [`mpmath`](https://github.com/mpmath/mpmath).
|
|
41
|
+
|
|
42
|
+
```shell
|
|
43
|
+
from mpmath import mp
|
|
44
|
+
mp.dps = 55 # set the desired precision
|
|
45
|
+
y = xp.astype(x, xp.float64)
|
|
46
|
+
# MPArray([mpf('100000000000000000000000000000000000000000000000000.0'),
|
|
47
|
+
# mpf('100000000000000000000000000000000000000000000000001.0')],
|
|
48
|
+
# dtype=float64)
|
|
49
|
+
xp.exp(y)
|
|
50
|
+
# MPArray([mpf('4.535356657536074105363661926352788592352627616652483210606e+43429448190325182765112891891660508229439700580366'),
|
|
51
|
+
# mpf('1.232837758776106335937289338644047193491770865081040599056e+43429448190325182765112891891660508229439700580367')],
|
|
52
|
+
# dtype=float64)
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
Arbitrary-precision equivalents of some SciPy special functions are available:
|
|
56
|
+
|
|
57
|
+
```shell
|
|
58
|
+
xp.special.ndtr(-y)
|
|
59
|
+
# MPArray([mpf('8.998129248551223738202088675072137970394711995443364311445e-2171472409516259138255644594583025411471985029018332830572230087781144236110599295949020894164292471'),
|
|
60
|
+
# mpf('6.461085845597172042837958170074610008600868697345677551181e-2171472409516259138255644594583025411471985029018376260020420755455996799425617236333573042128793971')],
|
|
61
|
+
# dtype=float64)
|
|
62
|
+
```
|
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["uv_build>=0.10.8,<0.11.0"]
|
|
3
|
+
build-backend = "uv_build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "mparray"
|
|
7
|
+
authors = [{ name = "Matt Haberland", email = "matt.haberland@gmail.com" }]
|
|
8
|
+
readme = "README.md"
|
|
9
|
+
license = { file = "LICENSE" }
|
|
10
|
+
classifiers = [
|
|
11
|
+
"License :: OSI Approved :: MIT License",
|
|
12
|
+
"Programming Language :: Python :: 3.13",
|
|
13
|
+
]
|
|
14
|
+
version = "0.1.0"
|
|
15
|
+
requires-python = ">=3.13"
|
|
16
|
+
dependencies = ["numpy", "mpmath", "scipy"]
|
|
17
|
+
description = "Array API-compatible, arbitrary-precision arrays"
|
|
18
|
+
|
|
19
|
+
[project.optional-dependencies]
|
|
20
|
+
test = [
|
|
21
|
+
"pytest",
|
|
22
|
+
]
|
|
23
|
+
docs = ["jupyter-book", "ghp-import"]
|
|
24
|
+
|
|
25
|
+
[project.urls]
|
|
26
|
+
Home = "https://github.com/mdhaber/mparray"
|
|
27
|
+
|
|
28
|
+
[tool.ruff]
|
|
29
|
+
target-version = "py313"
|
|
30
|
+
builtins = ["ellipsis"]
|
|
31
|
+
exclude = [".git", ".eggs", "build", "dist", "__pycache__"]
|
|
32
|
+
line-length = 100
|
|
33
|
+
|
|
34
|
+
[tool.ruff.lint]
|
|
35
|
+
ignore = [
|
|
36
|
+
"E402", # module level import not at top of file
|
|
37
|
+
# "E501", # line too long - let black worry about that
|
|
38
|
+
"E731", # do not assign a lambda expression, use a def
|
|
39
|
+
"UP038", # type union instead of tuple for isinstance etc
|
|
40
|
+
]
|
|
41
|
+
select = [
|
|
42
|
+
"F", # Pyflakes
|
|
43
|
+
"E", # Pycodestyle
|
|
44
|
+
"I", # isort
|
|
45
|
+
"UP", # Pyupgrade
|
|
46
|
+
"TID", # flake8-tidy-imports
|
|
47
|
+
"W",
|
|
48
|
+
]
|
|
49
|
+
extend-safe-fixes = [
|
|
50
|
+
"TID252", # absolute imports
|
|
51
|
+
]
|
|
52
|
+
fixable = ["I", "TID252", "UP"]
|
|
53
|
+
|
|
54
|
+
[tool.ruff.lint.isort]
|
|
55
|
+
known-first-party = ["mparray"]
|
|
56
|
+
|
|
57
|
+
[tool.ruff.lint.flake8-tidy-imports]
|
|
58
|
+
# Disallow all relative imports.
|
|
59
|
+
ban-relative-imports = "all"
|
|
60
|
+
|
|
61
|
+
[tool.coverage.run]
|
|
62
|
+
source = ["mparray"]
|
|
63
|
+
branch = true
|
|
64
|
+
|
|
65
|
+
[tool.coverage.report]
|
|
66
|
+
show_missing = true
|
|
67
|
+
exclude_lines = ["pragma: no cover", "if TYPE_CHECKING"]
|
|
68
|
+
|
|
69
|
+
[tool.pytest.ini_options]
|
|
70
|
+
filterwarnings = [
|
|
71
|
+
"error",
|
|
72
|
+
"ignore:invalid value encountered:RuntimeWarning",
|
|
73
|
+
"ignore:divide by zero encountered:RuntimeWarning",
|
|
74
|
+
"ignore:overflow encountered:RuntimeWarning",
|
|
75
|
+
]
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
"""Array API compatible, arbitrary-precision arrays."""
|
|
2
|
+
from importlib.metadata import version as _get_version
|
|
3
|
+
__version__ = _get_version("mparray")
|
|
4
|
+
del _get_version
|
|
5
|
+
|
|
6
|
+
from ._mparray import *
|
|
7
|
+
from ._mparray import __array_api_version__, __array_namespace_info__
|
|
8
|
+
from . import special
|
|
9
|
+
bool = np.bool
|
|
@@ -0,0 +1,730 @@
|
|
|
1
|
+
import collections
|
|
2
|
+
import inspect
|
|
3
|
+
import sys
|
|
4
|
+
import numpy as np
|
|
5
|
+
import mpmath
|
|
6
|
+
from mpmath import mp
|
|
7
|
+
import functools
|
|
8
|
+
import warnings
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class MPArray:
|
|
12
|
+
|
|
13
|
+
def __init__(self, obj, *, dtype=None, device=None, copy=None):
|
|
14
|
+
# TODO: fast path if input is MPArray and nothing is changing
|
|
15
|
+
|
|
16
|
+
if isinstance(obj, MPArray):
|
|
17
|
+
data = np.asarray(obj._data, device=device, copy=copy)
|
|
18
|
+
elif isinstance(obj, bool):
|
|
19
|
+
data = np.asarray(obj)
|
|
20
|
+
elif isinstance(obj, (int, mp.mpf, mp.mpc)):
|
|
21
|
+
data = np.asarray(obj, dtype=object, device=device, copy=copy)
|
|
22
|
+
else:
|
|
23
|
+
data = np.asarray(obj, device=device, copy=copy)
|
|
24
|
+
dtype = data.dtype if dtype is None else dtype
|
|
25
|
+
|
|
26
|
+
dtype = _get_dtype(obj) if dtype is None else _get_dtype(dtype)
|
|
27
|
+
dtype = result_type(dtype)
|
|
28
|
+
shape = data.shape
|
|
29
|
+
|
|
30
|
+
dtype_ = object
|
|
31
|
+
if np.isdtype(dtype, 'bool'):
|
|
32
|
+
type_ = bool
|
|
33
|
+
dtype_ = np.bool
|
|
34
|
+
elif np.isdtype(dtype, 'integral'):
|
|
35
|
+
type_ = int
|
|
36
|
+
elif np.isdtype(dtype, 'real floating'):
|
|
37
|
+
def type_(x):
|
|
38
|
+
try:
|
|
39
|
+
return mp.mpf(x) if isinstance(x, (int, mp.mpf, mp.mpc)) else mp.mpf(float(x))
|
|
40
|
+
except TypeError:
|
|
41
|
+
return mp.mpf(np.nan)
|
|
42
|
+
elif np.isdtype(dtype, 'complex floating'):
|
|
43
|
+
type_ = lambda x: mp.mpc(x) if isinstance(x, (mp.mpf, mp.mpc)) else mp.mpc(complex(x))
|
|
44
|
+
else:
|
|
45
|
+
type_ = mp.mpf
|
|
46
|
+
dtype = result_type(float64)
|
|
47
|
+
|
|
48
|
+
try:
|
|
49
|
+
data = np.asarray([type_(el) for el in data.ravel()], dtype=dtype_)
|
|
50
|
+
except TypeError:
|
|
51
|
+
type_ = mp.mpc
|
|
52
|
+
dtype = result_type(complex128)
|
|
53
|
+
data = np.asarray([type_(el) for el in data.ravel()], dtype=dtype_)
|
|
54
|
+
|
|
55
|
+
data = np.reshape(data, shape)
|
|
56
|
+
|
|
57
|
+
self._data = data
|
|
58
|
+
self._dtype = dtype
|
|
59
|
+
self._device = data.device
|
|
60
|
+
self._ndim = data.ndim
|
|
61
|
+
self._shape = data.shape
|
|
62
|
+
self._size = data.size
|
|
63
|
+
|
|
64
|
+
__array_priority__ = 1 # make reflected operators work with NumPy
|
|
65
|
+
|
|
66
|
+
@property
|
|
67
|
+
def dtype(self):
|
|
68
|
+
return self._dtype
|
|
69
|
+
|
|
70
|
+
@property
|
|
71
|
+
def device(self):
|
|
72
|
+
return self._device
|
|
73
|
+
|
|
74
|
+
@property
|
|
75
|
+
def ndim(self):
|
|
76
|
+
return self._ndim
|
|
77
|
+
|
|
78
|
+
@property
|
|
79
|
+
def shape(self):
|
|
80
|
+
return self._shape
|
|
81
|
+
|
|
82
|
+
@property
|
|
83
|
+
def size(self):
|
|
84
|
+
return self._size
|
|
85
|
+
|
|
86
|
+
def __array_namespace__(self, api_version=None):
|
|
87
|
+
if api_version is None or api_version == '2024.12':
|
|
88
|
+
import mparray as xp
|
|
89
|
+
return xp
|
|
90
|
+
else:
|
|
91
|
+
message = (f"MPArray interface for Array API version '{api_version}' "
|
|
92
|
+
"is not implemented.")
|
|
93
|
+
raise NotImplementedError(message)
|
|
94
|
+
|
|
95
|
+
def _call_super_method(self, method_name, *args, **kwargs):
|
|
96
|
+
method = getattr(self._data, method_name)
|
|
97
|
+
args = [_get_data(arg) for arg in args]
|
|
98
|
+
return method(*args, **kwargs)
|
|
99
|
+
|
|
100
|
+
def _validate_key(self, key):
|
|
101
|
+
if isinstance(key, tuple):
|
|
102
|
+
return tuple(self._validate_key(key_i) for key_i in key)
|
|
103
|
+
|
|
104
|
+
if isinstance(key, MPArray):
|
|
105
|
+
if np.isdtype(key.dtype, 'integral'):
|
|
106
|
+
return key._data.tolist()
|
|
107
|
+
return key._data
|
|
108
|
+
|
|
109
|
+
return key
|
|
110
|
+
|
|
111
|
+
# ## Indexing ##
|
|
112
|
+
|
|
113
|
+
def __getitem__(self, key):
|
|
114
|
+
key = self._validate_key(key)
|
|
115
|
+
return asarray(self._data[key], dtype=self.dtype, device=self.device)
|
|
116
|
+
|
|
117
|
+
def __setitem__(self, key, other):
|
|
118
|
+
key = self._validate_key(key)
|
|
119
|
+
other = asarray(other, dtype=self.dtype, device=self.device)
|
|
120
|
+
self._data.__setitem__(key, other._data[()])
|
|
121
|
+
|
|
122
|
+
def __deepcopy__(self, memo=None):
|
|
123
|
+
return asarray(self, copy=True)
|
|
124
|
+
|
|
125
|
+
# ## Visualization ##
|
|
126
|
+
def __repr__(self):
|
|
127
|
+
s = repr(self._data)
|
|
128
|
+
s = s.replace('array', 'MPArray')
|
|
129
|
+
return s.replace("dtype=object", f"dtype={self.dtype}")
|
|
130
|
+
|
|
131
|
+
def __str__(self):
|
|
132
|
+
# TODO: refine to show full precision?
|
|
133
|
+
return str(np.asarray(self._data, dtype=self.dtype))
|
|
134
|
+
|
|
135
|
+
# ## Linear Algebra Methods ##
|
|
136
|
+
def __matmul__(self, other):
|
|
137
|
+
return matmul(self, other)
|
|
138
|
+
|
|
139
|
+
def __imatmul__(self, other):
|
|
140
|
+
res = matmul(self, other)
|
|
141
|
+
self._data[...] = res.data[...]
|
|
142
|
+
return
|
|
143
|
+
|
|
144
|
+
def __rmatmul__(self, other):
|
|
145
|
+
other = asarray(other)
|
|
146
|
+
return matmul(other, self)
|
|
147
|
+
|
|
148
|
+
## Attributes ##
|
|
149
|
+
|
|
150
|
+
@property
|
|
151
|
+
def T(self):
|
|
152
|
+
return asarray(self._data.T, dtype=self.dtype)
|
|
153
|
+
|
|
154
|
+
@property
|
|
155
|
+
def mT(self):
|
|
156
|
+
return matrix_transpose(self)
|
|
157
|
+
|
|
158
|
+
# dlpack
|
|
159
|
+
def __dlpack_device__(self):
|
|
160
|
+
return self._data.__dlpack_device__()
|
|
161
|
+
|
|
162
|
+
def __dlpack__(self, *, stream, max_version, dl_device, copy):
|
|
163
|
+
# really not sure how to define this
|
|
164
|
+
return self._data.__dlpack__(stream=stream, max_version=max_version,
|
|
165
|
+
dl_device=dl_device, copy=copy)
|
|
166
|
+
|
|
167
|
+
def to_device(self, device, /, *, stream=None):
|
|
168
|
+
self._data = self._data.to_device(device, stream=stream)
|
|
169
|
+
|
|
170
|
+
def __index__(self):
|
|
171
|
+
if self.shape == () and np.isdtype(self.dtype, 'integral'):
|
|
172
|
+
return self._data[()]
|
|
173
|
+
else:
|
|
174
|
+
message = "Only integer scalar arrays can be converted to a scalar index."
|
|
175
|
+
raise ValueError(message)
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
## Methods ##
|
|
179
|
+
|
|
180
|
+
# Methods that return the result of a unary operation as an array
|
|
181
|
+
unary_names = (['__abs__', '__invert__', '__neg__', '__pos__'])
|
|
182
|
+
for name in unary_names:
|
|
183
|
+
def fun(self, name=name):
|
|
184
|
+
data = self._call_super_method(name)
|
|
185
|
+
dtype = self.dtype
|
|
186
|
+
dtype = ((np.float64 if "128" in str(dtype) else np.float32)
|
|
187
|
+
if ((name == '__abs__') and (dtype in [np.complex64, np.complex128]))
|
|
188
|
+
else dtype)
|
|
189
|
+
return asarray(data, dtype=dtype)
|
|
190
|
+
setattr(MPArray, name, fun)
|
|
191
|
+
|
|
192
|
+
# Methods that return the result of a unary operation as a Python scalar
|
|
193
|
+
unary_names_py = ['__bool__', '__complex__', '__float__', '__int__']
|
|
194
|
+
for name in unary_names_py:
|
|
195
|
+
def fun(self, name=name):
|
|
196
|
+
return self._call_super_method(name)
|
|
197
|
+
setattr(MPArray, name, fun)
|
|
198
|
+
|
|
199
|
+
# Methods that return the result of an elementwise binary operation
|
|
200
|
+
binary_names = ['__add__', '__sub__', '__and__', '__eq__', '__ge__', '__gt__',
|
|
201
|
+
'__le__', '__lshift__', '__lt__', '__mod__', '__mul__', '__ne__',
|
|
202
|
+
'__or__', '__rshift__', '__sub__', '__xor__']
|
|
203
|
+
# Methods that return the result of an elementwise binary operation (reflected)
|
|
204
|
+
rbinary_names = ['__radd__', '__rand__',
|
|
205
|
+
'__rlshift__', '__rmul__', '__ror__', '__rpow__',
|
|
206
|
+
'__rrshift__', '__rsub__', '__rxor__']
|
|
207
|
+
ensure_output_dtype = ['__eq__', '__ge__', '__gt__', '__le__', '__lt__',
|
|
208
|
+
'__ne__', '__rand__', '__ror__', '__rxor__']
|
|
209
|
+
for name in binary_names + rbinary_names:
|
|
210
|
+
def fun(self, other, name=name):
|
|
211
|
+
self, other = _promote(self, other)
|
|
212
|
+
dtype = None if name in ensure_output_dtype else self.dtype
|
|
213
|
+
data = self._call_super_method(name, other)
|
|
214
|
+
return asarray(data, dtype=dtype)
|
|
215
|
+
setattr(MPArray, name, fun)
|
|
216
|
+
|
|
217
|
+
|
|
218
|
+
# In-place methods
|
|
219
|
+
desired_names = ['__iadd__', '__iand__', '__ilshift__',
|
|
220
|
+
'__imul__', '__ior__', '__irshift__',
|
|
221
|
+
'__isub__', '__ixor__']
|
|
222
|
+
for name in desired_names:
|
|
223
|
+
def fun(self, other, name=name, **kwargs):
|
|
224
|
+
other = astype(other, self.dtype)
|
|
225
|
+
self._call_super_method(name, other)
|
|
226
|
+
return self
|
|
227
|
+
setattr(MPArray, name, fun)
|
|
228
|
+
|
|
229
|
+
mod = sys.modules[__name__].__dict__
|
|
230
|
+
|
|
231
|
+
## Constants ##
|
|
232
|
+
constant_names = ['e', 'inf', 'nan', 'pi']
|
|
233
|
+
for name in constant_names:
|
|
234
|
+
mod[name] = mp.mpf(getattr(mpmath, name))
|
|
235
|
+
newaxis = np.newaxis
|
|
236
|
+
|
|
237
|
+
|
|
238
|
+
## Creation Functions ##
|
|
239
|
+
def asarray(obj, /, *, dtype=None, device=None, copy=None):
|
|
240
|
+
if (isinstance(obj, MPArray) and not copy
|
|
241
|
+
and ((device is None) or (obj.device == device))
|
|
242
|
+
and ((dtype is None) or (obj.dtype == dtype))):
|
|
243
|
+
return obj
|
|
244
|
+
return MPArray(obj, dtype=dtype, device=device, copy=copy)
|
|
245
|
+
|
|
246
|
+
|
|
247
|
+
creation_functions = ['arange', 'empty', 'eye', 'from_dlpack',
|
|
248
|
+
'linspace', 'ones', 'zeros']
|
|
249
|
+
creation_functions_like = ['empty_like', 'ones_like', 'zeros_like']
|
|
250
|
+
# `full` and `full_like` created separately
|
|
251
|
+
# 'tril', 'triu', 'meshgrid' handled with array manipulation functions
|
|
252
|
+
for name in creation_functions:
|
|
253
|
+
def fun(*args, name=name, **kwargs):
|
|
254
|
+
data = getattr(np, name)(*args, **kwargs)
|
|
255
|
+
return asarray(data)
|
|
256
|
+
mod[name] = fun
|
|
257
|
+
|
|
258
|
+
for name in creation_functions_like:
|
|
259
|
+
def fun(x, /, name=name, **kwargs):
|
|
260
|
+
name = name.split("_")[0]
|
|
261
|
+
kwds = dict(shape=x.shape,
|
|
262
|
+
dtype=kwargs.get('dtype', None) or x.dtype,
|
|
263
|
+
device=kwargs.get('device', None) or x.device)
|
|
264
|
+
data = getattr(np, name)(**kwds)
|
|
265
|
+
return asarray(data)
|
|
266
|
+
mod[name] = fun
|
|
267
|
+
|
|
268
|
+
|
|
269
|
+
def full(shape, fill_value, *, dtype=None, device=None):
|
|
270
|
+
dtype = result_type(fill_value) if dtype is None else dtype
|
|
271
|
+
res = ones(shape, dtype=dtype, device=device)
|
|
272
|
+
res[...] = fill_value
|
|
273
|
+
return res
|
|
274
|
+
|
|
275
|
+
|
|
276
|
+
def full_like(x, /, fill_value, **kwargs):
|
|
277
|
+
dtype=kwargs.get('dtype', None) or x.dtype
|
|
278
|
+
device=kwargs.get('device', None) or x.device
|
|
279
|
+
return full(x.shape, fill_value, dtype=dtype, device=device)
|
|
280
|
+
|
|
281
|
+
|
|
282
|
+
## Data Type Functions and Data Types ##
|
|
283
|
+
def result_type(*args):
|
|
284
|
+
return np.result_type(*(_get_dtype(arg) for arg in args if arg is not None))
|
|
285
|
+
|
|
286
|
+
|
|
287
|
+
dtype_fun_names = ['can_cast', 'finfo', 'iinfo']
|
|
288
|
+
for name in dtype_fun_names:
|
|
289
|
+
# TODO: consider these more carefully
|
|
290
|
+
def fun(*args, name=name, **kwargs):
|
|
291
|
+
args = [_get_dtype(arg) for arg in args]
|
|
292
|
+
return getattr(np, name)(*args, **kwargs)
|
|
293
|
+
mod[name] = fun
|
|
294
|
+
|
|
295
|
+
dtype_names = ['int8', 'int16', 'int32', 'int64', 'uint8', 'uint16',
|
|
296
|
+
'uint32', 'uint64', 'float32', 'float64', 'complex64', 'complex128',
|
|
297
|
+
'isdtype'] # not really a dtype, but OK to treat it like one here
|
|
298
|
+
inspection_fun_names = ['__array_namespace_info__'] # TODO: replace this?
|
|
299
|
+
version_attribute_names = ['__array_api_version__']
|
|
300
|
+
for name in (dtype_names + inspection_fun_names + version_attribute_names):
|
|
301
|
+
mod[name] = getattr(np, name)
|
|
302
|
+
|
|
303
|
+
def astype(x, dtype, /, *, copy=True, device=None):
|
|
304
|
+
if device is None and not copy and dtype == x.dtype:
|
|
305
|
+
return x
|
|
306
|
+
# TODO: take care of copy=False error if impossible to satisfy
|
|
307
|
+
return asarray(x, dtype=dtype, device=device, copy=copy)
|
|
308
|
+
|
|
309
|
+
|
|
310
|
+
## Elementwise Functions ##
|
|
311
|
+
# TODO: fix `logical_` functions for non-boolean dtype
|
|
312
|
+
elementwise_numpy = ['equal', 'greater', 'greater_equal', 'less', 'less_equal',
|
|
313
|
+
'logical_and', 'logical_not', 'logical_or', 'logical_xor',
|
|
314
|
+
'not_equal']
|
|
315
|
+
for name in elementwise_numpy:
|
|
316
|
+
def fun(*args, name=name, **kwargs):
|
|
317
|
+
args = (_get_data(arg) for arg in args)
|
|
318
|
+
return asarray(getattr(np, name)(*args, **kwargs))
|
|
319
|
+
mod[name] = fun
|
|
320
|
+
|
|
321
|
+
# TODO: fix `bitwise_` functions for inappropriate dtypes
|
|
322
|
+
elementwise_no_dtype = ['abs', 'bitwise_and', 'bitwise_left_shift', 'bitwise_invert',
|
|
323
|
+
'bitwise_or', 'bitwise_right_shift', 'bitwise_xor', 'negative',
|
|
324
|
+
'positive', 'square']
|
|
325
|
+
elementwise_promote_numpy = ['add', 'remainder', 'multiply',
|
|
326
|
+
'maximum', 'minimum', 'subtract']
|
|
327
|
+
for name in elementwise_no_dtype + elementwise_promote_numpy:
|
|
328
|
+
def fun(*args, name=name, **kwargs):
|
|
329
|
+
args = _promote(*args)
|
|
330
|
+
dtype = args[0].dtype
|
|
331
|
+
args = tuple(_get_data(arg) for arg in args)
|
|
332
|
+
dtype = ((np.float64 if "128" in str(dtype) else np.float32)
|
|
333
|
+
if ((name == 'abs') and (dtype in [np.complex64, np.complex128]))
|
|
334
|
+
else dtype)
|
|
335
|
+
return asarray(getattr(np, name)(*args, **kwargs), dtype=dtype)
|
|
336
|
+
mod[name] = fun
|
|
337
|
+
|
|
338
|
+
|
|
339
|
+
def _dividelike_special_case(x1, x2, *, op):
|
|
340
|
+
return op(np.astype(x1._data, x1.dtype), np.astype(x2._data, x2.dtype))
|
|
341
|
+
|
|
342
|
+
|
|
343
|
+
def _dividelike(x1, x2, *, op):
|
|
344
|
+
# mpmath division by zero raises,
|
|
345
|
+
x1, x2 = _promote(x1, x2)
|
|
346
|
+
x1, x2 = broadcast_arrays(x1, x2)
|
|
347
|
+
res = empty(x1.shape, dtype=x1.dtype)
|
|
348
|
+
i = (x2 != 0)._data & isfinite(x2)._data
|
|
349
|
+
res._data[i] = op(x1[i]._data, x2[i]._data)
|
|
350
|
+
res._data[~i] = _dividelike_special_case(x1[~i], x2[~i], op=op)
|
|
351
|
+
return res
|
|
352
|
+
|
|
353
|
+
|
|
354
|
+
def divide(x1, x2, /):
|
|
355
|
+
return _dividelike(x1, x2, op=lambda x1, x2: x1 / x2)
|
|
356
|
+
|
|
357
|
+
|
|
358
|
+
def floor_divide(x1, x2, /):
|
|
359
|
+
return _dividelike(x1, x2, op=lambda x1, x2: x1 // x2)
|
|
360
|
+
|
|
361
|
+
|
|
362
|
+
def remainder(x1, x2, /):
|
|
363
|
+
return _dividelike(x1, x2, op=lambda x1, x2: x1 % x2)
|
|
364
|
+
|
|
365
|
+
|
|
366
|
+
def pow(x1, x2, /):
|
|
367
|
+
x1, x2 = _promote(x1, x2)
|
|
368
|
+
x1, x2 = broadcast_arrays(x1, x2)
|
|
369
|
+
res = empty(x1.shape, dtype=x1.dtype)
|
|
370
|
+
i = (x1 != 0)._data
|
|
371
|
+
res._data[i] = x1[i]._data ** x2[i]._data
|
|
372
|
+
res._data[~i] = np.astype(x1[~i]._data, x1.dtype) ** np.astype(x2[~i]._data, x2.dtype)
|
|
373
|
+
return res
|
|
374
|
+
|
|
375
|
+
|
|
376
|
+
setattr(MPArray, "__truediv__", divide)
|
|
377
|
+
setattr(MPArray, "__floordiv__", floor_divide)
|
|
378
|
+
setattr(MPArray, "__mod__", remainder)
|
|
379
|
+
setattr(MPArray, "__rtruediv__", lambda x, y: divide(y, x))
|
|
380
|
+
setattr(MPArray, "__rfloordiv__", lambda x, y: floor_divide(y, x))
|
|
381
|
+
setattr(MPArray, "__rmod__", lambda x, y: remainder(y, x))
|
|
382
|
+
setattr(MPArray, "__pow__", pow)
|
|
383
|
+
reciprocal = lambda x, /: divide(1, x)
|
|
384
|
+
|
|
385
|
+
mp.logaddexp = lambda x, y: mp.log(mp.exp(x) + mp.exp(y))
|
|
386
|
+
mp.imag = lambda x: x.imag
|
|
387
|
+
mp.real = lambda x: x.real
|
|
388
|
+
mp.round = lambda x: mp.nint(x)
|
|
389
|
+
mp.trunc = lambda x: mp.floor(x) if x > 0 else mp.ceil(x)
|
|
390
|
+
elementwise_mp = ['acos', 'acosh', 'asin', 'asinh', 'atan', 'atan2', 'atanh', 'cos',
|
|
391
|
+
'cosh', 'exp', 'expm1', 'hypot', 'log', 'log1p', 'log2',
|
|
392
|
+
'log10', 'logaddexp', 'sin', 'sinh', 'sqrt', 'tan', 'tanh']
|
|
393
|
+
elementwise_mp_float = ['ceil', 'conj', 'floor', 'imag', 'real', 'round', 'trunc']
|
|
394
|
+
for name in elementwise_mp + elementwise_mp_float:
|
|
395
|
+
def fun(*args, name=name, **kwargs):
|
|
396
|
+
atleast = bool if name in elementwise_mp_float else 1.0
|
|
397
|
+
args = _promote(*args, atleast=atleast)
|
|
398
|
+
|
|
399
|
+
if name in elementwise_mp_float and np.isdtype(args[0].dtype, ('bool', 'integral')):
|
|
400
|
+
return args[0] # TODO: fix imag
|
|
401
|
+
|
|
402
|
+
data = (_get_data(arg) for arg in args)
|
|
403
|
+
out = np.vectorize(getattr(mp, name), otypes=[object])(*data, **kwargs)
|
|
404
|
+
# TODO: preserve complex output dtype for funcs like acos
|
|
405
|
+
|
|
406
|
+
dtype = (np.float64 if "128" in str(args[0].dtype) else np.float32) if name in {'imag', 'real'} else args[0].dtype
|
|
407
|
+
return asarray(out, dtype=dtype)
|
|
408
|
+
mod[name] = fun
|
|
409
|
+
|
|
410
|
+
|
|
411
|
+
def _mimum(x1, x2, /, op):
|
|
412
|
+
res = op(x1, x2)
|
|
413
|
+
i = isnan(x1) | isnan(x2)
|
|
414
|
+
res[i] = nan
|
|
415
|
+
return res
|
|
416
|
+
|
|
417
|
+
|
|
418
|
+
def maximum(x1, x2, maximum=maximum):
|
|
419
|
+
return _mimum(x1, x2, op=maximum)
|
|
420
|
+
|
|
421
|
+
|
|
422
|
+
def minimum(x1, x2, minimum=minimum):
|
|
423
|
+
return _mimum(x1, x2, op=minimum)
|
|
424
|
+
|
|
425
|
+
|
|
426
|
+
elementwise_is = ['isfinite', 'isinf', 'isnan']
|
|
427
|
+
for name in elementwise_is:
|
|
428
|
+
def fun(arg, name=name, **kwargs):
|
|
429
|
+
arg = asarray(arg)
|
|
430
|
+
if np.isdtype(arg.dtype, ('bool', 'integral')):
|
|
431
|
+
return full_like(arg, name=='isfinite', dtype=np.bool)
|
|
432
|
+
|
|
433
|
+
out = np.vectorize(getattr(mp, name), otypes=[bool])(_get_data(arg), **kwargs)
|
|
434
|
+
return asarray(out, dtype=np.bool)
|
|
435
|
+
mod[name] = fun
|
|
436
|
+
|
|
437
|
+
|
|
438
|
+
def sign(x, /):
|
|
439
|
+
x = asarray(x)
|
|
440
|
+
if isdtype(x.dtype, ('bool', 'integral')):
|
|
441
|
+
return asarray(np.sign(x._data), dtype=x.dtype)
|
|
442
|
+
return asarray(_vectorize(mp.sign)(x), dtype=x.dtype)
|
|
443
|
+
|
|
444
|
+
|
|
445
|
+
def signbit(x, /):
|
|
446
|
+
return x < 0 # Python int/mp.mpf don't have -0 or signed NaNs
|
|
447
|
+
|
|
448
|
+
|
|
449
|
+
def copysign(x1, x2, /):
|
|
450
|
+
dtype = result_type(x1, x2)
|
|
451
|
+
return asarray(abs(x1) * sign(x2), dtype=dtype)
|
|
452
|
+
|
|
453
|
+
|
|
454
|
+
def nextafter(x1, x2, /):
|
|
455
|
+
x1, x2 = _promote(x1, x2)
|
|
456
|
+
inc = 10 ** (floor(log10(x1)) - mp.dps) # TODO: defined for other types
|
|
457
|
+
return x1 + sign(x2 - x1)*inc
|
|
458
|
+
|
|
459
|
+
|
|
460
|
+
def clip(x, /, min=None, max=None):
|
|
461
|
+
x, min, max = _promote(x, min, max)
|
|
462
|
+
dtype = x.dtype
|
|
463
|
+
x, min, max = _get_data(x, min, max)
|
|
464
|
+
out = np.clip(x, min=min, max=max)
|
|
465
|
+
return asarray(out, dtype=dtype)
|
|
466
|
+
|
|
467
|
+
## Indexing Functions
|
|
468
|
+
take_like = ['take', 'take_along_axis']
|
|
469
|
+
for name in take_like:
|
|
470
|
+
def fun(x, indices, name=name, **kwargs):
|
|
471
|
+
x = asarray(x)
|
|
472
|
+
dtype = x.dtype
|
|
473
|
+
x, indices = _get_data(x, indices)
|
|
474
|
+
indices = np.astype(indices, np.int64)
|
|
475
|
+
return asarray(getattr(np, name)(x, indices, **kwargs), dtype=dtype)
|
|
476
|
+
mod[name] = fun
|
|
477
|
+
|
|
478
|
+
## Inspection ##
|
|
479
|
+
# Included with dtype functions above
|
|
480
|
+
|
|
481
|
+
## Linear Algebra Functions ##
|
|
482
|
+
linalg_names = ['matmul', 'tensordot', 'vecdot']
|
|
483
|
+
for name in linalg_names:
|
|
484
|
+
def fun(x1, x2, /, name=name, **kwargs):
|
|
485
|
+
x1, x2 = _promote(x1, x2)
|
|
486
|
+
dtype = x1.dtype
|
|
487
|
+
x1, x2 = _get_data(x1, x2)
|
|
488
|
+
out = getattr(np, name)(x1, x2, **kwargs)
|
|
489
|
+
out = (getattr(np, name)(x1.astype(int), x2.astype(int), **kwargs)
|
|
490
|
+
if np.any(out == None) else out) # see gh-31019
|
|
491
|
+
return asarray(out, dtype=dtype)
|
|
492
|
+
mod[name] = fun
|
|
493
|
+
|
|
494
|
+
matrix_transpose = lambda x: asarray(_get_data(x).mT, dtype=x.dtype)
|
|
495
|
+
|
|
496
|
+
## Manipulation Functions ##
|
|
497
|
+
output_arrays = {'broadcast_arrays', 'unstack', 'meshgrid'}
|
|
498
|
+
|
|
499
|
+
manip_array_in_out = ['broadcast_arrays', 'meshgrid']
|
|
500
|
+
for name in manip_array_in_out:
|
|
501
|
+
def fun(*args, name=name, **kwargs):
|
|
502
|
+
res = getattr(np, name)(*(_get_data(arg) for arg in args), **kwargs)
|
|
503
|
+
return tuple(asarray(resi, dtype=_get_dtype(arg))
|
|
504
|
+
for resi, arg in zip(res, args))
|
|
505
|
+
mod[name] = fun
|
|
506
|
+
|
|
507
|
+
manip_tuple_in = ['concat', 'stack']
|
|
508
|
+
for name in manip_tuple_in:
|
|
509
|
+
def fun(args, name=name, **kwargs):
|
|
510
|
+
args = tuple(_promote(*args))
|
|
511
|
+
dtype = args[0].dtype
|
|
512
|
+
res = getattr(np, name)(tuple(_get_data(arg) for arg in args), **kwargs)
|
|
513
|
+
return asarray(res, dtype=dtype)
|
|
514
|
+
mod[name] = fun
|
|
515
|
+
|
|
516
|
+
manip_names = ['broadcast_to', 'expand_dims', 'flip', 'moveaxis', 'permute_dims',
|
|
517
|
+
'reshape', 'roll', 'squeeze', 'tile', 'tril', 'triu']
|
|
518
|
+
for name in manip_names:
|
|
519
|
+
def fun(x, *args, name=name, **kwargs):
|
|
520
|
+
x = asarray(x)
|
|
521
|
+
res = getattr(np, name)(_get_data(x), *args, **kwargs)
|
|
522
|
+
return asarray(res, dtype=x.dtype)
|
|
523
|
+
mod[name] = fun
|
|
524
|
+
|
|
525
|
+
|
|
526
|
+
def repeat(x, repeats, /, *, axis=None):
|
|
527
|
+
x = asarray(x)
|
|
528
|
+
repeats = np.asarray(_get_data(repeats), dtype=np.int64)
|
|
529
|
+
res = np.repeat(x._data, repeats, axis=axis)
|
|
530
|
+
return asarray(res, dtype=x.dtype)
|
|
531
|
+
|
|
532
|
+
|
|
533
|
+
def unstack(x, /, *, axis=0):
|
|
534
|
+
x = asarray(x)
|
|
535
|
+
res = np.unstack(x._data, axis=axis)
|
|
536
|
+
return tuple(asarray(resi, dtype=x.dtype) for resi in res)
|
|
537
|
+
|
|
538
|
+
|
|
539
|
+
broadcast_shapes = np.broadcast_shapes
|
|
540
|
+
|
|
541
|
+
## Searching Functions
|
|
542
|
+
def searchsorted(x1, x2, /, *, side='left', sorter=None):
|
|
543
|
+
x1, x2 = _promote(x1, x2)
|
|
544
|
+
x1, x2 = _get_data(x1, x2)
|
|
545
|
+
j = np.searchsorted(x1, x2, side=side,
|
|
546
|
+
sorter=_get_data(sorter).astype(int)
|
|
547
|
+
if sorter is not None else sorter)
|
|
548
|
+
return asarray(j)
|
|
549
|
+
|
|
550
|
+
|
|
551
|
+
def nonzero(x, /):
|
|
552
|
+
x = asarray(x)
|
|
553
|
+
res = np.nonzero(x._data)
|
|
554
|
+
return tuple(asarray(resi) for resi in res)
|
|
555
|
+
|
|
556
|
+
|
|
557
|
+
def where(condition, x1, x2, /):
|
|
558
|
+
condition = asarray(condition)
|
|
559
|
+
x1, x2 = _promote(x1, x2)
|
|
560
|
+
data = np.where(condition._data, x1._data, x2._data)
|
|
561
|
+
return asarray(data, dtype=x1.dtype)
|
|
562
|
+
|
|
563
|
+
# Defined below, in Statistical Functions
|
|
564
|
+
# argmax
|
|
565
|
+
# argmin
|
|
566
|
+
# count_nonzero
|
|
567
|
+
|
|
568
|
+
## Set Functions ##
|
|
569
|
+
unique_names = ['unique_values', 'unique_counts', 'unique_inverse', 'unique_all']
|
|
570
|
+
for name in unique_names:
|
|
571
|
+
def fun(x, /, name=name):
|
|
572
|
+
x = asarray(x)
|
|
573
|
+
res = getattr(np, name)(x._data)
|
|
574
|
+
if name == 'unique_values':
|
|
575
|
+
return asarray(res, dtype=x.dtype)
|
|
576
|
+
|
|
577
|
+
fields = res._fields
|
|
578
|
+
name_tuple = res.__class__.__name__
|
|
579
|
+
result_class = collections.namedtuple(name_tuple, fields)
|
|
580
|
+
|
|
581
|
+
result_list = []
|
|
582
|
+
for res_i, field_i in zip(res, fields):
|
|
583
|
+
dtype = x._dtype if field_i == 'values' else None
|
|
584
|
+
result_list.append(asarray(res_i, dtype=dtype))
|
|
585
|
+
return result_class(*result_list)
|
|
586
|
+
mod[name] = fun
|
|
587
|
+
|
|
588
|
+
|
|
589
|
+
def isin(x1, x2, /, *, invert=False):
|
|
590
|
+
x1, x2 = _promote(x1, x2)
|
|
591
|
+
data = np.isin(x1._data, x2._data, invert=invert)
|
|
592
|
+
return asarray(data, dtype=bool)
|
|
593
|
+
|
|
594
|
+
|
|
595
|
+
## Sorting Functions ##
|
|
596
|
+
sort_names = ['sort', 'argsort']
|
|
597
|
+
for name in sort_names:
|
|
598
|
+
def fun(x, /, *, name=name, axis=-1, descending=False, stable=True):
|
|
599
|
+
x = asarray(x)
|
|
600
|
+
x = -x if descending else x
|
|
601
|
+
res = getattr(np, name)(x._data, axis=axis, stable=stable)
|
|
602
|
+
res = -res if (descending and name == 'sort') else res
|
|
603
|
+
return asarray(res, dtype=x.dtype if name == 'sort' else None)
|
|
604
|
+
mod[name] = fun
|
|
605
|
+
|
|
606
|
+
## Statistical Functions and Utility Functions ##
|
|
607
|
+
statistical_names_float = ['mean', 'var', 'std']
|
|
608
|
+
statistical_names_dtype = ['max', 'min', 'sum', 'prod',
|
|
609
|
+
'cumulative_sum', 'cumulative_prod']
|
|
610
|
+
statistical_names_none = ['argmax', 'argmin', 'count_nonzero', 'all', 'any']
|
|
611
|
+
for name in statistical_names_float + statistical_names_dtype + statistical_names_none:
|
|
612
|
+
def fun(x, *args, name=name, **kwargs):
|
|
613
|
+
dtype = kwargs.pop('dtype', float if name in statistical_names_float else bool)
|
|
614
|
+
x, = _promote(x, atleast=dtype) # TODO: follow standard precisely?
|
|
615
|
+
with warnings.catch_warnings():
|
|
616
|
+
warnings.simplefilter("ignore")
|
|
617
|
+
res = getattr(np, name)(x._data, *args, **kwargs)
|
|
618
|
+
return asarray(res, dtype=None if name in statistical_names_none else x.dtype)
|
|
619
|
+
mod[name] = fun
|
|
620
|
+
|
|
621
|
+
|
|
622
|
+
def _minmax(x1, /, *, axis=None, keepdims=False, op):
|
|
623
|
+
res = op(x1, axis=axis, keepdims=keepdims)
|
|
624
|
+
if isdtype(x1.dtype, ("real floating", "complex floating")):
|
|
625
|
+
i = any(isnan(x1), axis=axis, keepdims=keepdims)
|
|
626
|
+
res[i] = nan
|
|
627
|
+
return res
|
|
628
|
+
|
|
629
|
+
|
|
630
|
+
def min(x1, /, *, axis=None, keepdims=False, min=min):
|
|
631
|
+
return _minmax(x1, axis=axis, keepdims=keepdims, op=min)
|
|
632
|
+
|
|
633
|
+
|
|
634
|
+
def max(x1, /, *, axis=None, keepdims=False, max=max):
|
|
635
|
+
return _minmax(x1, axis=axis, keepdims=keepdims, op=max)
|
|
636
|
+
|
|
637
|
+
|
|
638
|
+
def diff(x, /, *, axis=-1, n=1, prepend=None, append=None):
|
|
639
|
+
x, prepend, append = _promote(x, prepend, append)
|
|
640
|
+
prepend = prepend._data if prepend is not None else np._NoValue
|
|
641
|
+
append = append._data if append is not None else np._NoValue
|
|
642
|
+
res = np.diff(x._data, axis=axis, n=n, prepend=prepend, append=append)
|
|
643
|
+
return asarray(res, dtype=x.dtype)
|
|
644
|
+
|
|
645
|
+
|
|
646
|
+
_dont_mod_signature = {'clip', 'sort', 'argsort'}
|
|
647
|
+
preface = ["The following is the documentation for the corresponding "
|
|
648
|
+
f"attribute of NumPy.",
|
|
649
|
+
"MPArray behavior is the same except that the calculation is "
|
|
650
|
+
"carried out in the appropriate precision.\n\n"]
|
|
651
|
+
preface = "\n".join(preface)
|
|
652
|
+
mod_keys = list(mod.keys())
|
|
653
|
+
for attribute in mod_keys:
|
|
654
|
+
# Add documentation if it is not already present
|
|
655
|
+
if mod[attribute].__doc__:
|
|
656
|
+
continue
|
|
657
|
+
|
|
658
|
+
np_attr = getattr(np, attribute, None)
|
|
659
|
+
mod_attr = mod.get(attribute, None)
|
|
660
|
+
if np_attr is not None and mod_attr is not None:
|
|
661
|
+
|
|
662
|
+
if hasattr(np_attr, "__doc__"):
|
|
663
|
+
try:
|
|
664
|
+
np_doc = getattr(np_attr, "__doc__")
|
|
665
|
+
mod[attribute].__doc__ = preface + np_doc
|
|
666
|
+
except (AttributeError, TypeError):
|
|
667
|
+
pass
|
|
668
|
+
|
|
669
|
+
if attribute not in _dont_mod_signature:
|
|
670
|
+
try:
|
|
671
|
+
mod_attr.__signature__ = inspect.signature(np_attr)
|
|
672
|
+
except (ValueError, TypeError):
|
|
673
|
+
pass
|
|
674
|
+
|
|
675
|
+
try:
|
|
676
|
+
mod_attr.__name__ = np_attr.__name__
|
|
677
|
+
except (AttributeError, TypeError):
|
|
678
|
+
pass
|
|
679
|
+
|
|
680
|
+
|
|
681
|
+
def _xinfo(x):
|
|
682
|
+
np = x._np
|
|
683
|
+
if np.isdtype(x.dtype, 'integral'):
|
|
684
|
+
return np.iinfo(x.dtype)
|
|
685
|
+
elif np.isdtype(x.dtype, 'bool'):
|
|
686
|
+
binfo = dataclasses.make_dataclass("binfo", ['min', 'max'])
|
|
687
|
+
return binfo(min=False, max=True)
|
|
688
|
+
else:
|
|
689
|
+
return np.finfo(x.dtype)
|
|
690
|
+
|
|
691
|
+
|
|
692
|
+
def _get_data(*args):
|
|
693
|
+
if len(args) == 1:
|
|
694
|
+
x = args[0]
|
|
695
|
+
return x._data if isinstance(x, MPArray) else x
|
|
696
|
+
return tuple(_get_data(arg) for arg in args)
|
|
697
|
+
|
|
698
|
+
|
|
699
|
+
def _get_dtype(x):
|
|
700
|
+
if isinstance(x, bool) or x is np.bool:
|
|
701
|
+
return np.bool
|
|
702
|
+
elif isinstance(x, mp.mpf):
|
|
703
|
+
return np.float64
|
|
704
|
+
elif isinstance(x, mp.mpc):
|
|
705
|
+
return np.complex128
|
|
706
|
+
elif isinstance(x, type) and not ((x is bool) or (x is int)
|
|
707
|
+
or (x is float) or (x is complex)):
|
|
708
|
+
return x(0).dtype
|
|
709
|
+
elif (x is int) or (x is float) or (x is complex):
|
|
710
|
+
return x(1)
|
|
711
|
+
|
|
712
|
+
return getattr(x, "dtype", x)
|
|
713
|
+
|
|
714
|
+
|
|
715
|
+
def _promote(*args, atleast=bool):
|
|
716
|
+
dtype = result_type(*args, atleast)
|
|
717
|
+
return tuple((astype(arg, dtype) if arg is not None else arg) for arg in args)
|
|
718
|
+
|
|
719
|
+
|
|
720
|
+
def _vectorize(f):
|
|
721
|
+
|
|
722
|
+
@functools.wraps(f)
|
|
723
|
+
def wrapped(*args, **kwargs):
|
|
724
|
+
args = list(_promote(*args, atleast=float))
|
|
725
|
+
data = (_get_data(arg) for arg in args)
|
|
726
|
+
out = np.vectorize(f, otypes=[object])(*data, **kwargs)
|
|
727
|
+
# TODO: preserve complex output dtype for funcs like acos
|
|
728
|
+
return asarray(out, dtype=args[0].dtype)
|
|
729
|
+
|
|
730
|
+
return wrapped
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
from ._special import *
|
|
@@ -0,0 +1,230 @@
|
|
|
1
|
+
import sys as sys
|
|
2
|
+
import numpy as np
|
|
3
|
+
from mpmath import mp
|
|
4
|
+
import mparray as xp
|
|
5
|
+
from mparray._mparray import _vectorize as vectorize
|
|
6
|
+
from scipy import special
|
|
7
|
+
|
|
8
|
+
# add imported names to `_imports` to avoid altering their documentation and exposing
|
|
9
|
+
# as public members of `mparray.special`.
|
|
10
|
+
_imports = {'sys', 'np', 'mp', 'xp', 'vectorize', 'special'}
|
|
11
|
+
|
|
12
|
+
expm1 = vectorize(mp.expm1)
|
|
13
|
+
log1p = vectorize(mp.log1p)
|
|
14
|
+
factorial2 = vectorize(mp.fac2)
|
|
15
|
+
psi = vectorize(mp.digamma)
|
|
16
|
+
digamma = psi
|
|
17
|
+
ndtr = vectorize(mp.ncdf)
|
|
18
|
+
gamma = vectorize(mp.gamma)
|
|
19
|
+
gammaln = vectorize(mp.loggamma)
|
|
20
|
+
erf = vectorize(mp.erf)
|
|
21
|
+
erfc = vectorize(mp.erfc)
|
|
22
|
+
zeta = vectorize(mp.zeta)
|
|
23
|
+
poch = vectorize(mp.rf)
|
|
24
|
+
binom = vectorize(mp.binomial)
|
|
25
|
+
comb = binom
|
|
26
|
+
powm1 = vectorize(mp.powm1)
|
|
27
|
+
hyp1f1 = vectorize(mp.hyp1f1)
|
|
28
|
+
hyp2f1 = vectorize(mp.hyp2f1)
|
|
29
|
+
iv = vectorize(mp.besseli)
|
|
30
|
+
kv = vectorize(mp.besselk)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@vectorize
|
|
34
|
+
def gammainc(a, x):
|
|
35
|
+
return mp.gammainc(a, a=0, b=x, regularized=True)
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
@vectorize
|
|
39
|
+
def gammaincc(a, x):
|
|
40
|
+
return mp.gammainc(a, a=x, b=mp.inf, regularized=True)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
@vectorize
|
|
44
|
+
def ndtri(x):
|
|
45
|
+
if x == 0:
|
|
46
|
+
return -mp.inf
|
|
47
|
+
if x == 1:
|
|
48
|
+
return mp.inf
|
|
49
|
+
if x < 0 or x > 1:
|
|
50
|
+
return mp.nan
|
|
51
|
+
|
|
52
|
+
extra_dps = int(mp.ceil(-mp.log10(x)))
|
|
53
|
+
with mp.workdps(mp.dps + extra_dps):
|
|
54
|
+
return mp.sqrt(2) * mp.erfinv(2 * x - mp.one)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
@vectorize
|
|
58
|
+
def log_ndtr(x):
|
|
59
|
+
if x <= 0:
|
|
60
|
+
return mp.log(mp.ncdf(x))
|
|
61
|
+
else:
|
|
62
|
+
complement = mp.ncdf(-x)
|
|
63
|
+
return mp.log1p(-complement)
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
@vectorize
|
|
67
|
+
def betaln(x, y):
|
|
68
|
+
return mp.log(mp.beta(x, y))
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
@vectorize
|
|
72
|
+
def betainc(a, b, x):
|
|
73
|
+
return mp.betainc(a, b, 0, x, regularized=True)
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
@vectorize
|
|
77
|
+
def fdtr(dn, dd, x):
|
|
78
|
+
return mp.betainc(dn/2, dd/2, 0, x*dn/(dd + x*dn), regularized=True)
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
@vectorize
|
|
82
|
+
def fdtrc(dn, dd, x):
|
|
83
|
+
return mp.betainc(dn/2, dd/2, x*dn/(dd + x*dn), 1, regularized=True)
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
@vectorize
|
|
87
|
+
def xlogy(x, y): # needs accuracy review
|
|
88
|
+
return x*mp.log(y)
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
@vectorize
|
|
92
|
+
def xlog1py(x, y): # needs accuracy review
|
|
93
|
+
return x*mp.log1p(y)
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
@vectorize
|
|
97
|
+
def cosm1(x):
|
|
98
|
+
if x == 0:
|
|
99
|
+
# Handle this case separately to avoid blow up in extra_dps calculation.
|
|
100
|
+
return mp.zero
|
|
101
|
+
# second term in cosine series is x**2/2
|
|
102
|
+
# catastrophic cancellation also occurs near nonzero multiples of 2*pi,
|
|
103
|
+
# but doubling precision is enough here. We are being conservative by
|
|
104
|
+
# always at least doubling the precision.
|
|
105
|
+
extra_dps = max(mp.dps, 2*int(mp.ceil(-mp.log10(x))) + 1)
|
|
106
|
+
with mp.workdps(mp.dps + extra_dps):
|
|
107
|
+
return mp.cos(x) - mp.one
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
@vectorize
|
|
111
|
+
def logit(x): # needs accuracy review
|
|
112
|
+
res = mp.log(x) - mp.log1p(-x)
|
|
113
|
+
return res
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
@vectorize
|
|
117
|
+
def expit(x): # needs accuracy review
|
|
118
|
+
return mp.exp(x - mp.log1p(mp.exp(x)))
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def _boxcox_scalar(x, lmbda):
|
|
122
|
+
"""
|
|
123
|
+
y = (x**lmbda - 1) / lmbda if lmbda != 0
|
|
124
|
+
log(x) if lmbda == 0
|
|
125
|
+
"""
|
|
126
|
+
if x < 0:
|
|
127
|
+
return mp.nan
|
|
128
|
+
if lmbda != 0:
|
|
129
|
+
return mp.powm1(x, lmbda) / lmbda
|
|
130
|
+
else:
|
|
131
|
+
return mp.log(x)
|
|
132
|
+
|
|
133
|
+
|
|
134
|
+
@vectorize
|
|
135
|
+
def boxcox(x, lmbda):
|
|
136
|
+
return _boxcox_scalar(x, lmbda)
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
@vectorize
|
|
140
|
+
def boxcox1p(x, lmbda):
|
|
141
|
+
if x == 0:
|
|
142
|
+
# Handle x = 0 separately to avoid blow up in extra_dps calculation.
|
|
143
|
+
return mp.zero
|
|
144
|
+
extra_dps = max(0, int(mp.ceil(-mp.log10(abs(x)))))
|
|
145
|
+
with mp.workdps(mp.dps + extra_dps):
|
|
146
|
+
return _boxcox_scalar(mp.one + x, lmbda)
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
def logsumexp(a, axis=None, b=None):
|
|
150
|
+
# As far as I know, logsumexp is to avoid overflow, not to improve precision.
|
|
151
|
+
# mpmath doesn't overflow, so naive implementation should be OK.
|
|
152
|
+
return xp.log(xp.sum(b*xp.exp(a), axis=axis))
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
@vectorize
|
|
156
|
+
def ive(v, z):
|
|
157
|
+
return iv(v, z) * xp.exp(-xp.abs(xp.real(z)))
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
@vectorize
|
|
161
|
+
def i0e(x):
|
|
162
|
+
return ive(0, x)
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
@vectorize
|
|
166
|
+
def i1e(x):
|
|
167
|
+
return ive(1, x)
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
@vectorize
|
|
171
|
+
def kve(v, z):
|
|
172
|
+
return kv(v, z) * xp.exp(z)
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
@vectorize
|
|
176
|
+
def k0e(x):
|
|
177
|
+
return kve(0, x)
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
@vectorize
|
|
181
|
+
def k1e(x):
|
|
182
|
+
return kve(1, x)
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
@vectorize
|
|
186
|
+
def chdtr(v, x):
|
|
187
|
+
return gammainc(v / 2, x / 2)
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
@vectorize
|
|
191
|
+
def chdtrc(v, x):
|
|
192
|
+
return gammaincc(v / 2, x / 2)
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
@vectorize
|
|
196
|
+
def stdtr(df, t):
|
|
197
|
+
x = df / (t**2 + df)
|
|
198
|
+
p = betainc(df/2, mp.one/2, x)/2
|
|
199
|
+
return np.where(t < 0, p, mp.one - p)
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
# others to be added
|
|
203
|
+
# gammaincinv
|
|
204
|
+
# gammainccinv
|
|
205
|
+
# chdtri
|
|
206
|
+
# chndtr
|
|
207
|
+
# chndtrix
|
|
208
|
+
# stdtrit
|
|
209
|
+
# ndtri_exp
|
|
210
|
+
# tklmbda
|
|
211
|
+
# inv_boxcox
|
|
212
|
+
# inv_boxcox1p
|
|
213
|
+
# kolmogorov, smirnov
|
|
214
|
+
# erfcinv
|
|
215
|
+
# erfinv
|
|
216
|
+
# lambertw
|
|
217
|
+
|
|
218
|
+
|
|
219
|
+
# generate rough documentation
|
|
220
|
+
_preface = ["The following is the documentation for the corresponding "
|
|
221
|
+
f"attribute of `scipy.special`.",
|
|
222
|
+
"MPArray behavior is the same except that the calculation is "
|
|
223
|
+
"carried out in the appropriate precision.\n\n"]
|
|
224
|
+
_preface = "\n".join(_preface)
|
|
225
|
+
function_names = list(sys.modules[__name__].__dict__.keys())
|
|
226
|
+
for key in function_names:
|
|
227
|
+
if key in _imports or '_' in key:
|
|
228
|
+
continue
|
|
229
|
+
special_doc = getattr(special, key).__doc__
|
|
230
|
+
sys.modules[__name__].__dict__[key].__doc__ = _preface + special_doc
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
from ._testing import *
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import mparray as xp
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
def assert_allclose(res, ref, *args, strict=True, **kwargs):
|
|
6
|
+
if isinstance(res, xp.MPArray):
|
|
7
|
+
res = np.asarray(res._data, dtype=res.dtype)
|
|
8
|
+
return np.testing.assert_allclose(res, ref, *args, strict, **kwargs)
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def assert_equal(res, ref, *args, strict=True, **kwargs):
|
|
12
|
+
if isinstance(res, xp.MPArray):
|
|
13
|
+
res = np.asarray(res._data, dtype=res.dtype)
|
|
14
|
+
return np.testing.assert_equal(res, ref, *args, strict, **kwargs)
|