def result_type(space, __args__):
args_w, kw_w = __args__.unpack()
if kw_w:
raise oefmt(space.w_TypeError,
"result_type() takes no keyword arguments")
if not args_w:
raise oefmt(space.w_ValueError,
"at least one array or dtype is required")
result = None
for w_arg in args_w:
if isinstance(w_arg, W_NDimArray):
dtype = w_arg.get_dtype()
elif isinstance(w_arg, W_GenericBox) or (
space.isinstance_w(w_arg, space.w_int)
or space.isinstance_w(w_arg, space.w_float)
or space.isinstance_w(w_arg, space.w_complex)
or space.isinstance_w(w_arg, space.w_long)
or space.isinstance_w(w_arg, space.w_bool)):
dtype = ufuncs.find_dtype_for_scalar(space, w_arg)
else:
dtype = space.interp_w(
descriptor.W_Dtype,
space.call_function(space.gettypefor(descriptor.W_Dtype),
w_arg))
result = ufuncs.find_binop_result_dtype(space, result, dtype)
return result
def dtype_agreement(space, w_arr_list, shape, out=None):
""" agree on dtype from a list of arrays. if out is allocated,
use it's dtype, otherwise allocate a new one with agreed dtype
"""
from pypy.module.micronumpy.ufuncs import find_binop_result_dtype
if not space.is_none(out):
return out
dtype = None
for w_arr in w_arr_list:
if not space.is_none(w_arr):
dtype = find_binop_result_dtype(space, dtype, w_arr.get_dtype())
assert dtype is not None
out = base.W_NDimArray.from_shape(space, shape, dtype)
return out
def dtype_agreement(space, w_arr_list, shape, out=None):
""" agree on dtype from a list of arrays. if out is allocated,
use it's dtype, otherwise allocate a new one with agreed dtype
"""
from pypy.module.micronumpy.ufuncs import find_binop_result_dtype
if not space.is_none(out):
return out
dtype = None
for w_arr in w_arr_list:
if not space.is_none(w_arr):
dtype = find_binop_result_dtype(space, dtype, w_arr.get_dtype())
assert dtype is not None
out = base.W_NDimArray.from_shape(space, shape, dtype)
return out
def result_type(space, __args__):
args_w, kw_w = __args__.unpack()
if kw_w:
raise oefmt(space.w_TypeError, "result_type() takes no keyword arguments")
if not args_w:
raise oefmt(space.w_ValueError, "at least one array or dtype is required")
result = None
for w_arg in args_w:
if isinstance(w_arg, W_NDimArray):
dtype = w_arg.get_dtype()
elif isinstance(w_arg, W_GenericBox) or (
space.isinstance_w(w_arg, space.w_int) or
space.isinstance_w(w_arg, space.w_float) or
space.isinstance_w(w_arg, space.w_complex) or
space.isinstance_w(w_arg, space.w_long) or
space.isinstance_w(w_arg, space.w_bool)):
dtype = ufuncs.find_dtype_for_scalar(space, w_arg)
else:
dtype = space.interp_w(descriptor.W_Dtype,
space.call_function(space.gettypefor(descriptor.W_Dtype), w_arg))
result = ufuncs.find_binop_result_dtype(space, result, dtype)
return result
def where(space, w_arr, w_x=None, w_y=None):
"""where(condition, [x, y])
Return elements, either from `x` or `y`, depending on `condition`.
If only `condition` is given, return ``condition.nonzero()``.
Parameters
----------
condition : array_like, bool
When True, yield `x`, otherwise yield `y`.
x, y : array_like, optional
Values from which to choose. `x` and `y` need to have the same
shape as `condition`.
Returns
-------
out : ndarray or tuple of ndarrays
If both `x` and `y` are specified, the output array contains
elements of `x` where `condition` is True, and elements from
`y` elsewhere.
If only `condition` is given, return the tuple
``condition.nonzero()``, the indices where `condition` is True.
See Also
--------
nonzero, choose
Notes
-----
If `x` and `y` are given and input arrays are 1-D, `where` is
equivalent to::
[xv if c else yv for (c,xv,yv) in zip(condition,x,y)]
Examples
--------
>>> np.where([[True, False], [True, True]],
... [[1, 2], [3, 4]],
... [[9, 8], [7, 6]])
array([[1, 8],
[3, 4]])
>>> np.where([[0, 1], [1, 0]])
(array([0, 1]), array([1, 0]))
>>> x = np.arange(9.).reshape(3, 3)
>>> np.where( x > 5 )
(array([2, 2, 2]), array([0, 1, 2]))
>>> x[np.where( x > 3.0 )] # Note: result is 1D.
array([ 4., 5., 6., 7., 8.])
>>> np.where(x < 5, x, -1) # Note: broadcasting.
array([[ 0., 1., 2.],
[ 3., 4., -1.],
[-1., -1., -1.]])
NOTE: support for not passing x and y is unsupported
"""
if space.is_none(w_y):
if space.is_none(w_x):
raise OperationError(
space.w_NotImplementedError,
space.wrap("1-arg where unsupported right now"))
raise OperationError(
space.w_ValueError,
space.wrap("Where should be called with either 1 or 3 arguments"))
if space.is_none(w_x):
raise OperationError(
space.w_ValueError,
space.wrap("Where should be called with either 1 or 3 arguments"))
arr = convert_to_array(space, w_arr)
x = convert_to_array(space, w_x)
y = convert_to_array(space, w_y)
if x.is_scalar() and y.is_scalar() and arr.is_scalar():
if arr.get_dtype().itemtype.bool(arr.get_scalar_value()):
return x
return y
dtype = ufuncs.find_binop_result_dtype(space, x.get_dtype(), y.get_dtype())
shape = shape_agreement(space, arr.get_shape(), x)
shape = shape_agreement(space, shape, y)
out = W_NDimArray.from_shape(space, shape, dtype)
return loop.where(space, out, shape, arr, x, y, dtype)
def concatenate(space, w_args, w_axis=None):
args_w = space.listview(w_args)
if len(args_w) == 0:
raise oefmt(space.w_ValueError,
"need at least one array to concatenate")
args_w = [convert_to_array(space, w_arg) for w_arg in args_w]
if w_axis is None:
w_axis = space.wrap(0)
if space.is_none(w_axis):
args_w = [
w_arg.reshape(space, space.newlist([w_arg.descr_get_size(space)]))
for w_arg in args_w
]
w_axis = space.wrap(0)
dtype = args_w[0].get_dtype()
shape = args_w[0].get_shape()[:]
ndim = len(shape)
if ndim == 0:
raise oefmt(space.w_ValueError,
"zero-dimensional arrays cannot be concatenated")
axis = space.int_w(w_axis)
orig_axis = axis
if axis < 0:
axis = ndim + axis
if ndim == 1 and axis != 0:
axis = 0
if axis < 0 or axis >= ndim:
raise oefmt(space.w_IndexError, "axis %d out of bounds [0, %d)",
orig_axis, ndim)
for arr in args_w[1:]:
if len(arr.get_shape()) != ndim:
raise OperationError(
space.w_ValueError,
space.wrap(
"all the input arrays must have same number of dimensions")
)
for i, axis_size in enumerate(arr.get_shape()):
if i == axis:
shape[i] += axis_size
elif axis_size != shape[i]:
raise OperationError(
space.w_ValueError,
space.wrap("all the input array dimensions except for the "
"concatenation axis must match exactly"))
a_dt = arr.get_dtype()
if dtype.is_record() and a_dt.is_record():
# Record types must match
for f in dtype.fields:
if f not in a_dt.fields or \
dtype.fields[f] != a_dt.fields[f]:
raise OperationError(space.w_TypeError,
space.wrap("invalid type promotion"))
elif dtype.is_record() or a_dt.is_record():
raise OperationError(space.w_TypeError,
space.wrap("invalid type promotion"))
dtype = ufuncs.find_binop_result_dtype(space, dtype, arr.get_dtype())
# concatenate does not handle ndarray subtypes, it always returns a ndarray
res = W_NDimArray.from_shape(space, shape, dtype, 'C')
chunks = [Chunk(0, i, 1, i) for i in shape]
axis_start = 0
for arr in args_w:
if arr.get_shape()[axis] == 0:
continue
chunks[axis] = Chunk(axis_start, axis_start + arr.get_shape()[axis], 1,
arr.get_shape()[axis])
Chunks(chunks).apply(space, res).implementation.setslice(space, arr)
axis_start += arr.get_shape()[axis]
return res
def __init__(self, space, w_seq, w_flags, w_op_flags, w_op_dtypes,
w_casting, w_op_axes, w_itershape, buffersize=0, order='K'):
from pypy.module.micronumpy.ufuncs import find_binop_result_dtype
self.order = order
self.external_loop = False
self.buffered = False
self.tracked_index = ''
self.common_dtype = False
self.delay_bufalloc = False
self.grow_inner = False
self.ranged = False
self.refs_ok = False
self.reduce_ok = False
self.zerosize_ok = False
self.index_iter = None
self.done = False
self.first_next = True
self.op_axes = []
# convert w_seq operands to a list of W_NDimArray
if space.isinstance_w(w_seq, space.w_tuple) or \
space.isinstance_w(w_seq, space.w_list):
w_seq_as_list = space.listview(w_seq)
self.seq = [convert_to_array(space, w_elem)
if not space.is_none(w_elem) else None
for w_elem in w_seq_as_list]
else:
self.seq = [convert_to_array(space, w_seq)]
parse_func_flags(space, self, w_flags)
self.op_flags = parse_op_arg(space, 'op_flags', w_op_flags,
len(self.seq), parse_op_flag)
# handle w_op_axes
oa_ndim = -1
if not space.is_none(w_op_axes):
oa_ndim = self.set_op_axes(space, w_op_axes)
self.ndim = calculate_ndim(self.seq, oa_ndim)
# handle w_op_dtypes part 1: creating self.dtypes list from input
if not space.is_none(w_op_dtypes):
w_seq_as_list = space.listview(w_op_dtypes)
self.dtypes = [decode_w_dtype(space, w_elem) for w_elem in w_seq_as_list]
if len(self.dtypes) != len(self.seq):
raise oefmt(space.w_ValueError,
"op_dtypes must be a tuple/list matching the number of ops")
else:
self.dtypes = []
# handle None or writable operands, calculate my shape
outargs = [i for i in range(len(self.seq))
if self.seq[i] is None or self.op_flags[i].rw == 'w']
if len(outargs) > 0:
out_shape = shape_agreement_multiple(space, [self.seq[i] for i in outargs])
else:
out_shape = None
if space.isinstance_w(w_itershape, space.w_tuple) or \
space.isinstance_w(w_itershape, space.w_list):
self.shape = [space.int_w(i) for i in space.listview(w_itershape)]
else:
self.shape = shape_agreement_multiple(space, self.seq,
shape=out_shape)
if len(outargs) > 0:
# Make None operands writeonly and flagged for allocation
if len(self.dtypes) > 0:
out_dtype = self.dtypes[outargs[0]]
else:
out_dtype = None
for i in range(len(self.seq)):
if self.seq[i] is None:
self.op_flags[i].allocate = True
continue
if self.op_flags[i].rw == 'w':
continue
out_dtype = find_binop_result_dtype(
space, self.seq[i].get_dtype(), out_dtype)
for i in outargs:
if self.seq[i] is None:
# XXX can we postpone allocation to later?
self.seq[i] = W_NDimArray.from_shape(space, self.shape, out_dtype)
else:
if not self.op_flags[i].broadcast:
# Raises if ooutput cannot be broadcast
shape_agreement(space, self.shape, self.seq[i], False)
if self.tracked_index != "":
if self.order == "K":
self.order = self.seq[0].implementation.order
if self.tracked_index == "multi":
backward = False
else:
backward = self.order != self.tracked_index
self.index_iter = IndexIterator(self.shape, backward=backward)
# handle w_op_dtypes part 2: copy where needed if possible
if len(self.dtypes) > 0:
for i in range(len(self.seq)):
selfd = self.dtypes[i]
seq_d = self.seq[i].get_dtype()
if not selfd:
self.dtypes[i] = seq_d
elif selfd != seq_d:
if not 'r' in self.op_flags[i].tmp_copy:
raise oefmt(space.w_TypeError,
"Iterator operand required copying or "
"buffering for operand %d", i)
impl = self.seq[i].implementation
new_impl = impl.astype(space, selfd)
self.seq[i] = W_NDimArray(new_impl)
else:
#copy them from seq
self.dtypes = [s.get_dtype() for s in self.seq]
# create an iterator for each operand
self.iters = []
for i in range(len(self.seq)):
it = get_iter(space, self.order, self.seq[i], self.shape,
self.dtypes[i], self.op_flags[i], self)
it.contiguous = False
self.iters.append((it, it.reset()))
if self.external_loop:
coalesce_axes(self, space)
def __init__(self,
space,
w_seq,
w_flags,
w_op_flags,
w_op_dtypes,
w_casting,
w_op_axes,
w_itershape,
buffersize=0,
order='K'):
from pypy.module.micronumpy.ufuncs import find_binop_result_dtype
self.order = order
self.external_loop = False
self.buffered = False
self.tracked_index = ''
self.common_dtype = False
self.delay_bufalloc = False
self.grow_inner = False
self.ranged = False
self.refs_ok = False
self.reduce_ok = False
self.zerosize_ok = False
self.index_iter = None
self.done = False
self.first_next = True
self.op_axes = []
# convert w_seq operands to a list of W_NDimArray
if space.isinstance_w(w_seq, space.w_tuple) or \
space.isinstance_w(w_seq, space.w_list):
w_seq_as_list = space.listview(w_seq)
self.seq = [
convert_to_array(space, w_elem)
if not space.is_none(w_elem) else None
for w_elem in w_seq_as_list
]
else:
self.seq = [convert_to_array(space, w_seq)]
parse_func_flags(space, self, w_flags)
self.op_flags = parse_op_arg(space, 'op_flags', w_op_flags,
len(self.seq), parse_op_flag)
# handle w_op_axes
oa_ndim = -1
if not space.is_none(w_op_axes):
oa_ndim = self.set_op_axes(space, w_op_axes)
self.ndim = calculate_ndim(self.seq, oa_ndim)
# handle w_op_dtypes part 1: creating self.dtypes list from input
if not space.is_none(w_op_dtypes):
w_seq_as_list = space.listview(w_op_dtypes)
self.dtypes = [
decode_w_dtype(space, w_elem) for w_elem in w_seq_as_list
]
if len(self.dtypes) != len(self.seq):
raise oefmt(
space.w_ValueError,
"op_dtypes must be a tuple/list matching the number of ops"
)
else:
self.dtypes = []
# handle None or writable operands, calculate my shape
outargs = [
i for i in range(len(self.seq))
if self.seq[i] is None or self.op_flags[i].rw == 'w'
]
if len(outargs) > 0:
out_shape = shape_agreement_multiple(
space, [self.seq[i] for i in outargs])
else:
out_shape = None
if space.isinstance_w(w_itershape, space.w_tuple) or \
space.isinstance_w(w_itershape, space.w_list):
self.shape = [space.int_w(i) for i in space.listview(w_itershape)]
else:
self.shape = shape_agreement_multiple(space,
self.seq,
shape=out_shape)
if len(outargs) > 0:
# Make None operands writeonly and flagged for allocation
if len(self.dtypes) > 0:
out_dtype = self.dtypes[outargs[0]]
else:
out_dtype = None
for i in range(len(self.seq)):
if self.seq[i] is None:
self.op_flags[i].allocate = True
continue
if self.op_flags[i].rw == 'w':
continue
out_dtype = find_binop_result_dtype(
space, self.seq[i].get_dtype(), out_dtype)
for i in outargs:
if self.seq[i] is None:
# XXX can we postpone allocation to later?
self.seq[i] = W_NDimArray.from_shape(
space, self.shape, out_dtype)
else:
if not self.op_flags[i].broadcast:
# Raises if ooutput cannot be broadcast
shape_agreement(space, self.shape, self.seq[i], False)
if self.tracked_index != "":
if self.order == "K":
self.order = self.seq[0].implementation.order
if self.tracked_index == "multi":
backward = False
else:
backward = self.order != self.tracked_index
self.index_iter = IndexIterator(self.shape, backward=backward)
# handle w_op_dtypes part 2: copy where needed if possible
if len(self.dtypes) > 0:
for i in range(len(self.seq)):
selfd = self.dtypes[i]
seq_d = self.seq[i].get_dtype()
if not selfd:
self.dtypes[i] = seq_d
elif selfd != seq_d:
if not 'r' in self.op_flags[i].tmp_copy:
raise oefmt(
space.w_TypeError,
"Iterator operand required copying or "
"buffering for operand %d", i)
impl = self.seq[i].implementation
new_impl = impl.astype(space, selfd)
self.seq[i] = W_NDimArray(new_impl)
else:
#copy them from seq
self.dtypes = [s.get_dtype() for s in self.seq]
# create an iterator for each operand
self.iters = []
for i in range(len(self.seq)):
it = get_iter(space, self.order, self.seq[i], self.shape,
self.dtypes[i], self.op_flags[i], self)
it.contiguous = False
self.iters.append((it, it.reset()))
if self.external_loop:
coalesce_axes(self, space)
def concatenate(space, w_args, w_axis=None):
args_w = space.listview(w_args)
if len(args_w) == 0:
raise oefmt(space.w_ValueError, "need at least one array to concatenate")
args_w = [convert_to_array(space, w_arg) for w_arg in args_w]
if w_axis is None:
w_axis = space.wrap(0)
if space.is_none(w_axis):
args_w = [w_arg.reshape(space,
space.newlist([w_arg.descr_get_size(space)]))
for w_arg in args_w]
w_axis = space.wrap(0)
dtype = args_w[0].get_dtype()
shape = args_w[0].get_shape()[:]
ndim = len(shape)
if ndim == 0:
raise oefmt(space.w_ValueError,
"zero-dimensional arrays cannot be concatenated")
axis = space.int_w(w_axis)
orig_axis = axis
if axis < 0:
axis = ndim + axis
if ndim == 1 and axis != 0:
axis = 0
if axis < 0 or axis >= ndim:
raise oefmt(space.w_IndexError, "axis %d out of bounds [0, %d)",
orig_axis, ndim)
for arr in args_w[1:]:
if len(arr.get_shape()) != ndim:
raise OperationError(space.w_ValueError, space.wrap(
"all the input arrays must have same number of dimensions"))
for i, axis_size in enumerate(arr.get_shape()):
if i == axis:
shape[i] += axis_size
elif axis_size != shape[i]:
raise OperationError(space.w_ValueError, space.wrap(
"all the input array dimensions except for the "
"concatenation axis must match exactly"))
a_dt = arr.get_dtype()
if dtype.is_record() and a_dt.is_record():
# Record types must match
for f in dtype.fields:
if f not in a_dt.fields or \
dtype.fields[f] != a_dt.fields[f]:
raise OperationError(space.w_TypeError,
space.wrap("invalid type promotion"))
elif dtype.is_record() or a_dt.is_record():
raise OperationError(space.w_TypeError,
space.wrap("invalid type promotion"))
dtype = ufuncs.find_binop_result_dtype(space, dtype,
arr.get_dtype())
# concatenate does not handle ndarray subtypes, it always returns a ndarray
res = W_NDimArray.from_shape(space, shape, dtype, 'C')
chunks = [Chunk(0, i, 1, i) for i in shape]
axis_start = 0
for arr in args_w:
if arr.get_shape()[axis] == 0:
continue
chunks[axis] = Chunk(axis_start, axis_start + arr.get_shape()[axis], 1,
arr.get_shape()[axis])
Chunks(chunks).apply(space, res).implementation.setslice(space, arr)
axis_start += arr.get_shape()[axis]
return res
def where(space, w_arr, w_x=None, w_y=None):
"""where(condition, [x, y])
Return elements, either from `x` or `y`, depending on `condition`.
If only `condition` is given, return ``condition.nonzero()``.
Parameters
----------
condition : array_like, bool
When True, yield `x`, otherwise yield `y`.
x, y : array_like, optional
Values from which to choose. `x` and `y` need to have the same
shape as `condition`.
Returns
-------
out : ndarray or tuple of ndarrays
If both `x` and `y` are specified, the output array contains
elements of `x` where `condition` is True, and elements from
`y` elsewhere.
If only `condition` is given, return the tuple
``condition.nonzero()``, the indices where `condition` is True.
See Also
--------
nonzero, choose
Notes
-----
If `x` and `y` are given and input arrays are 1-D, `where` is
equivalent to::
[xv if c else yv for (c,xv,yv) in zip(condition,x,y)]
Examples
--------
>>> np.where([[True, False], [True, True]],
... [[1, 2], [3, 4]],
... [[9, 8], [7, 6]])
array([[1, 8],
[3, 4]])
>>> np.where([[0, 1], [1, 0]])
(array([0, 1]), array([1, 0]))
>>> x = np.arange(9.).reshape(3, 3)
>>> np.where( x > 5 )
(array([2, 2, 2]), array([0, 1, 2]))
>>> x[np.where( x > 3.0 )] # Note: result is 1D.
array([ 4., 5., 6., 7., 8.])
>>> np.where(x < 5, x, -1) # Note: broadcasting.
array([[ 0., 1., 2.],
[ 3., 4., -1.],
[-1., -1., -1.]])
NOTE: support for not passing x and y is unsupported
"""
if space.is_none(w_y):
if space.is_none(w_x):
raise OperationError(space.w_NotImplementedError, space.wrap(
"1-arg where unsupported right now"))
raise OperationError(space.w_ValueError, space.wrap(
"Where should be called with either 1 or 3 arguments"))
if space.is_none(w_x):
raise OperationError(space.w_ValueError, space.wrap(
"Where should be called with either 1 or 3 arguments"))
arr = convert_to_array(space, w_arr)
x = convert_to_array(space, w_x)
y = convert_to_array(space, w_y)
if x.is_scalar() and y.is_scalar() and arr.is_scalar():
if arr.get_dtype().itemtype.bool(arr.get_scalar_value()):
return x
return y
dtype = ufuncs.find_binop_result_dtype(space, x.get_dtype(),
y.get_dtype())
shape = shape_agreement(space, arr.get_shape(), x)
shape = shape_agreement(space, shape, y)
out = W_NDimArray.from_shape(space, shape, dtype)
return loop.where(space, out, shape, arr, x, y, dtype)
def test_binops(self, space):
bool_dtype = get_dtype_cache(space).w_booldtype
int8_dtype = get_dtype_cache(space).w_int8dtype
int32_dtype = get_dtype_cache(space).w_int32dtype
float64_dtype = get_dtype_cache(space).w_float64dtype
c64_dtype = get_dtype_cache(space).w_complex64dtype
c128_dtype = get_dtype_cache(space).w_complex128dtype
cld_dtype = get_dtype_cache(space).w_complexlongdtype
fld_dtype = get_dtype_cache(space).w_floatlongdtype
# Basic pairing
assert find_binop_result_dtype(space, bool_dtype, bool_dtype) is bool_dtype
assert find_binop_result_dtype(space, bool_dtype, float64_dtype) is float64_dtype
assert find_binop_result_dtype(space, float64_dtype, bool_dtype) is float64_dtype
assert find_binop_result_dtype(space, int32_dtype, int8_dtype) is int32_dtype
assert find_binop_result_dtype(space, int32_dtype, bool_dtype) is int32_dtype
assert find_binop_result_dtype(space, c64_dtype, float64_dtype) is c128_dtype
assert find_binop_result_dtype(space, c64_dtype, fld_dtype) is cld_dtype
assert find_binop_result_dtype(space, c128_dtype, fld_dtype) is cld_dtype
# With promote bool (happens on div), the result is that the op should
# promote bools to int8
assert find_binop_result_dtype(space, bool_dtype, bool_dtype, promote_bools=True) is int8_dtype
assert find_binop_result_dtype(space, bool_dtype, float64_dtype, promote_bools=True) is float64_dtype
# Coerce to floats
assert find_binop_result_dtype(space, bool_dtype, float64_dtype, promote_to_float=True) is float64_dtype
