Commit ⋅ probabilistic-numerics/probnum

from abc import ABC, abstractmethod
from typing import Generic, Tuple, TypeVar

from ._stopping_criterion import StoppingCriterion

ProblemType = TypeVar("ProblemType")
BeliefType = TypeVar("BeliefType")
StateType = TypeVar("StateType")

@@ -17,8 +19,8 @@


    Parameters
    ----------
    prior :
        Prior knowledge about quantities of interest for the numerical problem.
    stopping_criterion
        Stopping criterion determining when a desired terminal condition is met.

    References
    ----------

@@ -28,6 +30,11 @@

    .. [2] Cockayne, J., Oates, C., Sullivan Tim J. and Girolami M., Bayesian
       probabilistic numerical methods. *SIAM Review*, 61(4):756–789, 2019

    See Also
    --------
    ~probnum.linalg.solvers.ProbabilisticLinearSolver : Compose a custom
        probabilistic linear solver.

    Notes
    -----
    All PN methods should subclass this base class. Typically convenience functions

@@ -35,18 +42,19 @@

    derived subclass.
    """

    def __init__(self, prior: BeliefType):
        self.prior = prior
    def __init__(self, stopping_criterion: StoppingCriterion):
        self.stopping_criterion = stopping_criterion

    @abstractmethod
    def solve(
        self,
        problem: ProblemType,
        self, prior: BeliefType, problem: ProblemType, **kwargs
    ) -> Tuple[BeliefType, StateType]:
        """Solve the given numerical problem.

        Parameters
        ----------
        prior :
            Prior knowledge about quantities of interest of the numerical problem.
        problem :
            Numerical problem to be solved.
        """

@@ -1,10 +1,10 @@

"""Linear Algebra.

This package implements probabilistic numerical methods for the solution
of problems arising in linear algebra, such as the solution of linear
systems :math:`Ax=b`.
This package implements probabilistic numerical methods for the solution of problems
arising in linear algebra, such as the solution of linear systems :math:`Ax=b`.
"""
from probnum.linalg._problinsolve import bayescg, problinsolve
from probnum.linalg._bayescg import bayescg
from probnum.linalg._problinsolve import problinsolve

from . import solvers


@@ -3,8 +3,25 @@

Iterative probabilistic numerical methods solving linear systems :math:`Ax = b`.
"""

from typing import Generator, Optional, Tuple

class ProbabilisticLinearSolver:
import numpy as np

from probnum import ProbabilisticNumericalMethod, problems
from probnum.linalg.solvers import (
    belief_updates,
    beliefs,
    information_ops,
    policies,
    stopping_criteria,
)

from ._state import LinearSolverState


class ProbabilisticLinearSolver(
    ProbabilisticNumericalMethod[problems.LinearSystem, beliefs.LinearSystemBelief]
):
    r"""Compose a custom probabilistic linear solver.

    Class implementing probabilistic linear solvers. Such (iterative) solvers infer

@@ -20,7 +37,15 @@


    Parameters
    ----------

    policy
        Policy returning actions taken by the solver.
    information_op
        Information operator defining how information about the linear system is
        obtained given an action.
    belief_update
        Belief update defining how to update the QoI beliefs given new observations.
    stopping_criterion
        Stopping criterion determining when a desired terminal condition is met.

    References
    ----------

@@ -35,9 +60,283 @@


    See Also
    --------
    problinsolve : Solve linear systems in a Bayesian framework.
    bayescg : Solve linear systems with prior information on the solution.
    ~probnum.linalg.problinsolve : Solve linear systems in a Bayesian framework.
    ~probnum.linalg.bayescg : Solve linear systems with prior information on the solution.

    Examples
    --------
    Define a linear system.

    >>> import numpy as np
    >>> from probnum.problems import LinearSystem
    >>> from probnum.problems.zoo.linalg import random_spd_matrix

    >>> rng = np.random.default_rng(42)
    >>> n = 100
    >>> A = random_spd_matrix(rng=rng, dim=n)
    >>> b = rng.standard_normal(size=(n,))
    >>> linsys = LinearSystem(A=A, b=b)

    Create a custom probabilistic linear solver from pre-defined components.

    >>> from probnum.linalg.solvers import (
    ...     ProbabilisticLinearSolver,
    ...     belief_updates,
    ...     beliefs,
    ...     information_ops,
    ...     policies,
    ...     stopping_criteria,
    ... )

    >>> pls = ProbabilisticLinearSolver(
    ...     policy=policies.ConjugateGradientPolicy(),
    ...     information_op=information_ops.ProjectedRHSInformationOp(),
    ...     belief_update=belief_updates.solution_based.SolutionBasedProjectedRHSBeliefUpdate(),
    ...     stopping_criterion=(
    ...         stopping_criteria.MaxIterationsStoppingCriterion(100)
    ...         | stopping_criteria.ResidualNormStoppingCriterion(atol=1e-5, rtol=1e-5)
    ...     ),
    ... )

    Define a prior over the solution.

    >>> from probnum import linops, randvars
    >>> prior = beliefs.LinearSystemBelief(
    ...     x=randvars.Normal(
    ...         mean=np.zeros((n,)),
    ...         cov=np.eye(n),
    ...     ),
    ... )

    Solve the linear system using the custom solver.

    >>> belief, solver_state = pls.solve(prior=prior, problem=linsys)
    >>> np.linalg.norm(linsys.A @ belief.x.mean - linsys.b) / np.linalg.norm(linsys.b)
    7.1886e-06
    """

    def __init__(
        self,
        policy: policies.LinearSolverPolicy,
        information_op: information_ops.LinearSolverInformationOp,
        belief_update: belief_updates.LinearSolverBeliefUpdate,
        stopping_criterion: stopping_criteria.LinearSolverStoppingCriterion,
    ):
        self.policy = policy
        self.information_op = information_op
        self.belief_update = belief_update
        super().__init__(stopping_criterion=stopping_criterion)

    def solve_iterator(
        self,
        prior: beliefs.LinearSystemBelief,
        problem: problems.LinearSystem,
        rng: Optional[np.random.Generator] = None,
    ) -> Generator[LinearSolverState, None, None]:
        """Generator implementing the solver iteration.

        This function allows stepping through the solver iteration one step at a time
        and exposes the internal solver state.

        Parameters
        ----------
        prior
            Prior belief about the quantities of interest :math:`(x, A, A^{-1}, b)` of the linear system.
        problem
            Linear system to be solved.
        rng
            Random number generator.

        Yields
        ------
        solver_state
            State of the probabilistic linear solver.
        """
        solver_state = LinearSolverState(problem=problem, prior=prior, rng=rng)

        while True:

            yield solver_state

            # Check stopping criterion
            if self.stopping_criterion(solver_state=solver_state):
                break

            # Compute action via policy
            solver_state.action = self.policy(solver_state=solver_state)

            # Make observation via information operator
            solver_state.observation = self.information_op(solver_state=solver_state)

            # Update belief about the quantities of interest
            solver_state.belief = self.belief_update(solver_state=solver_state)

            # Advance state to next step and invalidate caches
            solver_state.next_step()

    def solve(
        self,
        prior: beliefs.LinearSystemBelief,
        problem: problems.LinearSystem,
        rng: Optional[np.random.Generator] = None,
    ) -> Tuple[beliefs.LinearSystemBelief, LinearSolverState]:
        r"""Solve the linear system.

        Parameters
        ----------
        prior
            Prior belief about the quantities of interest :math:`(x, A, A^{-1}, b)` of the linear system.
        problem
            Linear system to be solved.
        rng
            Random number generator.

        Returns
        -------
        belief
            Posterior belief :math:`(\mathsf{x}, \mathsf{A}, \mathsf{H}, \mathsf{b})`
            over the solution :math:`x`, the system matrix :math:`A`, its (pseudo-)inverse :math:`H=A^\dagger` and the right hand side :math:`b`.
        solver_state
            Final state of the solver.
        """
        solver_state = None

        for solver_state in self.solve_iterator(prior=prior, problem=problem, rng=rng):
            pass

        return solver_state.belief, solver_state


