def repeat(space, w_arr, repeats, w_axis):
arr = convert_to_array(space, w_arr)
if space.is_none(w_axis):
arr = arr.descr_flatten(space)
orig_size = arr.get_shape()[0]
shape = [arr.get_shape()[0] * repeats]
w_res = W_NDimArray.from_shape(space,
shape,
arr.get_dtype(),
w_instance=arr)
for i in range(repeats):
Chunks([Chunk(i, shape[0] - repeats + i, repeats, orig_size)
]).apply(space, w_res).implementation.setslice(space, arr)
else:
axis = space.int_w(w_axis)
shape = arr.get_shape()[:]
chunks = [Chunk(0, i, 1, i) for i in shape]
orig_size = shape[axis]
shape[axis] *= repeats
w_res = W_NDimArray.from_shape(space,
shape,
arr.get_dtype(),
w_instance=arr)
for i in range(repeats):
chunks[axis] = Chunk(i, shape[axis] - repeats + i, repeats,
orig_size)
Chunks(chunks).apply(space,
w_res).implementation.setslice(space, arr)
return w_res
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"))
dtype = find_result_type(space, args_w, [])
# 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 _prepare_slice_args(self, space, w_idx):
if space.isinstance_w(w_idx, space.w_str):
idx = space.str_w(w_idx)
dtype = self.dtype
if not dtype.is_record():
raise oefmt(
space.w_IndexError, "only integers, slices (`:`), "
"ellipsis (`...`), numpy.newaxis (`None`) and integer or "
"boolean arrays are valid indices")
elif idx not in dtype.fields:
raise oefmt(space.w_ValueError, "field named %s not found",
idx)
return RecordChunk(idx)
elif (space.isinstance_w(w_idx, space.w_int)
or space.isinstance_w(w_idx, space.w_slice)):
if len(self.get_shape()) == 0:
raise oefmt(space.w_ValueError, "cannot slice a 0-d array")
return Chunks(
[Chunk(*space.decode_index4(w_idx,
self.get_shape()[0]))])
elif isinstance(w_idx, W_NDimArray) and w_idx.is_scalar():
w_idx = w_idx.get_scalar_value().item(space)
if not space.isinstance_w(w_idx, space.w_int) and \
not space.isinstance_w(w_idx, space.w_bool):
raise OperationError(
space.w_IndexError,
space.wrap(
"arrays used as indices must be of integer (or boolean) type"
))
return Chunks(
[Chunk(*space.decode_index4(w_idx,
self.get_shape()[0]))])
elif space.is_w(w_idx, space.w_None):
return Chunks([NewAxisChunk()])
result = []
i = 0
for w_item in space.fixedview(w_idx):
if space.is_w(w_item, space.w_None):
result.append(NewAxisChunk())
else:
result.append(
Chunk(*space.decode_index4(w_item,
self.get_shape()[i])))
i += 1
return Chunks(result)