def multiply(a: LowRank, b: LowRank):
assert_compatible(a, b)
if structured(a.left, a.right, b.left, b.right):
warn_upmodule(
f"Multiplying {a} and {b}: converting factors to dense.",
category=ToDenseWarning,
)
al, am, ar = B.dense(a.left), B.dense(a.middle), B.dense(a.right)
bl, bm, br = B.dense(b.left), B.dense(b.middle), B.dense(b.right)
# Pick apart the matrices.
al, ar = B.unstack(al, axis=1), B.unstack(ar, axis=1)
bl, br = B.unstack(bl, axis=1), B.unstack(br, axis=1)
am = [B.unstack(x, axis=0) for x in B.unstack(am, axis=0)]
bm = [B.unstack(x, axis=0) for x in B.unstack(bm, axis=0)]
# Construct the factors.
left = B.stack(*[B.multiply(ali, blk) for ali in al for blk in bl], axis=1)
right = B.stack(*[B.multiply(arj, brl) for arj in ar for brl in br],
axis=1)
middle = B.stack(
*[
B.stack(*[amij * bmkl for amij in ami for bmkl in bmk], axis=0)
for ami in am for bmk in bm
],
axis=0,
)
return LowRank(left, right, middle)
def block(*rows):
"""Construct a matrix from its blocks, preserving structure when possible.
Assumes that every row has an equal number of blocks and that the sizes
of the blocks align to form a grid.
Args:
*rows (list): Rows of the block matrix.
Returns:
matrix: Assembled matrix with as much structured as possible.
"""
if len(rows) == 1 and len(rows[0]) == 1:
# There is just one block. Return it.
return rows[0][0]
res = _attempt_zero(rows)
if res is not None:
return res
res = _attempt_diagonal(rows)
if res is not None:
return res
# Could not preserve any structure. Simply concatenate them all densely.
warn_upmodule(
"Could not preserve structure in block matrix: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.concat2d(*[[B.dense(x) for x in row] for row in rows]))
def triangular_solve(a: UpperTriangular, b: AbstractMatrix, lower_a=True):
if lower_a:
warn_upmodule(
f'Solving against {a}, but "lower_a" is set to "True": ignoring flag.',
category=UserWarning,
)
return B.solve(a, b)
def power(a: AbstractMatrix, b: B.Numeric):
if structured(a):
warn_upmodule(
f"Taking an element-wise power of {a}: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.power(B.dense(a), b))
def matmul(a: AbstractMatrix, b: LowerTriangular, tr_a=False, tr_b=False):
if structured(a):
warn_upmodule(
f"Matrix-multiplying {a} and {b}: converting to dense.",
category=ToDenseWarning,
)
return B.matmul(a, B.dense(b), tr_a=tr_a, tr_b=tr_b)
def sqrt(a: AbstractMatrix):
if structured(a):
warn_upmodule(
f"Taking an element-wise square root of {a}: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.sqrt(B.dense(a)))
def matmul(a: Diagonal, b: Kronecker, tr_a=False, tr_b=False):
warn_upmodule(
f"Cannot efficiently matrix-multiply {a} by {b}: "
f"converting the Kronecker product to dense.",
category=ToDenseWarning,
)
return B.matmul(a, B.dense(b), tr_a=tr_a, tr_b=tr_b)
def cholesky(a: Woodbury):
if a.cholesky is None:
warn_upmodule(
f"Converting {a} to dense to compute its Cholesky decomposition.",
category=ToDenseWarning,
)
a.cholesky = LowerTriangular(B.cholesky(B.reg(B.dense(a))))
return a.cholesky
def concat(*elements: AbstractMatrix, axis=0):
if structured(*elements):
elements_str = ", ".join(map(str, elements[:3]))
if len(elements) > 3:
elements_str += "..."
warn_upmodule(
f"Concatenating {elements_str}: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.concat(*(B.dense(el) for el in elements), axis=axis))
def test_warn_upmodule(monkeypatch):
orig_warn = warnings.warn
def mock_warn(*args, **kw_args):
assert kw_args["stacklevel"] > 2
orig_warn(*args, **kw_args)
monkeypatch.setattr(warnings, "warn", mock_warn)
with pytest.warns(UserWarning, match="Test warning"):
warn_upmodule("Test warning", category=UserWarning)
def diag(a: LowRank):
if structured(a.left, a.right):
warn_upmodule(
f"Getting the diagonal of {a}: converting the factors to dense.",
category=ToDenseWarning,
)
diag_len = _diag_len(a)
left_mul = B.matmul(a.left, a.middle)
return B.sum(
B.multiply(
B.dense(left_mul)[:diag_len, :],
B.dense(a.right)[:diag_len, :]),
axis=1,
)
def diag(a, b):
# We could merge this with `block`, but `block` has a lot of overhead. It
# seems advantageous to optimise this common case.
warn_upmodule(
f"Constructing a dense block-diagonal matrix from "
f"{a} and {b}: converting to dense.",
category=ToDenseWarning,
)
a = B.dense(a)
b = B.dense(b)
dtype = B.dtype(a)
ar, ac = B.shape(a)
br, bc = B.shape(b)
return Dense(
B.concat2d([a, B.zeros(dtype, ar, bc)], [B.zeros(dtype, br, ac), b]))
def reshape(a: AbstractMatrix, rows: B.Int, cols: B.Int):
warn_upmodule(f"Converting {a} to dense for reshaping.",
category=ToDenseWarning)
return Dense(B.reshape(B.dense(a), rows, cols))
def root(a: Union[LowRank, Woodbury]):
warn_upmodule(f"Converting {a} to dense to compute its square root.",
category=ToDenseWarning)
return Dense(B.root(B.dense(a)))
def divide(a: AbstractMatrix, b: AbstractMatrix):
if structured(a, b):
warn_upmodule(
f"Dividing {a} by {b}: converting to dense.", category=ToDenseWarning
)
return Dense(B.divide(B.dense(a), B.dense(b)))
def __getitem__(self, item):
if structured(self):
warn_upmodule(f"Indexing into {self}: converting to dense.",
category=ToDenseWarning)
return B.dense(self)[item]
def isnan(a: AbstractMatrix):
if structured(a):
warn_upmodule(f'Applying "isnan" to {a}: converting to dense.',
category=ToDenseWarning)
return B.isnan(B.dense(a))
def take(a: AbstractMatrix, indices_or_mask, axis=0):
if structured(a):
warn_upmodule(f"Taking from {a}: converting to dense.", category=ToDenseWarning)
return B.take(B.dense(a), indices_or_mask, axis=axis)
def multiply(a: Constant, b: AbstractMatrix):
assert_compatible(a, b)
if structured(b):
warn_upmodule(f"Multiplying {a} and {b}: converting to dense.",
category=ToDenseWarning)
return Dense(a.const * B.dense(b))
def multiply(a: AbstractMatrix, b: AbstractMatrix):
if structured(a, b):
warn_upmodule(f"Multiplying {a} and {b}: converting to dense.",
category=ToDenseWarning)
return Dense(B.multiply(B.dense(a), B.dense(b)))
def add(a: AbstractMatrix, b: AbstractMatrix):
if structured(a) and structured(b):
warn_upmodule(
f"Adding {a} and {b}: converting to dense.", category=ToDenseWarning
)
return Dense(B.add(B.dense(a), B.dense(b)))
def solve(a: UpperTriangular, b: AbstractMatrix):
if structured(b):
warn_upmodule(f"Solving {a} x = {b}: converting to dense.",
category=ToDenseWarning)
return Dense(B.trisolve(a.mat, B.dense(b), lower_a=False))
def add(a: Constant, b: AbstractMatrix):
if structured(b):
warn_upmodule(
f"Adding {a} and {b}: converting to dense.", category=ToDenseWarning
)
return Dense(a.const + B.dense(b))
def squeeze(a: AbstractMatrix):
if structured(a):
warn_upmodule(f"Squeezing {a}: converting to dense.",
category=ToDenseWarning)
return B.squeeze(B.dense(a))
def solve(a: AbstractMatrix, b: AbstractMatrix):
if structured(a, b):
warn_upmodule(f"Solving {a} x = {b}: converting to dense.",
category=ToDenseWarning)
return B.solve(B.dense(a), B.dense(b))