class BayesCG(ProbabilisticLinearSolver):
    r"""Bayesian conjugate gradient method.

    Probabilistic linear solver taking prior information about the solution and
    choosing :math:`A`-conjugate actions to gain information about the solution
    by projecting the current residual.

    This code implements the method described in Cockayne et al. [1]_.

    Parameters
    ----------
    stopping_criterion
        Stopping criterion determining when a desired terminal condition is met.

    References
    ----------
    .. [1] Cockayne, J. et al., A Bayesian Conjugate Gradient Method, *Bayesian
       Analysis*, 2019
    """

    def __init__(
        self,
        stopping_criterion: stopping_criteria.LinearSolverStoppingCriterion = stopping_criteria.MaxIterationsStoppingCriterion()
        | stopping_criteria.ResidualNormStoppingCriterion(atol=1e-5, rtol=1e-5),
    ):
        super().__init__(
            policy=policies.ConjugateGradientPolicy(),
            information_op=information_ops.ProjectedRHSInformationOp(),
            belief_update=belief_updates.solution_based.SolutionBasedProjectedRHSBeliefUpdate(),
            stopping_criterion=stopping_criterion,
        )


class ProbabilisticKaczmarz(ProbabilisticLinearSolver):
    r"""Probabilistic Kaczmarz method.

    Probabilistic analogue of the (randomized) Kaczmarz method [1]_ [2]_, taking prior
    information about the solution and randomly choosing rows of the matrix :math:`A_i`
    and entries :math:`b_i` of the right-hand-side to obtain information about the solution.

    Parameters
    ----------
    stopping_criterion
        Stopping criterion determining when a desired terminal condition is met.

    References
    ----------
    .. [1] Kaczmarz, Stefan, Angenäherte Auflösung von Systemen linearer Gleichungen,
        *Bulletin International de l'Académie Polonaise des Sciences et des Lettres. Classe des Sciences Mathématiques et Naturelles. Série A, Sciences Mathématiques*, 1937
    .. [2] Strohmer, Thomas; Vershynin, Roman, A randomized Kaczmarz algorithm for
        linear systems with exponential convergence, *Journal of Fourier Analysis and Applications*, 2009
    """

    def __init__(
        self,
        stopping_criterion: stopping_criteria.LinearSolverStoppingCriterion = stopping_criteria.MaxIterationsStoppingCriterion()
        | stopping_criteria.ResidualNormStoppingCriterion(atol=1e-5, rtol=1e-5),
    ):
        super().__init__(
            policy=policies.RandomUnitVectorPolicy(),
            information_op=information_ops.ProjectedRHSInformationOp(),
            belief_update=belief_updates.solution_based.SolutionBasedProjectedRHSBeliefUpdate(),
            stopping_criterion=stopping_criterion,
        )


