def assign_references(self, dst_sel: BasicSelection, value):
# TODO (hme): This seems overly complicated, but correct. Double check it.
# Also, revisit some of the variable names. They will likely
# be confusing in the future.
# The destination has same block shape as value,
# but the destination selection may not have the same shape as value.
# May need to broadcast value to destination selection output shape.
dst_offset = dst_sel.position().value // np.array(
self._source.block_shape, dtype=np.int)
# Do we need to broadcast?
if (isinstance(value, ArrayView) and
(dst_sel.get_output_shape() != value.sel.get_output_shape())):
value = value.create()
if isinstance(value, ArrayView):
# This is the best case.
# We don't need to create value to perform the reference copy.
# No broadcasting required, so this should be okay.
src_offset = value.sel.position().value // np.array(
value._source.block_shape, dtype=np.int)
src_inflated_shape = dst_sel.get_broadcastable_shape()
src_inflated_block_shape = dst_sel.get_broadcastable_block_shape(
value.block_shape)
src_inflated_grid: ArrayGrid = ArrayGrid(src_inflated_shape,
src_inflated_block_shape,
self.grid.dtype.__name__)
for src_grid_entry_inflated in src_inflated_grid.get_entry_iterator(
):
# Num axes in value grid may be too small.
dst_grid_entry = tuple(
(np.array(src_grid_entry_inflated, dtype=np.int) +
dst_offset).tolist())
src_grid_entry = tuple(
(np.array(src_grid_entry_inflated, dtype=np.int) +
src_offset).tolist())
self._source.blocks[dst_grid_entry] = value._source.blocks[
src_grid_entry].copy()
elif isinstance(value, BlockArrayBase):
# The value has already been created, so just leverage value's existing grid iterator.
if value.shape != dst_sel.get_output_shape():
# Need to broadcast.
src_ba: BlockArrayBase = value.broadcast_to(
dst_sel.get_output_shape())
else:
src_ba: BlockArrayBase = value
src_inflated_shape = dst_sel.get_broadcastable_shape()
src_inflated_block_shape = dst_sel.get_broadcastable_block_shape(
src_ba.block_shape)
src_inflated_grid: ArrayGrid = ArrayGrid(src_inflated_shape,
src_inflated_block_shape,
self.grid.dtype.__name__)
src_grid_entry_iterator = list(src_ba.grid.get_entry_iterator())
for src_index, src_grid_entry_inflated in \
enumerate(src_inflated_grid.get_entry_iterator()):
src_grid_entry = src_grid_entry_iterator[src_index]
dst_grid_entry = tuple(
(np.array(src_grid_entry_inflated, dtype=np.int) +
dst_offset).tolist())
self._source.blocks[dst_grid_entry] = src_ba.blocks[
src_grid_entry].copy()
def basic_assign_single_step(self, dst_sel: BasicSelection, value):
assert isinstance(value, (ArrayView, BlockArrayBase))
dst_ba: BlockArrayBase = self._source
dst_sel_arr: np.ndarray = selection.BasicSelection.block_selection(
dst_ba.shape, dst_ba.block_shape)
dst_sel_clipped: np.ndarray = dst_sel_arr & dst_sel
assert dst_sel_clipped.shape == self._source.grid.grid_shape
# We create value's block array, in case we need to broadcast.
# This may not be necessary, but alternative solutions are extremely tedious.
# The result is a block array with replicated blocks,
# which match the output shape of dst_sel.
if isinstance(value, ArrayView):
src_ba_bc: BlockArrayBase = value.create().broadcast_to(
dst_sel.get_output_shape())
elif isinstance(value, BlockArrayBase):
src_ba_bc: BlockArrayBase = value.broadcast_to(
dst_sel.get_output_shape())
else:
raise Exception("Unexpected value type %s." % type(value))
# Different lengths occur when an index is used to perform
# a selection on an axis. Numpy semantics drops such axes. To allow operations
# between source and destination selections, dropped axes are restored with dimension 1
# so that selections are of equal length.
# We restore the dropped dimensions of the destination selection, because
# the source selection must be broadcastable to the destination selection
# for the assignment to be valid.
src_inflated_shape = dst_sel.get_broadcastable_shape()
# The block shapes need not be equal, but the broadcast source block shape must
# match the block shape we obtain below, so that there's a 1-to-1 correspondence
# between the grid entries.
src_inflated_block_shape = dst_sel.get_broadcastable_block_shape(
src_ba_bc.block_shape)
src_inflated_grid: ArrayGrid = ArrayGrid(src_inflated_shape,
src_inflated_block_shape,
self.grid.dtype.__name__)
src_sel_arr: np.ndarray = selection.BasicSelection.block_selection(
src_inflated_shape, src_inflated_block_shape)
src_sel_offset: np.ndarray = src_sel_arr + dst_sel.position()
# The enumeration of grid entries is identical if the broadcast source grid and
# inflated grid have the same number of blocks.
src_grid_entry_iterator = list(src_ba_bc.grid.get_entry_iterator())
for dst_grid_entry in dst_ba.grid.get_entry_iterator():
dst_sel_block: BasicSelection = dst_sel_arr[dst_grid_entry]
dst_sel_block_clipped: BasicSelection = dst_sel_clipped[
dst_grid_entry]
if dst_sel_block_clipped.is_empty():
continue
src_intersection_arr = src_sel_offset & dst_sel_block_clipped
src_oids = []
src_params = []
dst_params = []
dst_block: Block = dst_ba.blocks[dst_grid_entry]
for src_index, src_grid_entry_bc in enumerate(
src_inflated_grid.get_entry_iterator()):
src_intersection_block: BasicSelection = src_intersection_arr[
src_grid_entry_bc]
if src_intersection_block.is_empty():
continue
src_grid_entry = src_grid_entry_iterator[src_index]
src_block: Block = src_ba_bc.blocks[src_grid_entry]
src_oids.append(src_block.oid)
src_sel_block_offset: BasicSelection = src_sel_offset[
src_grid_entry_bc]
src_dep_sel_loc = src_intersection_block - src_sel_block_offset.position(
)
src_params.append((src_dep_sel_loc.selector(),
src_sel_block_offset.get_output_shape(),
src_block.transposed))
# We're looking at intersection of dst block and src block, so the
# location to which we assign must be offset by dst_sel_block.
dst_block_sel_loc: BasicSelection = (src_intersection_block -
dst_sel_block.position())
dst_params.append(
(dst_block_sel_loc.selector(), dst_block.transposed))
if len(src_oids) == 0:
continue
dst_block.oid = self._cm.update_block(dst_block.oid,
*src_oids,
src_params=src_params,
dst_params=dst_params,
syskwargs={
"grid_entry":
dst_block.grid_entry,
"grid_shape":
dst_block.grid_shape
})
