mirror of https://github.com/exaloop/codon.git
1653 lines
52 KiB
Python
1653 lines
52 KiB
Python
# Copyright (C) 2022-2025 Exaloop Inc. <https://exaloop.io>
|
|
|
|
import operator
|
|
import util
|
|
import zmath
|
|
from .ndarray import ndarray, flagsobj, newaxis
|
|
from .routines import _check_out, asarray, empty_like, empty, concatenate, moveaxis, broadcast_shapes, zeros, broadcast_to, atleast_1d
|
|
from .ufunc import UnaryUFunc, UnaryUFunc2, BinaryUFunc
|
|
from .npdatetime import datetime64, timedelta64
|
|
from .const import pi
|
|
|
|
def cumprod(a, axis=None, dtype: type = NoneType, out=None):
|
|
a = asarray(a)
|
|
|
|
if dtype is NoneType:
|
|
return cumprod(a, axis=axis, dtype=a.dtype, out=out)
|
|
|
|
if axis is None:
|
|
prod = util.cast(1, dtype)
|
|
if out is None:
|
|
result = empty((a.size, ), dtype=dtype)
|
|
else:
|
|
_check_out(out, (a.size, ))
|
|
result = out
|
|
i = 0
|
|
for idx in util.multirange(a.shape):
|
|
prod *= util.cast(a._ptr(idx)[0], dtype)
|
|
result._ptr((i, ))[0] = util.cast(prod, result.dtype)
|
|
i += 1
|
|
return result
|
|
else:
|
|
axis = util.normalize_axis_index(axis, a.ndim)
|
|
|
|
if out is None:
|
|
result = empty_like(a, dtype=dtype)
|
|
else:
|
|
_check_out(out, a.shape)
|
|
result = out
|
|
|
|
n = a.shape[axis]
|
|
|
|
for idx in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
accum = util.cast(1, dtype)
|
|
|
|
for i in range(n):
|
|
full_idx = util.tuple_insert(idx, axis, i)
|
|
p = a._ptr(full_idx)
|
|
q = result._ptr(full_idx)
|
|
accum *= util.cast(p[0], dtype)
|
|
q[0] = util.cast(accum, result.dtype)
|
|
|
|
return result
|
|
|
|
def cumsum(a, axis=None, dtype: type = NoneType, out=None):
|
|
a = asarray(a)
|
|
|
|
if dtype is NoneType:
|
|
return cumsum(a, axis=axis, dtype=a.dtype, out=out)
|
|
|
|
if axis is None:
|
|
sum = util.cast(0, dtype)
|
|
if out is None:
|
|
result = empty((a.size, ), dtype=dtype)
|
|
else:
|
|
_check_out(out, (a.size, ))
|
|
result = out
|
|
|
|
i = 0
|
|
for idx in util.multirange(a.shape):
|
|
sum += util.cast(a._ptr(idx)[0], dtype)
|
|
result._ptr((i, ))[0] = util.cast(sum, result.dtype)
|
|
i += 1
|
|
return result
|
|
|
|
else:
|
|
axis = util.normalize_axis_index(axis, a.ndim)
|
|
|
|
if out is None:
|
|
result = empty_like(a, dtype=dtype)
|
|
else:
|
|
_check_out(out, a.shape)
|
|
result = out
|
|
|
|
n = a.shape[axis]
|
|
|
|
for idx in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
accum = util.cast(0, dtype)
|
|
|
|
for i in range(n):
|
|
full_idx = util.tuple_insert(idx, axis, i)
|
|
p = a._ptr(full_idx)
|
|
q = result._ptr(full_idx)
|
|
accum += util.cast(p[0], dtype)
|
|
q[0] = util.cast(accum, result.dtype)
|
|
|
|
return result
|
|
|
|
def diff(a, axis: int = -1, n: int = 1, append=None, prepend=None):
|
|
a = asarray(a)
|
|
|
|
if append is None and prepend is None:
|
|
# If neither append nor prepend is provided, return a
|
|
pass
|
|
elif append is not None and prepend is None:
|
|
# If only append is provided, perform the append operation
|
|
a = append(a, append, axis=axis)
|
|
elif append is None and prepend is not None:
|
|
# If only prepend is provided, perform the prepend operation
|
|
a = append(prepend, a, axis=axis)
|
|
else:
|
|
# If both append and prepend are provided, concatenate them along the specified axis
|
|
a = concatenate((prepend, a, append), axis=axis)
|
|
|
|
if isinstance(a.dtype, datetime64):
|
|
g = a.view(timedelta64[a.dtype.base, a.dtype.num])
|
|
return diff(g, axis=axis, n=n, append=None, prepend=None)
|
|
|
|
axis = util.normalize_axis_index(axis, a.ndim)
|
|
axis_len = a.shape[axis]
|
|
|
|
new_shape = a.shape
|
|
value = new_shape[axis] - n
|
|
|
|
value = max(value,0)
|
|
|
|
new_shape = util.tuple_set(new_shape, axis, value)
|
|
dtype = a.dtype
|
|
|
|
result = empty(new_shape, dtype=dtype)
|
|
|
|
if result.size == 0:
|
|
return result
|
|
|
|
if n == 1:
|
|
# Iterate over all indices except the specified axis
|
|
for idx0 in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
|
|
for i in range(axis_len - 1):
|
|
idx = util.tuple_insert(idx0, axis, i)
|
|
x0 = a._ptr(idx)[0]
|
|
|
|
idx1 = util.tuple_insert(idx0, axis, i + 1)
|
|
x1 = a._ptr(idx1)[0]
|
|
|
|
result._ptr(idx)[0] = (x1 - x0)
|
|
|
|
return result
|
|
|
|
buf = Ptr[a.dtype](axis_len)
|
|
|
|
# Iterate over all indices except the specified axis
|
|
for idx0 in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
for i in range(axis_len):
|
|
idx = util.tuple_insert(idx0, axis, i)
|
|
e = a._ptr(idx)[0]
|
|
buf[i] = e
|
|
|
|
# Compute differences in-place on buf
|
|
for i in range(n):
|
|
for j in range(axis_len - i - 1):
|
|
buf[j] = buf[j + 1] - buf[j]
|
|
|
|
for i in range(axis_len - n):
|
|
idx = util.tuple_insert(idx0, axis, i)
|
|
result._ptr(idx)[0] = buf[i]
|
|
|
|
util.free(buf)
|
|
return result
|
|
|
|
def _gradient_helper_scaler(f, dx, dtype: type, axis: int, edge_order: int):
|
|
if isinstance(dx, datetime64):
|
|
compile_error("spacing datatype cannot be of type datetime")
|
|
|
|
if not (isinstance(f.dtype, timedelta64) or isinstance(dx, timedelta64)):
|
|
vtype = type(util.coerce(dtype, type(dx)))
|
|
if isinstance(dx, vtype):
|
|
dxx = util.cast(dx, dtype)
|
|
else:
|
|
dxx = util.cast(dx, vtype)
|
|
p = util.cast(-1.5, dtype)
|
|
q = util.cast(2.0, dtype)
|
|
r = util.cast(-0.5, dtype)
|
|
else:
|
|
dxx = util.cast(dx, float)
|
|
p = util.cast(-1.5, float)
|
|
q = util.cast(2.0, float)
|
|
r = util.cast(-0.5, float)
|
|
|
|
if f.shape[axis] < edge_order + 1:
|
|
raise ValueError(
|
|
"Shape of array too small to calculate a numerical gradient, "
|
|
"at least (edge_order + 1) elements are required.")
|
|
|
|
axis_len = f.shape[axis]
|
|
result = empty(f.shape, dtype=dtype)
|
|
for idx0 in util.multirange(util.tuple_delete(f.shape, axis)):
|
|
if edge_order == 1:
|
|
#for the first element
|
|
idx = util.tuple_insert(idx0, axis, 0)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, 1)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
result._ptr(idx)[0] = util.cast((x1 - x0) / dxx, dtype)
|
|
|
|
#for the last element
|
|
idx = util.tuple_insert(idx0, axis, axis_len - 1)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, axis_len - 2)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
result._ptr(idx)[0] = util.cast((x0 - x1) / dxx, dtype)
|
|
else:
|
|
#for the first element
|
|
idx = util.tuple_insert(idx0, axis, 0)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, 1)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
idx3 = util.tuple_insert(idx0, axis, 2)
|
|
x2 = util.cast(f._ptr(idx3)[0], dtype)
|
|
d = p / dxx
|
|
b = q / dxx
|
|
c = r / dxx
|
|
result._ptr(idx)[0] = util.cast(d * x0 + b * x1 + c * x2, dtype)
|
|
|
|
#for the last element
|
|
idx = util.tuple_insert(idx0, axis, axis_len - 1)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, axis_len - 2)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
idx3 = util.tuple_insert(idx0, axis, axis_len - 3)
|
|
x2 = util.cast(f._ptr(idx3)[0], dtype)
|
|
result._ptr(idx)[0] = util.cast((-c) * x2 + (-b) * x1 + (-d) * x0,
|
|
dtype)
|
|
|
|
for i in range(0, axis_len - 2):
|
|
idx = util.tuple_insert(idx0, axis, i + 1)
|
|
idx1 = util.tuple_insert(idx0, axis, i)
|
|
x0 = util.cast(f._ptr(idx1)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, i + 2)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
result._ptr(idx)[0] = util.cast(
|
|
(x1 - x0) / (util.cast(2, dtype) * dxx), dtype)
|
|
|
|
return result
|
|
|
|
def _gradient_helper_linear(f, dx, dtype: type, axis: int, edge_order: int):
|
|
if isinstance(dx, ndarray):
|
|
if dx.ndim == 0:
|
|
return _gradient_helper_scaler(f, dx.item(), dtype, axis,
|
|
edge_order)
|
|
dx = asarray(dx)
|
|
|
|
if isinstance(dx.dtype, datetime64):
|
|
compile_error("spacing datatype cannot be of type datetime")
|
|
|
|
if not (isinstance(f.dtype, timedelta64)
|
|
or isinstance(dx.dtype, timedelta64)):
|
|
vtype = type(util.coerce(dtype, dx.dtype))
|
|
if isinstance(vtype, dx.dtype):
|
|
dxx = dx.astype(dtype)
|
|
else:
|
|
dxx = dx.astype(vtype)
|
|
const = util.cast(2, dtype)
|
|
else:
|
|
dxx = dx.astype(float)
|
|
const = util.cast(2, float)
|
|
|
|
if dxx.ndim != 1:
|
|
compile_error("distances must be either scalars or 1d")
|
|
elif len(dxx) != f.shape[axis]:
|
|
raise ValueError("when 1d, distances must match "
|
|
"the length of the corresponding dimension")
|
|
|
|
if f.shape[axis] < edge_order + 1:
|
|
raise ValueError(
|
|
"Shape of array too small to calculate a numerical gradient, "
|
|
"at least (edge_order + 1) elements are required.")
|
|
|
|
val = len(dxx)
|
|
|
|
uniform_spacing = True
|
|
m = dxx._ptr((1, ))[0] - dxx._ptr((0, ))[0]
|
|
for i in range(val - 1):
|
|
if (m != (dxx._ptr((i + 1, ))[0] - dxx._ptr((i, ))[0])):
|
|
uniform_spacing = False
|
|
break
|
|
|
|
if uniform_spacing:
|
|
dval = m
|
|
return _gradient_helper_scaler(f, dval, dtype, axis, edge_order)
|
|
|
|
axis_len = f.shape[axis]
|
|
result = empty(f.shape, dtype=dtype)
|
|
for idx0 in util.multirange(util.tuple_delete(f.shape, axis)):
|
|
if edge_order == 1:
|
|
idx = util.tuple_insert(idx0, axis, 0)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, 1)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
result._ptr(idx)[0] = util.cast((x1 - x0) / (dxx._ptr(
|
|
(1, ))[0] - dxx._ptr((0, ))[0]), dtype)
|
|
idx = util.tuple_insert(idx0, axis, axis_len - 1)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, axis_len - 2)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
result._ptr(idx)[0] = util.cast((x0 - x1) / (dxx._ptr(
|
|
(dxx.size - 1, ))[0] - dxx._ptr((dxx.size - 2, ))[0]), dtype)
|
|
else:
|
|
idx = util.tuple_insert(idx0, axis, 0)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, 1)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
idx3 = util.tuple_insert(idx0, axis, 2)
|
|
x2 = util.cast(f._ptr(idx3)[0], dtype)
|
|
dx1 = dxx._ptr((1, ))[0] - dxx._ptr((0, ))[0]
|
|
dx2 = dxx._ptr((2, ))[0] - dxx._ptr((1, ))[0]
|
|
d = -(const * dx1 + dx2) / (dx1 * (dx1 + dx2))
|
|
b = (dx1 + dx2) / (dx1 * dx2)
|
|
c = -dx1 / (dx2 * (dx1 + dx2))
|
|
result._ptr(idx)[0] = util.cast(d * x0 + b * x1 + c * x2, dtype)
|
|
|
|
idx = util.tuple_insert(idx0, axis, axis_len - 1)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, axis_len - 2)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
idx3 = util.tuple_insert(idx0, axis, axis_len - 3)
|
|
x2 = util.cast(f._ptr(idx3)[0], dtype)
|
|
dx1 = dxx._ptr((dxx.size - 2, ))[0] - dxx._ptr((dxx.size - 3, ))[0]
|
|
dx2 = dxx._ptr((dxx.size - 1, ))[0] - dxx._ptr((dxx.size - 2, ))[0]
|
|
d = (dx2) / (dx1 * (dx1 + dx2))
|
|
b = -(dx2 + dx1) / (dx1 * dx2)
|
|
c = (const * dx2 + dx1) / (dx2 * (dx1 + dx2))
|
|
result._ptr(idx)[0] = util.cast(d * x2 + b * x1 + c * x0, dtype)
|
|
|
|
for i in range(0, axis_len - 2):
|
|
idx = util.tuple_insert(idx0, axis, i)
|
|
x0 = util.cast(f._ptr(idx)[0], dtype)
|
|
idx2 = util.tuple_insert(idx0, axis, i + 1)
|
|
x1 = util.cast(f._ptr(idx2)[0], dtype)
|
|
idx3 = util.tuple_insert(idx0, axis, i + 2)
|
|
x2 = util.cast(f._ptr(idx3)[0], dtype)
|
|
dx1 = dxx._ptr((i + 1, ))[0] - dxx._ptr((i, ))[0]
|
|
dx2 = dxx._ptr((i + 2, ))[0] - dxx._ptr((i + 1, ))[0]
|
|
d = -(dx2) / (dx1 * (dx1 + dx2))
|
|
b = (dx2 - dx1) / (dx1 * dx2)
|
|
c = dx1 / (dx2 * (dx1 + dx2))
|
|
result._ptr(idx2)[0] = util.cast(d * x0 + b * x1 + c * x2, dtype)
|
|
|
|
return result
|
|
|
|
def gradient(f, *varargs, axis=None, edge_order: int = 1):
|
|
|
|
def gradient_dtype(dtype: type):
|
|
if (dtype is int or dtype is byte or isinstance(dtype, Int)
|
|
or isinstance(dtype, UInt)):
|
|
return float()
|
|
else:
|
|
return dtype()
|
|
|
|
f = asarray(f)
|
|
|
|
if isinstance(f.dtype, datetime64):
|
|
g = f.view(timedelta64[f.dtype.base, f.dtype.num])
|
|
return gradient(g, *varargs, axis=axis, edge_order=edge_order)
|
|
|
|
rdtype = type(gradient_dtype(f.dtype))
|
|
result = empty(f.shape, dtype=rdtype)
|
|
|
|
if edge_order > 2:
|
|
raise ValueError("'edge_order' greater than 2 not supported")
|
|
|
|
if axis is None:
|
|
outval = []
|
|
dval = util.cast(1, rdtype)
|
|
if staticlen(varargs) == 0:
|
|
for i in staticrange(f.ndim):
|
|
result = _gradient_helper_scaler(f,
|
|
dval,
|
|
rdtype,
|
|
axis=i,
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
elif staticlen(varargs) == 1:
|
|
if f.ndim != 1 and (isinstance(varargs[0], List)
|
|
or isinstance(varargs[0], Tuple)):
|
|
compile_error("invalid number of arguments")
|
|
|
|
if f.ndim != 1 and isinstance(varargs[0], ndarray):
|
|
if varargs[0].ndim != 0:
|
|
compile_error("invalid number of arguments")
|
|
|
|
for i in staticrange(f.ndim):
|
|
if (isinstance(varargs[0], List)
|
|
or isinstance(varargs[0], Tuple)
|
|
or isinstance(varargs[0], ndarray)):
|
|
result = _gradient_helper_linear(f,
|
|
varargs[0],
|
|
rdtype,
|
|
axis=i,
|
|
edge_order=edge_order)
|
|
else:
|
|
result = _gradient_helper_scaler(f,
|
|
varargs[0],
|
|
rdtype,
|
|
axis=i,
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
else:
|
|
if f.ndim == staticlen(varargs):
|
|
for i in staticrange(f.ndim):
|
|
if (isinstance(varargs[i], List)
|
|
or isinstance(varargs[i], Tuple)
|
|
or isinstance(varargs[i], ndarray)):
|
|
result = _gradient_helper_linear(f,
|
|
varargs[i],
|
|
rdtype,
|
|
axis=i,
|
|
edge_order=edge_order)
|
|
else:
|
|
result = _gradient_helper_scaler(f,
|
|
varargs[i],
|
|
rdtype,
|
|
axis=i,
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
else:
|
|
compile_error("invalid number of arguments")
|
|
if f.ndim == 1:
|
|
return outval[0]
|
|
return outval
|
|
else:
|
|
if isinstance(axis, Tuple):
|
|
ax = util.normalize_axis_tuple(axis, f.ndim)
|
|
outval = []
|
|
dval = util.cast(1, rdtype)
|
|
|
|
if staticlen(varargs) == 0:
|
|
for i in range(staticlen(ax)):
|
|
result = _gradient_helper_scaler(f,
|
|
dval,
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
elif staticlen(varargs) == 1:
|
|
if staticlen(ax) != 1 and (isinstance(varargs[0], List)
|
|
or isinstance(varargs[0], Tuple)):
|
|
compile_error("invalid number of arguments")
|
|
if staticlen(ax) != 1 and isinstance(varargs[0], ndarray):
|
|
if varargs[0].ndim != 0:
|
|
compile_error("invalid number of arguments")
|
|
for i in range(staticlen(ax)):
|
|
if (isinstance(varargs[0], List)
|
|
or isinstance(varargs[0], Tuple)
|
|
or isinstance(varargs[0], ndarray)):
|
|
result = _gradient_helper_linear(f,
|
|
varargs[0],
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
else:
|
|
result = _gradient_helper_scaler(f,
|
|
varargs[0],
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
else:
|
|
if staticlen(ax) == staticlen(varargs):
|
|
for i in range(staticlen(ax)):
|
|
if (isinstance(varargs[i], List)
|
|
or isinstance(varargs[i], Tuple)
|
|
or isinstance(varargs[i], ndarray)):
|
|
result = _gradient_helper_linear(
|
|
f,
|
|
varargs[i],
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
else:
|
|
result = _gradient_helper_scaler(
|
|
f,
|
|
varargs[i],
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
else:
|
|
compile_error("invalid number of arguments")
|
|
|
|
if staticlen(ax) == 1:
|
|
return outval[0]
|
|
return outval
|
|
elif isinstance(axis, List[int]):
|
|
ax = util.normalize_axis_tuple(axis, f.ndim)
|
|
outval = []
|
|
dval = util.cast(1, rdtype)
|
|
|
|
if staticlen(varargs) == 0:
|
|
for i in range(len(ax)):
|
|
result = _gradient_helper_scaler(f,
|
|
dval,
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
elif staticlen(varargs) == 1:
|
|
if len(ax) != 1 and (isinstance(varargs[0], List)
|
|
or isinstance(varargs[0], Tuple)):
|
|
raise ValueError("invalid number of arguments")
|
|
if len(ax) != 1 and isinstance(varargs[0], ndarray):
|
|
if varargs[0].ndim != 0:
|
|
compile_error("invalid number of arguments")
|
|
|
|
for i in range(len(ax)):
|
|
if (isinstance(varargs[0], List)
|
|
or isinstance(varargs[0], Tuple)
|
|
or isinstance(varargs[0], ndarray)):
|
|
result = _gradient_helper_linear(f,
|
|
varargs[0],
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
else:
|
|
result = _gradient_helper_scaler(f,
|
|
varargs[0],
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
else:
|
|
if len(ax) == staticlen(varargs):
|
|
for i in range(len(ax)):
|
|
if (isinstance(varargs[i], List)
|
|
or isinstance(varargs[i], Tuple)
|
|
or isinstance(varargs[i], ndarray)):
|
|
result = _gradient_helper_linear(
|
|
f,
|
|
varargs[i],
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
else:
|
|
result = _gradient_helper_scaler(
|
|
f,
|
|
varargs[i],
|
|
rdtype,
|
|
axis=ax[i],
|
|
edge_order=edge_order)
|
|
outval.append(result)
|
|
else:
|
|
raise ValueError("invalid number of arguments")
|
|
return outval
|
|
else:
|
|
axis = util.normalize_axis_index(axis, f.ndim)
|
|
|
|
if staticlen(varargs) == 0:
|
|
return _gradient_helper_scaler(f,
|
|
dx=1,
|
|
dtype=rdtype,
|
|
axis=axis,
|
|
edge_order=edge_order)
|
|
else:
|
|
if staticlen(varargs) == 1:
|
|
if (isinstance(varargs[0], List)
|
|
or isinstance(varargs[0], Tuple)
|
|
or isinstance(varargs[0], ndarray)):
|
|
return _gradient_helper_linear(f,
|
|
varargs[0],
|
|
rdtype,
|
|
axis=axis,
|
|
edge_order=edge_order)
|
|
else:
|
|
return _gradient_helper_scaler(f,
|
|
varargs[0],
|
|
rdtype,
|
|
axis=axis,
|
|
edge_order=edge_order)
|
|
else:
|
|
compile_error("invalid number of arguments")
|
|
|
|
def ediff1d(ary, to_begin=None, to_end=None):
|
|
a = asarray(ary).ravel()
|
|
n = a.size - 1
|
|
|
|
if to_begin is not None:
|
|
begin = asarray(to_begin).ravel()
|
|
n += begin.size
|
|
else:
|
|
begin = None
|
|
|
|
if to_end is not None:
|
|
end = asarray(to_end).ravel()
|
|
n += end.size
|
|
else:
|
|
end = None
|
|
|
|
ans = empty(n, a.dtype)
|
|
k = 0
|
|
|
|
if to_begin is not None:
|
|
for i in range(begin.size):
|
|
ans._ptr((i, ))[0] = util.cast(begin._ptr((i, ))[0], a.dtype)
|
|
k = begin.size
|
|
|
|
for i in range(a.size - 1):
|
|
ans._ptr((k + i, ))[0] = a._ptr((i + 1, ))[0] - a._ptr((i, ))[0]
|
|
k += a.size - 1
|
|
|
|
if to_end is not None:
|
|
for i in range(end.size):
|
|
ans._ptr((k + i, ))[0] = util.cast(end._ptr((i, ))[0], a.dtype)
|
|
|
|
return ans
|
|
|
|
def cross(a, b, axisa: int = -1, axisb: int = -1, axisc: int = -1, axis = None):
|
|
if axis is not None:
|
|
axisa, axisb, axisc = (axis, ) * 3
|
|
a = asarray(a)
|
|
b = asarray(b)
|
|
|
|
# Check axisa and axisb are within bounds
|
|
axisa = util.normalize_axis_index(axisa, a.ndim)
|
|
axisb = util.normalize_axis_index(axisb, b.ndim)
|
|
|
|
# Move working axis to the end of the shape
|
|
a = moveaxis(a, axisa, -1)
|
|
b = moveaxis(b, axisb, -1)
|
|
msg = ("incompatible dimensions for cross product\n"
|
|
"(dimension must be 2 or 3)")
|
|
if a.shape[-1] not in (2, 3) or b.shape[-1] not in (2, 3):
|
|
raise ValueError(msg)
|
|
|
|
# Create the output array
|
|
x = a[..., 0]
|
|
x = asarray(x)
|
|
y = b[..., 0]
|
|
y = asarray(y)
|
|
shape = broadcast_shapes(x.shape, y.shape)
|
|
|
|
if a.shape[-1] == 3 or b.shape[-1] == 3:
|
|
new_shape = shape + (3, )
|
|
# Check axisc is within bounds
|
|
axisc = util.normalize_axis_index(axisc, len(new_shape))
|
|
else:
|
|
new_shape = shape + (1, )
|
|
|
|
dtype = type(util.coerce(a.dtype, b.dtype))
|
|
|
|
cp = zeros(new_shape, dtype)
|
|
|
|
# recast arrays as dtype
|
|
a = a.astype(dtype)
|
|
b = b.astype(dtype)
|
|
|
|
# create local aliases for readability
|
|
a0 = a[..., 0]
|
|
a1 = a[..., 1]
|
|
if a.shape[-1] == 3:
|
|
a2 = a[..., 2]
|
|
b0 = b[..., 0]
|
|
b1 = b[..., 1]
|
|
if b.shape[-1] == 3:
|
|
b2 = b[..., 2]
|
|
if cp.ndim != 0 and cp.shape[-1] == 3:
|
|
|
|
cp0 = cp[..., 0]
|
|
cp1 = cp[..., 1]
|
|
cp2 = cp[..., 2]
|
|
|
|
if a.shape[-1] == 2:
|
|
if b.shape[-1] == 2:
|
|
# a0 * b1 - a1 * b0
|
|
cp[0] = a[..., 0] * b[..., 1] - a[..., 1] * b[
|
|
..., 0] # cp2 = a0 * b1 - a1 * b0
|
|
return cp
|
|
|
|
else:
|
|
assert b.shape[-1] == 3
|
|
|
|
# cp0 = 0 - a1 * b2 (a2 = 0)
|
|
cp[..., 0] = -a[..., 1] * b[..., 2]
|
|
|
|
# cp1 = a0 * b2 - 0 (a2 = 0)
|
|
cp[..., 1] = a[..., 0] * b[..., 2]
|
|
|
|
# cp2 = a0 * b1 - a1 * b0
|
|
cp[..., 2] = a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0]
|
|
|
|
else:
|
|
assert a.shape[-1] == 3
|
|
if b.shape[-1] == 3:
|
|
|
|
# cp0 = a1 * b2 - a2 * b1
|
|
cp[..., 0] = a[..., 1] * b[..., 2] - a[..., 2] * b[..., 1]
|
|
|
|
# cp1 = a2 * b0 - a0 * b2
|
|
cp[..., 1] = a[..., 2] * b[..., 0] - a[..., 0] * b[..., 2]
|
|
|
|
# cp2 = a0 * b1 - a1 * b0
|
|
cp[..., 2] = a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0]
|
|
|
|
else:
|
|
assert b.shape[-1] == 2
|
|
|
|
# cp0 = a2 * b1 - 0 (b2 = 0)
|
|
cp[..., 0] = a[..., 2] * b[..., 1]
|
|
|
|
# cp1 = 0 - a2 * b0 (b2 = 0)
|
|
cp[..., 1] = -a[..., 2] * b[..., 0]
|
|
|
|
# cp2 = a0 * b1 - a1 * b0
|
|
cp[..., 2] = a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0]
|
|
|
|
return moveaxis(cp, -1, axisc)
|
|
|
|
def trapz(y, x=None, dx=1.0, axis: int = -1):
|
|
def trapz_helper(a, axis, dx, dtype: type):
|
|
axis = util.normalize_axis_index(axis, a.ndim)
|
|
axis_len = a.shape[axis]
|
|
new_shape = util.tuple_delete(a.shape, axis)
|
|
|
|
# Create empty array
|
|
result = empty(new_shape, dtype=dtype)
|
|
|
|
# Iterate over all indices except the specified axis
|
|
for idx0 in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
ans = dtype()
|
|
|
|
for i in range(axis_len - 1):
|
|
idx = util.tuple_insert(idx0, axis, i)
|
|
y0 = a._ptr(idx)[0]
|
|
|
|
idx = util.tuple_insert(idx0, axis, i + 1)
|
|
y1 = a._ptr(idx)[0]
|
|
|
|
ans += y1 + y0
|
|
|
|
result._ptr(idx0)[0] = ans * dx / dtype(2)
|
|
return result
|
|
|
|
def trapz_helper_with_x(a, axis, x, dtype: type):
|
|
axis = util.normalize_axis_index(axis, a.ndim)
|
|
axis_len = a.shape[axis]
|
|
new_shape = util.tuple_delete(a.shape, axis)
|
|
|
|
# Create empty array
|
|
result = empty(new_shape, dtype=dtype)
|
|
|
|
x = broadcast_to(x, a.shape)
|
|
|
|
# Iterate over all indices except the specified axis
|
|
for idx0 in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
ans = dtype()
|
|
|
|
for i in range(axis_len - 1):
|
|
idx = util.tuple_insert(idx0, axis, i)
|
|
x0 = x._ptr(idx)[0]
|
|
y0 = a._ptr(idx)[0]
|
|
|
|
idx = util.tuple_insert(idx0, axis, i + 1)
|
|
x1 = x._ptr(idx)[0]
|
|
y1 = a._ptr(idx)[0]
|
|
|
|
ans += (y1 + y0) * (x1 - x0)
|
|
|
|
result._ptr(idx0)[0] = ans / dtype(2)
|
|
return result
|
|
|
|
def trapz_helper_with_x_1d(a, axis, x, dtype: type):
|
|
axis = util.normalize_axis_index(axis, a.ndim)
|
|
axis_len = a.shape[axis]
|
|
|
|
if len(x) != axis_len:
|
|
raise ValueError(
|
|
"length of 'x' does not match dimension along axis")
|
|
|
|
new_shape = util.tuple_delete(a.shape, axis)
|
|
|
|
# Create empty array
|
|
result = empty(new_shape, dtype=dtype)
|
|
|
|
# Iterate over all indices except the specified axis
|
|
for idx0 in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
ans = dtype()
|
|
|
|
for i in range(axis_len - 1):
|
|
idx = util.tuple_insert(idx0, axis, i)
|
|
x0 = x._ptr((i, ))[0]
|
|
y0 = a._ptr(idx)[0]
|
|
|
|
idx = util.tuple_insert(idx0, axis, i + 1)
|
|
x1 = x._ptr((i + 1, ))[0]
|
|
y1 = a._ptr(idx)[0]
|
|
|
|
ans += (y1 + y0) * (x1 - x0)
|
|
|
|
result._ptr(idx0)[0] = ans / dtype(2)
|
|
return result
|
|
|
|
def trapz_dtype(dtype: type):
|
|
if (dtype is int or dtype is byte or isinstance(dtype, Int)
|
|
or isinstance(dtype, UInt)):
|
|
return float()
|
|
else:
|
|
return dtype()
|
|
|
|
def zero_dim_to_scalar(arr: ndarray):
|
|
if arr.ndim == 0:
|
|
return arr.item()
|
|
else:
|
|
return arr
|
|
|
|
y = asarray(y)
|
|
dtype = type(trapz_dtype(y.dtype))
|
|
axis = util.normalize_axis_index(axis, y.ndim)
|
|
if x is None:
|
|
return zero_dim_to_scalar(
|
|
trapz_helper(y, axis=axis, dx=dx, dtype=dtype))
|
|
else:
|
|
x = asarray(x)
|
|
if x.ndim == 1:
|
|
return zero_dim_to_scalar(
|
|
trapz_helper_with_x_1d(y, axis=axis, x=x, dtype=dtype))
|
|
else:
|
|
return zero_dim_to_scalar(
|
|
trapz_helper_with_x(y, axis=axis, x=x, dtype=dtype))
|
|
|
|
def convolve(v, a, mode: str = "full"):
|
|
# Convert inputs to NumPy arrays
|
|
v = atleast_1d(asarray(v))
|
|
a = atleast_1d(asarray(a))
|
|
|
|
if len(v) < len(a):
|
|
return convolve(a, v, mode=mode)
|
|
|
|
res_dtype = type(util.coerce(v.dtype, a.dtype))
|
|
|
|
if v.size == 0 or a.size == 0:
|
|
raise ValueError("Input arrays cannot be empty")
|
|
|
|
if a.ndim > 1 or v.ndim > 1:
|
|
compile_error(
|
|
"incompatible dimensions for convolution\n(dimension must be 0 or 1)"
|
|
)
|
|
|
|
a = a[::-1]
|
|
|
|
length = len(v)
|
|
n = len(a)
|
|
n_left = 0
|
|
n_right = 0
|
|
|
|
if mode == "valid":
|
|
reslen = length - n + 1
|
|
elif mode == "same":
|
|
if n % 2 == 1:
|
|
n_left = n // 2
|
|
n_right = n // 2
|
|
desc = n - n_left - 1
|
|
else:
|
|
n_left = n // 2
|
|
n_right = n - n_left - 1
|
|
desc = n - n_left
|
|
reslen = length
|
|
|
|
elif mode == "full":
|
|
n_right = n - 1
|
|
n_left = n - 1
|
|
desc = n - 1
|
|
reslen = length + n - 1
|
|
else:
|
|
raise ValueError(
|
|
f"mode must be one of 'valid', 'same', or 'full' (got {repr(mode)})"
|
|
)
|
|
|
|
# Allocate result array
|
|
ret = zeros(reslen, dtype=res_dtype)
|
|
for i in range(n_left):
|
|
k = desc - i
|
|
l = 0
|
|
for j in range(k, n, 1):
|
|
ret._ptr((i, ))[0] += util.cast(v._ptr(
|
|
(l, ))[0], res_dtype) * util.cast(a._ptr((j, ))[0], res_dtype)
|
|
l += 1
|
|
|
|
b = reslen - n_right
|
|
for i in range(n_left, b):
|
|
k = i - n_left
|
|
for j in range(n):
|
|
ret._ptr((i, ))[0] += util.cast(v._ptr(
|
|
(k, ))[0], res_dtype) * util.cast(a._ptr((j, ))[0], res_dtype)
|
|
k += 1
|
|
|
|
c = length - n + 1
|
|
for i in range(n_right):
|
|
k = length - (n - i) + 1
|
|
for j in range(n - i - 1):
|
|
ret._ptr((n_left + c +
|
|
i, ))[0] += util.cast(v._ptr(
|
|
(k, ))[0], res_dtype) * util.cast(
|
|
a._ptr((j, ))[0], res_dtype)
|
|
k += 1
|
|
|
|
return ret
|
|
|
|
def _apply(x, f, f_complex = None):
|
|
if f_complex is not None and (isinstance(x, complex) or isinstance(x, complex64)):
|
|
return f_complex(x)
|
|
elif isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
|
|
return f(x)
|
|
else:
|
|
return f(util.to_float(x))
|
|
|
|
def _apply2(x, y, f, f_complex = None):
|
|
if type(x) is not type(y):
|
|
compile_error("[internal error] type mismatch")
|
|
|
|
if f_complex is not None and (isinstance(x, complex) or isinstance(x, complex64)):
|
|
return f_complex(x, y)
|
|
elif isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
|
|
return f(x, y)
|
|
else:
|
|
return f(util.to_float(x), util.to_float(y))
|
|
|
|
def _fabs(x):
|
|
return _apply(x, util.fabs)
|
|
|
|
def _rint(x):
|
|
def rint_complex(x):
|
|
C = type(x)
|
|
return C(util.rint(x.real), util.rint(x.imag))
|
|
|
|
return _apply(x, util.rint, rint_complex)
|
|
|
|
def _exp(x):
|
|
return _apply(x, util.exp, zmath.exp)
|
|
|
|
def _exp2(x):
|
|
return _apply(x, util.exp2, zmath.exp2)
|
|
|
|
def _expm1(x):
|
|
return _apply(x, util.expm1, zmath.expm1)
|
|
|
|
def _log(x):
|
|
return _apply(x, util.log, zmath.log)
|
|
|
|
def _log2(x):
|
|
return _apply(x, util.log2, zmath.log2)
|
|
|
|
def _log10(x):
|
|
return _apply(x, util.log10, zmath.log10)
|
|
|
|
def _log1p(x):
|
|
return _apply(x, util.log1p, zmath.log1p)
|
|
|
|
def _sqrt(x):
|
|
return _apply(x, util.sqrt, zmath.sqrt)
|
|
|
|
def _cbrt(x):
|
|
return _apply(x, util.cbrt)
|
|
|
|
def _square(x):
|
|
return x * x
|
|
|
|
def _sin(x):
|
|
return _apply(x, util.sin, zmath.sin)
|
|
|
|
def _cos(x):
|
|
return _apply(x, util.cos, zmath.cos)
|
|
|
|
def _tan(x):
|
|
return _apply(x, util.tan, zmath.tan)
|
|
|
|
def _arcsin(x):
|
|
return _apply(x, util.asin, zmath.asin)
|
|
|
|
def _arccos(x):
|
|
return _apply(x, util.acos, zmath.acos)
|
|
|
|
def _arctan(x):
|
|
return _apply(x, util.atan, zmath.atan)
|
|
|
|
def _sinh(x):
|
|
return _apply(x, util.sinh, zmath.sinh)
|
|
|
|
def _cosh(x):
|
|
return _apply(x, util.cosh, zmath.cosh)
|
|
|
|
def _tanh(x):
|
|
return _apply(x, util.tanh, zmath.tanh)
|
|
|
|
def _arcsinh(x):
|
|
return _apply(x, util.asinh, zmath.asinh)
|
|
|
|
def _arccosh(x):
|
|
return _apply(x, util.acosh, zmath.acosh)
|
|
|
|
def _arctanh(x):
|
|
return _apply(x, util.atanh, zmath.atanh)
|
|
|
|
def _rad2deg(x):
|
|
r2d = 180.0 / util.PI
|
|
x = util.to_float(x)
|
|
F = type(x)
|
|
return x * F(r2d)
|
|
|
|
def _deg2rad(x):
|
|
d2r = util.PI / 180.0
|
|
x = util.to_float(x)
|
|
F = type(x)
|
|
return x * F(d2r)
|
|
|
|
def _arctan2(x, y):
|
|
return _apply2(x, y, util.atan2)
|
|
|
|
def _hypot(x, y):
|
|
return _apply2(x, y, util.hypot)
|
|
|
|
def _logaddexp(x, y):
|
|
return _apply2(x, y, util.logaddexp)
|
|
|
|
def _logaddexp2(x, y):
|
|
return _apply2(x, y, util.logaddexp2)
|
|
|
|
def _isnan(x):
|
|
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
|
|
return util.isnan(x)
|
|
elif isinstance(x, complex) or isinstance(x, complex64):
|
|
return util.isnan(x.real) or util.isnan(x.imag)
|
|
elif isinstance(x, datetime64) or isinstance(x, timedelta64):
|
|
return x._nat
|
|
else:
|
|
return False
|
|
|
|
def _isinf(x):
|
|
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
|
|
return util.isinf(x)
|
|
elif isinstance(x, complex) or isinstance(x, complex64):
|
|
return util.isinf(x.real) or util.isinf(x.imag)
|
|
else:
|
|
return False
|
|
|
|
def _isfinite(x):
|
|
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
|
|
return util.isfinite(x)
|
|
elif isinstance(x, complex) or isinstance(x, complex64):
|
|
return util.isfinite(x.real) and util.isfinite(x.imag)
|
|
elif isinstance(x, datetime64) or isinstance(x, timedelta64):
|
|
return not x._nat
|
|
else:
|
|
return True
|
|
|
|
def _signbit(x):
|
|
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
|
|
return util.signbit(x)
|
|
else:
|
|
T = type(x)
|
|
return x < T()
|
|
|
|
def _copysign(x, y):
|
|
return _apply2(x, y, util.copysign)
|
|
|
|
def _nextafter(x, y):
|
|
return _apply2(x, y, util.nextafter)
|
|
|
|
def _ldexp(x, y):
|
|
y = util.cast(y, int)
|
|
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
|
|
return util.ldexp(x, y)
|
|
else:
|
|
return util.ldexp(util.to_float(x), y)
|
|
|
|
def _floor(x):
|
|
return _apply(x, util.floor)
|
|
|
|
def _ceil(x):
|
|
return _apply(x, util.ceil)
|
|
|
|
def _trunc(x):
|
|
return _apply(x, util.trunc)
|
|
|
|
def _sign(x):
|
|
def sign1(x):
|
|
T = type(x)
|
|
if x < T(0):
|
|
return T(-1)
|
|
elif x > T(0):
|
|
return T(1)
|
|
else:
|
|
return x
|
|
|
|
if isinstance(x, complex):
|
|
if _isnan(x):
|
|
return complex(util.nan64(), 0.0)
|
|
return complex(sign1(x.real), 0) if x.real else complex(sign1(x.imag), 0)
|
|
elif isinstance(x, complex64):
|
|
if _isnan(x):
|
|
return complex64(util.nan64(), 0.0)
|
|
return complex64(sign1(x.real), 0) if x.real else complex64(sign1(x.imag), 0)
|
|
elif isinstance(x, timedelta64):
|
|
return x._sign
|
|
else:
|
|
return sign1(x)
|
|
|
|
def _heaviside(x, y):
|
|
def heaviside(x, y):
|
|
if isinstance(x, float16) and isinstance(y, float16):
|
|
if x < float16(0):
|
|
return float16(0)
|
|
elif x > float16(0):
|
|
return float16(1)
|
|
elif x == float16(0):
|
|
return y
|
|
else:
|
|
return x
|
|
elif isinstance(x, float32) and isinstance(y, float32):
|
|
if x < float32(0):
|
|
return float32(0)
|
|
elif x > float32(0):
|
|
return float32(1)
|
|
elif x == float32(0):
|
|
return y
|
|
else:
|
|
return x
|
|
elif isinstance(x, float) and isinstance(y, float):
|
|
if x < 0:
|
|
return 0.0
|
|
elif x > 0:
|
|
return 1.0
|
|
elif x == 0.0:
|
|
return y
|
|
else:
|
|
return x
|
|
|
|
return _apply2(x, y, heaviside)
|
|
|
|
def _conj(x):
|
|
if isinstance(x, complex) or isinstance(x, complex64):
|
|
return x.conjugate()
|
|
else:
|
|
return x
|
|
|
|
def _gcd(x, y):
|
|
''' # fails with optionals
|
|
if not (
|
|
isinstance(x, int) or
|
|
isinstance(x, Int) or
|
|
isinstance(x, UInt) or
|
|
isinstance(x, byte)
|
|
):
|
|
compile_error("gcd/lcm can only be used on integral types")
|
|
'''
|
|
|
|
while x:
|
|
z = x
|
|
x = y % x
|
|
y = z
|
|
return y
|
|
|
|
def _lcm(x, y):
|
|
gcd = _gcd(x, y)
|
|
return x // gcd * y if gcd else 0
|
|
|
|
def _reciprocal(x: T, T: type):
|
|
if (
|
|
isinstance(x, int) or
|
|
isinstance(x, Int) or
|
|
isinstance(x, UInt) or
|
|
isinstance(x, byte)
|
|
):
|
|
return T(1) // x
|
|
else:
|
|
return T(1) / x
|
|
|
|
def _logical_and(x, y):
|
|
return bool(x) and bool(y)
|
|
|
|
def _logical_or(x, y):
|
|
return bool(x) or bool(y)
|
|
|
|
def _logical_xor(x, y):
|
|
return bool(x) ^ bool(y)
|
|
|
|
def _logical_not(x):
|
|
return not bool(x)
|
|
|
|
def _coerce_types_for_minmax(x, y):
|
|
if isinstance(x, complex):
|
|
if isinstance(y, complex64):
|
|
return x, complex(y)
|
|
elif not isinstance(y, complex):
|
|
return x, complex(util.cast(y, float))
|
|
elif isinstance(x, complex64):
|
|
if isinstance(y, complex):
|
|
return complex(x), y
|
|
elif not isinstance(y, complex64):
|
|
return complex(x), complex(util.cast(y, float))
|
|
|
|
if isinstance(y, complex):
|
|
if isinstance(x, complex64):
|
|
return complex(x), y
|
|
elif not isinstance(x, complex):
|
|
return complex(util.cast(x, float)), y
|
|
elif isinstance(y, complex64):
|
|
if isinstance(x, complex):
|
|
return x, complex(y)
|
|
elif not isinstance(x, complex64):
|
|
return complex(util.cast(x, float)), complex(y)
|
|
|
|
T = type(util.coerce(type(x), type(y)))
|
|
return util.cast(x, T), util.cast(y, T)
|
|
|
|
def _compare_le(x, y):
|
|
if isinstance(x, complex) or isinstance(x, complex64):
|
|
return (x.real, x.imag) <= (y.real, y.imag)
|
|
else:
|
|
return x <= y
|
|
|
|
def _compare_ge(x, y):
|
|
if isinstance(x, complex) or isinstance(x, complex64):
|
|
return (x.real, x.imag) >= (y.real, y.imag)
|
|
else:
|
|
return x >= y
|
|
|
|
def _maximum(x, y):
|
|
x, y = _coerce_types_for_minmax(x, y)
|
|
|
|
if _isnan(x):
|
|
return x
|
|
|
|
if _isnan(y):
|
|
return y
|
|
|
|
return x if _compare_ge(x, y) else y
|
|
|
|
def _minimum(x, y):
|
|
x, y = _coerce_types_for_minmax(x, y)
|
|
|
|
if _isnan(x):
|
|
return x
|
|
|
|
if _isnan(y):
|
|
return y
|
|
|
|
return x if _compare_le(x, y) else y
|
|
|
|
def _fmax(x, y):
|
|
x, y = _coerce_types_for_minmax(x, y)
|
|
|
|
if _isnan(y):
|
|
return x
|
|
|
|
if _isnan(x):
|
|
return y
|
|
|
|
return x if _compare_ge(x, y) else y
|
|
|
|
def _fmin(x, y):
|
|
x, y = _coerce_types_for_minmax(x, y)
|
|
|
|
if _isnan(y):
|
|
return x
|
|
|
|
if _isnan(x):
|
|
return y
|
|
|
|
return x if _compare_le(x, y) else y
|
|
|
|
def _divmod_float(x, y):
|
|
F = type(x)
|
|
mod = util.cmod(x, y)
|
|
if not y:
|
|
return util.cdiv(x, y), mod
|
|
|
|
div = util.cdiv(x - mod, y)
|
|
if mod:
|
|
if (y < F(0)) != (mod < F(0)):
|
|
mod += y
|
|
div -= F(1)
|
|
else:
|
|
mod = util.copysign(F(0), y)
|
|
|
|
floordiv = F()
|
|
if div:
|
|
floordiv = util.floor(div)
|
|
if div - floordiv > F(0.5):
|
|
floordiv += F(1)
|
|
else:
|
|
floordiv = util.copysign(F(0), util.cdiv(x, y))
|
|
|
|
return floordiv, mod
|
|
|
|
def _divmod(x, y):
|
|
if isinstance(x, float16) and isinstance(y, float16):
|
|
return _divmod_float(x, y)
|
|
|
|
if isinstance(x, float32) and isinstance(y, float32):
|
|
return _divmod_float(x, y)
|
|
|
|
if isinstance(x, float) or isinstance(y, float):
|
|
return _divmod_float(util.cast(x, float), util.cast(y, float))
|
|
|
|
return (x // y, x % y)
|
|
|
|
def _modf(x):
|
|
return _apply(x, util.modf)
|
|
|
|
def _frexp(x):
|
|
def frexp(x):
|
|
a, b = util.frexp(x)
|
|
return a, i32(b)
|
|
|
|
return _apply(x, frexp)
|
|
|
|
def _spacing16(h: float16):
|
|
h_u16 = util.bitcast(h, u16)
|
|
h_exp = h_u16 & u16(0x7c00)
|
|
h_sig = h_u16 & u16(0x03ff)
|
|
|
|
if h_exp == u16(0x7c00):
|
|
return util.nan16()
|
|
elif h_u16 == u16(0x7bff):
|
|
return util.inf16()
|
|
elif (h_u16 & u16(0x8000)) and not h_sig:
|
|
if h_exp > u16(0x2c00):
|
|
return util.bitcast(h_exp - u16(0x2c00), float16)
|
|
elif h_exp > u16(0x0400):
|
|
return util.bitcast(u16(1) << ((h_exp >> u16(10)) - u16(2)), float16)
|
|
else:
|
|
return util.bitcast(u16(0x0001), float16)
|
|
elif h_exp > u16(0x2800):
|
|
return util.bitcast(h_exp - u16(0x2800), float16)
|
|
elif h_exp > u16(0x0400):
|
|
return util.bitcast(u16(1) << ((h_exp >> u16(10)) - u16(1)), float16)
|
|
else:
|
|
return util.bitcast(u16(0x0001), float16)
|
|
|
|
def _spacing(x):
|
|
if isinstance(x, float16):
|
|
return _spacing16(x)
|
|
elif isinstance(x, float32):
|
|
if util.isinf32(x):
|
|
return util.nan32()
|
|
p = util.inf32() if x >= float32(0) else -util.inf32()
|
|
return util.nextafter32(x, util.inf32()) - x
|
|
elif isinstance(x, float):
|
|
x = util.cast(x, float)
|
|
if util.isinf64(x):
|
|
return util.nan64()
|
|
p = util.inf64() if x >= 0 else -util.inf64()
|
|
return util.nextafter64(x, p) - x
|
|
else:
|
|
return _spacing(util.to_float(x))
|
|
|
|
def _isnat(x):
|
|
if isinstance(x, datetime64) or isinstance(x, timedelta64):
|
|
return x._nat
|
|
else:
|
|
compile_error("ufunc 'isnat' is only defined for np.datetime64 and np.timedelta64.")
|
|
|
|
def _floor_divide(x, y):
|
|
X = type(x)
|
|
Y = type(y)
|
|
if isinstance(X, Int) and isinstance(Y, Int):
|
|
return util.pydiv(x, y)
|
|
else:
|
|
return x // y
|
|
|
|
def _remainder(x, y):
|
|
X = type(x)
|
|
Y = type(y)
|
|
if isinstance(X, Int) and isinstance(Y, Int):
|
|
return util.pymod(x, y)
|
|
elif ((X is float and Y is float) or
|
|
(X is float32 and Y is float32) or
|
|
(X is float16 and Y is float16)):
|
|
return util.pyfmod(x, y)
|
|
else:
|
|
return x % y
|
|
|
|
def _fmod(x, y):
|
|
X = type(x)
|
|
Y = type(y)
|
|
if isinstance(X, Int) and isinstance(Y, Int):
|
|
return util.cmod_int(x, y)
|
|
elif ((X is float and Y is float) or
|
|
(X is float32 and Y is float32) or
|
|
(X is float16 and Y is float16)):
|
|
return util.cmod(x, y)
|
|
else:
|
|
return x % y
|
|
|
|
add = BinaryUFunc(operator.add, 'add', 0)
|
|
subtract = BinaryUFunc(operator.sub, 'subtract')
|
|
multiply = BinaryUFunc(operator.mul, 'multiply', 1)
|
|
divide = BinaryUFunc(operator.truediv, 'divide')
|
|
true_divide = divide
|
|
logaddexp = BinaryUFunc(_logaddexp, 'logaddexp')
|
|
logaddexp2 = BinaryUFunc(_logaddexp2, 'logaddexp2')
|
|
floor_divide = BinaryUFunc(_floor_divide, 'floor_divide')
|
|
negative = UnaryUFunc(operator.neg, 'negative')
|
|
positive = UnaryUFunc(operator.pos, 'positive')
|
|
power = BinaryUFunc(operator.pow, 'power')
|
|
remainder = BinaryUFunc(_remainder, 'remainder')
|
|
mod = remainder
|
|
fmod = BinaryUFunc(_fmod, 'fmod')
|
|
# divmod
|
|
absolute = UnaryUFunc(operator.abs, 'absolute')
|
|
abs = absolute
|
|
fabs = UnaryUFunc(_fabs, 'fabs')
|
|
rint = UnaryUFunc(_rint, 'rint')
|
|
sign = UnaryUFunc(_sign, 'sign')
|
|
heaviside = BinaryUFunc(_heaviside, 'heaviside')
|
|
conjugate = UnaryUFunc(_conj, 'conjugate')
|
|
conj = conjugate
|
|
exp = UnaryUFunc(_exp, 'exp')
|
|
exp2 = UnaryUFunc(_exp2, 'exp2')
|
|
log = UnaryUFunc(_log, 'log')
|
|
log2 = UnaryUFunc(_log2, 'log2')
|
|
log10 = UnaryUFunc(_log10, 'log10')
|
|
expm1 = UnaryUFunc(_expm1, 'expm1')
|
|
log1p = UnaryUFunc(_log1p, 'log1p')
|
|
sqrt = UnaryUFunc(_sqrt, 'sqrt')
|
|
square = UnaryUFunc(_square, 'square')
|
|
cbrt = UnaryUFunc(_cbrt, 'cbrt')
|
|
reciprocal = UnaryUFunc(_reciprocal, 'reciprocal')
|
|
gcd = BinaryUFunc(_gcd, 'gcd', 0)
|
|
lcm = BinaryUFunc(_lcm, 'lcm')
|
|
|
|
sin = UnaryUFunc(_sin, 'sin')
|
|
cos = UnaryUFunc(_cos, 'cos')
|
|
tan = UnaryUFunc(_tan, 'tan')
|
|
arcsin = UnaryUFunc(_arcsin, 'arcsin')
|
|
arccos = UnaryUFunc(_arccos, 'arccos')
|
|
arctan = UnaryUFunc(_arctan, 'arctan')
|
|
arctan2 = BinaryUFunc(_arctan2, 'arctan2')
|
|
hypot = BinaryUFunc(_hypot, 'hypot')
|
|
sinh = UnaryUFunc(_sinh, 'sinh')
|
|
cosh = UnaryUFunc(_cosh, 'cosh')
|
|
tanh = UnaryUFunc(_tanh, 'tanh')
|
|
arcsinh = UnaryUFunc(_arcsinh, 'arcsinh')
|
|
arccosh = UnaryUFunc(_arccosh, 'arccosh')
|
|
arctanh = UnaryUFunc(_arctanh, 'arctanh')
|
|
deg2rad = UnaryUFunc(_deg2rad, 'deg2rad')
|
|
radians = UnaryUFunc(_deg2rad, 'radians')
|
|
rad2deg = UnaryUFunc(_rad2deg, 'rad2deg')
|
|
degrees = UnaryUFunc(_rad2deg, 'degrees')
|
|
|
|
bitwise_and = BinaryUFunc(operator.and_, 'bitwise_and')
|
|
bitwise_or = BinaryUFunc(operator.or_, 'bitwise_or')
|
|
bitwise_xor = BinaryUFunc(operator.xor, 'bitwise_xor')
|
|
invert = UnaryUFunc(operator.invert, 'invert')
|
|
left_shift = BinaryUFunc(operator.lshift, 'left_shift')
|
|
right_shift = BinaryUFunc(operator.rshift, 'right_shift')
|
|
|
|
greater = BinaryUFunc(operator.gt, 'greater')
|
|
greater_equal = BinaryUFunc(operator.ge, 'greater_equal')
|
|
less = BinaryUFunc(operator.lt, 'less')
|
|
less_equal = BinaryUFunc(operator.le, 'less_equal')
|
|
not_equal = BinaryUFunc(operator.ne, 'not_equal')
|
|
equal = BinaryUFunc(operator.eq, 'equal')
|
|
|
|
logical_and = BinaryUFunc(_logical_and, 'logical_and')
|
|
logical_or = BinaryUFunc(_logical_or, 'logical_or')
|
|
logical_xor = BinaryUFunc(_logical_xor, 'logical_xor')
|
|
logical_not = UnaryUFunc(_logical_not, 'logical_not')
|
|
|
|
maximum = BinaryUFunc(_maximum, 'maximum')
|
|
minimum = BinaryUFunc(_minimum, 'minimum')
|
|
fmax = BinaryUFunc(_fmax, 'fmax')
|
|
fmin = BinaryUFunc(_fmin, 'fmin')
|
|
|
|
isfinite = UnaryUFunc(_isfinite, 'isfinite')
|
|
isinf = UnaryUFunc(_isinf, 'isinf')
|
|
isnan = UnaryUFunc(_isnan, 'isnan')
|
|
isnat = UnaryUFunc(_isnat, 'isnat')
|
|
signbit = UnaryUFunc(_signbit, 'signbit')
|
|
copysign = BinaryUFunc(_copysign, 'copysign')
|
|
nextafter = BinaryUFunc(_nextafter, 'nextafter')
|
|
spacing = UnaryUFunc(_spacing, 'spacing')
|
|
modf = UnaryUFunc2(_modf, 'modf')
|
|
ldexp = BinaryUFunc(_ldexp, 'ldexp')
|
|
frexp = UnaryUFunc2(_frexp, 'frexp')
|
|
floor = UnaryUFunc(_floor, 'floor')
|
|
ceil = UnaryUFunc(_ceil, 'ceil')
|
|
trunc = UnaryUFunc(_trunc, 'trunc')
|
|
|
|
def i0(x):
|
|
def i0_helper(x):
|
|
if isinstance(x, float) or isinstance(x, float32) or isinstance(x, float16):
|
|
return util.i0(x)
|
|
else:
|
|
return util.i0(util.cast(x, float))
|
|
|
|
x = asarray(x)
|
|
if x.dtype is complex or x.dtype is complex64:
|
|
compile_error("i0 is not supported for complex values")
|
|
return x.map(i0_helper)
|
|
|
|
def sinc(x):
|
|
def sinc_helper(x):
|
|
if isinstance(x, float32):
|
|
if not x:
|
|
x = float32(1e-20)
|
|
x *= float32(pi)
|
|
return util.sin32(x) / x
|
|
elif isinstance(x, float16):
|
|
if not x:
|
|
x = float16(1e-7)
|
|
x *= float16(pi)
|
|
return util.sin16(x) / x
|
|
elif isinstance(x, complex):
|
|
if not x:
|
|
x = complex(1e-20)
|
|
x *= pi
|
|
return zmath.sin(x) / x
|
|
elif isinstance(x, complex64):
|
|
if not x:
|
|
x = complex64(1e-20)
|
|
x *= float32(pi)
|
|
return zmath.sin(x) / x
|
|
else:
|
|
x = util.cast(x, float)
|
|
if not x:
|
|
x = 1e-20
|
|
x *= pi
|
|
return util.sin64(x) / x
|
|
|
|
ans = asarray(x).map(sinc_helper)
|
|
if ans.ndim == 0:
|
|
return ans.item()
|
|
else:
|
|
return ans
|
|
|
|
def nancumprod(a, axis=None, dtype: type = NoneType, out=None):
|
|
a = asarray(a)
|
|
|
|
if dtype is NoneType:
|
|
return nancumprod(a, axis=axis, dtype=a.dtype, out=out)
|
|
|
|
if axis is None:
|
|
prod = util.cast(1, dtype)
|
|
if out is None:
|
|
result = empty((a.size, ), dtype=dtype)
|
|
else:
|
|
_check_out(out, (a.size, ))
|
|
result = out
|
|
|
|
i = 0
|
|
for idx in util.multirange(a.shape):
|
|
if not isnan(a._ptr(idx)[0]):
|
|
prod *= util.cast(a._ptr(idx)[0], dtype)
|
|
result._ptr((i, ))[0] = util.cast(prod, result.dtype)
|
|
i += 1
|
|
return result
|
|
else:
|
|
axis = util.normalize_axis_index(axis, a.ndim)
|
|
|
|
if out is None:
|
|
result = empty_like(a, dtype=dtype)
|
|
else:
|
|
_check_out(out, a.shape)
|
|
result = out
|
|
|
|
n = a.shape[axis]
|
|
|
|
for idx in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
accum = util.cast(1, dtype)
|
|
|
|
for i in range(n):
|
|
full_idx = util.tuple_insert(idx, axis, i)
|
|
p = a._ptr(full_idx)
|
|
q = result._ptr(full_idx)
|
|
if not isnan(p[0]):
|
|
accum *= util.cast(p[0], dtype)
|
|
q[0] = util.cast(accum, result.dtype)
|
|
|
|
return result
|
|
|
|
def nancumsum(a, axis=None, dtype: type = NoneType, out=None):
|
|
a = asarray(a)
|
|
|
|
if dtype is NoneType:
|
|
return nancumsum(a, axis=axis, dtype=a.dtype, out=out)
|
|
|
|
if axis is None:
|
|
sum = util.cast(0, dtype)
|
|
if out is None:
|
|
result = empty((a.size, ), dtype=dtype)
|
|
else:
|
|
_check_out(out, (a.size, ))
|
|
result = out
|
|
|
|
i = 0
|
|
for idx in util.multirange(a.shape):
|
|
if not isnan(a._ptr(idx)[0]):
|
|
sum += util.cast(a._ptr(idx)[0], dtype)
|
|
result._ptr((i, ))[0] = util.cast(sum, result.dtype)
|
|
i += 1
|
|
return result
|
|
else:
|
|
axis = util.normalize_axis_index(axis, a.ndim)
|
|
|
|
if out is None:
|
|
result = empty_like(a, dtype=dtype)
|
|
else:
|
|
_check_out(out, a.shape)
|
|
result = out
|
|
|
|
n = a.shape[axis]
|
|
|
|
for idx in util.multirange(util.tuple_delete(a.shape, axis)):
|
|
accum = util.cast(0, dtype)
|
|
|
|
for i in range(n):
|
|
full_idx = util.tuple_insert(idx, axis, i)
|
|
p = a._ptr(full_idx)
|
|
q = result._ptr(full_idx)
|
|
if not isnan(p[0]):
|
|
accum += util.cast(p[0], dtype)
|
|
q[0] = util.cast(accum, result.dtype)
|
|
|
|
return result
|
|
|
|
@extend
|
|
class ndarray:
|
|
|
|
def cumprod(
|
|
self,
|
|
axis=None,
|
|
dtype: type = NoneType,
|
|
out=None
|
|
):
|
|
return cumprod(
|
|
self,
|
|
axis=axis,
|
|
dtype=dtype,
|
|
out=out
|
|
)
|
|
|
|
def cumsum(
|
|
self,
|
|
axis=None,
|
|
dtype: type = NoneType,
|
|
out=None
|
|
):
|
|
return cumsum(
|
|
self,
|
|
axis=axis,
|
|
dtype=dtype,
|
|
out=out
|
|
)
|