class MatrixBasedPLS(ProbabilisticLinearSolver):
    r"""Matrix-based probabilistic linear solver.

    Probabilistic linear solver updating beliefs over the system matrix and its
    inverse. The solver makes use of prior information and iteratively infers the matrix and its inverse by matrix-vector multiplication.

    This code implements the method described in Wenger et al. [1]_.

    Parameters
    ----------
    policy
        Policy returning actions taken by the solver.
    stopping_criterion
        Stopping criterion determining when a desired terminal condition is met.

    References
    ----------
    .. [1] Wenger, J. and Hennig, P., Probabilistic Linear Solvers for Machine Learning,
       *Advances in Neural Information Processing Systems (NeurIPS)*, 2020
    """

    def __init__(
        self,
        policy: policies.LinearSolverPolicy = policies.ConjugateGradientPolicy(),
        stopping_criterion: stopping_criteria.LinearSolverStoppingCriterion = stopping_criteria.MaxIterationsStoppingCriterion()
        | stopping_criteria.ResidualNormStoppingCriterion(atol=1e-5, rtol=1e-5),
    ):
        super().__init__(
            policy=policy,
            information_op=information_ops.MatVecInformationOp(),
            belief_update=belief_updates.matrix_based.MatrixBasedLinearBeliefUpdate(),
            stopping_criterion=stopping_criterion,
        )


class SymMatrixBasedPLS(ProbabilisticLinearSolver):
    r"""Symmetric matrix-based probabilistic linear solver.

    Probabilistic linear solver updating beliefs over the symmetric system matrix and its inverse. The solver makes use of prior information and iteratively infers the matrix and its inverse by matrix-vector multiplication.

    This code implements the method described in Wenger et al. [1]_.

    Parameters
    ----------
    policy
        Policy returning actions taken by the solver.
    stopping_criterion
        Stopping criterion determining when a desired terminal condition is met.

    References
    ----------
    .. [1] Wenger, J. and Hennig, P., Probabilistic Linear Solvers for Machine Learning,
       *Advances in Neural Information Processing Systems (NeurIPS)*, 2020
    """

    def __init__(
        self,
        policy: policies.LinearSolverPolicy = policies.ConjugateGradientPolicy(),
        stopping_criterion: stopping_criteria.LinearSolverStoppingCriterion = stopping_criteria.MaxIterationsStoppingCriterion()
        | stopping_criteria.ResidualNormStoppingCriterion(atol=1e-5, rtol=1e-5),
    ):
        super().__init__(
            policy=policy,
            information_op=information_ops.MatVecInformationOp(),
            belief_update=belief_updates.matrix_based.SymmetricMatrixBasedLinearBeliefUpdate(),
            stopping_criterion=stopping_criterion,
        )

@@ -1,5 +1,7 @@

"""Belief update in a solution-based inference view where the information is given by
projecting the current residual to a subspace."""
import numpy as np

import probnum  # pylint: disable="unused-import"
from probnum import randvars
from probnum.linalg.solvers.beliefs import LinearSystemBelief

