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 })