def diag(self, X: BlockArray) -> BlockArray:
if len(X.shape) == 1:
shape = X.shape[0], X.shape[0]
block_shape = X.block_shape[0], X.block_shape[0]
grid = ArrayGrid(shape, block_shape, X.dtype.__name__)
grid_meta = grid.to_meta()
rarr = BlockArray(grid, self.cm)
for grid_entry in grid.get_entry_iterator():
syskwargs = {"grid_entry": grid_entry, "grid_shape": grid.grid_shape}
if np.all(np.diff(grid_entry) == 0):
# This is a diagonal block.
rarr.blocks[grid_entry].oid = self.cm.diag(
X.blocks[grid_entry[0]].oid, 0, syskwargs=syskwargs
)
else:
rarr.blocks[grid_entry].oid = self.cm.new_block(
"zeros", grid_entry, grid_meta, syskwargs=syskwargs
)
elif len(X.shape) == 2:
out_shape = (min(X.shape),)
out_block_shape = (min(X.block_shape),)
# Obtain the block indices which contain the diagonal of the matrix.
diag_meta = array_utils.find_diag_output_blocks(X.blocks, out_shape[0])
output_block_arrays = []
out_grid_shape = (len(diag_meta),)
count = 0
# Obtain the diagonals.
for block_indices, offset, total_elements in diag_meta:
syskwargs = {"grid_entry": (count,), "grid_shape": out_grid_shape}
result_block_shape = (total_elements,)
block_grid = ArrayGrid(
result_block_shape,
result_block_shape,
X.blocks[block_indices].dtype.__name__,
)
block_array = BlockArray(block_grid, self.cm)
block_array.blocks[0].oid = self.cm.diag(
X.blocks[block_indices].oid, offset, syskwargs=syskwargs
)
output_block_arrays.append(block_array)
count += 1
if len(output_block_arrays) > 1:
# If there are multiple blocks, concatenate them.
return self.concatenate(
output_block_arrays, axis=0, axis_block_size=out_block_shape[0]
)
return output_block_arrays[0]
else:
raise ValueError("X must have 1 or 2 axes.")
return rarr
def empty(cls, shape, block_shape, dtype, cm: ComputeManager):
grid = ArrayGrid(shape=shape,
block_shape=block_shape,
dtype=dtype.__name__)
grid_meta = grid.to_meta()
arr = BlockArray(grid, cm)
for grid_entry in grid.get_entry_iterator():
arr.blocks[grid_entry].oid = cm.empty(grid_entry,
grid_meta,
syskwargs={
"grid_entry": grid_entry,
"grid_shape":
grid.grid_shape
})
return arr
def _new_array(
self, op_name: str, shape: tuple, block_shape: tuple, dtype: np.dtype = None
):
assert len(shape) == len(block_shape)
if dtype is None:
dtype = np.float64
grid = ArrayGrid(shape, block_shape, dtype.__name__)
grid_meta = grid.to_meta()
rarr = BlockArray(grid, self.cm)
for grid_entry in grid.get_entry_iterator():
rarr.blocks[grid_entry].oid = self.cm.new_block(
op_name,
grid_entry,
grid_meta,
syskwargs={"grid_entry": grid_entry, "grid_shape": grid.grid_shape},
)
return rarr
def eye(self, shape: tuple, block_shape: tuple, dtype: np.dtype = None):
assert len(shape) == len(block_shape) == 2
if dtype is None:
dtype = np.float64
grid = ArrayGrid(shape, block_shape, dtype.__name__)
grid_meta = grid.to_meta()
rarr = BlockArray(grid, self.cm)
for grid_entry in grid.get_entry_iterator():
syskwargs = {"grid_entry": grid_entry, "grid_shape": grid.grid_shape}
if np.all(np.diff(grid_entry) == 0):
# This is a diagonal block.
rarr.blocks[grid_entry].oid = self.cm.new_block(
"eye", grid_entry, grid_meta, syskwargs=syskwargs
)
else:
rarr.blocks[grid_entry].oid = self.cm.new_block(
"zeros", grid_entry, grid_meta, syskwargs=syskwargs
)
return rarr
def diag(self, X: BlockArray) -> BlockArray:
if len(X.shape) == 1:
shape = X.shape[0], X.shape[0]
block_shape = X.block_shape[0], X.block_shape[0]
grid = ArrayGrid(shape, block_shape, X.dtype.__name__)
grid_meta = grid.to_meta()
rarr = BlockArray(grid, self.cm)
for grid_entry in grid.get_entry_iterator():
syskwargs = {
"grid_entry": grid_entry,
"grid_shape": grid.grid_shape
}
if np.all(np.diff(grid_entry) == 0):
# This is a diagonal block.
rarr.blocks[grid_entry].oid = self.cm.diag(
X.blocks[grid_entry[0]].oid, syskwargs=syskwargs)
else:
rarr.blocks[grid_entry].oid = self.cm.new_block(
"zeros", grid_entry, grid_meta, syskwargs=syskwargs)
elif len(X.shape) == 2:
assert X.shape[0] == X.shape[1], "X must be a square array."
assert X.block_shape[0] == X.block_shape[
1], "block_shape must be square."
shape = X.shape[0],
block_shape = X.block_shape[0],
grid = ArrayGrid(shape, block_shape, X.dtype.__name__)
rarr = BlockArray(grid, self.cm)
for grid_entry in X.grid.get_entry_iterator():
out_grid_entry = grid_entry[:1]
out_grid_shape = grid.grid_shape[:1]
syskwargs = {
"grid_entry": out_grid_entry,
"grid_shape": out_grid_shape
}
if np.all(np.diff(grid_entry) == 0):
# This is a diagonal block.
rarr.blocks[out_grid_entry].oid = self.cm.diag(
X.blocks[grid_entry].oid, syskwargs=syskwargs)
else:
raise ValueError("X must have 1 or 2 axes.")
return rarr
def diag(self, X: BlockArray) -> BlockArray:
def find_diag_output_blocks(X: BlockArray, total_elements: int):
# The i,j entry corresponding to a block in X_blocks.
block_i, block_j = 0, 0
# The i,j entry within the current block.
element_i, element_j = 0, 0
# Keep track of the no of elements found so far.
count = 0
# Start at block 0,0.
block = X.blocks[(0, 0)]
# Each element contains block indices, diag offset,
# and the total elements required from the block.
diag_meta = []
while count < total_elements:
if element_i > block.shape[0] - 1:
block_i = block_i + 1
element_i = 0
if element_j > block.shape[1] - 1:
block_j = block_j + 1
element_j = 0
block = X.blocks[(block_i, block_j)]
block_rows, block_cols = block.shape[0], block.shape[1]
offset = -element_i if element_i > element_j else element_j
total_elements_block = (min(block_rows - 1 - element_i,
block_cols - 1 - element_j) + 1)
diag_meta.append(
((block_i, block_j), offset, total_elements_block))
count, element_i = (
count + total_elements_block,
element_i + total_elements_block,
)
element_j = element_j + total_elements_block
return diag_meta
if len(X.shape) == 1:
shape = X.shape[0], X.shape[0]
block_shape = X.block_shape[0], X.block_shape[0]
grid = ArrayGrid(shape, block_shape, X.dtype.__name__)
grid_meta = grid.to_meta()
rarr = BlockArray(grid, self.cm)
for grid_entry in grid.get_entry_iterator():
syskwargs = {
"grid_entry": grid_entry,
"grid_shape": grid.grid_shape
}
if np.all(np.diff(grid_entry) == 0):
# This is a diagonal block.
rarr.blocks[grid_entry].oid = self.cm.diag(
X.blocks[grid_entry[0]].oid, 0, syskwargs=syskwargs)
else:
rarr.blocks[grid_entry].oid = self.cm.new_block(
"zeros", grid_entry, grid_meta, syskwargs=syskwargs)
elif len(X.shape) == 2:
out_shape = (min(X.shape), )
out_block_shape = (min(X.block_shape), )
# Obtain the block indices which contain the diagonal of the matrix.
diag_meta = find_diag_output_blocks(X, out_shape[0])
output_block_arrays = []
out_grid_shape = (len(diag_meta), )
count = 0
# Obtain the diagonals.
for block_indices, offset, total_elements in diag_meta:
syskwargs = {
"grid_entry": (count, ),
"grid_shape": out_grid_shape
}
result_block_shape = (total_elements, )
block_grid = ArrayGrid(
result_block_shape,
result_block_shape,
X.blocks[block_indices].dtype.__name__,
)
block_array = BlockArray(block_grid, self.cm)
block_array.blocks[0].oid = self.cm.diag(
X.blocks[block_indices].oid, offset, syskwargs=syskwargs)
output_block_arrays.append(block_array)
count += 1
if len(output_block_arrays) > 1:
# If there are multiple blocks, concatenate them.
return self.concatenate(output_block_arrays,
axis=0,
axis_block_size=out_block_shape[0])
return output_block_arrays[0]
else:
raise ValueError("X must have 1 or 2 axes.")
return rarr