@@ -49,13 +51,16 @@

        gram = action_A @ cov_xy + self._noise_var
        gram_pinv = 1.0 / gram if gram > 0.0 else 0.0
        gain = cov_xy * gram_pinv
        cov_update = gain @ cov_xy.T
        cov_update = np.outer(gain, cov_xy)

        x = randvars.Normal(
            mean=solver_state.belief.x.mean + gain * proj_resid,
            cov=solver_state.belief.x.cov - cov_update,
        )
        Ainv = solver_state.belief.Ainv + cov_update
        if solver_state.belief.Ainv is None:
            Ainv = randvars.Constant(cov_update)
        else:
            Ainv = solver_state.belief.Ainv + cov_update

        return LinearSystemBelief(
            x=x, A=solver_state.belief.A, Ainv=Ainv, b=solver_state.belief.b

@@ -1,8 +1,7 @@

"""Linear system belief.

Class defining a belief about the quantities of interest of a linear
system such as its solution or the matrix inverse and any associated
hyperparameters.
Class defining a belief about the quantities of interest of a linear system such as its
solution or the matrix inverse and any associated hyperparameters.
"""

from typing import Mapping, Optional

@@ -23,9 +22,12 @@


    Random variables :math:`(\mathsf{x}, \mathsf{A}, \mathsf{H}, \mathsf{b})`
    modelling the solution :math:`x`, the system matrix :math:`A`, its (pseudo-)inverse
    :math:`H=A^{-1}` and the right hand side :math:`b` of a linear system :math:`Ax=b`, as well as any associated hyperparameters.
    :math:`H=A^{\dagger}` and the right hand side :math:`b` of a linear system :math:`Ax=b`,
    as well as any associated hyperparameters.

    For instantiation either a belief about the solution or the inverse and right hand side must be provided. Note that if both are specified, their consistency is not checked and depending on the algorithm either may be used.
    For instantiation either a belief about the solution or the inverse and right hand side
    must be provided. Note that if both are specified, their consistency is not checked and
    depending on the algorithm either may be used.

    Parameters
    ----------

@@ -5,22 +5,30 @@

linear systems.
"""

from probnum.linalg.solvers.matrixbased import (
    MatrixBasedSolver,
    SymmetricMatrixBasedSolver,
)

from . import belief_updates, beliefs, information_ops, policies, stopping_criteria
from ._probabilistic_linear_solver import ProbabilisticLinearSolver
from ._probabilistic_linear_solver import (
    BayesCG,
    MatrixBasedPLS,
    ProbabilisticKaczmarz,
    ProbabilisticLinearSolver,
    SymMatrixBasedPLS,
)
from ._state import LinearSolverState

# Public classes and functions. Order is reflected in documentation.
__all__ = [
    "ProbabilisticLinearSolver",
    "MatrixBasedSolver",
    "SymmetricMatrixBasedSolver",
    "LinearSolverState",
    "BayesCG",
    "ProbabilisticKaczmarz",
    "MatrixBasedPLS",
    "SymMatrixBasedPLS",
]

# Set correct module paths. Corrects links and module paths in documentation.
ProbabilisticLinearSolver.__module__ = "probnum.linalg.solvers"
LinearSolverState.__module__ = "probnum.linalg.solvers"
BayesCG.__module__ = "probnum.linalg.solvers"
ProbabilisticKaczmarz.__module__ = "probnum.linalg.solvers"
MatrixBasedPLS.__module__ = "probnum.linalg.solvers"
SymMatrixBasedPLS.__module__ = "probnum.linalg.solvers"

@@ -28,7 +28,12 @@


        Parameters
        ----------
        solver_state :
        solver_state
            Current state of the linear solver.

        Returns
        -------
        action
            Next action to take.
        """
        raise NotImplementedError

@@ -0,0 +1,61 @@

1	+	"""A Bayesian conjugate gradient method."""
2	+
3	+	from typing import Callable, Optional
4	+
5	+
6	+	def bayescg(
7	+	A,
8	+	b,
9	+	x0=None,
10	+	maxiter: Optional[int] = None,
11	+	atol: float = 1e-5,
12	+	rtol: float = 1e-5,
13	+	callback: Optional[Callable] = None,
14	+	):
15	+	r"""Bayesian Conjugate Gradient Method.
16	+
17	+	In the setting where :math:`A` is a symmetric positive-definite matrix, this solver
18	+	takes prior information on the solution and outputs a posterior belief over
19	+	:math:`x`. This code implements the method described in Cockayne et al. [1]_.
20	+
21	+	Note that the solution-based view of BayesCG and the matrix-based view of
22	+	:meth:`problinsolve` correspond [2]_.
23	+
24	+	Parameters
25	+	----------
26	+	A
27	+	shape=(n, n) -- A symmetric positive definite matrix (or linear operator). Only matrix-vector products :math:`Av` are used internally.
28	+	b
29	+	shape=(n, ) -- Right-hand side vector.
30	+	x0
31	+	shape=(n, ) -- Prior belief for the solution of the linear system.
32	+	maxiter
33	+	Maximum number of iterations. Defaults to :math:`10n`, where :math:`n` is the
34	+	dimension of :math:`A`.
35	+	atol
36	+	Absolute residual tolerance. If :math:`\lVert r_i \rVert = \lVert Ax_i - b
37	+	\rVert < \text{atol}`, the iteration terminates.
38	+	rtol
39	+	Relative residual tolerance. If :math:`\lVert r_i \rVert < \text{rtol}
40	+	\lVert b \rVert`, the iteration terminates.
41	+	callback
42	+	User-supplied function called after each iteration of the linear solver. It is
43	+	called as ``callback(xk, sk, yk, alphak, resid, **kwargs)`` and can be used to
44	+	return quantities from the iteration. Note that depending on the function
45	+	supplied, this can slow down the solver.
46	+
47	+	Returns
48	+	-------
49	+
50	+	References
51	+	----------
52	+	.. [1] Cockayne, J. et al., A Bayesian Conjugate Gradient Method, *Bayesian
53	+	Analysis*, 2019, 14, 937-1012
54	+	.. [2] Bartels, S. et al., Probabilistic Linear Solvers: A Unifying View,
55	+	Statistics and Computing, 2019
56	+
57	+	See Also
58	+	--------
59	+	problinsolve : Solve linear systems in a Bayesian framework.
60	+	"""
61	+	raise NotImplementedError

@@ -46,7 +46,13 @@

            action=solver_state.observation,
            observ=solver_state.action,
        )
        return LinearSystemBelief(A=A, Ainv=Ainv, x=None, b=solver_state.belief.b)

        if solver_state.belief.b is None:
            b = randvars.Constant(solver_state.problem.b)
        else:
            b = solver_state.belief.b

        return LinearSystemBelief(A=A, Ainv=Ainv, x=None, b=b)

    def _matrix_based_update(
        self, matrix: randvars.Normal, action: np.ndarray, observ: np.ndarray

@@ -1,9 +1,9 @@

"""Probabilistic numerical methods for solving linear systems.

This module provides routines to solve linear systems of equations in a
Bayesian framework. This means that a prior distribution over elements
of the linear system can be provided and is updated with information
collected by the solvers to return a posterior distribution.
This module provides routines to solve linear systems of equations in a Bayesian
framework. This means that a prior distribution over elements of the linear system can
be provided and is updated with information collected by the solvers to return a
posterior distribution.
"""

import warnings

@@ -230,54 +230,6 @@

    return x, A0, Ainv0, info


def bayescg(A, b, x0=None, maxiter=None, atol=None, rtol=None, callback=None):
    """Conjugate Gradients using prior information on the solution of the linear system.

    In the setting where :math:`A` is a symmetric positive-definite matrix, this solver
    takes prior information on the solution and outputs a posterior belief over
    :math:`x`. This code implements the method described in Cockayne et al. [1]_.

    Note that the solution-based view of BayesCG and the matrix-based view of
    :meth:`problinsolve` correspond [2]_.

    Parameters
    ----------
    A : array-like or LinearOperator, shape=(n,n)
        A square linear operator (or matrix). Only matrix-vector products :math:`Av` are
        used internally.
    b : array_like, shape=(n,) or (n, nrhs)
        Right-hand side vector or matrix in :math:`A x = b`.
    x0 : array-like or RandomVariable, shape=(n,) or or (n, nrhs)
        Prior belief over the solution of the linear system.
    maxiter : int
        Maximum number of iterations. Defaults to :math:`10n`, where :math:`n` is the
        dimension of :math:`A`.
    atol : float, optional
        Absolute residual tolerance. If :math:`\\lVert r_i \\rVert = \\lVert Ax_i - b
        \\rVert < \\text{atol}`, the iteration terminates.
    rtol : float, optional
        Relative residual tolerance. If :math:`\\lVert r_i \\rVert  < \\text{rtol}
        \\lVert b \\rVert`, the iteration terminates.
    callback : function, optional
        User-supplied function called after each iteration of the linear solver. It is
        called as ``callback(xk, sk, yk, alphak, resid, **kwargs)`` and can be used to
        return quantities from the iteration. Note that depending on the function
        supplied, this can slow down the solver.

    References
    ----------
    .. [1] Cockayne, J. et al., A Bayesian Conjugate Gradient Method, *Bayesian
       Analysis*, 2019, 14, 937-1012
    .. [2] Bartels, S. et al., Probabilistic Linear Solvers: A Unifying View,
       *Statistics and Computing*, 2019

    See Also
    --------
    problinsolve : Solve linear systems in a Bayesian framework.
    """
    raise NotImplementedError


def _check_linear_system(A, b, A0=None, Ainv0=None, x0=None):
    """Check linear system compatibility.


@@ -66,8 +66,8 @@

    See Also
    --------
    LambdaStoppingCriterion : Stopping criterion defined via an anonymous function.
    LinearSolverStoppingCriterion : Stopping criterion of a probabilistic linear solver.
    FiltSmoothStoppingCriterion : Stopping criterion of filters and smoothers.
    ~probnum.linalg.solvers.stopping_criteria.LinearSolverStoppingCriterion : Stopping criterion of a probabilistic linear solver.
    ~probnum.filtsmooth.optim.FiltSmoothStoppingCriterion : Stopping criterion of filters and smoothers.
    """

    @abc.abstractmethod

@@ -47,7 +47,13 @@

            action=solver_state.observation,
            observ=solver_state.action,
        )
        return LinearSystemBelief(A=A, Ainv=Ainv, x=None, b=solver_state.belief.b)

        if solver_state.belief.b is None:
            b = randvars.Constant(solver_state.problem.b)
        else:
            b = solver_state.belief.b

        return LinearSystemBelief(A=A, Ainv=Ainv, x=None, b=b)

    def _symmetric_matrix_based_update(
        self, matrix: randvars.Normal, action: np.ndarray, observ: np.ndarray

@@ -64,7 +70,9 @@

        gram = action.T @ covfactor_Ms
        gram_pinv = 1.0 / gram if gram > 0.0 else 0.0
        gain = covfactor_Ms * gram_pinv
        covfactor_update = gain @ covfactor_Ms.T
        covfactor_update = linops.aslinop(gain[:, None]) @ linops.aslinop(
            covfactor_Ms[None, :]
        )
        resid_gain = linops.aslinop(resid[:, None]) @ linops.aslinop(gain[None, :])

        return randvars.Normal(

@@ -31,10 +31,12 @@

        prior: "probnum.linalg.solvers.beliefs.LinearSystemBelief",
        rng: Optional[np.random.Generator] = None,
    ):
        self.problem = problem
        self.belief = prior
        self.problem: problems.LinearSystem = problem

        self.step = 0
        self.prior: "probnum.linalg.solvers.beliefs.LinearSystemBelief" = prior
        self._belief: "probnum.linalg.solvers.beliefs.LinearSystemBelief" = prior

        self._step: int = 0

        self._actions: List[np.ndarray] = [None]
        self._observations: List[Any] = [None]

@@ -42,11 +44,27 @@

            self.problem.A @ self.belief.x.mean - self.problem.b,
            None,
        ]
        self.rng = rng
        self.rng: Optional[np.random.Generator] = rng

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(step={self.step})"

    @property
    def step(self) -> int:
        """Current step of the solver."""
        return self._step

    @property
    def belief(self) -> "probnum.linalg.solvers.beliefs.LinearSystemBelief":
        """Belief over the quantities of interest of the linear system."""
        return self._belief

    @belief.setter
    def belief(
        self, belief: "probnum.linalg.solvers.beliefs.LinearSystemBelief"
    ) -> None:
        self._belief = belief

    @property
    def action(self) -> Optional[np.ndarray]:
        """Action of the solver for the current step.

@@ -101,10 +119,10 @@

    def next_step(self) -> None:
        """Advance the solver state to the next solver step.

        Called after a completed step / iteration of the probabilistic linear solver.
        Called after a completed step / iteration of the linear solver.
        """
        self._actions.append(None)
        self._observations.append(None)
        self._residuals.append(None)

        self.step += 1
        self._step += 1

Files		•	•	•	Coverage
src/probnum	6,668	5,960	229	479	89.38%
Project Totals (180 files)	6,668	5,960	229	479	89.38%

View CI Build Download

coverage:
  precision: 2
  status:
    project:
      default:
        target: auto
        threshold: 1%
    patch:
      default:
        target: 90%
        threshold: 1%

comment:
  layout: "reach, diff, files"
  behavior: default
  require_changes: true

Read the docs

Interact now

The inner-most circle is the entire project, moving away from the center are folders then, finally, a single file. The size and color of each slice is representing the number of statements and the coverage, respectively.

SVG

The top section represents the entire project. Proceeding with folders and finally individual files. The size and color of each slice is representing the number of statements and the coverage, respectively.

SVG

Each block represents a single file in the project. The size and color of each block is represented by the number of statements and the coverage, respectively.

SVG

Learn more about the coverage sunburst graph

230	230		return x, A0, Ainv0, info
231	231
232	232
233		-	def bayescg(A, b, x0=None, maxiter=None, atol=None, rtol=None, callback=None):
234		-	"""Conjugate Gradients using prior information on the solution of the linear system.
235		-
236		-	In the setting where :math:`A` is a symmetric positive-definite matrix, this solver
237		-	takes prior information on the solution and outputs a posterior belief over
238		-	:math:`x`. This code implements the method described in Cockayne et al. [1]_.
239		-
240		-	Note that the solution-based view of BayesCG and the matrix-based view of
241		-	:meth:`problinsolve` correspond [2]_.
242		-
243		-	Parameters
244		-	----------
245		-	A : array-like or LinearOperator, shape=(n,n)
246		-	A square linear operator (or matrix). Only matrix-vector products :math:`Av` are
247		-	used internally.
248		-	b : array_like, shape=(n,) or (n, nrhs)
249		-	Right-hand side vector or matrix in :math:`A x = b`.
250		-	x0 : array-like or RandomVariable, shape=(n,) or or (n, nrhs)
251		-	Prior belief over the solution of the linear system.
252		-	maxiter : int
253		-	Maximum number of iterations. Defaults to :math:`10n`, where :math:`n` is the
254		-	dimension of :math:`A`.
255		-	atol : float, optional
256		-	Absolute residual tolerance. If :math:`\\lVert r_i \\rVert = \\lVert Ax_i - b
257		-	\\rVert < \\text{atol}`, the iteration terminates.
258		-	rtol : float, optional
259		-	Relative residual tolerance. If :math:`\\lVert r_i \\rVert < \\text{rtol}
260		-	\\lVert b \\rVert`, the iteration terminates.
261		-	callback : function, optional
262		-	User-supplied function called after each iteration of the linear solver. It is
263		-	called as ``callback(xk, sk, yk, alphak, resid, **kwargs)`` and can be used to
264		-	return quantities from the iteration. Note that depending on the function
265		-	supplied, this can slow down the solver.
266		-
267		-	References
268		-	----------
269		-	.. [1] Cockayne, J. et al., A Bayesian Conjugate Gradient Method, *Bayesian
270		-	Analysis*, 2019, 14, 937-1012
271		-	.. [2] Bartels, S. et al., Probabilistic Linear Solvers: A Unifying View,
272		-	Statistics and Computing, 2019
273		-
274		-	See Also
275		-	--------
276		-	problinsolve : Solve linear systems in a Bayesian framework.
277		-	"""
278		-	raise NotImplementedError
279		-
280		-
281	233		def _check_linear_system(A, b, A0=None, Ainv0=None, x0=None):
282	234		"""Check linear system compatibility.
283	235

