def empty(cls, shape, block_shape, dtype, system):
grid = ArrayGrid(shape=shape,
block_shape=block_shape,
dtype=dtype.__name__)
grid_meta = grid.to_meta()
arr = BlockArray(grid, system)
for grid_entry in grid.get_entry_iterator():
arr.blocks[grid_entry].oid = system.empty(grid_entry, grid_meta,
syskwargs={
"grid_entry": grid_entry,
"grid_shape": grid.grid_shape
})
return arr
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.system)
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.system.new_block(
"eye", grid_entry, grid_meta, syskwargs=syskwargs)
else:
rarr.blocks[grid_entry].oid = self.system.new_block(
"zeros", grid_entry, grid_meta, syskwargs=syskwargs)
return rarr
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._system)
for grid_entry in grid.get_entry_iterator():
rarr.blocks[grid_entry].oid = self._system.new_block(
op_name,
grid_entry,
grid_meta,
syskwargs={
"grid_entry": grid_entry,
"grid_shape": grid.grid_shape
})
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.system)
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.system.diag(
X.blocks[grid_entry[0]].oid, syskwargs=syskwargs)
else:
rarr.blocks[grid_entry].oid = self.system.new_block(
"zeros", grid_entry, grid_meta, syskwargs=syskwargs)
elif len(X.shape) == 2:
assert X.shape[0] == X.shape[1]
assert X.block_shape[0] == X.block_shape[1]
shape = X.shape[0],
block_shape = X.block_shape[0],
grid = ArrayGrid(shape, block_shape, X.dtype.__name__)
rarr = BlockArray(grid, self.system)
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.system.diag(
X.blocks[grid_entry].oid, syskwargs=syskwargs)
else:
raise ValueError("X must have 1 or 2 axes.")
return rarr
def reshape(self, shape=None, block_shape=None):
# TODO (hme): Add support for arbitrary reshape.
if shape is None:
shape = self.shape
if block_shape is None:
block_shape = self.block_shape
if shape == self.shape and block_shape == self.block_shape:
return self
temp_shape = shape
temp_block_shape = block_shape
shape = []
block_shape = []
negative_one = False
for i, dim in enumerate(temp_shape):
if dim == -1:
assert len(self.shape) == 1
if negative_one:
raise Exception("Only one -1 permitted in reshape.")
negative_one = True
shape.append(self.shape[i])
assert temp_block_shape[i] == -1
block_shape.append(self.block_shape[0])
else:
shape.append(dim)
block_shape.append(temp_block_shape[i])
del temp_shape
shape = tuple(shape)
block_shape = tuple(block_shape)
assert np.product(shape) == np.product(self.shape)
# Make sure the difference is either a preceding or succeeding one.
if len(shape) > len(self.shape):
if shape[0] == 1:
grid_entry_op = "shift"
assert shape[1:] == self.shape
elif shape[-1] == 1:
grid_entry_op = "pop"
assert shape[:-1] == self.shape
else:
raise Exception()
elif len(shape) < len(self.shape):
if self.shape[0] == 1:
grid_entry_op = "prep"
assert self.shape[1:] == shape
elif self.shape[-1] == 1:
grid_entry_op = "app"
assert self.shape[:-1] == shape
else:
raise Exception()
else:
grid_entry_op = "none"
assert self.shape == shape
grid = ArrayGrid(shape=shape,
block_shape=block_shape,
dtype=self.grid.dtype.__name__)
grid_meta = grid.to_meta()
rarr = BlockArray(grid, self.system)
for grid_entry in grid.get_entry_iterator():
rarr.blocks[grid_entry].oid = self.system.empty(grid_entry, grid_meta,
syskwargs={
"grid_entry": grid_entry,
"grid_shape": grid.grid_shape
})
grid_entry_slice = grid.get_slice(grid_entry)
if grid_entry_op == "shift":
grid_entry_slice = tuple([0] + list(grid_entry_slice)[1:])
self_grid_entry_slice = self.grid.get_slice(grid_entry[1:])
elif grid_entry_op == "pop":
grid_entry_slice = tuple(list(grid_entry_slice)[:-1] + [0])
self_grid_entry_slice = self.grid.get_slice(grid_entry[:-1])
elif grid_entry_op == "prep":
self_grid_entry_slice = self.grid.get_slice(tuple([0] + list(grid_entry)))
elif grid_entry_op == "prep":
self_grid_entry_slice = self.grid.get_slice(tuple(list(grid_entry) + [0]))
else:
assert grid_entry_op == "none"
self_grid_entry_slice = grid_entry_slice
# TODO (hme): This is costly.
rarr[grid_entry_slice] = self[self_grid_entry_slice]
return rarr
