def __init__(self, source):
self.dtype = source.dtype
self.ndim = len(source.shape)
self.shape = source.shape
self.size = product(self.shape)
self._coord_index = np.array(
[product(source.shape[axis + 1:]) for axis in range(self.ndim)])
self._source = source
def sort_coords(coords, shape):
coord_index = np.array(
[product(shape[axis + 1:]) for axis in range(len(shape))])
shape = np.array(shape)
blocks = []
num_coords = 0
while True:
block = []
while len(block) < PER_BLOCK:
try:
coord = next(coords)
block.append(coord)
except StopIteration:
break
if block:
num_coords += len(block)
offsets = np.sum(block * coord_index, 1)
blocks.append(offsets)
else:
break
for block in merge_sort(blocks):
coords = (np.expand_dims(block, 1) // coord_index) % shape
yield from coords
def constant(shape, dtype, value):
assert isinstance(value, dtype)
size = product(shape)
blocks = [
np.ones([PER_BLOCK], dtype) * value
for _ in range(size // PER_BLOCK)
]
if size % PER_BLOCK > 0:
blocks.append(np.ones([size % PER_BLOCK], dtype) * value)
return BlockListBase(shape, dtype, blocks)
def expand(self, new_shape):
if not isinstance(self.accessor, SparseTable):
raise NotImplementedError
if len(new_shape) != len(self.shape):
return ValueError
elif not (np.array(new_shape) >= np.array(self.shape)).all():
return ValueError
self.accessor.shape = new_shape
self.accessor.size = product(new_shape)
self.shape = self.accessor.shape
self.size = self.accessor.size
def __setitem__(self, match, value):
match = validate_match(match, self.shape)
if len(match) == self.ndim and all(isinstance(c, int) for c in match):
index = sum(np.array(match) * self._coord_index)
self._blocks[index // PER_BLOCK][index %
PER_BLOCK] = self.dtype(value)
else:
update_shape = shape_of(self.shape, match)
if isinstance(value, Tensor):
if not isinstance(value, DenseTensor):
value = DenseTensor.from_sparse(value)
if value.shape != update_shape:
value = value.broadcast(update_shape)
value = list(iter(value))
else:
value = list(iter(value for _ in range(product(update_shape))))
assert len(value) == product(update_shape)
affected_coords = itertools.product(*affected(match, self.shape))
for coords in chunk_iter(affected_coords, PER_BLOCK):
indices = np.sum(coords * self._coord_index, 1)
block_ids = indices // PER_BLOCK
offsets = indices % PER_BLOCK
for block_id in np.unique(block_ids):
num_values = np.sum(block_ids == block_id)
block_offsets = offsets[:num_values]
block_values = value[:num_values]
self._blocks[block_id][block_offsets] = np.array(
block_values)
offsets = offsets[num_values:]
value = value[num_values:]
def product(self, axis=None):
if axis is None or (axis == 0 and self.ndim == 1):
if self.all():
return product(value for _, value in self.filled())
else:
return self.dtype(0)
assert axis < self.ndim
shape = list(self.shape)
del shape[axis]
multiplied = SparseTensor(shape, self.dtype)
if axis == 0:
for coord in self.filled_at(list(range(1, self.ndim))):
multiplied[coord] = self[(slice(None), ) + coord].product()
else:
for prefix in self.filled_at(list(range(axis))):
multiplied[prefix] = self[prefix].product(0)
return multiplied
def __init__(self, shape, dtype, blocks=None):
BlockList.__init__(self, shape, dtype)
if blocks is None:
blocks = [
np.zeros([PER_BLOCK], dtype)
for _ in range(self.size // PER_BLOCK)
]
if self.size % PER_BLOCK > 0:
blocks.append(np.zeros([self.size % PER_BLOCK], dtype))
else:
assert len(blocks) == math.ceil(self.size / PER_BLOCK)
for block in blocks[:-1]:
assert block.shape == (PER_BLOCK, )
if self.size % PER_BLOCK:
assert blocks[-1].shape == (self.size % PER_BLOCK, )
self._blocks = blocks
self._coord_index = np.array(
[product(shape[axis + 1:]) for axis in range(self.ndim)])
def filled(self, match=None):
assert list(self._left.filled(match)) == sorted(
self._left.filled(match))
assert list(self._right.filled(match)) == sorted(
self._right.filled(match))
left = self._left.filled(match)
right = self._right.filled(match)
left_done = False
right_done = False
left_next = None
right_next = None
coord_index = np.array([
product(self.shape[axis + 1:]) for axis in range(len(self.shape))
])
while True:
if left_next is None:
try:
left_next = next(left)
except StopIteration:
left_done = True
if right_next is None:
try:
right_next = next(right)
except StopIteration:
right_done = True
if left_done and right_next:
value = self._combinator(self._left.dtype(0), right_next[1])
if value:
yield (right_next[0], value)
right_next = None
if right_done and left_next:
value = self._combinator(left_next[1], self._right.dtype(0))
if value:
yield (left_next[0], value)
left_next = None
if left_done or right_done:
break
else:
(left_coord, left_value) = left_next
(right_coord, right_value) = right_next
left_index = np.sum(np.array(left_coord) * coord_index)
right_index = np.sum(np.array(right_coord) * coord_index)
if left_index == right_index:
left_next = None
right_next = None
value = self._combinator(left_value, right_value)
if value:
yield (left_coord, value)
elif left_index < right_index:
left_next = None
value = self._combinator(left_value, self._right.dtype(0))
if value:
yield (left_coord, value)
elif right_index < left_index:
right_next = None
value = self._combinator(self._left.dtype(0), right_value)
if value:
yield (right_coord, value)
else:
raise RuntimeError
if left_done and not right_done:
for coord, right_value in right:
value = self._combinator(self._left.dtype(0), right_value)
if value:
yield (coord, value)
elif right_done and not left_done:
for coord, left_value in left:
value = self._combinator(left_value, self._right.dtype(0))
if value:
yield (coord, value)
def __init__(self, shape, dtype):
self.dtype = dtype
self.ndim = len(shape)
self.shape = shape
self.size = product(shape)