101	119		def next_step(self) -> None:
102	120		"""Advance the solver state to the next solver step.
103	121
104		-	Called after a completed step / iteration of the probabilistic linear solver.
	122	+	Called after a completed step / iteration of the linear solver.
105	123		"""
106	124		self._actions.append(None)
107	125		self._observations.append(None)
108	126		self._residuals.append(None)
109	127
110		-	self.step += 1
	128	+	self._step += 1

Navigation	Overlay
`t` Navigate files	`h` Toggle hits
`y` Change url to tip of branch	`m` Toggle misses
`b / v` Jump to prev/next hit line	`p` Toggle partial
`z / x` Jump to prev/next missed or partial line	`1..9` Toggle flags
`shift + o` Open current page in GitHub	`a` Toggle all on
`/ or ?` Show keyboard shortcuts dialog	`c` Toggle context lines or commits

3	3		from abc import ABC, abstractmethod
4	4		from typing import Generic, Tuple, TypeVar
5	5
	6	+	from ._stopping_criterion import StoppingCriterion
	7	+
6	8		ProblemType = TypeVar("ProblemType")
7	9		BeliefType = TypeVar("BeliefType")
8	10		StateType = TypeVar("StateType")

17	19
18	20		Parameters
19	21		----------
20		-	prior :
21		-	Prior knowledge about quantities of interest for the numerical problem.
	22	+	stopping_criterion
	23	+	Stopping criterion determining when a desired terminal condition is met.
