# Copyright (C) 2022-2025 Exaloop Inc. <https://exaloop.io>
import util
from .npdatetime import datetime64, timedelta64, busdaycalendar, _apply_busines_day_offset, \
_apply_busines_day_count, _apply_is_business_day, _BUSDAY_FORWARD, \
_BUSDAY_FOLLOWING, _BUSDAY_BACKWARD, _BUSDAY_PRECEDING, \
_BUSDAY_MODIFIEDFOLLOWING, _BUSDAY_MODIFIEDPRECEDING, _BUSDAY_NAT, \
_BUSDAY_RAISE
from .ndarray import ndarray, flagsobj
def _check_out(out, shape):
if not isinstance(out, ndarray):
compile_error("output must be an array")
if out.ndim != staticlen(shape):
compile_error("output parameter has the wrong number of dimensions")
if out.shape != shape:
raise ValueError("output parameter has incorrect shape")
############
# Creation #
############
def _inner_type(t):
if isinstance(t, ndarray):
return t.dtype()
elif isinstance(t, List) or isinstance(t, Tuple):
return _inner_type(t[0])
else:
return t
def _extract_shape(t):
if isinstance(t, ndarray):
return t.shape
elif isinstance(t, List) or isinstance(t, Tuple):
rest = _extract_shape(t[0])
return (len(t), *rest)
else:
return ()
def _flatten(t, shape, D: type):
def helper(t, p, start: int, D: type):
for s in t:
if isinstance(s, ndarray):
cc, fc = s._contig
if cc:
str.memcpy((p + start).as_byte(), s.data.as_byte(), s.nbytes)
else:
for idx in util.multirange(s.shape):
p[start] = s._ptr(idx)[0]
start += 1
start += s.size
elif isinstance(s, List) or isinstance(s, Tuple):
start = helper(s, p, start, D)
else:
p[start] = util.cast(s, D)
start += 1
return start
p = Ptr[D](util.count(shape))
helper(t, p, 0, D)
return p
def _validate_shape(t, shape):
def error():
raise ValueError('array dimensions mismatch (jagged arrays are not allowed)')
if staticlen(shape) == 0:
return
else:
if not hasattr(type(t), "__getitem__"):
error()
else:
if isinstance(t, ndarray):
if not util.tuple_equal(t.shape, shape):
error()
else:
if len(t) != shape[0]:
error()
for s in t:
_validate_shape(s, shape[1:])
def _array(a, dtype: type = NoneType, copy: bool = True, order: str = 'K'):
if dtype is NoneType:
if isinstance(a, ndarray):
return _array(a, a.dtype, copy, order)
else:
return _array(a, type(_inner_type(a)), copy, order)
ndarray._check_order(order)
if isinstance(a, ndarray):
if copy:
return a.astype(dtype, order=order, copy=True)
else:
if order == 'K' or order == 'A':
return a.astype(dtype, copy=False)
cc, fc = a._contig
if order == 'C':
if cc:
return a.astype(dtype, copy=False)
else:
return a.astype(dtype, order='C', copy=False)
if order == 'F':
if fc:
return a.astype(dtype, copy=False)
else:
return a.astype(dtype, order='F', copy=False)
elif isinstance(a, List) or isinstance(a, Tuple):
shape = _extract_shape(a)
data = _flatten(a, shape, dtype)
_validate_shape(a, shape)
result = ndarray(shape, data)
if order == 'F':
return result.astype(dtype, order='F')
else:
return result
else:
shape = ()
data = Ptr[dtype](1)
data[0] = util.cast(a, dtype)
return ndarray(shape, data)
def array(a, dtype: type = NoneType, copy: bool = True, order: str = 'K', ndmin: Static[int] = 0):
result = _array(a, dtype=dtype, copy=copy, order=order)
if staticlen(result.shape) < ndmin:
o = (1,) * (ndmin - staticlen(result.shape))
return result.reshape((*o, *result.shape))
return result
def asarray(a, dtype: type = NoneType, order: str = 'K'):
return array(a, dtype=dtype, copy=False, order=order)
def asanyarray(a, dtype: type = NoneType, order: str = 'K'):
return asarray(a, dtype=dtype, order=order)
def asarray_chkfinite(a, dtype: type = NoneType, order: str = 'K'):
# Note: this is c/p from ndmath
def isfinite(x):
if isinstance(x, float) or isinstance(x, float32):
return util.isfinite(x)
elif isinstance(x, complex) or isinstance(x, complex64):
return util.isfinite(x.real) and util.isfinite(x.imag)
else:
return True
a = asarray(a, dtype=dtype, order=order)
for idx in util.multirange(a.shape):
if not isfinite(a._ptr(idx)[0]):
raise ValueError("array must not contain infs or NaNs")
return a
def empty(shape, dtype: type = float, order: str = 'C'):
if isinstance(shape, int):
return empty((shape,), dtype, order)
for s in shape:
if s < 0:
raise ValueError('negative dimensions are not allowed')
ccontig = (order == 'C')
fcontig = (order == 'F')
if not (ccontig or fcontig):
raise ValueError("'order' must be 'C' or 'F'")
data = Ptr[dtype](util.count(shape))
return ndarray(shape, data, fcontig=fcontig)
def empty_like(prototype, dtype: type = NoneType, order: str = 'K', shape = None):
prototype = asarray(prototype)
if dtype is NoneType:
return empty_like(prototype, dtype=prototype.dtype, order=order, shape=None)
if shape is None:
return empty_like(prototype, dtype=dtype, order=order, shape=prototype.shape)
cc, fc = prototype._contig
if order == 'A':
order = 'F' if fc else 'C'
elif order == 'K':
if staticlen(shape) == prototype.ndim:
if cc or prototype.ndim <= 1:
order = 'C'
elif fc:
order = 'F'
else:
strides = (0,) * staticlen(prototype.strides)
pstrides = Ptr[int](__ptr__(strides).as_byte())
r = util.tuple_range(staticlen(shape))
perm, strides_sorted = util.sort_by_stride(r, prototype.strides)
stride = util.sizeof(dtype)
ndim = prototype.ndim
for idim in range(ndim - 1, -1, -1):
iperm = perm[idim]
pstrides[iperm] = stride
stride *= shape[iperm]
data = Ptr[dtype](util.count(shape))
return ndarray(shape, strides, data)
else:
order = 'C'
return empty(shape, dtype, order)
def zeros(shape, dtype: type = float, order: str = 'C'):
result = empty(shape, dtype, order)
str.memset(result.data.as_byte(), byte(0), result.nbytes)
return result
def zeros_like(prototype, dtype: type = NoneType, order: str = 'K'):
result = empty_like(prototype, dtype, order)
str.memset(result.data.as_byte(), byte(0), result.nbytes)
return result
def ones(shape, dtype: type = float, order: str = 'C'):
result = empty(shape, dtype, order)
result.map(lambda x: util.cast(1, result.dtype), inplace=True)
return result
def ones_like(prototype, dtype: type = NoneType, order: str = 'K'):
result = empty_like(prototype, dtype, order)
result.map(lambda x: util.cast(1, result.dtype), inplace=True)
return result
def identity(n: int, dtype: type = float):
result = zeros((n, n), dtype)
p = result.data
for i in range(n):
p[i * (n + 1)] = dtype(1)
return result
def _diag_count(k: int, n: int, m: int):
count = 0
if k < m and k > -n:
count = min(n, m)
if k > 0:
d = max(m - n, 0)
if k > d:
count -= (k - d)
elif k < 0:
d = max(n - m, 0)
if k < -d:
count -= (-k - d)
return count
def eye(N: int, M: Optional[int] = None, k: int = 0, dtype: type = float, order: str = 'C'):
n: int = N
m: int = M if M is not None else n
result = zeros((n, m), dtype=dtype, order=order)
p = result.data
for i in range(_diag_count(k, n, m)):
if k >= 0:
result[i, i + k] = dtype(1)
else:
j = n - i - 1
result[i - k, i] = dtype(1)
return result
def diag(v, k: int = 0):
v = asarray(v)
data = v.data
T = type(data[0])
if staticlen(v.shape) == 1:
count = v.shape[0]
n = count + abs(k)
result = zeros((n, n), dtype=T)
p = result.data
if k > 0:
p += k
elif k < 0:
p += (-k) * n
for i in range(count):
q = v._ptr((i,))
p[0] = q[0]
p += n + 1
return result
elif staticlen(v.shape) == 2:
n, m = v.shape
sn, sm = v.strides
new_shape = (_diag_count(k, n, m),)
new_strides = (sn + sm,)
new_data = data
if new_shape[0] > 0:
if k >= 0:
new_data = v._ptr((0, k))
else:
new_data = v._ptr((-k, 0))
return ndarray(new_shape, new_strides, new_data)
else:
compile_error('Input must be 1- or 2-d.')
def diagflat(v, k: int = 0):
return diag(asarray(v).flatten(), k)
def tri(N: int, M: Optional[int] = None, k: int = 0, dtype: type = float):
n: int = N
m: int = M if M is not None else n
result = zeros((n, m), dtype=dtype)
p = result.data
for i in range(n):
for j in range(min(i + k + 1, m)):
p[i*m + j] = dtype(1)
return result
def triu(x, k: int = 0):
x = asarray(x)
T = x.dtype
if staticlen(x.shape) == 0:
compile_error('Cannot call triu on 0-d array.')
elif staticlen(x.shape) == 1:
n = x.shape[0]
result = zeros((n, n), dtype=T)
p = result.data
for i in range(n):
for j in range(max(0, i + k), n):
p[i*n + j] = x[j]
return result
else:
y = x.copy()
n, m = x.shape[-2], x.shape[-1]
pre = (slice(None, None, None),) * (staticlen(x.shape) - 2)
for i in range(n):
for j in range(min(i + k, m)):
y[(*pre, i, j)] = T(0)
return y
def tril(x, k: int = 0):
x = asarray(x)
T = x.dtype
if staticlen(x.shape) == 0:
compile_error('Cannot call tril on 0-d array.')
elif staticlen(x.shape) == 1:
n = x.shape[0]
result = zeros((n, n), dtype=T)
p = result.data
for i in range(n):
for j in range(min(i + k + 1, n)):
p[i*n + j] = x[j]
return result
else:
y = x.copy()
n, m = x.shape[-2], x.shape[-1]
pre = (slice(None, None, None),) * (staticlen(x.shape) - 2)
for i in range(n):
for j in range(max(0, i + k + 1), m):
y[(*pre, i, j)] = T(0)
return y
def vander(x, N: Optional[int] = None, increasing: bool = False):
x = asarray(x)
if staticlen(x.shape) != 1:
compile_error('x must be a one-dimensional array or sequence.')
T = x.dtype
n: int = x.shape[0]
m: int = N if N is not None else n
result = zeros((n, m), dtype=T)
p = result.data
for i in range(n):
base = x._ptr((i,))[0]
for j in range(m):
power = j if increasing else m - j - 1
p[i*m + j] = base ** power
return result
@pure
@llvm
def _signbit(x: float) -> bool:
%y = bitcast double %x to i64
%z = icmp slt i64 %y, 0
%b = zext i1 %z to i8
ret i8 %b
_imin: Static[int] = -9223372036854775808
_imax: Static[int] = 9223372036854775807
def _safe_ceil(value: float):
ivalue = util.ceil64(value)
if util.isnan64(ivalue):
raise ValueError('arange: cannot compute array length')
if not (float(_imin) <= ivalue <= float(_imax)):
raise OverflowError('arange: overflow while computing length')
return int(ivalue)
def _datetime_arange(start, stop, step, dtype: type):
def convert(x, dtype: type):
if isinstance(x, datetime64) or isinstance(x, timedelta64):
return x.value
elif isinstance(x, int):
return x
elif isinstance(x, str):
return dtype(x, dtype.base).value
else:
compile_error("datetime_arange inputs must be datetime64, timedelta64 or int")
def nat(x: int):
return timedelta64(x, 'generic')._nat
if not (isinstance(dtype, datetime64) or isinstance(dtype, timedelta64)):
compile_error("datetime_arange was given a non-datetime dtype")
a = convert(start, dtype)
b = convert(stop, dtype)
c = 0
if step is None:
c = 1
elif isinstance(step, datetime64) or isinstance(step, str):
compile_error("cannot use a datetime as a step in arange")
else:
c = convert(step, dtype)
if ((isinstance(start, datetime64) or isinstance(start, str)) and not
(isinstance(stop, datetime64) or isinstance(stop, str))):
b += a
if nat(a) or nat(b) or nat(c):
raise ValueError("arange: cannot use NaT (not-a-time) datetime values")
length = len(range(a, b, c))
ans = empty(length, dtype)
for i in range(length):
ans.data[i] = dtype(a, dtype.base)
a += c
return ans
def arange(start: float, stop: float, step: float, dtype: type = float):
if step == 0.0:
raise ValueError('step cannot be zero')
delta = stop - start
tmp_len = delta / step
length = 0
if tmp_len == 0.0 and delta != 0.0:
if _signbit(tmp_len):
length = 0
else:
length = 1
else:
length = _safe_ceil(tmp_len)
if length <= 0:
return empty(0, dtype=dtype)
result = empty(length, dtype=dtype)
p = result.data
i = start
j = 0
while (i < stop) if step > 0.0 else (i > stop):
p[j] = util.cast(i, dtype)
j += 1
i += step
return result
@overload
def arange(stop: float, step: float, dtype: type = float):
return arange(0.0, stop, step, dtype)
@overload
def arange(start: float, stop: float, dtype: type = float):
return arange(start, stop, 1.0, dtype)
@overload
def arange(stop: float, dtype: type = float):
return arange(0.0, stop, 1.0, dtype)
@overload
def arange(start: int, stop: int, step: int, dtype: type = int):
if isinstance(dtype, datetime64) or isinstance(dtype, timedelta64):
return _datetime_arange(start, stop, step, dtype)
length = len(range(start, stop, step))
result = empty(length, dtype=dtype)
p = result.data
j = 0
for i in range(start, stop, step):
p[j] = util.cast(i, dtype)
j += 1
return result
@overload
def arange(stop: int, step: int, dtype: type = int):
return arange(0, stop, step, dtype)
@overload
def arange(start: int, stop: int, dtype: type = int):
return arange(start, stop, 1, dtype)
@overload
def arange(stop: int, dtype: type = int):
return arange(0, stop, 1, dtype)
@overload
def arange(start: datetime64, stop, step = None, dtype: type = NoneType):
if dtype is NoneType:
return _datetime_arange(start, stop, step, type(start))
else:
return _datetime_arange(start, stop, step, dtype)
@overload
def arange(start: timedelta64, stop, step, dtype: type = NoneType):
if dtype is NoneType:
return _datetime_arange(start, stop, step, type(start))
else:
return _datetime_arange(start, stop, step, dtype)
@overload
def arange(start: timedelta64, stop, dtype: type = NoneType):
if dtype is NoneType:
return _datetime_arange(start, stop, 1, type(start))
else:
return _datetime_arange(start, stop, 1, dtype)
@overload
def arange(stop: timedelta64, dtype: type = NoneType):
if dtype is NoneType:
return _datetime_arange(0, stop, 1, type(stop))
else:
return _datetime_arange(0, stop, 1, dtype)
@overload
def arange(start: str, stop, step = None, dtype: type = datetime64['D', 1]):
return _datetime_arange(start, stop, step, dtype)
def linspace(start: float, stop: float, num: int = 50,
endpoint: bool = True, retstep: Static[int] = False,
dtype: type = float):
if num < 0:
raise ValueError(f'Number of samples, {num}, must be non-negative.')
delta = stop - start
div = (num - 1) if endpoint else num
step = delta / div
result = empty(num, dtype=dtype)
p = result.data
if div > 0:
if step == 0:
for i in range(num):
p[i] = util.cast(((i / div) * delta) + start, dtype)
else:
for i in range(num):
p[i] = util.cast((i * step) + start, dtype)
else:
for i in range(num):
p[i] = util.cast((i * delta) + start, dtype)
step = util.nan64()
if endpoint and num > 1:
p[num - 1] = stop
if retstep:
return result, step
else:
return result
def _linlogspace(start: float, stop: float, num: int = 50, base: float = 10.0,
out_sign: int = 1, endpoint: bool = True, retstep: Static[int] = False,
dtype: type = float, log: Static[int] = False):
if num < 0:
raise ValueError(f'Number of samples, {num}, must be non-negative.')
delta = stop - start
div = (num - 1) if endpoint else num
step = delta / div
result = empty(num, dtype=dtype)
p = result.data
if div > 0:
if step == 0:
for i in range(num):
y = ((i / div) * delta) + start
if log:
y = util.pow64(base, y)
y *= out_sign
p[i] = util.cast(y, dtype)
else:
for i in range(num):
y = (i * step) + start
if log:
y = util.pow64(base, y)
y *= out_sign
p[i] = util.cast(y, dtype)
else:
for i in range(num):
y = (i * delta) + start
if log:
y = util.pow64(base, y)
y *= out_sign
p[i] = util.cast(y, dtype)
step = util.nan64()
if endpoint and num > 1:
y = stop
if log:
y = util.pow64(base, y)
y *= out_sign
p[num - 1] = util.cast(y, dtype)
if retstep:
return result, step
else:
return result
def linspace(start: float, stop: float, num: int = 50,
endpoint: bool = True, retstep: Static[int] = False,
dtype: type = float):
return _linlogspace(start=start, stop=stop, num=num,
endpoint=endpoint, retstep=retstep,
dtype=dtype, log=False)
def logspace(start: float, stop: float, num: int = 50,
endpoint: bool = True, base: float = 10.0,
retstep: Static[int] = False,
dtype: type = float):
return _linlogspace(start=start, stop=stop, num=num,
endpoint=endpoint, retstep=retstep,
dtype=dtype, base=base, log=True)
def geomspace(start: float, stop: float, num: int = 50,
endpoint: bool = True, dtype: type = float):
if start == 0 or stop == 0:
raise ValueError('Geometric sequence cannot include zero')
out_sign = 1
if start < 0 and stop < 0:
start, stop = -start, -stop
out_sign = -out_sign
start = start + (stop - stop)
stop = stop + (start - start)
log_start = util.log10_64(start)
log_stop = util.log10_64(stop)
return _linlogspace(start=log_start, stop=log_stop, num=num,
endpoint=endpoint, retstep=False,
dtype=dtype, base=10.0, out_sign=out_sign,
log=True)
def fromfunction(function, shape, dtype: type = float, **kwargs):
result_dtype = type(
function(
*tuple(util.zero(dtype)
for _ in staticrange(staticlen(shape))),
**kwargs))
result = empty(shape, dtype=result_dtype)
for idx in util.multirange(shape):
p = result._ptr(idx)
args = tuple(util.cast(i, dtype) for i in idx)
p[0] = function(*args, **kwargs)
return result
def fromiter(iterable, dtype: type, count: int = -1):
if count < 0:
return array([a for a in iterable], dtype=dtype)
else:
result = empty((count,), dtype=dtype)
if count:
p = result.data
i = 0
for a in iterable:
p[i] = util.cast(a, dtype)
i += 1
if i == count:
break
if i != count:
raise ValueError(f'iterator too short: Expected {count} but iterator had only {i} items.')
return result
def frombuffer(buffer: str, dtype: type = float, count: int = -1, offset: int = 0):
if count < 0:
count = len(buffer) // util.sizeof(dtype)
p = Ptr[dtype](buffer.ptr + offset)
return ndarray((count,), p)
################
# Broadcasting #
################
def broadcast_shapes(*args):
def _largest(args):
if staticlen(args) == 1:
return args[0]
a = args[0]
b = _largest(args[1:])
if staticlen(b) > staticlen(a):
return b
else:
return a
def _ensure_tuple(x):
if isinstance(x, Tuple):
return x
else:
return (x,)
if staticlen(args) == 0:
return ()
args = tuple(_ensure_tuple(a) for a in args)
for a in args:
for i in a:
if i < 0:
raise ValueError('negative dimensions are not allowed')
t = _largest(args)
N: Static[int] = staticlen(t)
ans = (0,) * N
p = Ptr[int](__ptr__(ans).as_byte())
for i in staticrange(N):
p[i] = t[i]
for a in args:
for i in staticrange(staticlen(a)):
x = a[len(a) - 1 - i]
q = p + (len(t) - 1 - i)
y = q[0]
if y == 1:
q[0] = x
elif x != 1 and x != y:
raise ValueError('shape mismatch: objects cannot be broadcast to a single shape')
return ans
def broadcast_to(x, shape):
x = asarray(x)
if not isinstance(shape, Tuple):
return broadcast_to(x, (shape,))
N: Static[int] = staticlen(x.shape)
if staticlen(shape) < N:
compile_error('input operand has more dimensions than allowed by the axis remapping')
shape1, shape2 = shape[:-N], shape[-N:]
substrides = (0,) * N
p = Ptr[int](__ptr__(substrides).as_byte())
if N > 0:
for i in range(N):
a = x.shape[i]
b = shape2[i]
if a == b:
p[i] = x.strides[i]
elif a == 1:
p[i] = 0
else:
raise ValueError(f'cannot broadcast array of shape {x.shape} to shape {shape}')
z = (0,) * (staticlen(shape) - N)
new_strides = (*z, *substrides)
return ndarray(shape, new_strides, x.data)
def broadcast_arrays(*args):
def _ensure_array(x):
if isinstance(x, ndarray):
return x
else:
return array(x)
args = tuple(_ensure_array(a) for a in args)
shapes = tuple(a.shape for a in args)
bshape = broadcast_shapes(*shapes)
return [broadcast_to(a, bshape) for a in args]
def meshgrid(*xi, copy: bool = True, sparse: Static[int] = False, indexing: Static[str] = 'xy'):
def make_shape(i, ndim: Static[int]):
t = (1,) * ndim
p = Ptr[int](__ptr__(t).as_byte())
p[i] = -1
return t
def build_output(xi, i: int = 0, ndim: Static[int]):
if staticlen(xi) == 0:
return ()
x = xi[0]
y = array(x).reshape(make_shape(i, ndim))
rest = build_output(xi[1:], i + 1, ndim)
return (y, *rest)
if indexing != 'xy' and indexing != 'ij':
compile_error("Valid values for `indexing` are 'xy' and 'ij'.")
ndim: Static[int] = staticlen(xi)
s0 = (1,) * ndim
output = build_output(xi, ndim=ndim)
if indexing == 'xy' and ndim > 1:
# switch first and second axis
output0 = output[0].reshape(1, -1, *s0[2:])
output1 = output[1].reshape(-1, 1, *s0[2:])
output = (output0, output1, *output[2:])
if not sparse:
# Return the full N-D matrix (not only the 1-D vector)
return [a for a in broadcast_arrays(*output)]
if copy:
return [a.copy() for a in output]
return [a for a in output]
class _broadcast[A, S]:
_arrays: A
_shape: S
_index: int
def __init__(self, arrays: A, shape: S):
self._arrays = arrays
self._shape = shape
self._index = 0
@property
def shape(self):
return self._shape
@property
def index(self):
return self._index
@property
def size(self):
return util.count(self.shape)
@property
def ndim(self):
return staticlen(self.shape)
@property
def nd(self):
return self.ndim
@property
def numiter(self):
return staticlen(self._arrays)
@property
def iters(self):
def get_flat(arr, index, bshape):
f = broadcast_to(arr, bshape).flat
f.index = index
return f
return tuple(get_flat(a, self.index, self.shape) for a in self._arrays)
def __iter__(self):
arrays = self._arrays
n = self.size
while self.index < n:
idx = util.index_to_coords(self.index, self.shape)
self._index += 1
yield tuple(a._ptr(idx, broadcast=True)[0] for a in arrays)
def reset(self):
self._index = 0
def broadcast(*args):
arrays = tuple(asarray(a) for a in args)
shape = broadcast_shapes(*tuple(a.shape for a in arrays))
return _broadcast(arrays, shape)
def full(shape, fill_value, dtype: type = NoneType, order: str = 'C'):
if isinstance(shape, int):
sh = (shape,)
else:
sh = shape
fv = asarray(fill_value)
if fv.ndim != 0:
broadcast_to(fv, shape) # error check
if dtype is NoneType:
result = empty(shape, fv.dtype, order)
if fv.ndim == 0:
e = fv.item()
result.map(lambda x: e, inplace=True)
else:
for idx in util.multirange(shape):
result._ptr(idx)[0] = fv._ptr(idx, broadcast=True)[0]
return result
else:
result = empty(shape, dtype, order)
if fv.ndim == 0:
e = fv.item()
result.map(lambda x: util.cast(e, result.dtype), inplace=True)
else:
for idx in util.multirange(shape):
result._ptr(idx)[0] = util.cast(fv._ptr(idx, broadcast=True)[0], result.dtype)
return result
def full_like(prototype, fill_value, dtype: type = NoneType, order: str = 'K'):
prototype = asarray(prototype)
shape = prototype.shape
fv = asarray(fill_value)
if fv.ndim != 0:
broadcast_to(fv, shape) # error check
if dtype is NoneType:
result = empty_like(prototype, fv.dtype, order)
if fv.ndim == 0:
e = fv.item()
result.map(lambda x: e, inplace=True)
else:
for idx in util.multirange(shape):
result._ptr(idx)[0] = fv._ptr(idx, broadcast=True)[0]
return result
else:
result = empty_like(prototype, dtype, order)
if fv.ndim == 0:
e = fv.item()
result.map(lambda x: util.cast(e, result.dtype), inplace=True)
else:
for idx in util.multirange(shape):
result._ptr(idx)[0] = util.cast(fv._ptr(idx, broadcast=True)[0], result.dtype)
return result
################
# Manipulation #
################
def copyto(dst: ndarray, src, where = True):
src = asarray(src)
dst_dtype = dst.dtype
src_dtype = src.dtype
if isinstance(where, bool):
if not where:
return
if dst_dtype is src_dtype and src._contig_match(dst):
str.memcpy(dst.data.as_byte(), src.data.as_byte(), dst.nbytes)
return
src = broadcast_to(src, dst.shape)
where = broadcast_to(asarray(where), dst.shape)
for idx in util.multirange(dst.shape):
w = where._ptr(idx)
if w[0]:
p = src._ptr(idx)
q = dst._ptr(idx)
q[0] = util.cast(p[0], dst_dtype)
def ndim(a):
return asarray(a).ndim
def size(a, axis: Optional[int] = None):
a = asarray(a)
if axis is None:
return a.size
else:
return a.shape[axis]
def shape(a):
if isinstance(a, ndarray):
return a.shape
else:
shape = _extract_shape(a)
_validate_shape(a, shape)
return shape
def reshape(a, newshape, order: str = 'C'):
return asarray(a).reshape(newshape, order=order)
def transpose(a, axes=None):
return a.transpose(axes)
def ravel(a, order: str = 'C'):
return asarray(a).ravel(order=order)
def ascontiguousarray(a, dtype: type = NoneType):
return asarray(a, dtype=dtype, order='C')
def asfortranarray(a, dtype: type = NoneType):
return asarray(a, dtype=dtype, order='F')
def asfarray(a, dtype: type = float):
if (dtype is not float and
dtype is not float32 and
dtype is not complex and
dtype is not complex64):
return asfarray(a, float)
return asarray(a, dtype=dtype)
def moveaxis(a, source, destination):
a = asarray(a)
source = util.normalize_axis_tuple(source, a.ndim, 'source')
destination = util.normalize_axis_tuple(destination, a.ndim, 'destination')
if len(source) != len(destination):
raise ValueError('`source` and `destination` arguments must have '
'the same number of elements')
order = [n for n in range(a.ndim) if n not in source]
for dest, src in sorted(zip(destination, source)):
order.insert(dest, src)
order = Ptr[type(a.shape)](order.arr.ptr.as_byte())[0]
return a.transpose(order)
def swapaxes(a, axis1: int, axis2: int):
return asarray(a).swapaxes(axis1, axis2)
def atleast_1d(*arys):
def atl1d(a):
a = asarray(a)
if staticlen(a.shape) == 0:
return a.reshape(1)
else:
return a
if staticlen(arys) == 1:
return atl1d(arys[0])
else:
return tuple(atl1d(a) for a in arys)
def atleast_2d(*arys):
def atl2d(a):
a = asarray(a)
if staticlen(a.shape) == 0:
return a.reshape(1, 1)
elif staticlen(a.shape) == 1:
return a[None, :]
else:
return a
if staticlen(arys) == 1:
return atl2d(arys[0])
else:
return tuple(atl2d(a) for a in arys)
def atleast_3d(*arys):
def atl3d(a):
a = asarray(a)
if staticlen(a.shape) == 0:
return a.reshape(1, 1, 1)
elif staticlen(a.shape) == 1:
return a[None, :, None]
elif staticlen(a.shape) == 2:
return a[:, :, None]
else:
return a
if staticlen(arys) == 1:
return atl3d(arys[0])
else:
return tuple(atl3d(a) for a in arys)
def require(a, dtype: type = NoneType, requirements = None):
REQ_C: Static[int] = 1
REQ_F: Static[int] = 2
REQ_A: Static[int] = 4
REQ_W: Static[int] = 8
REQ_O: Static[int] = 16
REQ_E: Static[int] = 32
if requirements is None:
return asarray(a, dtype=dtype)
if not requirements:
return asarray(a, dtype=dtype)
req = 0
for x in requirements:
if x in ('C', 'C_CONTIGUOUS', 'CONTIGUOUS'):
req |= REQ_C
elif x in ('F', 'F_CONTIGUOUS', 'FORTRAN'):
req |= REQ_F
elif x in ('A', 'ALIGNED'):
req |= REQ_A
elif x in ('W', 'WRITEABLE'):
req |= REQ_W
elif x in ('O', 'OWNDATA'):
req |= REQ_O
elif x in ('E', 'ENSUREARRAY'):
req |= REQ_E
else:
raise ValueError("invalid requirement: " + repr(x))
order = 'A'
if (req & REQ_C) and (req & REQ_F):
raise ValueError('Cannot specify both "C" and "F" order')
elif req & REQ_F:
order = 'F'
req &= ~REQ_F
elif req & REQ_C:
order = 'C'
req &= ~REQ_C
# Codon-NumPy ignores other flags/properties currently
copy = ((req & REQ_O) != 0)
arr = array(a, dtype=dtype, order=order, copy=copy)
return arr
def _copy_data_c(dst: cobj, src: cobj, shape: Ptr[int],
strides: Ptr[int], dst_strides: Ptr[int],
ndim: Static[int], element_size: int,
block: int, block_dim: int):
if ndim < 0:
return
if ndim == block_dim:
str.memcpy(dst, src, block)
else:
src_stride = strides[0]
dst_stride = dst_strides[0]
for j in range(shape[0]):
_copy_data_c(dst + j * dst_stride, src + j * src_stride,
shape + 1, strides + 1, dst_strides + 1,
ndim - 1, element_size, block, block_dim)
def _copy_data_f(dst: cobj, src: cobj, shape: Ptr[int],
strides: Ptr[int], dst_strides: Ptr[int],
ndim: Static[int], element_size: int,
block: int, block_dim: int):
if ndim < 0:
return
if ndim == block_dim:
str.memcpy(dst, src, block)
else:
src_stride = strides[ndim - 1]
dst_stride = dst_strides[ndim - 1]
for j in range(shape[ndim - 1]):
_copy_data_f(dst + j * dst_stride, src + j * src_stride,
shape, strides, dst_strides,
ndim - 1, element_size, block, block_dim)
def concatenate(arrays, axis = 0, out = None, dtype: type = NoneType):
def check_array_dims_static(arrays, ndim: Static[int]):
if staticlen(arrays) > 0:
if staticlen(arrays[0].shape) != ndim:
compile_error("all the input arrays must have same number of dimensions")
check_array_dims_static(arrays[1:], ndim)
def find_out_type(arrays, dtype: type):
if staticlen(arrays) == 0:
return dtype()
else:
x = util.coerce(arrays[0].dtype, dtype)
return find_out_type(arrays[1:], type(x))
def concat_inner(arrays, axis: int, out, dtype: type):
if out is not None and dtype is NoneType:
return concat_inner(arrays, axis=axis, out=out, dtype=out.dtype)
if out is None and dtype is NoneType:
compile_error("[internal error] bad out and dtype given to concat_inner")
arr0 = asarray(arrays[0])
ndim: Static[int] = arrays[0].ndim
axis = util.normalize_axis_index(axis, ndim)
shape = arr0.shape
pshape = Ptr[int](__ptr__(shape).as_byte())
for i in range(1, len(arrays)):
arr = arrays[i]
arr_shape = arr.shape
if arr.ndim != ndim:
compile_error("all the input arrays must have same number of dimensions")
for idim in staticrange(ndim):
if idim == axis:
pshape[idim] += arr_shape[idim]
elif pshape[idim] != arr_shape[idim]:
raise ValueError("all the input array dimensions except for "
"the concatenation axis must match exactly")
corder = True
if out is not None:
if dtype is not NoneType:
compile_error("can only specify one of 'out' and 'dtype'")
_check_out(out, shape)
ret = out
else:
num_c = 0
num_f = 0
for array in arrays:
cc, fc = array._contig
if cc:
num_c += 1
if fc:
num_f += 1
corder = (num_c >= num_f)
ret = empty(shape, dtype, order=('C' if corder else 'F'))
element_size = util.sizeof(dtype)
offset = 0
dst_strides = ret.strides
dst_stride_axis = dst_strides[axis]
element_size = util.sizeof(dtype)
for array in arrays:
array_shape = array.shape
if array.dtype is dtype:
strides = array.strides
dst = ret.data.as_byte() + offset * dst_stride_axis
src = array.data.as_byte()
block = element_size
block_dim = ndim
if corder:
i = ndim - 1
while i >= 0:
if strides[i] == block and dst_strides[i] == block:
block *= array_shape[i]
else:
block_dim = ndim - 1 - i
break
i -= 1
_copy_data_c(dst, src, Ptr[int](__ptr__(array_shape).as_byte()),
Ptr[int](__ptr__(strides).as_byte()),
Ptr[int](__ptr__(dst_strides).as_byte()),
array.ndim, element_size, block, block_dim)
else:
i = 0
while i < ndim:
if strides[i] == block and dst_strides[i] == block:
block *= array_shape[i]
else:
block_dim = i
break
i += 1
_copy_data_f(dst, src, Ptr[int](__ptr__(array_shape).as_byte()),
Ptr[int](__ptr__(strides).as_byte()),
Ptr[int](__ptr__(dst_strides).as_byte()),
array.ndim, element_size, block, block_dim)
else:
for src_idx in util.multirange(array_shape):
dst_idx = util.tuple_add(src_idx, axis, offset)
ret._ptr(dst_idx)[0] = util.cast(array._ptr(src_idx)[0], dtype)
offset += array_shape[axis]
return ret
def concat_flatten(arrays, out):
dtype = out.dtype
out_ccontig = out._contig[0]
i = 0
for a in arrays:
if a.dtype is dtype and out_ccontig and a._contig[0]:
q = out._ptr((i,))
n = a.size
str.memcpy(q.as_byte(), a.data.as_byte(), n * util.sizeof(dtype))
i += n
else:
for idx in util.multirange(a.shape):
p = a._ptr(idx)
q = out._ptr((i,))
q[0] = util.cast(p[0], dtype)
i += 1
return out
def concat_tuple(arrays, axis = 0, out = None, dtype: type = NoneType):
if staticlen(arrays) == 0:
compile_error("need at least one array to concatenate")
arrays = tuple(asarray(arr) for arr in arrays)
if axis is None:
tot = 0
for a in arrays:
tot += a.size
shape = (tot,)
if out is None:
if dtype is NoneType:
x = find_out_type(arrays[1:], arrays[0].dtype)
return concat_flatten(arrays, empty(shape, dtype=type(x)))
else:
return concat_flatten(arrays, empty(shape, dtype=dtype))
else:
if staticlen(out.shape) != 1:
compile_error("Output array has wrong dimensionality")
if not util.tuple_equal(out.shape, shape):
raise ValueError("Output array is the wrong shape")
return concat_flatten(arrays, out)
else:
ndim: Static[int] = staticlen(arrays[0].shape)
if ndim == 0:
compile_error("zero-dimensional arrays cannot be concatenated")
check_array_dims_static(arrays[1:], ndim)
if out is None:
if dtype is NoneType:
x = find_out_type(arrays[1:], arrays[0].dtype)
return concat_inner(arrays, axis, out=None, dtype=type(x))
else:
return concat_inner(arrays, axis, out=None, dtype=dtype)
else:
return concat_inner(arrays, axis, out, dtype=out.dtype)
def concat_list(arrays, axis = 0, out = None, dtype: type = NoneType):
if len(arrays) == 0:
raise ValueError("need at least one array to concatenate")
if not isinstance(arrays[0], ndarray):
arrays = [asarray(arr) for arr in arrays]
if axis is None:
if axis is not None and len(arrays) == 1:
util.normalize_axis_index(axis, 1) # error check
tot = 0
for a in arrays:
tot += a.size
shape = (tot,)
if out is None:
if dtype is NoneType:
return concat_flatten(arrays, empty(shape, dtype=arrays[0].dtype))
else:
return concat_flatten(arrays, empty(shape, dtype=dtype))
else:
if staticlen(out.shape) != 1:
compile_error("Output array has wrong dimensionality")
if not util.tuple_equal(out.shape, shape):
raise ValueError("Output array is the wrong shape")
return concat_flatten(arrays, out)
else:
ndim: Static[int] = staticlen(arrays[0].shape)
if ndim == 0:
compile_error("zero-dimensional arrays cannot be concatenated")
for arr in arrays:
if arr.ndim != arrays[0].ndim:
raise ValueError("all the input arrays must have same number of dimensions")
shape = arrays[0].shape
pshape = Ptr[int](__ptr__(shape).as_byte())
axis = util.normalize_axis_index(axis, ndim)
for iarrays in range(1, len(arrays)):
arr_shape = arrays[iarrays].shape
for idim in range(ndim):
if idim == axis:
pshape[idim] += arr_shape[idim]
elif pshape[idim] != arr_shape[idim]:
raise ValueError("all the input array dimensions except for the "
"concatenation axis must match exactly")
if out is None:
if dtype is NoneType:
return concat_inner(arrays, axis, out=None, dtype=arrays[0].dtype)
else:
return concat_inner(arrays, axis, out=None, dtype=dtype)
else:
return concat_inner(arrays, axis, out=out, dtype=out.dtype)
if out is not None:
if dtype is not NoneType:
compile_error("concatenate() only takes `out` or `dtype` as an argument, but both were provided.")
if not isinstance(out, ndarray):
compile_error("'out' must be an array")
if isinstance(arrays, Tuple):
return concat_tuple(arrays, axis=axis, out=out, dtype=dtype)
else:
return concat_list(arrays, axis=axis, out=out, dtype=dtype)
def expand_dims(a, axis):
a = asarray(a)
old_ndims: Static[int] = staticlen(a.shape)
new_ndims: Static[int] = (old_ndims +
staticlen(util.normalize_axis_tuple(axis, 9999999)))
axis = util.normalize_axis_tuple(axis, new_ndims)
old_shape = a.shape
pold_shape = Ptr[int](__ptr__(old_shape).as_byte())
new_shape = (1,) * new_ndims
pnew_shape = Ptr[int](__ptr__(new_shape).as_byte())
j = 0
for i in range(new_ndims):
if i not in axis:
pnew_shape[i] = pold_shape[j]
j += 1
return a.reshape(new_shape)
def _apply(fn, seq):
if isinstance(seq, Tuple):
return tuple(fn(a) for a in seq)
elif isinstance(seq, List):
return [fn(a) for a in seq]
else:
compile_error("expected a tuple or a list as input")
def stack(arrays, axis: int = 0, out = None, dtype: type = NoneType):
if not (isinstance(arrays, Tuple) or
(isinstance(arrays, List) and
isinstance(arrays[0], ndarray))):
return asarray(arrays)
if len(arrays) == 0:
raise ValueError("need at least one array to stack")
arrays = _apply(asarray, arrays)
arr0 = arrays[0]
for arr in arrays:
if not util.tuple_equal(arr0.shape, arr.shape):
raise ValueError("all input arrays must have the same shape")
result_ndim = arrays[0].ndim + 1
axis = util.normalize_axis_index(axis, result_ndim)
expanded_arrays = _apply(lambda arr: expand_dims(arr, axis=axis), arrays)
return concatenate(expanded_arrays, axis=axis, out=out, dtype=dtype)
def vstack(tup, dtype: type = NoneType):
if not (isinstance(tup, Tuple) or
(isinstance(tup, List) and
isinstance(tup[0], ndarray))):
return asarray(tup)
arrs = _apply(atleast_2d, tup)
return concatenate(arrs, axis=0, dtype=dtype)
def hstack(tup, dtype: type = NoneType):
if not (isinstance(tup, Tuple) or
(isinstance(tup, List) and
isinstance(tup[0], ndarray))):
return asarray(tup)
arrs = _apply(atleast_1d, tup)
axis = 0 if (len(arrs) > 0 and arrs[0].ndim == 1) else 1
return concatenate(arrs, axis=axis, dtype=dtype)
def dstack(tup):
arrs = _apply(atleast_3d, tup)
return concatenate(arrs, axis=2)
row_stack = vstack
def column_stack(tup):
def fix_array(v):
arr = asarray(v)
if staticlen(arr.shape) < 2:
return array(arr, copy=False, ndmin=2).T
else:
return arr
arrs = _apply(fix_array, tup)
return concatenate(arrs, axis=1)
def repeat(a, repeats, axis = None):
def neg_rep_error():
raise ValueError("negative dimensions are not allowed")
a = asarray(a, order='C')
dtype = a.dtype
shape = a.shape
pshape = Ptr[int](__ptr__(shape).as_byte())
if isinstance(repeats, int):
if repeats < 0:
neg_rep_error()
else:
for rep in repeats:
if rep < 0:
neg_rep_error()
if isinstance(axis, int):
axis = util.normalize_axis_index(axis, a.ndim)
n = pshape[axis]
nel = 1
for i in range(axis + 1, a.ndim):
nel *= pshape[i]
chunk = nel * a.itemsize
n_outer = 1
for i in range(axis):
n_outer *= pshape[i]
rep_shape = shape
prep_shape = Ptr[int](__ptr__(rep_shape).as_byte())
if isinstance(repeats, int):
prep_shape[axis] *= repeats
repeated = empty(rep_shape, dtype=dtype)
old_data = a.data.as_byte()
new_data = repeated.data.as_byte()
for i in range(n_outer):
for j in range(n):
for k in range(repeats):
str.memcpy(new_data, old_data, chunk)
new_data += chunk
old_data += chunk
for src in util.multirange(shape):
e = a._ptr(src)[0]
off = src[axis]
for r in range(repeats):
dst = util.tuple_set(src, axis, off * repeats + r)
p = repeated._ptr(dst)
p[0] = e
return repeated
else:
axis_dim = prep_shape[axis]
if len(repeats) != axis_dim:
raise ValueError("length of 'repeats' does not match axis size")
prep_shape[axis] = 0
for rep in repeats:
prep_shape[axis] += rep
repeated = empty(rep_shape, dtype=dtype)
old_data = a.data.as_byte()
new_data = repeated.data.as_byte()
for i in range(n_outer):
for j in range(n):
for k in range(repeats[j]):
str.memcpy(new_data, old_data, chunk)
new_data += chunk
old_data += chunk
return repeated
elif axis is None:
if isinstance(repeats, int):
a_size = a.size
rep_size = a_size * repeats
repeated = empty((rep_size,), dtype=dtype)
p = a.data
q = repeated.data
off = 0
for i in range(a_size):
elem = p[i]
for _ in range(repeats):
q[off] = elem
off += 1
return repeated
else:
a_size = a.size
if len(repeats) != a_size:
raise ValueError("length of 'repeats' does not match array size")
rep_tot = 0
for rep in repeats:
rep_tot += rep
repeated = empty((rep_tot,), dtype=dtype)
p = a.data
q = repeated.data
rep_idx = 0
off = 0
for i in range(a_size):
elem = p[i]
for _ in range(repeats[i]):
q[off] = elem
off += 1
return repeated
else:
compile_error("'axis' must be None or an int")
def delete(arr, obj, axis = None):
arr = asarray(arr)
dtype = arr.dtype
shape = arr.shape
newshape = shape
pnewshape = Ptr[int](__ptr__(newshape).as_byte())
cc, fc = arr._contig
arrorder = 'F' if (fc and not cc) else 'C'
if isinstance(axis, int):
axis = util.normalize_axis_index(axis, arr.ndim)
shape_no_axis = util.tuple_delete(shape, axis)
N = arr.shape[axis]
if isinstance(obj, int):
pnewshape[axis] -= 1
new = empty(newshape, dtype, arrorder)
obj = util.normalize_index(obj, axis, N)
for i in range(obj):
for idx0 in util.multirange(shape_no_axis):
idx = util.tuple_insert(idx0, axis, i)
p = arr._ptr(idx)
q = new._ptr(idx)
q[0] = p[0]
for i in range(obj + 1, N):
for idx0 in util.multirange(shape_no_axis):
idx1 = util.tuple_insert(idx0, axis, i)
idx2 = util.tuple_insert(idx0, axis, i - 1)
p = arr._ptr(idx1)
q = new._ptr(idx2)
q[0] = p[0]
return new
elif isinstance(obj, slice):
start, stop, step = obj.adjust_indices(N)
xr = range(start, stop, step)
numtodel = len(xr)
if numtodel <= 0:
return arr.copy(order=arrorder)
if step < 0:
step = -step
start = xr[-1]
stop = xr[0] + 1
pnewshape[axis] -= numtodel
new = empty(newshape, dtype, arrorder)
if start:
for i in range(start):
for idx0 in util.multirange(shape_no_axis):
idx = util.tuple_insert(idx0, axis, i)
p = arr._ptr(idx)
q = new._ptr(idx)
q[0] = p[0]
if stop != N:
for i in range(stop, N):
for idx0 in util.multirange(shape_no_axis):
idx1 = util.tuple_insert(idx0, axis, i)
idx2 = util.tuple_insert(idx0, axis, i - numtodel)
p = arr._ptr(idx1)
q = new._ptr(idx2)
q[0] = p[0]
if step != 1:
off = start
for i in range(start, stop):
if i in xr:
continue
for idx0 in util.multirange(shape_no_axis):
idx1 = util.tuple_insert(idx0, axis, i)
idx2 = util.tuple_insert(idx0, axis, off)
p = arr._ptr(idx1)
q = new._ptr(idx2)
q[0] = p[0]
off += 1
return new
else:
if isinstance(obj[0], int):
remove = [util.normalize_index(r, axis, N) for r in obj]
remove.sort()
# remove duplicates
n = len(remove)
numtodel = 0
if n < 2:
numtodel = n
else:
j = 0
for i in range(n - 1):
if remove[i] != remove[i+1]:
remove[j] = remove[i]
j += 1
remove[j] = remove[n-1]
j += 1
numtodel = j
if numtodel == 0:
return arr.copy(order=arrorder)
pnewshape[axis] -= numtodel
new = empty(newshape, dtype, arrorder)
curr = 0
skip = remove[curr]
for i in range(N):
if i == skip:
curr += 1
if curr < numtodel:
skip = remove[curr]
continue
for idx0 in util.multirange(shape_no_axis):
idx1 = util.tuple_insert(idx0, axis, i)
idx2 = util.tuple_insert(idx0, axis, i - curr)
p = arr._ptr(idx1)
q = new._ptr(idx2)
q[0] = p[0]
return new
elif isinstance(obj[0], bool):
if len(obj) != N:
raise ValueError(
f"boolean array argument obj to delete must be one dimensional and match the axis length of {N}")
numtodel = 0
for r in obj:
if r:
numtodel += 1
pnewshape[axis] -= numtodel
new = empty(newshape, dtype, arrorder)
off = 0
for i in range(N):
if obj[i]:
continue
for idx0 in util.multirange(shape_no_axis):
idx1 = util.tuple_insert(idx0, axis, i)
idx2 = util.tuple_insert(idx0, axis, off)
p = arr._ptr(idx1)
q = new._ptr(idx2)
q[0] = p[0]
off += 1
return new
else:
compile_error("arrays used as indices must be of integer (or boolean) type")
elif axis is None:
return delete(arr.ravel(), obj, axis=0)
else:
compile_error("'axis' must be None or an int")
def append(arr, values, axis = None):
arr = asarray(arr)
values = asarray(values)
if staticlen(arr.shape) != staticlen(values.shape):
compile_error("'arr' and 'values' must have the same number of dimensions")
dtype1 = arr.dtype
dtype2 = values.dtype
dtype_out = type(util.coerce(dtype1, dtype2))
ndim: Static[int] = staticlen(arr.shape)
shape = arr.shape
val_shape = values.shape
arr_nbytes = arr.nbytes
val_nbytes = values.nbytes
newshape = shape
pnewshape = Ptr[int](__ptr__(newshape).as_byte())
cc, fc = arr._contig
arrorder = 'F' if (fc and not cc) else 'C'
if isinstance(axis, int):
axis = util.normalize_axis_index(axis, ndim)
if util.tuple_delete(val_shape, axis) != util.tuple_delete(shape, axis):
raise ValueError("'arr' and 'values' must have the same shape aside from specified axis")
pnewshape[axis] += val_shape[axis]
new = empty(newshape, dtype_out, arrorder)
off = shape[axis]
if dtype1 is dtype_out and ((axis == 0 and cc) or (axis == ndim - 1 and fc)):
str.memcpy(new.data.as_byte(), arr.data.as_byte(), arr_nbytes)
else:
for idx in util.multirange(shape):
p = arr._ptr(idx)
q = new._ptr(idx)
q[0] = util.cast(p[0], dtype_out)
if dtype2 is dtype_out and ((axis == 0 and cc) or (axis == ndim - 1 and fc)):
q = new.data.as_byte() + arr_nbytes
str.memcpy(q, values.data.as_byte(), val_nbytes)
else:
for idx1 in util.multirange(val_shape):
idx2 = util.tuple_add(idx1, axis, off)
p = values._ptr(idx1)
q = new._ptr(idx2)
q[0] = util.cast(p[0], dtype_out)
return new
elif axis is None:
new = empty((arr.size + values.size,), dtype_out, arrorder)
q = new.data
off = 0
val_cc = values._contig[0]
if dtype1 is dtype_out and cc:
str.memcpy(q.as_byte(), arr.data.as_byte(), arr_nbytes)
else:
for idx in util.multirange(shape):
p = arr._ptr(idx)
q[off] = util.cast(p[0], dtype_out)
off += 1
if dtype2 is dtype_out and val_cc:
str.memcpy(q.as_byte() + arr_nbytes, values.data.as_byte(), val_nbytes)
else:
for idx in util.multirange(val_shape):
p = values._ptr(idx)
q[off] = util.cast(p[0], dtype_out)
off += 1
return new
else:
compile_error("'axis' must be None or an int")
def insert(arr, obj, values, axis = None):
def normalize_special(idx: int, axis: int, n: int, numnew: int):
idx0 = idx
if idx < 0:
idx += n
if idx < 0 or idx >= n + numnew:
raise IndexError(f'index {idx0} is out of bounds for axis {axis} with size {n + numnew}')
return idx
arr = asarray(arr)
values = asarray(values)
dtype = arr.dtype
shape = arr.shape
newshape = shape
pnewshape = Ptr[int](__ptr__(newshape).as_byte())
cc, fc = arr._contig
arrorder = 'F' if (fc and not cc) else 'C'
ndim: Static[int] = staticlen(arr.shape)
if isinstance(axis, int):
axis = util.normalize_axis_index(axis, arr.ndim)
N = pnewshape[axis]
numnew = 0
indices: List[Tuple[int,int]]
if isinstance(obj, slice):
start, stop, step, length = obj.adjust_indices(N)
numnew = length
indices = [(normalize_special(a, axis, N, numnew), i)
for i, a in enumerate(range(start, stop, step))]
indices.sort()
elif isinstance(obj, int):
numnew = 1
indices = [(normalize_special(obj, axis, N, numnew), 0)]
else:
numnew = len(obj)
indices = [(normalize_special(a, axis, N, numnew), i)
for i, a in enumerate(obj)]
indices.sort()
numnew = len(indices)
for i in range(numnew):
idx, rank = indices[i]
idx += i
util.normalize_index(idx, axis, N + numnew) # error check
indices[i] = (idx, rank)
if numnew == 0:
return arr.copy(order=arrorder)
pnewshape[axis] += numnew
new = empty(newshape, dtype, arrorder)
newshape_no_axis = util.tuple_delete(newshape, axis)
curr = 0
off = 0
next_index, next_rank = indices[curr]
multiple_values = False
if staticlen(values.shape) > 0:
multiple_values = (len(values) == numnew)
for ai in range(pnewshape[axis]):
if ai == next_index:
for idx0 in util.multirange(newshape_no_axis):
idx1 = util.tuple_insert(idx0, axis, 0)
idx2 = util.tuple_insert(idx0, axis, ai)
q = new._ptr(idx2)
if staticlen(values.shape) == 0:
q[0] = util.cast(values.data[0], dtype)
else:
if multiple_values:
p = asarray(values[next_rank])._ptr(idx1, broadcast=True)
q[0] = util.cast(p[0], dtype)
else:
p = values._ptr(idx1, broadcast=True)
q[0] = util.cast(p[0], dtype)
curr += 1
if curr < numnew:
next_index, next_rank = indices[curr]
else:
for idx0 in util.multirange(newshape_no_axis):
idx1 = util.tuple_insert(idx0, axis, off)
idx2 = util.tuple_insert(idx0, axis, ai)
p = arr._ptr(idx1)
q = new._ptr(idx2)
q[0] = util.cast(p[0], dtype)
off += 1
return new
elif axis is None:
return insert(arr.ravel(), obj, values, axis=0)
else:
compile_error("'axis' must be None or an int")
def array_split(ary, indices_or_sections, axis: int = 0):
def correct(idx: int, n: int):
if idx > n:
return n
elif idx < -n:
return 0
elif idx < 0:
return idx + n
else:
return idx
def slice_axis(arr, axis: int, start: int, stop: int):
ndim: Static[int] = staticlen(ary.shape)
dtype = arr.dtype
shape = arr.shape
pshape = Ptr[int](__ptr__(shape).as_byte())
limit = pshape[axis]
start = correct(start, limit)
stop = correct(stop, limit)
base = (0,) * ndim
pbase = Ptr[int](__ptr__(base).as_byte())
pbase[axis] = start
pshape[axis] = stop - start if stop > start else 0
sub = arr._ptr(base) if stop > start else Ptr[dtype]()
return ndarray(shape, arr.strides, sub)
ary = asarray(ary)
axis = util.normalize_axis_index(axis, ary.ndim)
ntotal = ary.shape[axis]
if isinstance(indices_or_sections, int):
nsections = indices_or_sections
if nsections <= 0:
raise ValueError("number sections must be larger than 0.")
neach, extras = divmod(ntotal, nsections)
result = List(capacity=nsections)
start = 0
for i in range(nsections):
stop = start + neach + (1 if i < extras else 0)
result.append(slice_axis(ary, axis, start, stop))
start = stop
return result
else:
result = List(capacity=(len(indices_or_sections) + 1))
prev = 0
for s in indices_or_sections:
result.append(slice_axis(ary, axis, prev, s))
prev = s
result.append(slice_axis(ary, axis, prev, ntotal))
return result
def split(ary, indices_or_sections, axis: int = 0):
ary = asarray(ary)
if isinstance(indices_or_sections, int):
sections = indices_or_sections
axis = util.normalize_axis_index(axis, ary.ndim)
N = ary.shape[axis]
if N % sections:
raise ValueError("array split does not result in an equal division")
return array_split(ary, indices_or_sections, axis)
def vsplit(ary, indices_or_sections):
ary = asarray(ary)
if staticlen(ary.shape) < 2:
compile_error("vsplit only works on arrays of 2 or more dimensions")
return split(ary, indices_or_sections, 0)
def hsplit(ary, indices_or_sections):
ary = asarray(ary)
if staticlen(ary.shape) == 0:
compile_error("hsplit only works on arrays of 1 or more dimensions")
return split(ary, indices_or_sections, 1 if ary.ndim > 1 else 0)
def dsplit(ary, indices_or_sections):
ary = asarray(ary)
if staticlen(ary.shape) < 3:
compile_error("dsplit only works on arrays of 3 or more dimensions")
return split(ary, indices_or_sections, 2)
def trim_zeros(filt, trim: str = 'fb'):
filt = asarray(filt)
if staticlen(filt.shape) != 1:
compile_error("trim_zeros() only takes 1-dimensional arrays as input")
fb = (trim == 'fb' or trim == 'bf')
just_f = (trim == 'f')
just_b = (trim == 'b')
trim_front = fb or just_f
trim_back = fb or just_b
if trim_front:
n = len(filt)
i = 0
while i < n and not filt[i]:
i += 1
filt = filt[i:]
if trim_back:
n = len(filt)
i = n - 1
while i >= 0 and not filt[i]:
i -= 1
filt = filt[:i+1]
return filt
def flip(m, axis = None):
m = asarray(m)
ndim: Static[int] = staticlen(m.shape)
if axis is None:
return flip(m, util.tuple_range(ndim))
elif isinstance(axis, int):
return flip(m, (axis,))
axis = util.normalize_axis_tuple(axis, ndim)
dtype = m.dtype
offset = 0
shape = m.shape
strides = m.strides
pshape = Ptr[int](__ptr__(shape).as_byte())
pstrides = Ptr[int](__ptr__(strides).as_byte())
for i in staticrange(ndim):
if i in axis:
st = pstrides[i]
offset += (pshape[i] - 1) * st
pstrides[i] = -st
return ndarray(shape, strides, Ptr[dtype](m.data.as_byte() + offset))
def fliplr(m):
m = asarray(m)
if staticlen(m.shape) < 2:
compile_error("Input must be >= 2-d.")
return flip(m, axis=1)
def flipud(m):
m = asarray(m)
if staticlen(m.shape) < 1:
compile_error("Input must be >= 1-d.")
return flip(m, axis=0)
def rot90(m, k: int = 1, axes: Tuple[int,int] = (0, 1)):
m = asarray(m)
ndim: Static[int] = staticlen(m.shape)
if axes[0] == axes[1] or abs(axes[0] - axes[1]) == ndim:
raise ValueError("Axes must be different.")
if (axes[0] >= ndim or axes[0] < -ndim
or axes[1] >= ndim or axes[1] < -ndim):
raise ValueError(f"Axes={axes} out of range for array of ndim={m.ndim}.")
k %= 4
if k == 0:
return m[:]
if k == 2:
return flip(flip(m, axes[0]), axes[1])
axes_list = util.tuple_range(ndim)
axes_list = util.tuple_swap(axes_list, axes[0], axes[1])
if k == 1:
return flip(m, axes[1]).transpose(axes_list)
else:
# k == 3
return flip(m.transpose(axes_list), axes[1])
def resize(a, new_shape):
if isinstance(new_shape, int):
return resize(a, (new_shape,))
a = asarray(a)
new_size = 1
for dim_length in new_shape:
new_size *= dim_length
if dim_length < 0:
raise ValueError('all elements of `new_shape` must be non-negative')
if a.size == 0 or new_size == 0:
# First case must zero fill. The second would have repeats == 0.
return zeros(new_shape, dtype=a.dtype)
new = empty(new_shape, dtype=a.dtype)
off = 0
go = True
while go:
for idx in util.multirange(a.shape):
p = a._ptr(idx)
new.data[off] = p[0]
off += 1
if off == new_size:
go = False
break
return new
def tile(A, reps):
if isinstance(reps, int):
return tile(A, (reps,))
A = asarray(A)
dtype = A.dtype
ndim: Static[int] = staticlen(A.shape)
d: Static[int] = staticlen(reps)
if ndim < d:
new_shape = ((1,) * (d - ndim)) + A.shape
return tile(A.reshape(new_shape), reps)
if ndim > d:
new_reps = ((1,) * (ndim - d)) + reps
return tile(A, new_reps)
out_shape = util.tuple_apply(int.__mul__, A.shape, reps)
new = empty(out_shape, dtype=dtype)
for idx1 in util.multirange(out_shape):
idx2 = util.tuple_apply(int.__mod__, idx1, A.shape)
q = new._ptr(idx1)
p = A._ptr(idx2)
q[0] = p[0]
return new
def roll(a, shift, axis = None):
a = asarray(a)
ndim: Static[int] = staticlen(a.shape)
if axis is None:
return roll(a.ravel(), shift, axis=0).reshape(a.shape)
if isinstance(axis, int):
return roll(a, shift, (axis,))
if isinstance(shift, int):
return roll(a, (shift,), axis)
na: Static[int] = staticlen(axis)
ns: Static[int] = staticlen(shift)
if na == 0:
compile_error("empty tuple given for 'axis'")
if ns == 0:
compile_error("empty tuple given for 'shift'")
if na == 1 and ns != 1:
return roll(a, shift, axis * ns)
if na != 1 and ns == 1:
return roll(a, shift * na, axis)
if staticlen(axis) != staticlen(shift):
compile_error("'shift' and 'axis' must be tuples of the same size")
axis = util.normalize_axis_tuple(axis, a.ndim, allow_duplicates=True)
shifts = (0,) * ndim
pshifts = Ptr[int](__ptr__(shifts).as_byte())
for i in staticrange(staticlen(axis)):
pshifts[axis[i]] += shift[i]
new = empty_like(a)
@pure
@llvm
def cdiv(a: int, b: int) -> int:
%0 = sdiv i64 %a, %b
ret i64 %0
def pymod(a: int, b: int):
d = cdiv(a, b)
m = a - d * b
if m and ((b ^ m) < 0):
m += b
return m
for idx1 in util.multirange(a.shape):
idx2 = util.tuple_apply(int.__add__, idx1, shifts)
idx2 = util.tuple_apply(pymod, idx2, a.shape)
p = a._ptr(idx1)
q = new._ptr(idx2)
q[0] = p[0]
return new
def _atleast_nd(a, ndim: Static[int]):
return array(a, ndmin=ndim, copy=False)
def _accumulate(values):
import itertools
return list(itertools.accumulate(values))
def _longest_shape(a):
if isinstance(a, ndarray):
return (0,) * staticlen(a.shape)
elif isinstance(a, List):
if len(a) == 0:
raise ValueError("Lists passed to block cannot be empty")
return _longest_shape(a[0])
elif isinstance(a, Tuple):
if staticlen(a) == 0:
compile_error("Tuples passed to block cannot be empty")
if staticlen(a) == 1:
return _longest_shape(a[0])
ls1 = _longest_shape(a[0])
ls2 = _longest_shape(a[1:])
if staticlen(ls1) > staticlen(ls2):
return ls1
else:
return ls2
else:
return ()
def _bottom_index(a):
if isinstance(a, ndarray):
return ()
elif isinstance(a, List):
return (0,) + _bottom_index(a[0])
elif isinstance(a, Tuple):
bi0 = _bottom_index(a[0])
for i in staticrange(1, staticlen(a)):
if staticlen(_bottom_index(a[i])) != staticlen(bi0):
compile_error("Depths of block argument are mismatched")
return (0,) + bi0
else:
return ()
def _final_size(a):
if isinstance(a, ndarray):
return a.size
elif isinstance(a, List) or isinstance(a, Tuple):
ans = 0
for x in a:
ans += _final_size(x)
return ans
else:
return 1
def _block(arrays, max_depth: Static[int], result_ndim: Static[int], depth: Static[int] = 0):
if depth < max_depth:
arrs = [_block(arr, max_depth, result_ndim, depth+1)
for arr in arrays]
return concatenate(arrs, axis=-(max_depth-depth))
else:
return _atleast_nd(arrays, result_ndim)
def _block_concatenate(arrays, list_ndim: Static[int], result_ndim: Static[int]):
result = _block(arrays, list_ndim, result_ndim)
if list_ndim == 0:
result = result.copy()
return result
def _concatenate_shapes(shapes, axis: Static[int]):
# Cache a result that will be reused.
shape_at_axis = [shape[axis] for shape in shapes]
# Take a shape, any shape
first_shape = shapes[0]
first_shape_pre = first_shape[:axis]
first_shape_post = first_shape[axis+1:]
if any(shape[:axis] != first_shape_pre or
shape[axis+1:] != first_shape_post for shape in shapes):
raise ValueError(
f"Mismatched array shapes in block along axis {axis}.")
shape = (first_shape_pre + (sum(shape_at_axis),) + first_shape[axis+1:])
offsets_at_axis = _accumulate(shape_at_axis)
slice_prefixes = [(slice(start, end),)
for start, end in zip([0] + offsets_at_axis,
offsets_at_axis)]
return shape, slice_prefixes
def _block_info_recursion(arrs, max_depth: Static[int], result_ndim: Static[int], depth: Static[int] = 0):
if depth < max_depth:
shapes = []
slices = []
arrays = []
for arr in arrs:
sh, sl, ar = _block_info_recursion(arr, max_depth, result_ndim, depth+1)
shapes.append(sh)
slices.append(sl)
arrays.append(ar)
axis: Static[int] = result_ndim - max_depth + depth
shape, slice_prefixes = _concatenate_shapes(shapes, axis)
# Prepend the slice prefix and flatten the slices
slices = [slice_prefix + the_slice
for slice_prefix, inner_slices in zip(slice_prefixes, slices)
for the_slice in inner_slices]
# Flatten the array list
arrays_flat = []
for arr in arrays:
for a in arr:
arrays_flat.append(a)
return shape, slices, arrays_flat
else:
arr = _atleast_nd(arrs, result_ndim)
return arr.shape, [()], [arr]
def _block_slicing(arrays, list_ndim: Static[int], result_ndim: Static[int]):
shape, slices, arrays = _block_info_recursion(
arrays, list_ndim, result_ndim)
dtype = arrays[0].dtype
C_order = True
F_order = True
for arr in arrays:
cc, fc = arr._contig
C_order = C_order and cc
F_order = F_order and fc
if not C_order and not F_order:
break
order = 'F' if F_order and not C_order else 'C'
result = empty(shape=shape, dtype=dtype, order=order)
for the_slice, arr in zip(slices, arrays):
result[(Ellipsis,) + the_slice] = arr
return result
def block(arrays):
ls = _longest_shape(arrays)
bi = _bottom_index(arrays)
final_size = _final_size(arrays)
list_ndim: Static[int] = staticlen(bi)
arr_ndim: Static[int] = staticlen(ls)
result_ndim: Static[int] = arr_ndim if arr_ndim > list_ndim else list_ndim
# Note: This is just the heuristic NumPy uses. I have not tested how
# good it is for Codon.
if list_ndim * final_size > (2 * 512 * 512):
return _block_slicing(arrays, list_ndim, result_ndim)
else:
return _block_concatenate(arrays, list_ndim, result_ndim)
def _close(a, b, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False):
A = type(a)
B = type(b)
if A is not B:
C = type(util.coerce(A, B))
if (C is not float and
C is not float32 and
C is not float16 and
C is not complex and
C is not complex64):
return _close(util.to_float(a),
util.to_float(b),
rtol=rtol,
atol=atol,
equal_nan=equal_nan)
else:
return _close(util.cast(a, C),
util.cast(b, C),
rtol=rtol,
atol=atol,
equal_nan=equal_nan)
elif (A is not float and
A is not float32 and
A is not float16 and
A is not complex and
A is not complex64):
return _close(util.to_float(a),
util.to_float(b),
rtol=rtol,
atol=atol,
equal_nan=equal_nan)
if A is float or A is float32 or A is float16:
fin_a = util.isfinite(a)
fin_b = util.isfinite(b)
if fin_a and fin_b:
return abs(a - b) <= A(atol) + A(rtol) * abs(b)
else:
nan_a = util.isnan(a)
nan_b = util.isnan(b)
if nan_a or nan_b:
if equal_nan:
return nan_a and nan_b
else:
return False
else:
return a == b
else: # complex or complex64
R = type(a.real)
fin_a = util.isfinite(a.real) and util.isfinite(a.imag)
fin_b = util.isfinite(b.real) and util.isfinite(b.imag)
if fin_a and fin_b:
return abs(a - b) <= R(atol) + R(rtol) * abs(b)
else:
nan_a = util.isnan(a.real) or util.isnan(a.imag)
nan_b = util.isnan(b.real) or util.isnan(b.imag)
if nan_a or nan_b:
if equal_nan:
return nan_a and nan_b
else:
return False
else:
return a == b
def isclose(a, b, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False):
a = asarray(a)
b = asarray(b)
if a.ndim == 0 and b.ndim == 0:
return _close(a.item(), b.item(), rtol=rtol, atol=atol, equal_nan=equal_nan)
ans_shape = broadcast_shapes(a.shape, b.shape)
ans = empty(ans_shape, bool)
for idx in util.multirange(ans_shape):
xa = a._ptr(idx, broadcast=True)[0]
xb = b._ptr(idx, broadcast=True)[0]
xans = _close(xa, xb, rtol=rtol, atol=atol, equal_nan=equal_nan)
ans._ptr(idx)[0] = xans
return ans
def allclose(a, b, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False):
a = asarray(a)
b = asarray(b)
if a.ndim == 0 and b.ndim == 0:
return _close(a.item(), b.item(), rtol=rtol, atol=atol, equal_nan=equal_nan)
ans_shape = broadcast_shapes(a.shape, b.shape)
for idx in util.multirange(ans_shape):
xa = a._ptr(idx, broadcast=True)[0]
xb = b._ptr(idx, broadcast=True)[0]
if not _close(xa, xb, rtol=rtol, atol=atol, equal_nan=equal_nan):
return False
return True
def array_equal(a1, a2, equal_nan: bool = False):
from .ndmath import isnan
a1 = asarray(a1)
a2 = asarray(a2)
if a1.ndim != a2.ndim:
return False
if a1.shape != a2.shape:
return False
dtype = type(util.coerce(a1.dtype, a2.dtype))
if a1._contig_match(a2):
p1 = a1.data
p2 = a2.data
n = a1.size
for i in range(n):
e1 = util.cast(p1[i], dtype)
e2 = util.cast(p2[i], dtype)
if equal_nan:
if e1 != e2 and not (isnan(e1) and isnan(e2)):
return False
else:
if e1 != e2:
return False
else:
for idx in util.multirange(a1.shape):
e1 = util.cast(a1._ptr(idx)[0], dtype)
e2 = util.cast(a2._ptr(idx)[0], dtype)
if equal_nan:
if e1 != e2 and not (isnan(e1) and isnan(e2)):
return False
else:
if e1 != e2:
return False
return True
def array_equiv(a1, a2):
def can_broadcast(s1, s2):
if staticlen(s1) > staticlen(s2):
return can_broadcast(s1[-staticlen(s2):], s2)
elif staticlen(s1) < staticlen(s2):
return can_broadcast(s1, s2[-staticlen(s1):])
else:
for i in staticrange(staticlen(s1)):
d1 = s1[i]
d2 = s2[i]
if d1 != d2 and not (d1 == 1 or d2 == 1):
return False
return True
a1 = asarray(a1)
a2 = asarray(a2)
if not can_broadcast(a1.shape, a2.shape):
return False
bshape = broadcast_shapes(a1.shape, a2.shape)
dtype = type(util.coerce(a1.dtype, a2.dtype))
for idx in util.multirange(bshape):
e1 = util.cast(a1._ptr(idx, broadcast=True)[0], dtype)
e2 = util.cast(a2._ptr(idx, broadcast=True)[0], dtype)
if e1 != e2:
return False
return True
def squeeze(a, axis = None):
if axis is None:
compile_error("squeeze() must specify 'axis' argument in Codon-NumPy")
if isinstance(axis, int):
return squeeze(a, (axis,))
if not isinstance(axis, Tuple):
compile_error("'axis' must be an int or a tuple of ints")
a = asarray(a)
axis = util.normalize_axis_tuple(axis, a.ndim)
shape = a.shape
strides = a.strides
new_shape = (0,) * (staticlen(shape) - staticlen(axis))
new_strides = (0,) * (staticlen(shape) - staticlen(axis))
pnew_shape = Ptr[int](__ptr__(new_shape).as_byte())
pnew_strides = Ptr[int](__ptr__(new_strides).as_byte())
j = 0
for i in staticrange(staticlen(shape)):
if i in axis:
if shape[i] != 1:
raise ValueError("cannot select an axis to squeeze out which has size not equal to one")
else:
pnew_shape[j] = shape[i]
pnew_strides[j] = strides[i]
j += 1
return ndarray(new_shape, new_strides, a.data)
def pad(array, pad_width, mode = 'constant', **kwargs):
from .ndmath import isnan
def unpack_params(x, iaxis: int, name: Static[str]):
if isinstance(x, Tuple):
if staticlen(x) == 1:
if isinstance(x[0], Tuple):
if staticlen(x[0]) == 1:
return x[0][0], x[0][0]
elif staticlen(x[0]) == 2:
return x[0]
else:
compile_error("invalid parameter '" + name + "' given")
else:
return x[0], x[0]
elif staticlen(x) == 2:
if isinstance(x[0], Tuple):
return x[iaxis]
else:
return x
else:
return x[iaxis]
else:
return x, x
def unpack_stat_length(k: int, iaxis: int, kwargs):
stat_length = kwargs.get('stat_length', k)
s1, s2 = unpack_params(stat_length, iaxis, 'stat_length')
if s1 <= 0 or s2 <= 0:
raise ValueError("'stat_length' must contain positive values")
s1 = min(s1, k)
s2 = min(s2, k)
return s1, s2
def round_if_needed(x, dtype: type):
if dtype is int or isinstance(dtype, Int) or isinstance(dtype, UInt):
return util.cast(util.rint(x), dtype)
else:
return util.cast(x, dtype)
def pad_from_function(a: ndarray, pw, padding_func, kwargs, extra = None):
shape = a.shape
strides = a.strides
ndim: Static[int] = staticlen(shape)
for i in staticrange(ndim):
length = shape[i]
stride = strides[i]
for idx in util.multirange(util.tuple_delete(shape, i)):
p = a._ptr(util.tuple_insert(idx, i, 0))
vector = ndarray((length,), (stride,), p)
padding_func(vector, pw[i], i, kwargs, extra)
def pad_constant(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
p1, p2 = iaxis_pad_width
cval = kwargs.get('constant_values', 0)
dtype = vector.dtype
c1, c2 = unpack_params(cval, iaxis, 'constant_values')
n = vector.size
# Note we don't need to do range checks since padded
# array size will always ensure we're not out of bounds,
# as long as each `iaxis_pad_width` is non-negative.
for i in range(0, p1):
vector._ptr((i,))[0] = util.cast(c1, dtype)
for i in range(n - p2, n):
vector._ptr((i,))[0] = util.cast(c2, dtype)
def pad_wrap(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
p1, p2 = iaxis_pad_width
n = vector.size
k = n - (p1 + p2)
for i in range(n - p2, n):
vector._ptr((i,))[0] = vector._ptr((i - k,))[0]
for i in range(p1 - 1, -1, -1):
vector._ptr((i,))[0] = vector._ptr((i + k,))[0]
def pad_edge(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
p1, p2 = iaxis_pad_width
n = vector.size
c1 = vector._ptr((p1,))[0]
c2 = vector._ptr((n - 1 - p2,))[0]
for i in range(0, p1):
vector._ptr((i,))[0] = c1
for i in range(n - p2, n):
vector._ptr((i,))[0] = c2
def pad_max(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
p1, p2 = iaxis_pad_width
n = vector.size
k = n - (p1 + p2)
s1, s2 = unpack_stat_length(k, iaxis, kwargs)
m1 = vector._ptr((p1,))[0]
for i in range(p1 + 1, p1 + s1):
e = vector._ptr((i,))[0]
if e > m1:
m1 = e
if s1 == k and s2 == k:
m2 = m1
else:
m2 = vector._ptr((n - p2 - s2,))[0]
for i in range(n - p2 - s2 + 1, n - p2):
e = vector._ptr((i,))[0]
if e > m2:
m2 = e
for i in range(0, p1):
vector._ptr((i,))[0] = m1
for i in range(n - p2, n):
vector._ptr((i,))[0] = m2
def pad_min(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
p1, p2 = iaxis_pad_width
n = vector.size
k = n - (p1 + p2)
s1, s2 = unpack_stat_length(k, iaxis, kwargs)
m1 = vector._ptr((p1,))[0]
for i in range(p1 + 1, p1 + s1):
e = vector._ptr((i,))[0]
if e < m1:
m1 = e
if s1 == k and s2 == k:
m2 = m1
else:
m2 = vector._ptr((n - p2 - s2,))[0]
for i in range(n - p2 - s2 + 1, n - p2):
e = vector._ptr((i,))[0]
if e < m2:
m2 = e
for i in range(0, p1):
vector._ptr((i,))[0] = m1
for i in range(n - p2, n):
vector._ptr((i,))[0] = m2
def pad_mean(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
p1, p2 = iaxis_pad_width
n = vector.size
k = n - (p1 + p2)
s1, s2 = unpack_stat_length(k, iaxis, kwargs)
m1 = vector._ptr((p1,))[0]
for i in range(p1 + 1, p1 + s1):
m1 += vector._ptr((i,))[0]
if s1 == k and s2 == k:
m2 = m1
else:
m2 = vector._ptr((n - p2 - s2,))[0]
for i in range(n - p2 - s2 + 1, n - p2):
m2 += vector._ptr((i,))[0]
dtype = vector.dtype
av1 = round_if_needed(m1 / util.cast(s1, dtype), dtype)
av2 = round_if_needed(m2 / util.cast(s2, dtype), dtype)
for i in range(0, p1):
vector._ptr((i,))[0] = av1
for i in range(n - p2, n):
vector._ptr((i,))[0] = av2
def pad_median(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
p1, p2 = iaxis_pad_width
n = vector.size
k = n - (p1 + p2)
s1, s2 = unpack_stat_length(k, iaxis, kwargs)
buf = extra
nan_idx = -1
for i in range(p1, p1 + s1):
e = vector._ptr((i,))[0]
if isnan(e):
nan_idx = i
break
else:
buf[i - p1] = e
dtype = vector.dtype
if nan_idx >= 0:
m1 = vector._ptr((nan_idx,))[0]
else:
m1a, m1b = util.median(buf, s1)
if m1a == m1b:
m1 = m1a
else:
m1 = round_if_needed((m1a + m1b) / util.cast(2, dtype), dtype)
if s1 == k and s2 == k:
m2 = m1
else:
nan_idx = -1
base = n - p2 - s2
for i in range(base, n - p2):
e = vector._ptr((i,))[0]
if isnan(e):
nan_idx = i
break
else:
buf[i - base] = e
if nan_idx >= 0:
m2 = vector._ptr((nan_idx,))[0]
else:
m2a, m2b = util.median(buf, s2)
if m2a == m2b:
m2 = m2a
else:
m2 = round_if_needed((m2a + m2b) / util.cast(2, dtype), dtype)
for i in range(0, p1):
vector._ptr((i,))[0] = m1
for i in range(n - p2, n):
vector._ptr((i,))[0] = m2
def pad_linear_ramp(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
def fill_linear(vec: ndarray, offset: int, start: float,
stop: float, num: int, rev: bool):
if num == 0:
return
dtype = vec.dtype
delta = stop - start
step = delta / num
for i in range(num):
j = num - 1 - i if rev else i
p = vec._ptr((j + offset,))
e = ((i / num) * delta) + start if not step else (i * step) + start
p[0] = round_if_needed(e, dtype)
p1, p2 = iaxis_pad_width
n = vector.size
end_values = kwargs.get('end_values', 0)
start1, end2 = unpack_params(end_values, iaxis, 'end_values')
end1 = vector._ptr((p1,))[0]
start2 = vector._ptr((n - 1 - p2,))[0]
fill_linear(vector,
offset=0,
start=util.cast(start1, float),
stop=util.cast(end1, float),
num=p1,
rev=False)
fill_linear(vector,
offset=(n - p2),
start=util.cast(end2, float),
stop=util.cast(start2, float),
num=p2,
rev=True)
def pad_reflect_or_symmetric(vector: ndarray,
iaxis_pad_width: Tuple[int, int],
iaxis: int,
kwargs,
extra):
p1, p2 = iaxis_pad_width
n = vector.size
k = n - (p1 + p2)
diff = extra
if k == 1:
e = vector._ptr((p1,))[0]
for i in range(0, p1):
vector._ptr((i,))[0] = e
for i in range(n - p2, n):
vector._ptr((i,))[0] = e
return
even = kwargs.get('reflect_type', 'even') != 'odd'
# left side
i = p1 - 1
while i >= 0:
z = i
edge = vector._ptr((z + 1,))[0]
for j in range(k - diff):
e = vector._ptr((z + j + (1 + diff),))[0]
if not even:
e = (edge + edge) - e
vector._ptr((i,))[0] = e
i -= 1
if i < 0:
break
# right side
i = n - p2
while i < n:
z = i
edge = vector._ptr((z - 1,))[0]
for j in range(k - diff):
e = vector._ptr((z - j - (1 + diff),))[0]
if not even:
e = (edge + edge) - e
vector._ptr((i,))[0] = e
i += 1
if i >= n:
break
a = asarray(array)
s = a.shape
dtype = a.dtype
ndim: Static[int] = staticlen(s)
if ndim == 0:
return a
if isinstance(mode, str):
if a.size == 0 and mode not in ('empty', 'constant'):
raise ValueError("can't pad empty array using modes other than 'constant' or 'empty'")
if isinstance(pad_width, int):
pw = ((pad_width, pad_width),) * ndim
elif isinstance(pad_width, Tuple[int]):
pw = ((pad_width[0], pad_width[0]),) * ndim
elif isinstance(pad_width, Tuple[int, int]):
pw = (pad_width,) * ndim
elif isinstance(pad_width, Tuple[Tuple[int, int]]):
pw = (pad_width[0],) * ndim
elif staticlen(pad_width) == 1:
pw = pad_width * ndim
else:
pw = pad_width
if staticlen(pw) != ndim:
compile_error("invalid pad_width")
for p in pw:
if p[0] < 0 or p[1] < 0:
raise ValueError("padding can't be negative")
new_shape = tuple(s[i] + pw[i][0] + pw[i][1] for i in staticrange(ndim))
if isinstance(mode, str):
ans = empty(new_shape, dtype)
else:
ans = zeros(new_shape, dtype)
# copy in original array
for idx in util.multirange(s):
idx1 = tuple(idx[i] + pw[i][0] for i in staticrange(ndim))
p = a._ptr(idx)
q = ans._ptr(idx1)
q[0] = p[0]
if isinstance(mode, str):
if mode == 'empty':
pass # do nothing
elif mode == 'constant':
pad_from_function(ans, pw, pad_constant, kwargs)
elif mode == 'wrap':
pad_from_function(ans, pw, pad_wrap, kwargs)
elif mode == 'edge':
pad_from_function(ans, pw, pad_edge, kwargs)
elif mode == 'linear_ramp':
pad_from_function(ans, pw, pad_linear_ramp, kwargs)
elif mode == 'maximum':
pad_from_function(ans, pw, pad_max, kwargs)
elif mode == 'minimum':
pad_from_function(ans, pw, pad_min, kwargs)
elif mode == 'mean':
pad_from_function(ans, pw, pad_mean, kwargs)
elif mode == 'median':
buf_size = 0
for iaxis in range(ndim):
s1, s2 = unpack_stat_length(s[iaxis], iaxis, kwargs)
if s1 > buf_size:
buf_size = s1
if s2 > buf_size:
buf_size = s2
buf = Ptr[dtype](buf_size)
pad_from_function(ans, pw, pad_median, kwargs, extra=buf)
elif mode == 'reflect':
pad_from_function(ans, pw, pad_reflect_or_symmetric, kwargs, extra=1)
elif mode == 'symmetric':
pad_from_function(ans, pw, pad_reflect_or_symmetric, kwargs, extra=0)
else:
raise ValueError(f"mode {repr(mode)} is not supported")
else:
pad_from_function(ans, pw, mode, kwargs)
return ans
#############
# Searching #
#############
def nonzero(a):
a = asarray(a)
if staticlen(a.shape) == 0:
# Note: this is technically deprecated behavior
if bool(a.item()):
ans = empty(1, int)
ans.data[0] = 0
return (ans,)
else:
return (empty(0, int),)
else:
num_true = 0
if a._is_contig:
for i in range(a.size):
if bool(a.data[i]):
num_true += 1
else:
for idx in util.multirange(a.shape):
if bool(a._ptr(idx)[0]):
num_true += 1
ans = tuple(empty(num_true, int) for _ in a.shape)
k = 0
for idx in util.multirange(a.shape):
if bool(a._ptr(idx)[0]):
for i in staticrange(staticlen(ans)):
ans[i].data[k] = idx[i]
k += 1
return ans
def flatnonzero(a):
a = asarray(a)
if staticlen(a.shape) <= 1:
return nonzero(a)[0]
else:
sz = a.size
cc, fc = a._contig
num_true = 0
if cc or fc:
for i in range(sz):
if bool(a.data[i]):
num_true += 1
else:
for idx in util.multirange(a.shape):
if bool(a._ptr(idx)[0]):
num_true += 1
ans = empty(num_true, int)
k = 0
if cc:
for i in range(sz):
if bool(a.data[i]):
ans.data[k] = i
k += 1
else:
j = 0
for idx in util.multirange(a.shape):
if bool(a._ptr(idx)[0]):
ans.data[k] = j
k += 1
j += 1
return ans
def argwhere(a):
a = asarray(a)
if staticlen(a.shape) == 0:
m = 1 if bool(a.item()) else 0
return empty((m, 0), int)
else:
num_true = 0
if a._is_contig:
for i in range(a.size):
if bool(a.data[i]):
num_true += 1
else:
for idx in util.multirange(a.shape):
if bool(a._ptr(idx)[0]):
num_true += 1
ans = empty((num_true, a.ndim), int)
k = 0
for idx in util.multirange(a.shape):
if bool(a._ptr(idx)[0]):
for i in range(a.ndim):
ans._ptr((k, i))[0] = idx[i]
k += 1
return ans
def where(condition, x, y):
condition = asarray(condition)
x = asarray(x)
y = asarray(y)
T = type(util.coerce(x.dtype, y.dtype))
if condition._contig_match(x) and condition._contig_match(y):
cc, fc = condition._contig
ans = empty(condition.shape, dtype=T, order=('C' if cc else 'F'))
sz = condition.size
pc = condition.data
px = x.data
py = y.data
q = ans.data
for i in range(sz):
if pc[i]:
q[i] = util.cast(px[i], T)
else:
q[i] = util.cast(py[i], T)
return ans
else:
bshape = broadcast_shapes(condition.shape, x.shape, y.shape)
ans = empty(bshape, dtype=T)
for idx in util.multirange(bshape):
q = ans._ptr(idx)
if condition._ptr(idx, broadcast=True)[0]:
q[0] = util.cast(x._ptr(idx, broadcast=True)[0], T)
else:
q[0] = util.cast(y._ptr(idx, broadcast=True)[0], T)
return ans
@overload
def where(condition):
return nonzero(asarray(condition))
def extract(condition, arr):
condition = ravel(asarray(condition))
arr = ravel(asarray(arr))
num_true = 0
for i in range(condition.size):
if condition.data[i]:
num_true += 1
if num_true > arr.size:
raise IndexError(f"index {num_true} is out of bounds for axis 0 with size {arr.size}")
ans = empty(num_true, arr.dtype)
k = 0
i = 0
m = min(arr.size, condition.size)
while i < m and k < num_true:
if condition.data[i]:
ans.data[k] = arr.data[i]
k += 1
i += 1
return ans
def _less_than(a, b):
T = type(util.coerce(type(a), type(b)))
return util.cast(a, T) < util.cast(b, T)
def _bisect_left(p: Ptr[T], s: int, n: int, x: X, lo: int, hi: int, T: type, X: type):
while lo < hi:
mid = (lo + hi) >> 1
elem = Ptr[T](p.as_byte() + (mid * s))[0]
if _less_than(elem, x):
lo = mid + 1
else:
hi = mid
return lo, hi
def _bisect_right(p: Ptr[T], s: int, n: int, x: X, lo: int, hi: int, T: type, X: type):
while lo < hi:
mid = (lo + hi) >> 1
elem = Ptr[T](p.as_byte() + (mid * s))[0]
if _less_than(x, elem):
hi = mid
else:
lo = mid + 1
return lo, hi
def _bisect_left_perm(p: Ptr[T], s: int, n: int, x: X, perm: Ptr[int],
sp: int, lo: int, hi: int, T: type, X: type):
while lo < hi:
mid = (lo + hi) >> 1
mid_remap = Ptr[int](perm.as_byte() + (mid * sp))[0]
elem = Ptr[T](p.as_byte() + (mid_remap * s))[0]
if _less_than(elem, x):
lo = mid + 1
else:
hi = mid
return lo, hi
def _bisect_right_perm(p: Ptr[T], s: int, n: int, x: X, perm: Ptr[int],
sp: int, lo: int, hi: int, T: type, X: type):
while lo < hi:
mid = (lo + hi) >> 1
mid_remap = Ptr[int](perm.as_byte() + (mid * sp))[0]
elem = Ptr[T](p.as_byte() + (mid_remap * s))[0]
if _less_than(x, elem):
hi = mid
else:
lo = mid + 1
return lo, hi
def searchsorted(a, v, side: str = 'left', sorter = None):
a = asarray(a)
v = asarray(v)
if staticlen(a.shape) != 1:
compile_error("searchsorted requires 1-dimensional input array")
n = a.size
s = a.strides[0]
left = False
if side == 'left':
left = True
elif side == 'right':
left = False
else:
raise ValueError(f"side must be 'left' or 'right' (got {repr(side)})")
if sorter is None:
perm = None
sp = 0
else:
perm = asarray(sorter)
if staticlen(perm.shape) != 1:
compile_error("sorter argument must be 1-dimensional")
if perm.dtype is not int:
compile_error("sorter must only contain integers")
if perm.size != n:
raise ValueError("sorter.size must equal a.size")
for sorter_idx in perm:
if sorter_idx < 0 or sorter_idx >= n:
raise ValueError("Sorter index out of range.")
sp = perm.strides[0]
# copy for large needles to improve cache performance
if v.size > a.size:
a = ascontiguousarray(a)
s = a.strides[0]
if staticlen(v.shape) == 0:
x = v.item()
if sorter is None:
if left:
return _bisect_left(a.data, s, n, x, lo=0, hi=n)[0]
else:
return _bisect_right(a.data, s, n, x, lo=0, hi=n)[0]
else:
if left:
return _bisect_left_perm(a.data, s, n, x, perm.data, sp, lo=0, hi=n)[0]
else:
return _bisect_right_perm(a.data, s, n, x, perm.data, sp, lo=0, hi=n)[0]
else:
ans = empty(v.shape, int)
if v.size == 0:
return ans
last_x = v._ptr((0,) * staticlen(v.shape))[0]
lo = 0
hi = n
for idx in util.multirange(v.shape):
x = v._ptr(idx)[0]
if last_x < x:
hi = n
else:
lo = 0
hi = hi + 1 if hi < n else n
last_x = x
if sorter is None:
if left:
lo, hi = _bisect_left(a.data, s, n, x, lo=lo, hi=hi)
else:
lo, hi = _bisect_right(a.data, s, n, x, lo=lo, hi=hi)
else:
if left:
lo, hi = _bisect_left_perm(a.data, s, n, x, perm.data, sp, lo=lo, hi=hi)
else:
lo, hi = _bisect_right_perm(a.data, s, n, x, perm.data, sp, lo=lo, hi=hi)
ans._ptr(idx)[0] = lo
return ans
############
# Indexing #
############
def take(a, indices, axis = None, out = None, mode: str = 'raise'):
a = atleast_1d(asarray(a, order='C'))
if axis is None:
return take(a.ravel(), indices, axis=0, out=out, mode=mode)
indices = asarray(indices, order='C')
axis = util.normalize_axis_index(axis, a.ndim)
n = 1
m = 1
chunk = 1
nd: Static[int] = staticlen(a.shape) + staticlen(indices.shape) - 1
shape = (0,) * nd
ashape0 = a.shape
ishape0 = indices.shape
ashape = Ptr[int](__ptr__(ashape0).as_byte())
ishape = Ptr[int](__ptr__(ishape0).as_byte())
pshape = Ptr[int](__ptr__(shape).as_byte())
for i in staticrange(nd):
if i < axis:
pshape[i] = ashape[i]
n *= pshape[i]
else:
if i < axis + indices.ndim:
pshape[i] = ishape[i - axis]
m *= pshape[i]
else:
pshape[i] = ashape[i - indices.ndim + 1]
chunk *= pshape[i]
if out is None:
res = empty(shape, dtype=a.dtype)
elif isinstance(out, ndarray):
if staticlen(out.shape) != staticlen(shape) or out.dtype is not a.dtype:
compile_error("output array does not match result of ndarray.take")
if out.shape != shape:
raise ValueError("output array does not match result of ndarray.take")
res = asarray(out, order='C')
else:
compile_error("output must be an array")
max_item = a.shape[axis]
nelem = chunk
itemsize = res.itemsize
chunk *= itemsize
src = a.data.as_byte()
dest = res.data.as_byte()
indices_data = indices.data
if max_item == 0 and res.size != 0:
raise IndexError("cannot do a non-empty take from an empty axes.")
if mode == 'raise':
for i in range(n):
for j in range(m):
tmp = indices_data[j]
tmp = util.normalize_index(tmp, axis, max_item)
tmp_src = src + tmp * chunk
str.memcpy(dest, tmp_src, chunk)
dest += chunk
src += chunk * max_item
elif mode == 'wrap':
for i in range(n):
for j in range(m):
tmp = indices_data[j]
if tmp < 0:
while tmp < 0:
tmp += max_item
elif tmp >= max_item:
while tmp >= max_item:
tmp -= max_item
tmp_src = src + tmp * chunk
str.memcpy(dest, tmp_src, chunk)
dest += chunk
src += chunk * max_item
elif mode == 'clip':
for i in range(n):
for j in range(m):
tmp = indices_data[j]
if tmp < 0:
tmp = 0
elif tmp >= max_item:
tmp = max_item - 1
tmp_src = src + tmp * chunk
str.memcpy(dest, tmp_src, chunk)
dest += chunk
src += chunk * max_item
else:
raise ValueError(f"clipmode must be one of 'clip', 'raise', or 'wrap' (got {repr(mode)})")
if res.ndim == 0:
return res.item()
else:
return res
def indices(dimensions, dtype: type = int, sparse: Static[int] = False):
if not isinstance(dimensions, Tuple):
compile_error("dimensions must be a tuple of integers")
N: Static[int] = staticlen(dimensions)
shape = (1,) * N
if sparse:
return tuple(
arange(dimensions[i], dtype=dtype).reshape(
shape[:i] + (dimensions[i],) + shape[i+1:]) for i in staticrange(N))
res = empty((N,) + dimensions, dtype=dtype)
for i in staticrange(N):
dim = dimensions[i]
idx = arange(dim, dtype=dtype).reshape(shape[:i] + (dim,) + shape[i+1:])
res[i] = idx
return res
def ix_(*args):
def ix_one(new, nd: Static[int], k: Static[int]):
new = asarray(new)
if staticlen(new.shape) != 1:
compile_error("Cross index must be 1 dimensional")
if new.dtype is bool:
newx = new.nonzero()[0]
else:
newx = new
return newx.reshape((1,)*k + (newx.size,) + (1,)*(nd-k-1))
nd: Static[int] = staticlen(args)
return tuple(ix_one(args[k], nd, k) for k in staticrange(nd))
def ravel_multi_index(multi_index, dims, mode: str = 'raise', order: str = 'C'):
def fix_index(idx: int, axis: int, dim: int, mode: str):
if mode == 'raise':
if idx < 0 or idx >= dim:
raise ValueError("invalid entry in coordinates array")
elif mode == 'wrap':
if idx < 0:
while idx < 0:
idx += dim
elif idx >= dim:
while idx >= dim:
idx -= dim
elif mode == 'clip':
if idx < 0:
idx = 0
elif idx >= dim:
idx = dim - 1
return idx
def gather_non_ints(x):
if staticlen(x) == 0:
return ()
i = x[0]
rest = gather_non_ints(x[1:])
if isinstance(i, int):
return rest
else:
return (i,) + rest
if isinstance(dims, int):
return ravel_multi_index(multi_index, (dims,), mode=mode, order=order)
for d in dims:
if d <= 0:
raise ValueError("dimensions must be positive")
if mode not in ('raise', 'wrap', 'clip'):
raise ValueError(f"clipmode must be one of 'clip', 'raise', or 'wrap' (got {repr(mode)})")
corder = True
if order == 'C':
corder = True
elif order == 'F':
corder = False
else:
raise ValueError("only 'C' or 'F' order is permitted")
N: Static[int] = staticlen(dims)
if isinstance(multi_index, List[int]):
if len(multi_index) != N:
raise ValueError(f"parameter multi_index must be a sequence of length {N}")
idx = tuple(fix_index(multi_index[j], j, dims[j], mode) for j in staticrange(N))
return util.coords_to_index(idx, dims) if corder else util.coords_to_findex(idx, dims)
elif isinstance(multi_index, Tuple):
if staticlen(gather_non_ints(multi_index)) == 0:
if staticlen(multi_index) != staticlen(dims):
compile_error("parameter multi_index does not match dims in size")
idx = tuple(fix_index(multi_index[j], j, dims[j], mode) for j in staticrange(N))
return util.coords_to_index(idx, dims) if corder else util.coords_to_findex(idx, dims)
midx = vstack(multi_index)
else:
midx = asarray(multi_index)
if staticlen(midx.shape) != 2:
compile_error("multi_index must be 2 dimensional")
if len(multi_index) != N:
raise ValueError(f"parameter multi_index must be a sequence of length {N}")
ans = empty(midx.shape[1], int)
for i in range(ans.size):
idx = tuple(midx._ptr((j, i))[0] for j in staticrange(N))
idx = tuple(fix_index(idx[j], j, dims[j], mode) for j in staticrange(N))
res = util.coords_to_index(idx, dims) if corder else util.coords_to_findex(idx, dims)
ans.data[i] = res
return ans
def unravel_index(indices, shape, order: str = 'C'):
def check(idx: int, n: int):
if idx < 0 or idx >= n:
raise ValueError(f"index {idx} is out of bounds for array with size {n}")
return idx
if isinstance(shape, int):
return unravel_index(indices, (shape,), order=order)
corder = True
if order == 'C':
corder = True
elif order == 'F':
corder = False
else:
raise ValueError("only 'C' or 'F' order is permitted")
N: Static[int] = staticlen(shape)
if N == 0:
if not isinstance(indices, int):
compile_error("multiple indices are not supported for 0d arrays")
check(indices, 1)
return ()
n = 1
for s in shape:
if s < 0:
raise ValueError("dimensions must be non-negative")
n *= s
if isinstance(indices, int):
check(indices, n)
return util.index_to_coords(indices, shape) if corder else util.index_to_fcoords(indices, shape)
indices = asarray(indices)
N: Static[int] = staticlen(shape)
ans = tuple(empty(indices.shape, int) for _ in shape)
for idx in util.multirange(indices.shape):
index = indices._ptr(idx)[0]
check(index, n)
res = util.index_to_coords(index, shape) if corder else util.index_to_fcoords(index, shape)
for i in staticrange(N):
ans[i]._ptr(idx)[0] = res[i]
return ans
def diag_indices(n: int, ndim: int):
idx = arange(n)
return [idx for _ in range(ndim)]
@overload
def diag_indices(n: int, ndim: Static[int] = 2):
idx = arange(n)
return (idx,) * ndim
def diag_indices_from(arr):
arr = asarray(arr)
shape = arr.shape
ndim: Static[int] = staticlen(shape)
if ndim < 2:
compile_error("input array must be at least 2-d")
for i in staticrange(1, ndim):
if shape[i] != shape[0]:
raise ValueError("All dimensions of input must be of equal length")
return diag_indices(shape[0], ndim)
def mask_indices(n: int, mask_func, k = 0):
m = ones((n, n), int)
a = mask_func(m, k)
return nonzero(a != 0)
def tril_indices(n: int, k: int = 0, m: Optional[int] = None):
tri_ = tri(n, m, k=k, dtype=bool)
return tuple(broadcast_to(inds, tri_.shape)[tri_]
for inds in indices(tri_.shape, sparse=True))
def tril_indices_from(arr, k: int = 0):
arr = asarray(arr)
if staticlen(arr.shape) != 2:
compile_error("input array must be 2-d")
return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1])
def triu_indices(n: int, k: int = 0, m: Optional[int] = None):
tri_ = ~tri(n, m, k=k - 1, dtype=bool)
return tuple(broadcast_to(inds, tri_.shape)[tri_]
for inds in indices(tri_.shape, sparse=True))
def triu_indices_from(arr, k: int = 0):
arr = asarray(arr)
if staticlen(arr.shape) != 2:
compile_error("input array must be 2-d")
return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1])
def take_along_axis(arr, indices, axis):
def dim_mismatch():
compile_error("`indices` and `arr` must have the same number of dimensions")
arr = asarray(arr)
indices = asarray(indices)
if indices.dtype is not int:
compile_error("`indices` must be an integer array")
if axis is None:
if staticlen(indices.shape) != 1:
dim_mismatch()
out = empty(indices.size, arr.dtype)
for i in range(indices.size):
out.data[i] = arr._get_flat(indices._ptr((i,))[0], check=True)
return out
if staticlen(arr.shape) != staticlen(indices.shape):
dim_mismatch()
axis = util.normalize_axis_index(axis, arr.ndim)
M = arr.shape[axis]
J = indices.shape[axis]
bshape = broadcast_shapes(util.tuple_delete(arr.shape, axis),
util.tuple_delete(indices.shape, axis))
out_shape = util.tuple_insert(bshape, axis, J)
out = empty(out_shape, arr.dtype)
a_stride = arr.strides[axis]
indices_stride = indices.strides[axis]
out_stride = out.strides[axis]
for idx in util.multirange(bshape):
base_idx = util.tuple_insert(idx, axis, 0)
a_1d = ndarray((M,), (a_stride,), arr._ptr(base_idx, broadcast=True))
indices_1d = ndarray((J,), (indices_stride,), indices._ptr(base_idx, broadcast=True))
out_1d = ndarray((J,), (out_stride,), out._ptr(base_idx))
for j in range(J):
out_1d._ptr((j,))[0] = a_1d[indices_1d._ptr((j,))[0]]
return out
def choose(a, choices, out = None, mode: str = 'raise'):
MODE_RAISE: Static[int] = 0
MODE_WRAP : Static[int] = 1
MODE_CLIP : Static[int] = 2
a = asarray(a)
if a.dtype is not int:
compile_error("first argument must be an integer array")
xmode = 0
if mode == 'raise':
xmode = MODE_RAISE
elif mode == 'wrap':
xmode = MODE_WRAP
elif mode == 'clip':
xmode = MODE_CLIP
else:
raise ValueError(f"clipmode must be one of 'clip', 'raise', or 'wrap' (got {repr(mode)})")
n = len(choices)
if n == 0:
raise ValueError("0-length sequence.")
if isinstance(choices, Tuple):
xchoices1 = tuple(asarray(c) for c in choices)
bshape = broadcast_shapes(*(tuple(c.shape for c in xchoices1) + (a.shape,)))
xchoices = tuple(broadcast_to(c, bshape) for c in xchoices1)
elif isinstance(choices, List):
if not isinstance(choices[0], ndarray):
xchoices = [asarray(c) for c in choices]
else:
xchoices = choices
bshape = broadcast_shapes(xchoices[0].shape, a.shape)
elif isinstance(choices, ndarray):
xchoices = choices
bshape = broadcast_shapes(xchoices[0].shape, a.shape)
else:
compile_error("'choices' must be a list, tuple or array")
dtype = xchoices[0].dtype
if out is None:
ans = empty(bshape, dtype)
else:
if not isinstance(out, ndarray):
compile_error("'out' must be an array")
if out.dtype is not dtype:
compile_error("'out' has incorrect dtype")
if staticlen(bshape) != staticlen(out.shape):
compile_error("'out' has incorrect number of dimensions")
if bshape != out.shape:
raise ValueError("'out' has incorrect shape")
ans = out
for idx in util.multirange(bshape):
mi = a._ptr(idx, broadcast=True)[0]
if mi < 0 or mi >= n:
if xmode == MODE_RAISE:
raise ValueError("invalid entry in choice array")
elif xmode == MODE_WRAP:
if mi < 0:
while mi < 0:
mi += n;
else:
while mi >= n:
mi -= n
elif xmode == MODE_CLIP:
if mi < 0:
mi = 0
elif mi >= n:
mi = n - 1
choice = xchoices[mi]._ptr(idx, broadcast=True)[0]
ans._ptr(idx)[0] = choice
return ans
def compress(condition, a, axis = None, out = None):
condition = asarray(condition)
a = asarray(a)
if staticlen(condition.shape) != 1:
compile_error("condition must be a 1-d array")
if axis is None:
num_true = 0
i = 0
n = a.size
for c in condition:
if c:
if i >= n:
raise IndexError(f"index {i} is out of bounds for axis 0 with size {n}")
num_true += 1
i += 1
if out is None:
ans = empty(num_true, a.dtype)
ans_cc = True
else:
if not isinstance(out, ndarray):
compile_error("'out' must be an array")
if staticlen(out.shape) != 1:
compile_error("'out' must be 1-dimensional")
if out.size != num_true:
raise ValueError("'out' has incorrect length")
ans = out
ans_cc = ans._contig[0]
a_cc = a._contig[0]
k = 0
if a_cc and ans_cc:
for i in range(min(condition.size, n)):
if condition._ptr((i,))[0]:
ans.data[k] = util.cast(a.data[i], ans.dtype)
k += 1
else:
i = 0
for idx in util.multirange(a.shape):
if i >= condition.size:
break
if condition[i]:
ans._ptr((k,))[0] = util.cast(a._ptr(idx)[0], ans.dtype)
k += 1
i += 1
return ans
axis = util.normalize_axis_index(axis, a.ndim)
num_true = 0
i = 0
n = a.shape[axis]
for c in condition:
if c:
if i >= n:
raise IndexError(f"index {i} is out of bounds for axis {axis} with size {n}")
num_true += 1
i += 1
ans_shape = util.tuple_set(a.shape, axis, num_true)
if out is None:
ans = empty(ans_shape, a.dtype)
else:
if not isinstance(out, ndarray):
compile_error("'out' must be an array")
if staticlen(out.shape) != staticlen(ans_shape):
compile_error("'out' has incorrect number of dimensions")
if out.shape != ans_shape:
raise ValueError("'out' has incorrect shape")
ans = out
sub_shape = util.tuple_delete(ans_shape, axis)
k = 0
for i in range(min(condition.size, n)):
if condition._ptr((i,))[0]:
for idx in util.multirange(sub_shape):
idx1 = util.tuple_insert(idx, axis, i)
idx2 = util.tuple_insert(idx, axis, k)
p = a._ptr(idx1)
q = ans._ptr(idx2)
q[0] = util.cast(p[0], ans.dtype)
k += 1
return ans
def diagonal(a, offset: int = 0, axis1: int = 0, axis2: int = 1):
a = asarray(a)
shape = a.shape
strides = a.strides
ndim: Static[int] = staticlen(shape)
if ndim < 2:
compile_error("diag requires an array of at least two dimensions")
axis1 = util.normalize_axis_index(axis1, ndim)
axis2 = util.normalize_axis_index(axis2, ndim)
if axis1 == axis2:
raise ValueError("axis1 and axis2 cannot be the same")
if axis1 > axis2:
axis1, axis2 = axis2, axis1
dim1 = shape[axis1]
dim2 = shape[axis2]
stride1 = strides[axis1]
stride2 = strides[axis2]
data = a.data
if offset >= 0:
offset_stride = stride2
dim2 -= offset
else:
offset = -offset
offset_stride = stride1
dim1 -= offset
diag_size = dim2 if dim2 < dim1 else dim1
if diag_size < 0:
diag_size = 0
else:
data = Ptr[a.dtype](data.as_byte() + (offset * offset_stride))
ret_shape = util.tuple_delete(util.tuple_delete(shape, axis2), axis1) + (diag_size,)
ret_strides = util.tuple_delete(util.tuple_delete(strides, axis2), axis1) + (stride1 + stride2,)
return ndarray(ret_shape, ret_strides, data)
def select(condlist: List, choicelist: List, default = 0):
n = len(condlist)
if n != len(choicelist):
raise ValueError(
"list of cases must be same length as list of conditions")
if n == 0:
raise ValueError("select with an empty condition list is not possible")
condlist = [asarray(cond) for cond in condlist]
if condlist[0].dtype is not bool:
compile_error("condlist entries should be boolean ndarray")
cond_bshape = condlist[0].shape
for cond in condlist:
cond_bshape = broadcast_shapes(cond_bshape, cond.shape)
for i in range(n):
condlist[i] = broadcast_to(condlist[i], cond_bshape)
choicelist = [asarray(choice) for choice in choicelist]
choice_bshape = choicelist[0].shape
for choice in choicelist:
choice_bshape = broadcast_shapes(choice_bshape, choice.shape)
for i in range(n):
choicelist[i] = broadcast_to(choicelist[i], choice_bshape)
default = asarray(default)
ans_shape = broadcast_shapes(cond_bshape, choice_bshape, default.shape)
dtype = type(util.coerce(choicelist[0].dtype, default.dtype))
ans = empty(ans_shape, dtype)
for idx in util.multirange(ans_shape):
found = False
p = ans._ptr(idx)
for i in range(n):
if condlist[i]._ptr(idx, broadcast=True)[0]:
p[0] = util.cast(choicelist[i]._ptr(idx, broadcast=True)[0], dtype)
found = True
break
if not found:
p[0] = util.cast(default._ptr(idx, broadcast=True)[0], dtype)
return ans
def place(arr: ndarray, mask, vals):
mask = asarray(mask)
vals = asarray(vals)
ni = arr.size
nm = mask.size
nv = vals.size
if nm != ni:
raise ValueError("place: mask and data must be the same size")
if nv <= 0:
if mask.any():
raise ValueError("Cannot insert from an empty array!");
return
cc1, _ = arr._contig
cc2, _ = mask._contig
cc3, _ = vals._contig
j = 0
if cc1 and cc2 and cc3:
for i in range(ni):
if mask.data[i]:
if j >= nv:
j = 0
arr.data[i] = util.cast(vals.data[j], arr.dtype)
j += 1
else:
for i in range(ni):
if mask._get_flat(i, check=False):
if j >= nv:
j = 0
arr._set_flat(i, vals._get_flat(j, check=False), check=False)
j += 1
def put(a: ndarray, ind, v, mode: str = 'raise'):
def fix_index(idx: int, dim: int, mode: str):
if mode == 'raise':
if idx < 0 or idx >= dim:
raise ValueError("invalid entry in coordinates array")
elif mode == 'wrap':
if idx < 0:
while idx < 0:
idx += dim
elif idx >= dim:
while idx >= dim:
idx -= dim
elif mode == 'clip':
if idx < 0:
idx = 0
elif idx >= dim:
idx = dim - 1
return idx
if mode not in ('raise', 'wrap', 'clip'):
raise ValueError(f"clipmode must be one of 'clip', 'raise', or 'wrap' (got {repr(mode)})")
ind = asarray(ind)
v = asarray(v)
na = a.size
ni = ind.size
nv = v.size
if ni == 0 or nv == 0:
return
if na == 0 and (mode == 'wrap' or mode == 'clip'):
raise ValueError("empty array given to put")
cc1, _ = a._contig
cc2, _ = ind._contig
cc3, _ = v._contig
j = 0
if cc1 and cc2 and cc3:
for i in range(ni):
idx = util.cast(ind.data[i], int)
idx = fix_index(idx, na, mode)
if j >= nv:
j = 0
a.data[idx] = util.cast(v.data[j], a.dtype)
j += 1
else:
for i in range(ni):
idx = util.cast(ind._get_flat(i, check=False), int)
idx = fix_index(idx, na, mode)
if j >= nv:
j = 0
a._set_flat(idx, v._get_flat(j, check=False), check=True)
j += 1
def put_along_axis(arr: ndarray, indices, values, axis):
def dim_mismatch():
compile_error("`indices` and `arr` must have the same number of dimensions")
indices = asarray(indices)
values = asarray(values)
if indices.dtype is not int:
compile_error("`indices` must be an integer array")
if axis is None:
nv = values.size
j = 0
for i in range(indices.size):
if j >= nv:
j = 0
arr._set_flat(indices._get_flat(i, check=False),
values._get_flat(j, check=False),
check=True)
j += 1
return
bshape = broadcast_shapes(indices.shape, values.shape)
indices = broadcast_to(indices, bshape)
values = broadcast_to(values, bshape)
if staticlen(arr.shape) != staticlen(indices.shape):
dim_mismatch()
axis = util.normalize_axis_index(axis, arr.ndim)
M = arr.shape[axis]
J = indices.shape[axis]
bshape = broadcast_shapes(util.tuple_delete(arr.shape, axis),
util.tuple_delete(indices.shape, axis))
a_stride = arr.strides[axis]
indices_stride = indices.strides[axis]
values_stride = values.strides[axis]
for idx in util.multirange(bshape):
base_idx = util.tuple_insert(idx, axis, 0)
a_1d = ndarray((M,), (a_stride,), arr._ptr(base_idx, broadcast=True))
indices_1d = ndarray((J,), (indices_stride,), indices._ptr(base_idx, broadcast=True))
values_1d = ndarray((J,), (values_stride,), values._ptr(base_idx, broadcast=True))
for j in range(J):
a_1d[indices_1d._ptr((j,))[0]] = util.cast(values_1d._ptr((j,))[0], arr.dtype)
def putmask(a: ndarray, mask, values):
mask = asarray(mask)
values = asarray(values)
na = a.size
nv = values.size
if na != mask.size:
raise ValueError("putmask: mask and data must be the same size")
if na == 0 or nv == 0:
return
cc1, _ = a._contig
cc2, _ = mask._contig
cc3, _ = values._contig
j = 0
if cc1 or cc2 or cc3:
for i in range(na):
if mask.data[i]:
a.data[i] = values.data[j]
j += 1
if j >= nv:
j = 0
else:
for i in range(na):
if mask._get_flat(i, check=False):
a._set_flat(i, values._get_flat(j, check=False), check=False)
j += 1
if j >= nv:
j = 0
def fill_diagonal(a: ndarray, val, wrap: bool = False):
if staticlen(a.shape) < 2:
compile_error("array must be at least 2-d")
val = asarray(val)
nv = val.size
end = a.size
if nv == 0 or end == 0:
return
if staticlen(a.shape) == 2:
step = a.shape[1] + 1
if not wrap:
end = min(end, a.shape[1] * a.shape[1])
else:
for i in range(1, a.ndim):
if a.shape[i] != a.shape[0]:
raise ValueError("All dimensions of input must be of equal length")
step = 1
acc = 1
for s in a.shape[:-1]:
acc *= s
step += acc
j = 0
for i in range(0, end, step):
a._set_flat(i, val._get_flat(j, check=False), check=False)
j += 1
if j >= nv:
j = 0
########
# Misc #
########
def _power_of_ten(n: int):
p10 = (1e0, 1e1, 1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8)
if n < 9:
return Ptr[float](__ptr__(p10).as_byte())[n]
else:
ret = 1e9
while n > 9:
ret *= 10.
n -= 1
return ret
def _round_int(x: T, decimals: int, T: type):
if decimals >= 0:
return x
else:
f = _power_of_ten(-decimals)
y = util.cast(x, float)
return T(int(util.rint(y / f) * f))
def _round(x: T, decimals: int, T: type):
if isinstance(x, complex) or isinstance(x, complex64):
return T(_round(x.real, decimals), _round(x.imag, decimals))
if (isinstance(x, int) or
isinstance(x, Int) or
isinstance(x, UInt) or
isinstance(x, byte) or
isinstance(x, bool)):
return _round_int(x, decimals)
if not (isinstance(x, float) or isinstance(x, float32)):
compile_error("don't know how to round type '" + T.__name__ + "'")
if decimals == 0:
return util.rint(x)
elif decimals > 0:
f = T(_power_of_ten(decimals))
return util.rint(x * f) / f
else:
f = T(_power_of_ten(-decimals))
return util.rint(x / f) * f
def round(a, decimals: int = 0, out = None):
a = asarray(a)
if staticlen(a.shape) == 0:
return _round(a.data[0], decimals)
cc, fc = a._contig
n = a.size
if out is None:
ans = empty_like(a, order=('F' if fc and not cc else 'C'))
else:
if not isinstance(out, ndarray):
compile_error("output must be an array")
if out.dtype is not a.dtype:
compile_error("output has wrong type")
if not util.tuple_equal(a.shape, out.shape):
raise ValueError("invalid output shape")
ans = out
cc1, fc1 = ans._contig
if (cc and cc1) or (fc and fc1):
p = a.data
q = ans.data
for i in range(n):
q[i] = _round(p[i], decimals)
else:
for idx in util.multirange(a.shape):
p = a._ptr(idx)
q = ans._ptr(idx)
q[0] = _round(p[0], decimals)
return ans
around = round
def _clip_full(a, a_min, a_max, out):
a = asarray(a)
a_min = asarray(a_min)
a_max = asarray(a_max)
bshape = broadcast_shapes(a.shape, a_min.shape, a_max.shape)
if out is None:
ans = empty(bshape, a.dtype)
elif isinstance(out, ndarray):
if staticlen(bshape) != staticlen(out.shape):
compile_error("'out' has incorrect number of dimensions")
if bshape != out.shape:
raise ValueError("'out' has incorrect shape")
ans = out
else:
compile_error("'out' must be an array")
maxf = lambda a, b: a if a > b else b
minf = lambda a, b: a if a < b else b
for idx in util.multirange(bshape):
x = a._ptr(idx, broadcast=True)[0]
xmin = a_min._ptr(idx, broadcast=True)[0]
xmax = a_max._ptr(idx, broadcast=True)[0]
x = util.cast(x, ans.dtype)
xmin = util.cast(xmin, ans.dtype)
xmax = util.cast(xmax, ans.dtype)
ans._ptr(idx)[0] = minf(xmax, maxf(x, xmin))
return ans
def _clip_min(a, a_min, out):
a = asarray(a)
a_min = asarray(a_min)
bshape = broadcast_shapes(a.shape, a_min.shape)
if out is None:
ans = empty(bshape, a.dtype)
elif isinstance(out, ndarray):
if staticlen(bshape) != staticlen(out.shape):
compile_error("'out' has incorrect number of dimensions")
if bshape != out.shape:
raise ValueError("'out' has incorrect shape")
ans = out
else:
compile_error("'out' must be an array")
maxf = lambda a, b: a if a > b else b
for idx in util.multirange(bshape):
x = a._ptr(idx, broadcast=True)[0]
xmin = a_min._ptr(idx, broadcast=True)[0]
x = util.cast(x, ans.dtype)
xmin = util.cast(xmin, ans.dtype)
ans._ptr(idx)[0] = maxf(x, xmin)
return ans
def _clip_max(a, a_max, out):
a = asarray(a)
a_max = asarray(a_max)
bshape = broadcast_shapes(a.shape, a_max.shape)
if out is None:
ans = empty(bshape, a.dtype)
elif isinstance(out, ndarray):
if staticlen(bshape) != staticlen(out.shape):
compile_error("'out' has incorrect number of dimensions")
if bshape != out.shape:
raise ValueError("'out' has incorrect shape")
ans = out
else:
compile_error("'out' must be an array")
minf = lambda a, b: a if a < b else b
for idx in util.multirange(bshape):
x = a._ptr(idx, broadcast=True)[0]
xmax = a_max._ptr(idx, broadcast=True)[0]
x = util.cast(x, ans.dtype)
xmax = util.cast(xmax, ans.dtype)
ans._ptr(idx)[0] = minf(xmax, x)
return ans
def clip(a, a_min, a_max, out = None):
if a_min is None and a_max is None:
compile_error("One of max or min must be given")
elif a_max is None:
return _clip_min(a, a_min, out)
elif a_min is None:
return _clip_max(a, a_max, out)
else:
return _clip_full(a, a_min, a_max, out)
def ndenumerate(arr):
arr = asarray(arr)
for idx in util.multirange(arr.shape):
yield (idx, arr._ptr(idx)[0])
def ndindex(*shape):
if staticlen(shape) == 1:
if isinstance(shape[0], Tuple):
return ndindex(*shape[0])
for s in shape:
if s < 0:
raise ValueError("negative dimensions are not allowed")
return util.multirange(shape)
def iterable(y):
if not hasattr(y, "__iter__"):
return False
elif isinstance(y, ndarray):
if staticlen(y.shape) == 0:
return False
else:
return True
else:
return True
def packbits(a, axis = None, bitorder: str = 'big'):
a = asarray(a)
if (a.dtype is not bool and
a.dtype is not int and
a.dtype is not byte and
not isinstance(a.dtype, Int) and
not isinstance(a.dtype, UInt)):
compile_error("Expected an input array of integer or boolean data type")
little_endian = False
if bitorder == 'big':
little_endian = False
elif bitorder == 'little':
little_endian = True
else:
raise ValueError("'order' must be either 'little' or 'big'")
if axis is None:
pack_size = ((a.size - 1) >> 3) + 1
ans = empty(pack_size, dtype=u8)
m = 0
k = 0
e = u8(0)
for idx in util.multirange(a.shape):
if m == 8:
ans.data[k] = e
k += 1
e = u8(0)
m = 0
b = u8(1 if a._ptr(idx)[0] else 0)
if little_endian:
e = (e >> u8(1)) | (b << u8(7))
else:
e = (e << u8(1)) | b
m += 1
if k < pack_size:
if little_endian:
e >>= u8(8 - m)
else:
e <<= u8(8 - m)
ans.data[k] = e
return ans
if not isinstance(axis, int):
compile_error("'axis' must be an int or None")
axis = util.normalize_axis_index(axis, a.ndim)
n = a.shape[axis]
pack_size = ((n - 1) >> 3) + 1
ans_shape = util.tuple_set(a.shape, axis, pack_size)
ans = empty(ans_shape, dtype=u8)
for idx0 in util.multirange(util.tuple_delete(a.shape, axis)):
m = 0
k = 0
e = u8(0)
for i in range(n):
if m == 8:
ans._ptr(util.tuple_insert(idx0, axis, k))[0] = e
k += 1
e = u8(0)
m = 0
b = u8(1 if a._ptr(util.tuple_insert(idx0, axis, i))[0] else 0)
if little_endian:
e = (e >> u8(1)) | (b << u8(7))
else:
e = (e << u8(1)) | b
m += 1
if k < pack_size:
if little_endian:
e >>= u8(8 - m)
else:
e <<= u8(8 - m)
ans._ptr(util.tuple_insert(idx0, axis, k))[0] = e
return ans
def unpackbits(a, axis = None, count = None, bitorder: str = 'big'):
a = asarray(a)
if a.dtype is not u8:
compile_error("Expected an input array of unsigned byte data type")
little_endian = False
if bitorder == 'big':
little_endian = False
elif bitorder == 'little':
little_endian = True
else:
raise ValueError("'order' must be either 'little' or 'big'")
if axis is None:
unpack_size = a.size * 8
if count is not None:
if count < 0:
if -count > unpack_size:
raise ValueError("-count larger than number of elements")
unpack_size += count
else:
unpack_size = count
ans = empty(unpack_size, dtype=u8)
k = 0
for idx in util.multirange(a.shape):
e = a._ptr(idx)[0]
for i in range(8):
if k >= unpack_size:
break
sh = u8(i if little_endian else 7 - i)
ans.data[k] = (e & (u8(1) << sh)) >> sh
k += 1
if k >= unpack_size:
break
while k < unpack_size:
ans.data[k] = u8(0)
k += 1
return ans
if not isinstance(axis, int):
compile_error("'axis' must be an int or None")
axis = util.normalize_axis_index(axis, a.ndim)
n = a.shape[axis]
unpack_size = n * 8
if count is not None:
if count < 0:
if -count > unpack_size:
raise ValueError("-count larger than number of elements")
unpack_size += count
else:
unpack_size = count
ans_shape = util.tuple_set(a.shape, axis, unpack_size)
ans = empty(ans_shape, dtype=u8)
if unpack_size == 0:
return ans
for idx0 in util.multirange(util.tuple_delete(a.shape, axis)):
k = 0
for m in range(n):
e = a._ptr(util.tuple_insert(idx0, axis, m))[0]
for i in range(8):
if k >= unpack_size:
break
sh = u8(i if little_endian else 7 - i)
ans._ptr(util.tuple_insert(idx0, axis, k))[0] = (e & (u8(1) << sh)) >> sh
k += 1
if k >= unpack_size:
break
while k < unpack_size:
ans._ptr(util.tuple_insert(idx0, axis, k))[0] = u8(0)
k += 1
return ans
def _is_pos_neg_inf(x, pos: bool, out = None):
from .ndmath import isinf, signbit
x = asarray(x)
if out is None:
ans = empty(x.shape, bool)
else:
if not isinstance(out, ndarray):
compile_error("'out' must be an array")
if out.ndim != x.ndim:
compile_error("'out' has incorrect number of dimensions")
if x.shape != out.shape:
raise ValueError("'out' has incorrect shape")
ans = out
for idx in util.multirange(x.shape):
e = x._ptr(idx)[0]
b = isinf(e)
if pos:
b = b and not signbit(e)
else:
b = b and signbit(e)
ans._ptr(idx)[0] = util.cast(b, ans.dtype)
if out is None and ans.ndim == 0:
return ans.item()
else:
return ans
def isposinf(x, out = None):
return _is_pos_neg_inf(x, pos=True, out=out)
def isneginf(x, out = None):
return _is_pos_neg_inf(x, pos=False, out=out)
def iscomplex(x):
x = asarray(x)
if x.dtype is complex or x.dtype is complex64:
ans = x.map(lambda c: bool(c.imag))
else:
ans = zeros(x.shape, bool)
if ans.ndim == 0:
return ans.item()
else:
return ans
def iscomplexobj(x):
if isinstance(x, ndarray):
return x.dtype is complex or x.dtype is complex64
else:
dtype = asarray(x).dtype
return dtype is complex or dtype is complex64
def isreal(x):
x = asarray(x)
if x.dtype is complex or x.dtype is complex64:
ans = x.map(lambda c: not bool(c.imag))
else:
ans = ones(x.shape, bool)
if ans.ndim == 0:
return ans.item()
else:
return ans
def isrealobj(x):
return not iscomplexobj(x)
def isfortran(a: ndarray):
cc, fc = a._contig
return fc and not cc
def isscalar(element):
T = type(element)
return (T is int or
T is float or
T is complex or
T is complex64 or
T is bool or
T is byte or
isinstance(T, Int) or
isinstance(T, UInt) or
T is str or
T is NoneType)
def _array_get_part(arr: ndarray, imag: Static[int]):
if arr.dtype is complex:
offset = util.sizeof(float) if imag else 0
data = Ptr[float](arr.data.as_byte() + offset)
return ndarray(arr.shape, arr.strides, data)
elif arr.dtype is complex64:
offset = util.sizeof(float32) if imag else 0
data = Ptr[float32](arr.data.as_byte() + offset)
return ndarray(arr.shape, arr.strides, data)
else:
if imag:
n = arr.size
data = Ptr[arr.dtype](n)
str.memset(data.as_byte(), byte(0), n * arr.itemsize)
return ndarray(arr.shape, data)
else:
return arr
def real(val):
return _array_get_part(asarray(val), imag=False)
def imag(val):
return _array_get_part(asarray(val), imag=True)
def _real_set(arr: ndarray, val):
val = broadcast_to(asarray(val), arr.shape)
for idx in util.multirange(arr.shape):
x = val._ptr(idx)[0]
p = arr._ptr(idx)
if arr.dtype is complex:
p[0] = complex(util.cast(x, float), p[0].imag)
elif arr.dtype is complex64:
p[0] = complex64(util.cast(x, float32), p[0].imag)
else:
p[0] = util.cast(x, arr.dtype)
def _imag_set(arr: ndarray, val):
val = broadcast_to(asarray(val), arr.shape)
for idx in util.multirange(arr.shape):
x = val._ptr(idx)[0]
p = arr._ptr(idx)
if arr.dtype is complex:
p[0] = complex(p[0].real, util.cast(x, float))
elif arr.dtype is complex64:
p[0] = complex64(p[0].imag, util.cast(x, float32))
else:
compile_error("array does not have imaginary part to set")
@extend
class ndarray:
def take(self, indices, axis = None, out = None, mode: str = 'raise'):
return take(self, indices, axis=axis, out=out, mode=mode)
def squeeze(self, axis = None):
return squeeze(self, axis)
def nonzero(self):
return nonzero(self)
def searchsorted(self, v, side: str = 'left', sorter = None):
return searchsorted(self, v=v, side=side, sorter=sorter)
def repeat(self, repeats, axis = None):
return repeat(self, repeats, axis=axis)
def compress(self, condition, axis = None, out = None):
return compress(condition, self, axis=axis, out=out)
def choose(self, choices, out = None, mode: str = 'raise'):
return choose(self, choices, out, mode)
def diagonal(self, offset: int = 0, axis1: int = 0, axis2: int = 1):
return diagonal(self, offset, axis1=axis1, axis2=axis2)
def put(self, ind, v, mode: str = 'raise'):
return put(self, ind, v, mode)
def round(self, decimals: int = 0, out = None):
return round(self, decimals, out=out)
def clip(self, min = None, max = None, out = None):
return clip(self, min, max, out)
@property
def real(self):
return real(self)
@real.setter
def real(self, val):
_real_set(self, val)
@property
def imag(self):
return imag(self)
@imag.setter
def imag(self, val):
_imag_set(self, val)
def conj(self):
if hasattr(dtype(), 'conjugate'):
return self._op_unary(lambda a: a.conjugate())
else:
return self
def conjugate(self):
return self.conj()
@flat.setter
def flat(self, val):
val = asarray(val)
j = 0
m = val.size
if m == 0:
return
for i in range(self.size):
p = self._ptr_flat(i, check=False)
if j == m:
j = 0
q = val._ptr_flat(j, check=False)
p[0] = util.cast(q[0], dtype)
j += 1
# Not supported:
# - rollaxis (use is discouraged in NumPy docs)
# - asmatrix (matrix class not supported)
############
# Datetime #
############
@extend
class busdaycalendar:
@property
def weekmask(self):
wm = empty(7, bool)
for i in range(7):
wm.data[i] = self._wm[i]
return wm
@property
def holidays(self):
hd = empty(self._nholidays, datetime64['D', 1])
for i in range(len(hd)):
hd.data[i] = self._holidays[i]
return hd
def _get_busdaycal(weekmask=None, holidays=None, busdaycal=None):
if busdaycal is None:
if weekmask is None:
w = "1111100"
else:
w = weekmask
return busdaycalendar(weekmask=w, holidays=holidays)
else:
if weekmask is not None or holidays is not None:
compile_error("Cannot supply both the weekmask/holidays and the "
"busdaycal parameters to busday_offset()")
if not isinstance(busdaycal, busdaycalendar):
compile_error("busdaycal parameter must be a busdaycalendar")
return busdaycal
def busday_offset(dates, offsets, roll: str = "raise", weekmask=None,
holidays=None, busdaycal=None, out=None):
cal = _get_busdaycal(weekmask=weekmask, holidays=holidays, busdaycal=busdaycal)
wm = cal._wm
busdays_in_weekmask = wm.count
dates = asarray(dates, dtype=datetime64['D', 1])
offsets = asarray(offsets)
if offsets.dtype is not int:
compile_error("offsets parameter must be an array of integers")
roll_code = 0
if roll == "forward":
roll_code = _BUSDAY_FORWARD
elif roll == "following":
roll_code = _BUSDAY_FOLLOWING
elif roll == "backward":
roll_code = _BUSDAY_BACKWARD
elif roll == "preceding":
roll_code = _BUSDAY_PRECEDING
elif roll == "modifiedfollowing":
roll_code = _BUSDAY_MODIFIEDFOLLOWING
elif roll == "modifiedpreceding":
roll_code = _BUSDAY_MODIFIEDPRECEDING
elif roll == "nat":
roll_code = _BUSDAY_NAT
elif roll == "raise":
roll_code = _BUSDAY_RAISE
else:
raise ValueError(f"Invalid business day roll parameter \"{roll}\"")
bshape = broadcast_shapes(dates.shape, offsets.shape)
if out is not None:
if not isinstance(out, ndarray):
compile_error("'out' must be an array")
if out.dtype is not datetime64['D', 1]:
compile_error("'out' must have dtype datetime64[D]")
if out.ndim != staticlen(bshape):
compile_error("'out' has incorrect number of dimensions")
if out.shape != bshape:
raise ValueError("'out' has incorrect shape")
ans = out
else:
ans = empty(bshape, datetime64['D', 1])
for idx in util.multirange(bshape):
d = dates._ptr(idx, broadcast=True)[0]
o = offsets._ptr(idx, broadcast=True)[0]
b = _apply_busines_day_offset(d, o, roll_code, wm, busdays_in_weekmask,
cal._holidays, cal._holidays + cal._nholidays)
ans._ptr(idx)[0] = b
if ans.ndim == 0 and out is None:
return ans.item()
else:
return ans
def busday_count(begindates, enddates, weekmask=None,
holidays=None, busdaycal=None, out=None):
cal = _get_busdaycal(weekmask=weekmask, holidays=holidays, busdaycal=busdaycal)
wm = cal._wm
busdays_in_weekmask = wm.count
begindates = asarray(begindates, dtype=datetime64['D', 1])
enddates = asarray(enddates, dtype=datetime64['D', 1])
if out is not None:
if not isinstance(out, ndarray):
compile_error("'out' must be an array")
if out.dtype is not int:
compile_error("'out' must have dtype int")
bshape = broadcast_shapes(begindates.shape, enddates.shape, out.shape)
ans = out
else:
bshape = broadcast_shapes(begindates.shape, enddates.shape)
ans = empty(bshape, int)
for idx in util.multirange(bshape):
d1 = begindates._ptr(idx, broadcast=True)[0]
d2 = enddates._ptr(idx, broadcast=True)[0]
c = _apply_busines_day_count(d1, d2, wm, busdays_in_weekmask,
cal._holidays, cal._holidays + cal._nholidays)
ans._ptr(idx, broadcast=(out is not None))[0] = c
if ans.ndim == 0 and out is None:
return ans.item()
else:
return ans
def is_busday(dates, weekmask=None, holidays=None, busdaycal=None, out=None):
cal = _get_busdaycal(weekmask=weekmask, holidays=holidays, busdaycal=busdaycal)
wm = cal._wm
busdays_in_weekmask = wm.count
dates = asarray(dates, dtype=datetime64['D', 1])
if out is not None:
if not isinstance(out, ndarray):
compile_error("'out' must be an array")
if out.dtype is not bool:
compile_error("'out' must have dtype bool")
if out.ndim != dates.ndim:
compile_error("'out' has incorrect number of dimensions")
if out.shape != dates.shape:
raise ValueError("'out' has incorrect shape")
ans = out
else:
ans = empty(dates.shape, bool)
for idx in util.multirange(dates.shape):
d = dates._ptr(idx)[0]
ans._ptr(idx)[0] = _apply_is_business_day(d, wm, cal._holidays, cal._holidays + cal._nholidays)
if ans.ndim == 0 and out is None:
return ans.item()
else:
return ans
def datetime_data(dtype: type):
if not isinstance(dtype, datetime64) and not isinstance(dtype, timedelta64):
compile_error("cannot get datetime metadata from non-datetime type")
return (dtype.base, dtype.num)
def datetime_as_string(arr, unit=None, timezone: str = "naive"):
arr = asarray(arr)
if not isinstance(arr.dtype, datetime64):
compile_error("input must have type NumPy datetime")
return arr.map(lambda d: d._as_string(unit=unit, timezone=timezone))