"""Finite-dimensional linear operators."""
from typing import Callable, Optional, Tuple, Union
import numpy as np
import scipy.linalg
import scipy.sparse.linalg
import probnum.utils
from probnum import config
from probnum.typing import DTypeLike, ScalarLike, ShapeLike
BinaryOperandType = Union[
"LinearOperator", ScalarLike, np.ndarray, scipy.sparse.spmatrix
]
# pylint: disable="too-many-lines"
class LinearOperator:
r"""Composite base class for finite-dimensional linear operators.
This class provides a way to define finite-dimensional linear operators without
explicitly constructing a matrix representation. Instead it suffices to define a
matrix-vector product and a shape attribute. This avoids unnecessary memory usage
and can often be more convenient to derive.
:class:`LinearOperator` instances can be multiplied, added and exponentiated. This
happens lazily: the result of these operations is a new, composite
:class:`LinearOperator`, that defers linear operations to the original operators and
combines the results.
To construct a concrete :class:`LinearOperator`, either pass appropriate callables to
the constructor of this class, or subclass it.
A subclass must implement either one of the methods ``_matvec`` and ``_matmat``, and
the attributes/properties ``shape`` (pair of integers) and ``dtype`` (may be
``None``). It may call the ``__init__`` on this class to have these attributes
validated. Implementing ``_matvec`` automatically implements ``_matmat`` (using a
naive algorithm) and vice-versa.
Optionally, a subclass may implement ``_rmatvec`` or ``_adjoint`` to implement the
Hermitian adjoint (conjugate transpose). As with ``_matvec`` and ``_matmat``,
implementing either ``_rmatvec`` or ``_adjoint`` implements the other automatically.
Implementing ``_adjoint`` is preferable; ``_rmatvec`` is mostly there for backwards
compatibility.
Parameters
----------
shape :
Matrix dimensions `(M, N)`.
dtype :
Data type of the operator.
matmul :
Callable which computes the matrix-matrix product :math:`y = A V`, where
:math:`A` is the linear operator and :math:`V` is an :math:`N \times K` matrix.
The callable must support broadcasted matrix products, i.e. the argument
:math:`V` might also be a stack of matrices in which case the broadcasting rules
of :func:`np.matmul` must apply.
Note that the argument to this callable is guaranteed to have at least two
dimensions.
rmatmul :
Callable which implements the matrix-matrix product, i.e. :math:`A @ V`, where
:math:`A` is the linear operator and :math:`V` is a matrix of shape `(N, K)`.
todense :
Callable which returns a dense matrix representation of the linear operator as a
:class:`np.ndarray`. The output of this function must be equivalent to the
output of :code:`A.matmat(np.eye(N, dtype=A.dtype))`.
rmatvec :
Callable which implements the matrix-vector product with the adjoint of the
operator, i.e. :math:`A^H v`, where :math:`A^H` is the conjugate transpose of
the linear operator :math:`A` and :math:`v` is a vector of shape `(N,)`.
This argument will be ignored if `adjoint` is given.
rmatmat :
Returns :math:`A^H V`, where :math:`V` is a dense matrix with dimensions (M, K).
See Also
--------
aslinop : Transform into a LinearOperator.
Examples
--------
>>> import numpy as np
>>> from probnum.linops import LinearOperator
>>> @LinearOperator.broadcast_matvec
... def mv(v):
... return np.array([2 * v[0] - v[1], 3 * v[1]])
>>> A = LinearOperator(shape=(2, 2), dtype=np.float_, matmul=mv)
>>> A
<LinearOperator with shape=(2, 2) and dtype=float64>
>>> A @ np.array([1., 2.])
array([0., 6.])
>>> A @ np.ones(2)
array([1., 3.])
"""
# pylint: disable=too-many-public-methods
def __init__(
self,
shape: ShapeLike,
dtype: DTypeLike,
*,
matmul: Callable[[np.ndarray], np.ndarray],
rmatmul: Optional[Callable[[np.ndarray], np.ndarray]] = None,
apply: Callable[[np.ndarray, int], np.ndarray] = None,
todense: Optional[Callable[[], np.ndarray]] = None,
transpose: Optional[Callable[[np.ndarray], "LinearOperator"]] = None,
inverse: Optional[Callable[[], "LinearOperator"]] = None,
rank: Optional[Callable[[], np.intp]] = None,
eigvals: Optional[Callable[[], np.ndarray]] = None,
cond: Optional[
Callable[[Optional[Union[None, int, str, np.floating]]], np.number]
] = None,
det: Optional[Callable[[], np.number]] = None,
logabsdet: Optional[Callable[[], np.flexible]] = None,
trace: Optional[Callable[[], np.number]] = None,
):
self.__shape = probnum.utils.as_shape(shape, ndim=2)
# DType
self.__dtype = np.dtype(dtype)
if not np.issubdtype(self.__dtype, np.number):
raise TypeError("The dtype of a linear operator must be numeric.")
if np.issubdtype(self.__dtype, np.complexfloating):
raise TypeError("Linear operators do not support complex dtypes.")
self.__matmul = matmul # (self @ x)
self.__rmatmul = rmatmul # (x @ self)
self.__apply = apply # __call__
self.__todense = todense
self.__transpose = transpose
self.__inverse = inverse
# Matrix properties
self.__rank = rank
self.__eigvals = eigvals
self.__cond = cond
self.__det = det
self.__logabsdet = logabsdet
self.__trace = trace
# Caches
self.__todense_cache = None
self.__rank_cache = None
self.__eigvals_cache = None
self.__cond_cache = {}
self.__det_cache = None
self.__logabsdet_cache = None
self.__trace_cache = None
@property
def shape(self) -> Tuple[int, int]:
"""Shape of the linear operator.
Defined as a tuple of the output and input dimension of
operator.
"""
return self.__shape
@property
def ndim(self) -> int:
"""Number of linear operator dimensions.
Defined analogously to :attr:`numpy.ndarray.ndim`.
"""
return 2
@property
def size(self) -> int:
return self.__shape[0] * self.__shape[1]
@property
def dtype(self) -> np.dtype:
"""Data type of the linear operator."""
return self.__dtype
@property
def is_square(self) -> bool:
"""Whether input dimension matches output dimension."""
return self.shape[0] == self.shape[1]
def __repr__(self) -> str:
return (
f"<{self.__class__.__name__} with "
f"shape={self.shape} and "
f"dtype={str(self.dtype)}>"
)
[docs] def __call__(self, x: np.ndarray, axis: Optional[int] = None) -> np.ndarray:
if axis is not None and (axis < -x.ndim or axis >= x.ndim):
raise np.AxisError(axis, ndim=x.ndim)
if x.ndim == 1:
return self @ x
elif x.ndim > 1:
if axis is None:
axis = -1
if axis < 0:
axis += x.ndim
if x.shape[axis] != self.__shape[1]:
raise ValueError(
f"Dimension mismatch. Expected array with {self.__shape[1]} "
f"entries along axis {axis}, but got array with shape {x.shape}."
)
if axis == (x.ndim - 1):
return (self @ x[..., np.newaxis])[..., 0]
elif axis == (x.ndim - 2):
return self @ x
else:
if self.__apply is None:
return np.moveaxis(self @ np.moveaxis(x, axis, -2), -2, axis)
return self.__apply(x, axis)
else:
raise ValueError("The operand must be at least one dimensional.")
[docs] def astype(
self,
dtype: DTypeLike,
order: str = "K",
casting: str = "unsafe",
subok: bool = True,
copy: bool = True,
) -> "LinearOperator":
"""Cast a linear operator to a different ``dtype``.
Parameters
----------
dtype:
Data type to which the linear operator is cast.
order:
Memory layout order of the result.
casting:
Controls what kind of data casting may occur.
subok:
If True, then sub-classes will be passed-through (default). False is currently not supported for linear operators.
copy:
Whether to return a new linear operator, even if ``dtype`` is the same.
"""
dtype = np.dtype(dtype)
if not np.can_cast(self.dtype, dtype, casting=casting):
raise TypeError(
f"Cannot cast linear operator from {self.dtype} to {dtype} "
f"according to the rule {casting}"
)
if not subok:
raise NotImplementedError(
"Setting `subok` to `False` is not supported for linear operators"
)
return self._astype(dtype, order, casting, copy)
def _astype(
self, dtype: np.dtype, order: str, casting: str, copy: bool
) -> "LinearOperator":
if self.dtype == dtype and not copy:
return self
else:
return _TypeCastLinearOperator(self, dtype, order, casting, copy)
[docs] def todense(self, cache: bool = True) -> np.ndarray:
"""Dense matrix representation of the linear operator.
This method can be computationally very costly depending on the shape of the
linear operator. Use with caution.
Returns
-------
matrix : np.ndarray
Matrix representation of the linear operator.
"""
if self.__todense_cache is None:
if self.__todense is not None:
dense = self.__todense()
else:
dense = self @ np.eye(self.shape[1], dtype=self.__dtype, order="F")
if not cache:
return dense
self.__todense_cache = dense
return self.__todense_cache
####################################################################################
# Derived Quantities
####################################################################################
[docs] def rank(self) -> np.intp:
"""Rank of the linear operator."""
if self.__rank_cache is None:
if self.__rank is not None:
self.__rank_cache = self.__rank()
else:
self.__rank_cache = np.linalg.matrix_rank(self.todense(cache=False))
return self.__rank_cache
[docs] def eigvals(self) -> np.ndarray:
"""Eigenvalue spectrum of the linear operator."""
if self.__eigvals_cache is None:
if not self.is_square:
raise np.linalg.LinAlgError(
"Eigenvalues are only defined for square operators"
)
if self.__eigvals is not None:
self.__eigvals_cache = self.__eigvals()
else:
self.__eigvals_cache = np.linalg.eigvals(self.todense(cache=False))
self.__eigvals_cache.setflags(write=False)
return self.__eigvals_cache
[docs] def cond(self, p=None) -> np.inexact:
"""Compute the condition number of the linear operator.
The condition number of the linear operator with respect to the ``p`` norm. It
measures how much the solution :math:`x` of the linear system :math:`Ax=b`
changes with respect to small changes in :math:`b`.
Parameters
----------
p : {None, 1, , 2, , inf, 'fro'}, optional
Order of the norm:
======= ============================
p norm for matrices
======= ============================
None 2-norm, computed directly via singular value decomposition
'fro' Frobenius norm
np.inf max(sum(abs(x), axis=1))
1 max(sum(abs(x), axis=0))
2 2-norm (largest sing. value)
======= ============================
Returns
-------
cond :
The condition number of the linear operator. May be infinite.
"""
if p not in self.__cond_cache:
if not self.is_square:
raise np.linalg.LinAlgError(
"The condition number is only defined for square operators"
)
if self.__cond is not None:
self.__cond_cache[p] = self.__cond(p)
else:
self.__cond_cache[p] = np.linalg.cond(self.todense(cache=False), p=p)
return self.__cond_cache[p]
[docs] def det(self) -> np.inexact:
"""Determinant of the linear operator."""
if self.__det_cache is None:
if not self.is_square:
raise np.linalg.LinAlgError(
"The determinant is only defined for square operators"
)
if self.__det is not None:
self.__det_cache = self.__det()
else:
self.__det_cache = np.linalg.det(self.todense(cache=False))
return self.__det_cache
[docs] def logabsdet(self) -> np.inexact:
"""Log absolute determinant of the linear operator."""
if self.__logabsdet_cache is None:
if not self.is_square:
raise np.linalg.LinAlgError(
"The determinant is only defined for square operators"
)
if self.__logabsdet is not None:
self.__logabsdet_cache = self.__logabsdet()
else:
self.__logabsdet_cache = self._logabsdet_fallback()
return self.__logabsdet_cache
def _logabsdet_fallback(self) -> np.inexact:
if self.det() == 0:
return probnum.utils.as_numpy_scalar(-np.inf, dtype=self._inexact_dtype)
else:
return np.log(np.abs(self.det()))
[docs] def trace(self) -> np.number:
r"""Trace of the linear operator.
Computes the trace of a square linear operator :math:`\text{tr}(A) =
\sum_{i-1}^n A_{ii}`.
Returns
-------
trace : float
Trace of the linear operator.
Raises
------
LinAlgError :
If :meth:`trace` is called on a non-square matrix.
"""
if self.__trace_cache is None:
if not self.is_square:
raise np.linalg.LinAlgError(
"The trace is only defined for square operators."
)
if self.__trace is not None:
self.__trace_cache = self.__trace()
else:
self.__trace_cache = self._trace_fallback()
return self.__trace_cache
def _trace_fallback(self) -> np.number:
vec = np.zeros(self.shape[1], dtype=self.dtype)
vec[0] = 1
trace = (self @ vec)[0]
vec[0] = 0
for i in range(1, self.shape[0]):
vec[i] = 1
trace += (self @ vec)[i]
vec[i] = 0
return trace
####################################################################################
# Unary Arithmetic
####################################################################################
def __neg__(self) -> "LinearOperator":
from ._arithmetic import ( # pylint: disable=import-outside-toplevel
NegatedLinearOperator,
)
return NegatedLinearOperator(self)
@property
def T(self) -> "LinearOperator":
if self.__transpose is None:
return self._transpose_fallback()
return self.__transpose()
[docs] def transpose(self, *axes: Union[int, Tuple[int]]) -> "LinearOperator":
"""Transpose this linear operator.
Can be abbreviated self.T instead of self.transpose().
"""
if len(axes) > 0:
if len(axes) == 1 and isinstance(axes[0], tuple):
axes = axes[0]
if len(axes) != 2:
raise ValueError(
f"The given axes {axes} don't match the linear operator with shape "
f"{self.shape}."
)
axes_int = []
for axis in axes:
axis = int(axis)
if not -2 <= axis <= 1:
raise np.AxisError(axis, ndim=2)
if axis < 0:
axis += 2
axes_int.append(axis)
axes = tuple(axes_int)
if axes == (0, 1):
return self
if axes == (1, 0):
return self.T
raise ValueError("Cannot transpose a linear operator along repeated axes.")
return self.T
def _transpose_fallback(self) -> "LinearOperator":
if self.__rmatmul is not None:
return TransposedLinearOperator(
self,
matmul=lambda x: self.__rmatmul(x[..., np.newaxis])[..., :],
)
return TransposedLinearOperator(self)
[docs] def inv(self) -> "LinearOperator":
"""Inverse of the linear operator."""
if self.__inverse is None:
return _InverseLinearOperator(self)
return self.__inverse()
####################################################################################
# Binary Arithmetic
####################################################################################
__array_ufunc__ = None
"""
This prevents numpy from calling elementwise arithmetic operations allowing
expressions like `y = np.array([1, 1]) + linop` to call the arithmetic operations
defined by `LinearOperator` instead of elementwise. Thus no array of
`LinearOperator`s but a `LinearOperator` with the correct shape is returned.
"""
def __add__(self, other: BinaryOperandType) -> "LinearOperator":
from ._arithmetic import add # pylint: disable=import-outside-toplevel
return add(self, other)
def __radd__(self, other: BinaryOperandType) -> "LinearOperator":
from ._arithmetic import add # pylint: disable=import-outside-toplevel
return add(other, self)
def __sub__(self, other: BinaryOperandType) -> "LinearOperator":
from ._arithmetic import sub # pylint: disable=import-outside-toplevel
return sub(self, other)
def __rsub__(self, other: BinaryOperandType) -> "LinearOperator":
from ._arithmetic import sub # pylint: disable=import-outside-toplevel
return sub(other, self)
def __mul__(self, other: BinaryOperandType) -> "LinearOperator":
from ._arithmetic import mul # pylint: disable=import-outside-toplevel
return mul(self, other)
def __rmul__(self, other: BinaryOperandType) -> "LinearOperator":
from ._arithmetic import mul # pylint: disable=import-outside-toplevel
return mul(other, self)
def _is_type_shape_dtype_equal(self, other: "LinearOperator") -> bool:
return (
isinstance(self, type(other))
and self.shape == other.shape
and self.dtype == other.dtype
)
def __matmul__(
self, other: BinaryOperandType
) -> Union["LinearOperator", np.ndarray]:
"""Matrix-vector multiplication.
Performs the operation `y = self @ x` where `self` is an MxN linear operator
and `x` is a 1-d array or random variable.
Parameters
----------
x :
An array or `RandomVariable` with shape `(N,)` or `(N, 1)`.
Returns
-------
y :
A `np.matrix` or `np.ndarray` or `RandomVariable` with shape `(M,)` or `(M, 1)`,
depending on the type and shape of the x argument.
Notes
-----
This matvec wraps the user-specified matvec routine or overridden
_matvec method to ensure that y has the correct shape and type.
"""
if isinstance(other, np.ndarray):
x = other
M, N = self.shape
if x.ndim == 1 and x.shape == (N,):
y = self.__matmul(x[:, np.newaxis])[:, 0]
assert y.ndim == 1
assert y.shape == (M,)
elif x.ndim > 1 and x.shape[-2] == N:
y = self.__matmul(x)
assert y.ndim > 1
assert y.shape == x.shape[:-2] + (M, x.shape[-1])
else:
raise ValueError(
f"Dimension mismatch. Expected operand of shape ({N},) or "
f"(..., {N}, K), but got {x.shape}."
)
return y
else:
from ._arithmetic import matmul # pylint: disable=import-outside-toplevel
return matmul(self, other)
def __rmatmul__(
self, other: BinaryOperandType
) -> Union["LinearOperator", np.ndarray]:
if isinstance(other, np.ndarray):
x = other
M, N = self.shape
if x.ndim >= 1 and x.shape[-1] == M:
if self.__rmatmul is not None:
if x.ndim == 1:
y = self.__rmatmul(x[np.newaxis, :])[0, :]
else:
y = self.__rmatmul(x)
else:
y = (self.T)(x, axis=-1)
else:
raise ValueError(
f"Dimension mismatch. Expected operand of shape (..., {M}), but "
f"got {x.shape}."
)
assert y.ndim >= 1
assert y.shape[-1] == N
assert y.shape[:-1] == x.shape[:-1]
return y
else:
from ._arithmetic import matmul # pylint: disable=import-outside-toplevel
return matmul(other, self)
####################################################################################
# Automatic `(r)mat{vec,mat}`` to `(r)matmul` Broadcasting
####################################################################################
[docs] @classmethod
def broadcast_matvec(
cls, matvec: Callable[[np.ndarray], np.ndarray]
) -> Callable[[np.ndarray], np.ndarray]:
"""Broadcasting for a (implicitly defined) matrix-vector product.
Convenience function / decorator to broadcast the definition of
a matrix-vector product. This can be used to easily construct a
new linear operator only from a matrix-vector product.
"""
def _matmul(x: np.ndarray) -> np.ndarray:
if x.ndim == 2 and x.shape[1] == 1:
return matvec(x[:, 0])[:, np.newaxis]
return np.apply_along_axis(matvec, -2, x)
return _matmul
[docs] @classmethod
def broadcast_matmat(
cls, matmat: Callable[[np.ndarray], np.ndarray]
) -> Callable[[np.ndarray], np.ndarray]:
"""Broadcasting for a (implicitly defined) matrix-matrix product.
Convenience function / decorator to broadcast the definition of
a matrix-matrix product to vectors. This can be used to easily
construct a new linear operator only from a matrix-matrix
product.
"""
def _matmul(x: np.ndarray) -> np.ndarray:
if x.ndim == 2:
return matmat(x)
return _apply_to_matrix_stack(matmat, x)
return _matmul
[docs] @classmethod
def broadcast_rmatvec(
cls, rmatvec: Callable[[np.ndarray], np.ndarray]
) -> Callable[[np.ndarray], np.ndarray]:
def _rmatmul(x: np.ndarray) -> np.ndarray:
if x.ndim == 2 and x.shape[0] == 1:
return rmatvec(x[0, :])[np.newaxis, :]
return np.apply_along_axis(rmatvec, -1, x)
return _rmatmul
[docs] @classmethod
def broadcast_rmatmat(
cls, rmatmat: Callable[[np.ndarray], np.ndarray]
) -> Callable[[np.ndarray], np.ndarray]:
def _rmatmul(x: np.ndarray) -> np.ndarray:
if x.ndim == 2:
return rmatmat(x)
return _apply_to_matrix_stack(rmatmat, x)
return _rmatmul
@property
def _inexact_dtype(self) -> np.dtype:
if np.issubdtype(self.dtype, np.inexact):
return self.dtype
else:
return np.double
def _apply_to_matrix_stack(
mat_fn: Callable[[np.ndarray], np.ndarray], x: np.ndarray
) -> np.ndarray:
idcs = np.ndindex(x.shape[:-2])
# Shape and dtype inference
idx0 = next(idcs)
y0 = mat_fn(x[idx0])
# Result buffer
y = np.empty(x.shape[:-2] + y0.shape, dtype=y0.dtype)
# Fill buffer
y[idx0] = y0
for idx in idcs:
y[idx] = mat_fn(x[idx])
return y
class TransposedLinearOperator(LinearOperator):
"""Transposition of a linear operator."""
def __init__(
self,
linop: LinearOperator,
matmul: Optional[Callable[[np.ndarray], np.ndarray]] = None,
):
self._linop = linop
if matmul is None:
matmul = lambda x: self.todense(cache=True) @ x
super().__init__(
shape=(self._linop.shape[1], self._linop.shape[0]),
dtype=self._linop.dtype,
matmul=matmul,
rmatmul=lambda x: self._linop(x, axis=-1),
todense=lambda: self._linop.todense(cache=False).T.copy(order="C"),
transpose=lambda: self._linop,
inverse=None, # lambda: self._linop.inv().T,
rank=self._linop.rank,
det=self._linop.det,
logabsdet=self._linop.logabsdet,
trace=self._linop.trace,
)
def _astype(
self, dtype: np.dtype, order: str, casting: str, copy: bool
) -> "LinearOperator":
return self._linop.astype(dtype, order=order, casting=casting, copy=copy).T
def __repr__(self) -> str:
return f"Transpose of {self._linop}"
class _InverseLinearOperator(LinearOperator):
def __init__(self, linop: LinearOperator):
if not linop.is_square:
raise np.linalg.LinAlgError("Only square operators can be inverted.")
self._linop = linop
self.__factorization = None
tmatmul = LinearOperator.broadcast_matmat(self._tmatmat)
super().__init__(
shape=self._linop.shape,
dtype=self._linop._inexact_dtype,
matmul=LinearOperator.broadcast_matmat(self._matmat),
rmatmul=lambda x: tmatmul(x[..., np.newaxis])[..., 0],
transpose=lambda: TransposedLinearOperator(self, matmul=tmatmul),
inverse=lambda: self._linop,
det=lambda: 1 / self._linop.det(),
logabsdet=lambda: -self._linop.logabsdet(),
)
def __repr__(self) -> str:
return f"Inverse of {self._linop}"
@property
def factorization(self):
if self.__factorization is None:
self.__factorization = scipy.linalg.lu_factor(
self._linop.todense(cache=False), overwrite_a=False
)
return self.__factorization
def _matmat(self, x: np.ndarray) -> np.ndarray:
return scipy.linalg.lu_solve(self.factorization, x, trans=0, overwrite_b=False)
def _tmatmat(self, x: np.ndarray) -> np.ndarray:
return scipy.linalg.lu_solve(self.factorization, x, trans=1, overwrite_b=False)
class _TypeCastLinearOperator(LinearOperator):
def __init__(
self,
linop: LinearOperator,
dtype: DTypeLike,
order: str = "K",
casting: str = "unsafe",
copy: bool = True,
):
self._linop = linop
dtype = np.dtype(dtype)
if not np.can_cast(self._linop.dtype, dtype, casting=casting):
raise TypeError(
f"Cannot cast linear operator from {self._linop.dtype} to {dtype} "
f"according to the rule {casting}"
)
super().__init__(
self._linop.shape,
dtype,
matmul=lambda x: (self._linop @ x).astype(
np.result_type(self.dtype, x.dtype), copy=False
),
rmatmul=lambda x: (x @ self._linop).astype(
np.result_type(x.dtype, self.dtype), copy=False
),
apply=lambda x, axis: self._linop(x, axis=axis).astype(
np.result_type(self.dtype, x.dtype), copy=False
),
todense=lambda: self._linop.todense(cache=False).astype(
dtype, order=order, copy=copy
),
transpose=lambda: self._linop.T.astype(dtype),
inverse=lambda: self._linop.inv().astype(self._inexact_dtype),
rank=self._linop.rank,
eigvals=lambda: self._linop.eigvals().astype(self._inexact_dtype),
cond=lambda p: self._linop.cond(p=p).astype(self._inexact_dtype),
det=lambda: self._linop.det().astype(self._inexact_dtype),
logabsdet=lambda: self._linop.logabsdet().astype(self._inexact_dtype),
trace=lambda: self._linop.trace().astype(dtype),
)
def _astype(
self, dtype: np.dtype, order: str, casting: str, copy: bool
) -> "LinearOperator":
if self.dtype == dtype and not copy:
return self
elif dtype == self._linop.dtype and not copy:
return self._linop
else:
return _TypeCastLinearOperator(self, dtype, order, casting, copy)
class Matrix(LinearOperator):
"""A linear operator defined via a matrix.
Parameters
----------
A : array-like or scipy.sparse.spmatrix
The explicit matrix.
"""
def __init__(
self,
A: Union[np.ndarray, scipy.sparse.spmatrix],
):
if isinstance(A, scipy.sparse.spmatrix):
self.A = A
shape = self.A.shape
dtype = self.A.dtype
matmul = LinearOperator.broadcast_matmat(lambda x: self.A @ x)
rmatmul = LinearOperator.broadcast_rmatmat(lambda x: x @ self.A)
todense = self.A.toarray
inverse = self._sparse_inv
trace = lambda: self.A.diagonal().sum()
else:
self.A = np.asarray(A)
shape = self.A.shape
dtype = self.A.dtype
matmul = lambda x: A @ x
rmatmul = lambda x: x @ A
todense = lambda: self.A
inverse = None
trace = lambda: np.trace(self.A)
transpose = lambda: Matrix(self.A.T)
super().__init__(
shape,
dtype,
matmul=matmul,
rmatmul=rmatmul,
todense=todense,
transpose=transpose,
inverse=inverse,
trace=trace,
)
def _astype(
self, dtype: np.dtype, order: str, casting: str, copy: bool
) -> "LinearOperator":
if isinstance(self.A, np.ndarray):
A_astype = self.A.astype(dtype, order=order, casting=casting, copy=copy)
else:
assert isinstance(self.A, scipy.sparse.spmatrix)
A_astype = self.A.astype(dtype, casting=casting, copy=copy)
if A_astype is self.A:
return self
return Matrix(A_astype)
def _sparse_inv(self) -> "Matrix":
try:
return Matrix(scipy.sparse.linalg.inv(self.A.tocsc()))
except RuntimeError as err:
raise np.linalg.LinAlgError(str(err)) from err
def _matmul_matrix(self, other: "Matrix") -> "Matrix":
if not config.lazy_matrix_matrix_matmul:
if not self.shape[1] == other.shape[0]:
raise ValueError(f"Matmul shape mismatch {self.shape} x {other.shape}")
return Matrix(A=self.A @ other.A)
return NotImplemented
def __eq__(self, other: LinearOperator) -> bool:
if not self._is_type_shape_dtype_equal(other):
return False
return np.all(self.A == other.A)
def __neg__(self) -> "Matrix":
return Matrix(-self.A)
class Identity(LinearOperator):
"""The identity operator.
Parameters
----------
shape :
Shape of the identity operator.
"""
def __init__(
self,
shape: ShapeLike,
dtype: DTypeLike = np.double,
):
shape = probnum.utils.as_shape(shape)
if len(shape) == 1:
shape = 2 * shape
elif len(shape) != 2:
raise ValueError("The shape of a linear operator must be two-dimensional.")
if shape[0] != shape[1]:
raise np.linalg.LinAlgError("An identity operator must be square.")
super().__init__(
shape,
dtype,
matmul=lambda x: x.astype(np.result_type(self.dtype, x.dtype), copy=False),
rmatmul=lambda x: x.astype(np.result_type(self.dtype, x.dtype), copy=False),
apply=lambda x, axis: x.astype(
np.result_type(self.dtype, x.dtype), copy=False
),
todense=lambda: np.identity(self.shape[0], dtype=dtype),
transpose=lambda: self,
inverse=lambda: self,
rank=lambda: np.intp(shape[0]),
eigvals=lambda: np.ones(shape[0], dtype=self._inexact_dtype),
cond=self._cond,
det=lambda: probnum.utils.as_numpy_scalar(1.0, dtype=self._inexact_dtype),
logabsdet=lambda: probnum.utils.as_numpy_scalar(
0.0, dtype=self._inexact_dtype
),
trace=lambda: probnum.utils.as_numpy_scalar(
self.shape[0], dtype=self.dtype
),
)
def _cond(self, p: Union[None, int, float, str]) -> np.inexact:
if p is None or p in (2, 1, np.inf, -2, -1, -np.inf):
return probnum.utils.as_numpy_scalar(1.0, dtype=self._inexact_dtype)
elif p == "fro":
return probnum.utils.as_numpy_scalar(
self.shape[0], dtype=self._inexact_dtype
)
else:
return np.linalg.cond(self.todense(cache=False), p=p)
def _astype(
self, dtype: np.dtype, order: str, casting: str, copy: bool
) -> "Identity":
if dtype == self.dtype and not copy:
return self
else:
return Identity(self.shape, dtype=dtype)
def __eq__(self, other: LinearOperator) -> bool:
return self._is_type_shape_dtype_equal(other)
class Selection(LinearOperator):
def __init__(self, indices, shape, dtype=np.double):
if np.ndim(indices) > 1:
raise ValueError(
"Selection LinOp expects an integer or (1D) iterable of "
f"integers. Received {type(indices)} with shape {np.shape(indices)}."
)
if shape[0] > shape[1]:
raise ValueError(
f"Invalid shape {shape} for Selection LinOp. If the "
"output-dimension (shape[0]) is larger than the input-dimension "
"(shape[1]), consider using `Embedding`."
)
self._indices = probnum.utils.as_shape(indices)
assert len(self._indices) == shape[0]
super().__init__(
dtype=dtype,
shape=shape,
matmul=lambda x: _selection_matmul(self.indices, x),
todense=self._todense,
transpose=lambda: Embedding(
take_indices=np.arange(len(self._indices)),
put_indices=self._indices,
shape=(self.shape[1], self.shape[0]),
),
)
@property
def indices(self):
return self._indices
def _todense(self):
res = np.eye(self.shape[1], self.shape[1])
return _selection_matmul(self.indices, res)
def _selection_matmul(indices, M):
return np.take(M, indices=indices, axis=-2)
class Embedding(LinearOperator):
def __init__(
self, take_indices, put_indices, shape, fill_value=0.0, dtype=np.double
):
if np.ndim(take_indices) > 1:
raise ValueError(
"Embedding LinOp expects an integer or (1D) iterable of "
f"integers. Received {type(take_indices)} with shape "
f"{np.shape(take_indices)}."
)
if np.ndim(put_indices) > 1:
raise ValueError(
"Embedding LinOp expects an integer or (1D) iterable of "
f"integers. Received {type(put_indices)} with shape "
f"{np.shape(put_indices)}."
)
if shape[0] < shape[1]:
raise ValueError(
f"Invalid shape {shape} for Embedding LinOp. If the "
"output-dimension (shape[0]) is smaller than the input-dimension "
"(shape[1]), consider using `Selection`."
)
self._take_indices = probnum.utils.as_shape(take_indices)
self._put_indices = probnum.utils.as_shape(put_indices)
self._fill_value = fill_value
super().__init__(
dtype=dtype,
shape=shape,
matmul=lambda x: _embedding_matmul(self, x),
todense=self._todense,
transpose=lambda: Selection(
indices=put_indices, shape=(self.shape[1], self.shape[0])
),
)
def _todense(self):
return self.T.todense().T
def _embedding_matmul(embedding, M):
res_shape = np.array(M.shape)
res_shape[-2] = embedding.shape[0]
res = np.full(shape=tuple(res_shape), fill_value=embedding._fill_value)
take_from_M = M[..., np.array(embedding._take_indices), :]
res[..., np.array(embedding._put_indices), :] = take_from_M
return res