# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import itertools
import numpy as np
from ... import opcodes as OperandDef
from ...operands import OperandStage
from ...serialize import KeyField, StringField, AnyField, Int64Field, Int32Field
from ...config import options
from ...utils import check_chunks_unknown_shape
from ...tiles import TilesError
from ..operands import TensorOperand, TensorOperandMixin
from ..core import TENSOR_TYPE, TensorOrder
from ..datasource.array import tensor as astensor
from ..array_utils import as_same_device, device
class TensorSearchsorted(TensorOperand, TensorOperandMixin):
_op_type_ = OperandDef.SEARCHSORTED
_input = KeyField('input')
_values = AnyField('values')
_side = StringField('side')
_combine_size = Int32Field('combine_size')
# offset is used only for map stage
_offset = Int64Field('offset')
def __init__(self, values=None, side=None, dtype=None, gpu=None, combine_size=None,
stage=None, offset=None, **kw):
super().__init__(_values=values, _side=side, _dtype=dtype, _gpu=gpu,
_combine_size=combine_size, _stage=stage, _offset=offset, **kw)
def _set_inputs(self, inputs):
super()._set_inputs(inputs)
self._input = self._inputs[0]
if len(self._inputs) == 2:
self._values = self._inputs[1]
@property
def input(self):
return self._input
@property
def values(self):
return self._values
@property
def side(self):
return self._side
@property
def offset(self):
return self._offset
@property
def combine_size(self):
return self._combine_size
def __call__(self, a, v):
inputs = [a]
if isinstance(v, TENSOR_TYPE):
inputs.append(v)
shape = v.shape
else:
shape = ()
return self.new_tensor(inputs, shape=shape, order=TensorOrder.C_ORDER)
@classmethod
def _tile_one_chunk(cls, op, a, v, out):
chunks = []
if len(op.inputs) == 1:
v_chunks = [v]
else:
v_chunks = v.chunks
for v_chunk in v_chunks:
chunk_op = op.copy().reset_key()
in_chunks = [a.chunks[0]]
if len(op.inputs) == 2:
in_chunks.append(v_chunk)
v_shape = v_chunk.shape if hasattr(v_chunk, 'shape') else ()
chunk_idx = v_chunk.index if len(op.inputs) == 2 else (0,)
chunk = chunk_op.new_chunk(in_chunks, shape=v_shape,
index=chunk_idx, order=out.order)
chunks.append(chunk)
new_op = op.copy().reset_key()
nsplits = ((s,) for s in out.shape) if len(op.inputs) == 1 else v.nsplits
return new_op.new_tensors(op.inputs, out.shape,
chunks=chunks, nsplits=nsplits)
@classmethod
def _combine_chunks(cls, to_combine, op, stage, v, idx):
from ..merge import TensorStack
v_shape = v.shape if hasattr(v, 'shape') else ()
combine_op = TensorStack(axis=0, dtype=op.outputs[0].dtype)
combine_chunk = combine_op.new_chunk(to_combine, shape=v_shape)
chunk_op = op.copy().reset_key()
chunk_op._stage = stage
in_chunks = [combine_chunk]
if len(op.inputs) == 2:
in_chunks.append(v)
return chunk_op.new_chunk(in_chunks, shape=v_shape, index=idx,
order=op.outputs[0].order)
@classmethod
def _tile_tree_reduction(cls, op, a, v, out):
check_chunks_unknown_shape(op.inputs, TilesError)
combine_size = op.combine_size or options.combine_size
input_len = len(op.inputs)
v_chunks = [v] if input_len == 1 else v.chunks
out_chunks = []
for v_chunk in v_chunks:
offsets = [0] + np.cumsum(a.nsplits[0]).tolist()[:-1]
v_shape = v_chunk.shape if hasattr(v_chunk, 'shape') else ()
v_index = v_chunk.index if hasattr(v_chunk, 'index') else (0,)
chunks = []
for i, c in enumerate(a.chunks):
chunk_op = op.copy().reset_key()
chunk_op._stage = OperandStage.map
chunk_op._offset = offsets[i]
in_chunks = [c]
if input_len == 2:
in_chunks.append(v_chunk)
chunks.append(chunk_op.new_chunk(in_chunks, shape=v_shape,
index=c.index, order=out.order))
while len(chunks) > combine_size:
new_chunks = []
it = itertools.count(0)
while True:
j = next(it)
to_combine = chunks[j * combine_size: (j + 1) * combine_size]
if len(to_combine) == 0:
break
new_chunks.append(
cls._combine_chunks(to_combine, op, OperandStage.combine, v_chunk, (j,)))
chunks = new_chunks
chunk = cls._combine_chunks(chunks, op, OperandStage.reduce, v_chunk, v_index)
out_chunks.append(chunk)
new_op = op.copy().reset_key()
nsplits = ((s,) for s in out.shape) if len(op.inputs) == 1 else v.nsplits
return new_op.new_tensors(op.inputs, out.shape,
chunks=out_chunks, nsplits=nsplits)
@classmethod
def tile(cls, op):
a = op.inputs[0]
out = op.outputs[0]
input_len = len(op.inputs)
if input_len == 1:
v = op.values
else:
v = op.inputs[1]
if len(a.chunks) == 1:
return cls._tile_one_chunk(op, a, v, out)
return cls._tile_tree_reduction(op, a, v, out)
@classmethod
def _execute_without_stage(cls, xp, a, v, op):
return xp.searchsorted(a, v, side=op.side)
@classmethod
def _execute_map(cls, xp, a, v, op):
# in the map phase, calculate the indices and positions
# for instance, a=[1, 4, 6], v=5, return will be (2, 6)
indices = xp.atleast_1d(xp.searchsorted(a, v, side=op.side))
data_indices = indices.copy()
# if the value is larger than all data
# for instance, a=[1, 4, 6], v=7
# return will be (2, 6), not (3, 6), thus needs to subtract 1
data_indices = xp.subtract(data_indices, 1, out=data_indices,
where=data_indices >= len(a))
data = a[data_indices]
if op.offset > 0:
indices = xp.add(indices, op.offset, out=indices)
if np.isscalar(v):
indices, data = indices[0], data[0]
return indices, data
@classmethod
def _execute_combine(cls, xp, a, v, op):
inp_indices, inp_data = a
if np.isscalar(v):
ind = xp.searchsorted(inp_data, v, side=op.side)
if ind >= len(inp_data):
ind -= 1
return inp_indices[ind], inp_data[ind]
else:
ret_indices = np.empty(v.shape, dtype=np.intp)
ret_data = np.empty(v.shape, dtype=inp_data.dtype)
for idx in itertools.product(*(range(s) for s in v.shape)):
ind = xp.searchsorted(inp_data[(slice(None),) + idx], v[idx], side=op.side)
if ind >= len(inp_indices):
ind -= 1
ret_indices[idx] = inp_indices[(ind,) + idx]
ret_data[idx] = inp_data[(ind,) + idx]
return ret_indices, ret_data
@classmethod
def _execute_reduce(cls, xp, a, v, op):
inp_indices, inp_data = a
if np.isscalar(v):
ind = xp.searchsorted(inp_data, v, side=op.side)
if ind >= len(inp_indices):
ind -= 1
return inp_indices[ind]
else:
indices = np.empty(v.shape, dtype=np.intp)
for idx in itertools.product(*(range(s) for s in v.shape)):
ind = xp.searchsorted(inp_data[(slice(None),) + idx], v[idx], side=op.side)
if ind >= len(inp_indices):
ind -= 1
indices[idx] = inp_indices[(ind,) + idx]
return indices
@classmethod
def execute(cls, ctx, op):
a = ctx[op.inputs[0].key]
v = ctx[op.inputs[1].key] if len(op.inputs) == 2 else op.values
data = []
if isinstance(a, tuple):
data.extend(a)
else:
data.append(a)
if len(op.inputs) == 2:
data.append(v)
data, device_id, xp = as_same_device(
data, device=op.device, ret_extra=True)
if isinstance(a, tuple):
a = data[:2]
else:
a = data[0]
if len(op.inputs) == 2:
v = data[-1]
with device(device_id):
if op.stage is None:
ret = cls._execute_without_stage(xp, a, v, op)
elif op.stage == OperandStage.map:
ret = cls._execute_map(xp, a, v, op)
elif op.stage == OperandStage.combine:
ret = cls._execute_combine(xp, a, v, op)
else:
ret = cls._execute_reduce(xp, a, v, op)
ctx[op.outputs[0].key] = ret
[docs]def searchsorted(a, v, side='left', sorter=None, combine_size=None):
"""
Find indices where elements should be inserted to maintain order.
Find the indices into a sorted tensor `a` such that, if the
corresponding elements in `v` were inserted before the indices, the
order of `a` would be preserved.
Assuming that `a` is sorted:
====== ============================
`side` returned index `i` satisfies
====== ============================
left ``a[i-1] < v <= a[i]``
right ``a[i-1] <= v < a[i]``
====== ============================
Parameters
----------
a : 1-D array_like
Input tensor. If `sorter` is None, then it must be sorted in
ascending order, otherwise `sorter` must be an array of indices
that sort it.
v : array_like
Values to insert into `a`.
side : {'left', 'right'}, optional
If 'left', the index of the first suitable location found is given.
If 'right', return the last such index. If there is no suitable
index, return either 0 or N (where N is the length of `a`).
sorter : 1-D array_like, optional
Optional tensor of integer indices that sort array a into ascending
order. They are typically the result of argsort.
combine_size: int, optional
The number of chunks to combine.
Returns
-------
indices : tensor of ints
Array of insertion points with the same shape as `v`.
See Also
--------
sort : Return a sorted copy of a tensor.
histogram : Produce histogram from 1-D data.
Notes
-----
Binary search is used to find the required insertion points.
This function is a faster version of the builtin python `bisect.bisect_left`
(``side='left'``) and `bisect.bisect_right` (``side='right'``) functions,
which is also vectorized in the `v` argument.
Examples
--------
>>> import mars.tensor as mt
>>> mt.searchsorted([1,2,3,4,5], 3).execute()
2
>>> mt.searchsorted([1,2,3,4,5], 3, side='right').execute()
3
>>> mt.searchsorted([1,2,3,4,5], [-10, 10, 2, 3]).execute()
array([0, 5, 1, 2])
"""
if not isinstance(a, TENSOR_TYPE) and sorter is not None and \
not isinstance(sorter, TENSOR_TYPE):
a = astensor(np.asarray(a)[sorter])
else:
a = astensor(a)
if sorter is not None:
a = a[sorter]
if a.ndim != 1:
raise ValueError('`a` should be 1-d tensor')
if a.issparse():
# does not support sparse tensor
raise ValueError('`a` should be a dense tensor')
if side not in {'left', 'right'}:
raise ValueError(f"'{side}' is an invalid value for keyword 'side'")
if not np.isscalar(v):
v = astensor(v)
op = TensorSearchsorted(values=v, side=side, dtype=np.dtype(np.intp),
combine_size=combine_size)
return op(a, v)