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