22	24
23	25		References
24	26		----------

28	30		.. [2] Cockayne, J., Oates, C., Sullivan Tim J. and Girolami M., Bayesian
29	31		probabilistic numerical methods. SIAM Review, 61(4):756–789, 2019
30	32
	33	+	See Also
	34	+	--------
	35	+	~probnum.linalg.solvers.ProbabilisticLinearSolver : Compose a custom
	36	+	probabilistic linear solver.
	37	+
31	38		Notes
32	39		-----
33	40		All PN methods should subclass this base class. Typically convenience functions

35	42		derived subclass.
36	43		"""
37	44
38		-	def __init__(self, prior: BeliefType):
39		-	self.prior = prior
	45	+	def __init__(self, stopping_criterion: StoppingCriterion):
	46	+	self.stopping_criterion = stopping_criterion
40	47
41	48		@abstractmethod
42	49		def solve(
43		-	self,
44		-	problem: ProblemType,
	50	+	self, prior: BeliefType, problem: ProblemType, **kwargs
45	51		) -> Tuple[BeliefType, StateType]:
46	52		"""Solve the given numerical problem.
47	53
48	54		Parameters
49	55		----------
	56	+	prior :
	57	+	Prior knowledge about quantities of interest of the numerical problem.
50	58		problem :
51	59		Numerical problem to be solved.
52	60		"""

1	1		"""Linear Algebra.
2	2
3		-	This package implements probabilistic numerical methods for the solution
4		-	of problems arising in linear algebra, such as the solution of linear
5		-	systems :math:`Ax=b`.
	3	+	This package implements probabilistic numerical methods for the solution of problems
	4	+	arising in linear algebra, such as the solution of linear systems :math:`Ax=b`.
