codekingpro/portable-devtools

1"""
2Array methods which are called by both the C-code for the method
3and the Python code for the NumPy-namespace function
4
5"""
6import os
7import pickle
8import warnings
9from contextlib import nullcontext
10
11import numpy as np
12from numpy._core import multiarray as mu, numerictypes as nt, umath as um
13from numpy._core.multiarray import asanyarray
14from numpy._globals import _NoValue
15
16# save those O(100) nanoseconds!
17bool_dt = mu.dtype("bool")
18umr_maximum = um.maximum.reduce
19umr_minimum = um.minimum.reduce
20umr_sum = um.add.reduce
21umr_prod = um.multiply.reduce
22umr_bitwise_count = um.bitwise_count
23umr_any = um.logical_or.reduce
24umr_all = um.logical_and.reduce
25
26# Complex types to -> (2,)float view for fast-path computation in _var()
27_complex_to_float = {
28    nt.dtype(nt.csingle): nt.dtype(nt.single),
29    nt.dtype(nt.cdouble): nt.dtype(nt.double),
30}
31# Special case for windows: ensure double takes precedence
32if nt.dtype(nt.longdouble) != nt.dtype(nt.double):
33    _complex_to_float.update({
34        nt.dtype(nt.clongdouble): nt.dtype(nt.longdouble),
35    })
36
37# avoid keyword arguments to speed up parsing, saves about 15%-20% for very
38# small reductions
39def _amax(a, axis=None, out=None, keepdims=False,
40          initial=_NoValue, where=True):
41    return umr_maximum(a, axis, None, out, keepdims, initial, where)
42
43def _amin(a, axis=None, out=None, keepdims=False,
44          initial=_NoValue, where=True):
45    return umr_minimum(a, axis, None, out, keepdims, initial, where)
46
47def _sum(a, axis=None, dtype=None, out=None, keepdims=False,
48         initial=_NoValue, where=True):
49    return umr_sum(a, axis, dtype, out, keepdims, initial, where)
50
51def _prod(a, axis=None, dtype=None, out=None, keepdims=False,
52          initial=_NoValue, where=True):
53    return umr_prod(a, axis, dtype, out, keepdims, initial, where)
54
55def _any(a, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
56    # By default, return a boolean for any and all
57    if dtype is None:
58        dtype = bool_dt
59    # Parsing keyword arguments is currently fairly slow, so avoid it for now
60    if where is True:
61        return umr_any(a, axis, dtype, out, keepdims)
62    return umr_any(a, axis, dtype, out, keepdims, where=where)
63
64def _all(a, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
65    # By default, return a boolean for any and all
66    if dtype is None:
67        dtype = bool_dt
68    # Parsing keyword arguments is currently fairly slow, so avoid it for now
69    if where is True:
70        return umr_all(a, axis, dtype, out, keepdims)
71    return umr_all(a, axis, dtype, out, keepdims, where=where)
72
73def _count_reduce_items(arr, axis, keepdims=False, where=True):
74    # fast-path for the default case
75    if where is True:
76        # no boolean mask given, calculate items according to axis
77        if axis is None:
78            axis = tuple(range(arr.ndim))
79        elif not isinstance(axis, tuple):
80            axis = (axis,)
81        items = 1
82        for ax in axis:
83            items *= arr.shape[mu.normalize_axis_index(ax, arr.ndim)]
84        items = nt.intp(items)
85    else:
86        # TODO: Optimize case when `where` is broadcast along a non-reduction
87        # axis and full sum is more excessive than needed.
88
89        # guarded to protect circular imports
90        from numpy.lib._stride_tricks_impl import broadcast_to
91        # count True values in (potentially broadcasted) boolean mask
92        items = umr_sum(broadcast_to(where, arr.shape), axis, nt.intp, None,
93                        keepdims)
94    return items
95
96def _clip(a, min=None, max=None, out=None, **kwargs):
97    if a.dtype.kind in "iu":
98        # If min/max is a Python integer, deal with out-of-bound values here.
99        # (This enforces NEP 50 rules as no value based promotion is done.)
100        if type(min) is int and min <= np.iinfo(a.dtype).min:
101            min = None
102        if type(max) is int and max >= np.iinfo(a.dtype).max:
103            max = None
104
105    if min is None and max is None:
106        # return identity
107        return um.positive(a, out=out, **kwargs)
108    elif min is None:
109        return um.minimum(a, max, out=out, **kwargs)
110    elif max is None:
111        return um.maximum(a, min, out=out, **kwargs)
112    else:
113        return um.clip(a, min, max, out=out, **kwargs)
114
115def _mean(a, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
116    arr = asanyarray(a)
117
118    is_float16_result = False
119
120    rcount = _count_reduce_items(arr, axis, keepdims=keepdims, where=where)
121    if rcount == 0 if where is True else umr_any(rcount == 0, axis=None):
122        warnings.warn("Mean of empty slice", RuntimeWarning, stacklevel=2)
123
124    # Cast bool, unsigned int, and int to float64 by default
125    if dtype is None:
126        if issubclass(arr.dtype.type, (nt.integer, nt.bool)):
127            dtype = mu.dtype('f8')
128        elif issubclass(arr.dtype.type, nt.float16):
129            dtype = mu.dtype('f4')
130            is_float16_result = True
131
132    ret = umr_sum(arr, axis, dtype, out, keepdims, where=where)
133    if isinstance(ret, mu.ndarray):
134        ret = um.true_divide(
135                ret, rcount, out=ret, casting='unsafe', subok=False)
136        if is_float16_result and out is None:
137            ret = arr.dtype.type(ret)
138    elif hasattr(ret, 'dtype'):
139        if is_float16_result:
140            ret = arr.dtype.type(ret / rcount)
141        else:
142            ret = ret.dtype.type(ret / rcount)
143    else:
144        ret = ret / rcount
145
146    return ret
147
148def _var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False, *,
149         where=True, mean=None):
150    arr = asanyarray(a)
151
152    rcount = _count_reduce_items(arr, axis, keepdims=keepdims, where=where)
153    # Make this warning show up on top.
154    if ddof >= rcount if where is True else umr_any(ddof >= rcount, axis=None):
155        warnings.warn("Degrees of freedom <= 0 for slice", RuntimeWarning,
156                      stacklevel=2)
157
158    # Cast bool, unsigned int, and int to float64 by default
159    if dtype is None and issubclass(arr.dtype.type, (nt.integer, nt.bool)):
160        dtype = mu.dtype('f8')
161
162    if mean is not None:
163        arrmean = mean
164    else:
165        # Compute the mean.
166        # Note that if dtype is not of inexact type then arraymean will
167        # not be either.
168        arrmean = umr_sum(arr, axis, dtype, keepdims=True, where=where)
169        # The shape of rcount has to match arrmean to not change the shape of
170        # out in broadcasting. Otherwise, it cannot be stored back to arrmean.
171        if rcount.ndim == 0:
172            # fast-path for default case when where is True
173            div = rcount
174        else:
175            # matching rcount to arrmean when where is specified as array
176            div = rcount.reshape(arrmean.shape)
177        if isinstance(arrmean, mu.ndarray):
178            arrmean = um.true_divide(arrmean, div, out=arrmean,
179                                     casting='unsafe', subok=False)
180        elif hasattr(arrmean, "dtype"):
181            arrmean = arrmean.dtype.type(arrmean / rcount)
182        else:
183            arrmean = arrmean / rcount
184
185    # Compute sum of squared deviations from mean
186    # Note that x may not be inexact and that we need it to be an array,
187    # not a scalar.
188    x = um.subtract(arr, arrmean, out=...)
189    if issubclass(arr.dtype.type, (nt.floating, nt.integer)):
190        x = um.square(x, out=x)
191    # Fast-paths for built-in complex types
192    elif (_float_dtype := _complex_to_float.get(x.dtype)) is not None:
193        xv = x.view(dtype=(_float_dtype, (2,)))
194        um.square(xv, out=xv)
195        x = um.add(xv[..., 0], xv[..., 1], out=x.real)
196    # Most general case; includes handling object arrays containing imaginary
197    # numbers and complex types with non-native byteorder
198    else:
199        x = um.multiply(x, um.conjugate(x), out=x).real
200
201    ret = umr_sum(x, axis, dtype, out, keepdims=keepdims, where=where)
202
203    # Compute degrees of freedom and make sure it is not negative.
204    rcount = um.maximum(rcount - ddof, 0)
205
206    # divide by degrees of freedom
207    if isinstance(ret, mu.ndarray):
208        ret = um.true_divide(
209                ret, rcount, out=ret, casting='unsafe', subok=False)
210    elif hasattr(ret, 'dtype'):
211        ret = ret.dtype.type(ret / rcount)
212    else:
213        ret = ret / rcount
214
215    return ret
216
217def _std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False, *,
218         where=True, mean=None):
219    ret = _var(a, axis=axis, dtype=dtype, out=out, ddof=ddof,
220               keepdims=keepdims, where=where, mean=mean)
221
222    if isinstance(ret, mu.ndarray):
223        ret = um.sqrt(ret, out=ret)
224    elif hasattr(ret, 'dtype'):
225        ret = ret.dtype.type(um.sqrt(ret))
226    else:
227        ret = um.sqrt(ret)
228
229    return ret
230
231def _ptp(a, axis=None, out=None, keepdims=False):
232    return um.subtract(
233        umr_maximum(a, axis, None, out, keepdims),
234        umr_minimum(a, axis, None, None, keepdims),
235        out
236    )
237
238def _dump(self, file, protocol=2):
239    if hasattr(file, 'write'):
240        ctx = nullcontext(file)
241    else:
242        ctx = open(os.fspath(file), "wb")
243    with ctx as f:
244        pickle.dump(self, f, protocol=protocol)
245
246def _dumps(self, protocol=2):
247    return pickle.dumps(self, protocol=protocol)
248
249def _bitwise_count(a, out=None, *, where=True, casting='same_kind',
250          order='K', dtype=None, subok=True):
251    return umr_bitwise_count(a, out, where=where, casting=casting,
252            order=order, dtype=dtype, subok=subok)
253