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 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
+ ```
@@ -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)