6	5		"""
7		-	from probnum.linalg._problinsolve import bayescg, problinsolve
	6	+	from probnum.linalg._bayescg import bayescg
	7	+	from probnum.linalg._problinsolve import problinsolve
8	8
9	9		from . import solvers
10	10

3	3		Iterative probabilistic numerical methods solving linear systems :math:`Ax = b`.
4	4		"""
5	5
	6	+	from typing import Generator, Optional, Tuple
6	7
7		-	class ProbabilisticLinearSolver:
	8	+	import numpy as np
	9	+
	10	+	from probnum import ProbabilisticNumericalMethod, problems
	11	+	from probnum.linalg.solvers import (
	12	+	belief_updates,
	13	+	beliefs,
	14	+	information_ops,
	15	+	policies,
	16	+	stopping_criteria,
	17	+	)
	18	+
	19	+	from ._state import LinearSolverState
	20	+
	21	+
	22	+	class ProbabilisticLinearSolver(
	23	+	ProbabilisticNumericalMethod[problems.LinearSystem, beliefs.LinearSystemBelief]
	24	+	):
8	25		r"""Compose a custom probabilistic linear solver.
9	26
10	27		Class implementing probabilistic linear solvers. Such (iterative) solvers infer

20	37
21	38		Parameters
22	39		----------
23		-
	40	+	policy
	41	+	Policy returning actions taken by the solver.
	42	+	information_op
	43	+	Information operator defining how information about the linear system is
	44	+	obtained given an action.
	45	+	belief_update
	46	+	Belief update defining how to update the QoI beliefs given new observations.
	47	+	stopping_criterion
	48	+	Stopping criterion determining when a desired terminal condition is met.
24	49
25	50		References
26	51		----------

1	1		"""Belief update in a solution-based inference view where the information is given by
2	2		projecting the current residual to a subspace."""
	3	+	import numpy as np
	4	+
3	5		import probnum # pylint: disable="unused-import"
4	6		from probnum import randvars
5	7		from probnum.linalg.solvers.beliefs import LinearSystemBelief

49	51		gram = action_A @ cov_xy + self._noise_var
50	52		gram_pinv = 1.0 / gram if gram > 0.0 else 0.0
51	53		gain = cov_xy * gram_pinv
52		-	cov_update = gain @ cov_xy.T
	54	+	cov_update = np.outer(gain, cov_xy)
53	55
54	56		x = randvars.Normal(
55	57		mean=solver_state.belief.x.mean + gain * proj_resid,
56	58		cov=solver_state.belief.x.cov - cov_update,
57	59		)
58		-	Ainv = solver_state.belief.Ainv + cov_update
	60	+	if solver_state.belief.Ainv is None:
	61	+	Ainv = randvars.Constant(cov_update)
	62	+	else:
	63	+	Ainv = solver_state.belief.Ainv + cov_update
59	64
60	65		return LinearSystemBelief(
61	66		x=x, A=solver_state.belief.A, Ainv=Ainv, b=solver_state.belief.b

1	1		"""Linear system belief.
2	2
3		-	Class defining a belief about the quantities of interest of a linear
4		-	system such as its solution or the matrix inverse and any associated
5		-	hyperparameters.
	3	+	Class defining a belief about the quantities of interest of a linear system such as its
	4	+	solution or the matrix inverse and any associated hyperparameters.
6	5		"""
7	6
8	7		from typing import Mapping, Optional

23	22
24	23		Random variables :math:`(\mathsf{x}, \mathsf{A}, \mathsf{H}, \mathsf{b})`
25	24		modelling the solution :math:`x`, the system matrix :math:`A`, its (pseudo-)inverse
26		-	:math:`H=A^{-1}` and the right hand side :math:`b` of a linear system :math:`Ax=b`, as well as any associated hyperparameters.
	25	+	:math:`H=A^{\dagger}` and the right hand side :math:`b` of a linear system :math:`Ax=b`,
	26	+	as well as any associated hyperparameters.
27	27
28		-	For instantiation either a belief about the solution or the inverse and right hand side must be provided. Note that if both are specified, their consistency is not checked and depending on the algorithm either may be used.
	28	+	For instantiation either a belief about the solution or the inverse and right hand side
	29	+	must be provided. Note that if both are specified, their consistency is not checked and
	30	+	depending on the algorithm either may be used.
29	31
30	32		Parameters
31	33		----------

5	5		linear systems.
6	6		"""
7	7
8		-	from probnum.linalg.solvers.matrixbased import (
9		-	MatrixBasedSolver,
10		-	SymmetricMatrixBasedSolver,
11		-	)
12		-
13	8		from . import belief_updates, beliefs, information_ops, policies, stopping_criteria
14		-	from ._probabilistic_linear_solver import ProbabilisticLinearSolver
	9	+	from ._probabilistic_linear_solver import (
	10	+	BayesCG,
	11	+	MatrixBasedPLS,
	12	+	ProbabilisticKaczmarz,
	13	+	ProbabilisticLinearSolver,
	14	+	SymMatrixBasedPLS,
	15	+	)
15	16		from ._state import LinearSolverState
16	17
17	18		# Public classes and functions. Order is reflected in documentation.
18	19		__all__ = [
19	20		"ProbabilisticLinearSolver",
20		-	"MatrixBasedSolver",
21		-	"SymmetricMatrixBasedSolver",
22	21		"LinearSolverState",
	22	+	"BayesCG",
	23	+	"ProbabilisticKaczmarz",
	24	+	"MatrixBasedPLS",
	25	+	"SymMatrixBasedPLS",
23	26		]
24	27
25	28		# Set correct module paths. Corrects links and module paths in documentation.
26	29		ProbabilisticLinearSolver.__module__ = "probnum.linalg.solvers"
	30	+	LinearSolverState.__module__ = "probnum.linalg.solvers"
	31	+	BayesCG.__module__ = "probnum.linalg.solvers"
	32	+	ProbabilisticKaczmarz.__module__ = "probnum.linalg.solvers"
	33	+	MatrixBasedPLS.__module__ = "probnum.linalg.solvers"
	34	+	SymMatrixBasedPLS.__module__ = "probnum.linalg.solvers"

28	28
29	29		Parameters
30	30		----------
31		-	solver_state :
	31	+	solver_state
32	32		Current state of the linear solver.
	33	+
	34	+	Returns
	35	+	-------
	36	+	action
	37	+	Next action to take.
33	38		"""
34	39		raise NotImplementedError

46	46		action=solver_state.observation,
47	47		observ=solver_state.action,
48	48		)
49		-	return LinearSystemBelief(A=A, Ainv=Ainv, x=None, b=solver_state.belief.b)
	49	+
	50	+	if solver_state.belief.b is None:
	51	+	b = randvars.Constant(solver_state.problem.b)
	52	+	else:
	53	+	b = solver_state.belief.b
	54	+
	55	+	return LinearSystemBelief(A=A, Ainv=Ainv, x=None, b=b)
50	56
51	57		def _matrix_based_update(
52	58		self, matrix: randvars.Normal, action: np.ndarray, observ: np.ndarray

1	1		"""Probabilistic numerical methods for solving linear systems.
2	2
3		-	This module provides routines to solve linear systems of equations in a
4		-	Bayesian framework. This means that a prior distribution over elements
5		-	of the linear system can be provided and is updated with information
6		-	collected by the solvers to return a posterior distribution.
	3	+	This module provides routines to solve linear systems of equations in a Bayesian
	4	+	framework. This means that a prior distribution over elements of the linear system can
	5	+	be provided and is updated with information collected by the solvers to return a
	6	+	posterior distribution.
7	7		"""
8	8
9	9		import warnings

66	66		See Also
67	67		--------
68	68		LambdaStoppingCriterion : Stopping criterion defined via an anonymous function.
69		-	LinearSolverStoppingCriterion : Stopping criterion of a probabilistic linear solver.
70		-	FiltSmoothStoppingCriterion : Stopping criterion of filters and smoothers.
	69	+	~probnum.linalg.solvers.stopping_criteria.LinearSolverStoppingCriterion : Stopping criterion of a probabilistic linear solver.
	70	+	~probnum.filtsmooth.optim.FiltSmoothStoppingCriterion : Stopping criterion of filters and smoothers.
71	71		"""
72	72
73	73		@abc.abstractmethod

64	70		gram = action.T @ covfactor_Ms
65	71		gram_pinv = 1.0 / gram if gram > 0.0 else 0.0
66	72		gain = covfactor_Ms * gram_pinv
67		-	covfactor_update = gain @ covfactor_Ms.T
	73	+	covfactor_update = linops.aslinop(gain[:, None]) @ linops.aslinop(
	74	+	covfactor_Ms[None, :]
	75	+	)
68	76		resid_gain = linops.aslinop(resid[:, None]) @ linops.aslinop(gain[None, :])
69	77
70	78		return randvars.Normal(

31	31		prior: "probnum.linalg.solvers.beliefs.LinearSystemBelief",
32	32		rng: Optional[np.random.Generator] = None,
33	33		):
34		-	self.problem = problem
35		-	self.belief = prior
	34	+	self.problem: problems.LinearSystem = problem
36	35
37		-	self.step = 0
	36	+	self.prior: "probnum.linalg.solvers.beliefs.LinearSystemBelief" = prior
	37	+	self._belief: "probnum.linalg.solvers.beliefs.LinearSystemBelief" = prior
	38	+
	39	+	self._step: int = 0
38	40
39	41		self._actions: List[np.ndarray] = [None]
40	42		self._observations: List[Any] = [None]

42	44		self.problem.A @ self.belief.x.mean - self.problem.b,
43	45		None,
44	46		]
45		-	self.rng = rng
	47	+	self.rng: Optional[np.random.Generator] = rng
46	48
47	49		def __repr__(self) -> str:
48	50		return f"{self.__class__.__name__}(step={self.step})"
49	51
	52	+	@property
	53	+	def step(self) -> int:
	54	+	"""Current step of the solver."""
	55	+	return self._step
	56	+
	57	+	@property
	58	+	def belief(self) -> "probnum.linalg.solvers.beliefs.LinearSystemBelief":
	59	+	"""Belief over the quantities of interest of the linear system."""
	60	+	return self._belief
	61	+
	62	+	@belief.setter
	63	+	def belief(
	64	+	self, belief: "probnum.linalg.solvers.beliefs.LinearSystemBelief"
	65	+	) -> None:
	66	+	self._belief = belief
	67	+
50	68		@property
51	69		def action(self) -> Optional[np.ndarray]:
52	70		"""Action of the solver for the current step.

1	coverage:
2	precision: 2
3	status:
4	project:
5	default:
6	target: auto
7	threshold: 1%
8	patch:
9	default:
10	target: 90%
11	threshold: 1%
12
13	comment:
14	layout: "reach, diff, files"
15	behavior: default
16	require_changes: true

Updating our web app