Reorthogonalization for a stable implementation of the probabilistic linear solver

src/probnum/linalg/solvers/policies/_conjugate_gradient.py

@@ -45,7 +45,7 @@ def __call__(
         self, solver_state: "probnum.linalg.solvers.LinearSolverState"
     ) -> np.ndarray:
 
-        action = -solver_state.residual.copy()
+        residual = solver_state.residual.copy()
 
         if self._reorthogonalization_fn_residual is not None and solver_state.step == 0:
             solver_state.cache["reorthogonalized_residuals"] = [solver_state.residual]
@@ -62,15 +62,17 @@ def __call__(
 
             # A-conjugacy correction (in exact arithmetic)
             beta = (np.linalg.norm(residual) / np.linalg.norm(prev_residual)) ** 2
-            action += beta * solver_state.actions[solver_state.step - 1]
+            action = -residual + beta * solver_state.actions[solver_state.step - 1]
 
             # Reorthogonalization of the resulting action
             if self._reorthogonalization_fn_action is not None:
                 return self._reorthogonalized_action(
                     action=action, solver_state=solver_state
                 )
 
-        return action
+            return action
+
+        return -residual
 
     def _reorthogonalized_residuals(
         self,
@@ -82,7 +84,7 @@ def _reorthogonalized_residuals(
             orthogonal_basis=np.asarray(
                 solver_state.cache["reorthogonalized_residuals"]
             ),
-            inprod=None,
+            inner_product=None,
         )
         solver_state.cache["reorthogonalized_residuals"].append(residual)
         prev_residual = solver_state.cache["reorthogonalized_residuals"][
@@ -110,5 +112,5 @@ def _reorthogonalized_action(
         return self._reorthogonalization_fn_action(
             v=action,
             orthogonal_basis=orthogonal_basis,
-            inprod=inprod_matrix,
+            inner_product=inprod_matrix,
         )

src/probnum/utils/linalg/__init__.py

@@ -1,12 +1,12 @@
 """Utility functions that involve numerical linear algebra."""
 
 from ._cholesky_updates import cholesky_update, tril_to_positive_tril
-from ._inner_product import euclidean_inprod, euclidean_norm
+from ._inner_product import induced_norm, inner_product
 from ._orthogonalize import double_gram_schmidt, gram_schmidt, modified_gram_schmidt
 
 __all__ = [
-    "euclidean_inprod",
-    "euclidean_norm",
+    "inner_product",
+    "induced_norm",
     "cholesky_update",
     "tril_to_positive_tril",
     "gram_schmidt",

src/probnum/utils/linalg/_inner_product.py

@@ -1,65 +1,76 @@
 """Functions defining useful inner products."""
+from __future__ import annotations
 
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
 
-from probnum import linops
+if TYPE_CHECKING:
+    from probnum import linops
 
 
-def euclidean_inprod(
+def inner_product(
     v: np.ndarray,
     w: np.ndarray,
     A: Optional[Union[np.ndarray, linops.LinearOperator]] = None,
 ) -> np.ndarray:
-    r"""(Modified) Euclidean inner product :math:`\langle v, w \rangle_A := v^T A w`.
+    r"""Inner product :math:`\langle v, w \rangle_A := v^T A w`.
+
+    For arrays the function computes the inner product over the last axes of the
+    two arrays ``v`` and ``w``.
 
     Parameters
     ----------
     v
-        First vector.
+        First array.
     w
-        Second vector.
+        Second array.
     A
         Symmetric positive (semi-)definite matrix defining the geometry.
 
     Returns
     -------
-    inprod
-        Inner product.
+    inprod :
+        *shape=(\*v.shape[:-1], \*w.shape[:-1])* -- Inner product of ``v`` and ``w``. If they are both 1-D arrays then a scalar is returned; otherwise an array is returned.
     """
-
-    v_T = v[..., None, :]
     w = w[..., :, None]
 
     if A is None:
-        vw_inprod = v_T @ w
+        vw_inprod = np.dot(v, w)
     else:
-        vw_inprod = v_T @ (A @ w)
+        vw_inprod = np.dot(v, A @ w)
 
-    return np.squeeze(vw_inprod, axis=(-2, -1))
+    return np.squeeze(vw_inprod, axis=(-1))
 
 
-def euclidean_norm(
+def induced_norm(
     v: np.ndarray,
     A: Optional[Union[np.ndarray, linops.LinearOperator]] = None,
+    axis: int = -1,
 ) -> np.ndarray:
-    r"""(Modified) Euclidean norm :math:`\lVert v \rVert_A := \sqrt{v^T A v}`.
+    r"""Induced norm :math:`\lVert v \rVert_A := \sqrt{v^T A v}`.
+
+    Computes the induced norm over the given axis of the array.
 
     Parameters
     ----------
     v
-        Vector.
+        Array.
     A
-        Symmetric positive (semi-)definite matrix defining the geometry.
+        Symmetric positive (semi-)definite linear operator defining the geometry.
+    axis
+        Specifies the axis along which to compute the vector norms.
 
     Returns
     -------
-    norm
-        Vector norm.
+    norm :
+        Vector norm of ``v`` along the given ``axis``.
     """
 
     if A is None:
-        return np.linalg.norm(v, ord=2, axis=-1, keepdims=False)
+        return np.linalg.norm(v, ord=2, axis=axis, keepdims=False)
+
+    v = np.moveaxis(v, axis, -1)
+    w = np.squeeze(A @ v[..., :, None], axis=-1)
 
-    return np.sqrt(euclidean_inprod(v, v, A))
+    return np.sqrt(np.sum(v * w, axis=-1))

src/probnum/utils/linalg/_orthogonalize.py

@@ -1,18 +1,20 @@
 """Orthogonalization of vectors."""
 
+from functools import partial
 from typing import Callable, Iterable, Optional, Union
 
 import numpy as np
 
 from probnum import linops
 
-from ._inner_product import euclidean_inprod, euclidean_norm
+from ._inner_product import induced_norm
+from ._inner_product import inner_product as inner_product_fn
 
 
 def gram_schmidt(
     v: np.ndarray,
     orthogonal_basis: Iterable[np.ndarray],
-    inprod: Optional[
+    inner_product: Optional[
         Union[
             np.ndarray,
             linops.LinearOperator,
@@ -32,24 +34,27 @@ def gram_schmidt(
         Vector to orthogonalize.
     orthogonal_basis
         Orthogonal basis.
-    inprod
-        Inner product.
+    inner_product
+        Inner product defining orthogonality. Can be either a :class`numpy.ndarray` or a :class:`Callable`
+        defining the inner product. Defaults to the euclidean inner product.
     normalize
         Normalize the output vector, s.t. :math:`\langle v', v' \rangle = 1`.
 
     Returns
     -------
-    v_orth
+    v_orth :
         Orthogonalized vector.
     """
-    if inprod is None:
-        inprod_fn = euclidean_inprod
-        norm_fn = euclidean_norm
-    elif isinstance(inprod, (np.ndarray, linops.LinearOperator)):
-        inprod_fn = lambda v, w: euclidean_inprod(v, w, A=inprod)
-        norm_fn = lambda v: euclidean_norm(v, A=inprod)
+    orthogonal_basis = np.atleast_2d(orthogonal_basis)
+
+    if inner_product is None:
+        inprod_fn = inner_product_fn
+        norm_fn = partial(induced_norm, axis=-1)
+    elif isinstance(inner_product, (np.ndarray, linops.LinearOperator)):
+        inprod_fn = lambda v, w: inner_product_fn(v, w, A=inner_product)
+        norm_fn = lambda v: induced_norm(v, A=inner_product, axis=-1)
     else:
-        inprod_fn = inprod
+        inprod_fn = inner_product
         norm_fn = lambda v: np.sqrt(inprod_fn(v, v))
 
     v_orth = v.copy()
@@ -66,7 +71,7 @@ def gram_schmidt(
 def modified_gram_schmidt(
     v: np.ndarray,
     orthogonal_basis: Iterable[np.ndarray],
-    inprod: Optional[
+    inner_product: Optional[
         Union[
             np.ndarray,
             linops.LinearOperator,
@@ -86,24 +91,27 @@ def modified_gram_schmidt(
         Vector to orthogonalize.
     orthogonal_basis
         Orthogonal basis.
-    inprod
-        Inner product.
+    inner_product
+        Inner product defining orthogonality. Can be either a :class:`numpy.ndarray` or a :class:`Callable`
+        defining the inner product. Defaults to the euclidean inner product.
     normalize
         Normalize the output vector, s.t. :math:`\langle v', v' \rangle = 1`.
 
     Returns
     -------
-    v_orth
+    v_orth :
         Orthogonalized vector.
     """
-    if inprod is None:
-        inprod_fn = euclidean_inprod
-        norm_fn = euclidean_norm
-    elif isinstance(inprod, (np.ndarray, linops.LinearOperator)):
-        inprod_fn = lambda v, w: euclidean_inprod(v, w, A=inprod)
-        norm_fn = lambda v: euclidean_norm(v, A=inprod)
+    orthogonal_basis = np.atleast_2d(orthogonal_basis)
+
+    if inner_product is None:
+        inprod_fn = inner_product_fn
+        norm_fn = induced_norm
+    elif isinstance(inner_product, (np.ndarray, linops.LinearOperator)):
+        inprod_fn = lambda v, w: inner_product_fn(v, w, A=inner_product)
+        norm_fn = lambda v: induced_norm(v, A=inner_product)
     else:
-        inprod_fn = inprod
+        inprod_fn = inner_product
         norm_fn = lambda v: np.sqrt(inprod_fn(v, v))
 
     v_orth = v.copy()
@@ -120,34 +128,40 @@ def modified_gram_schmidt(
 def double_gram_schmidt(
     v: np.ndarray,
     orthogonal_basis: Iterable[np.ndarray],
-    inprod: Optional[
+    inner_product: Optional[
         Union[
             np.ndarray,
             linops.LinearOperator,
             Callable[[np.ndarray, np.ndarray], np.ndarray],
         ]
     ] = None,
     normalize: bool = False,
+    gram_schmidt_fn: Callable = modified_gram_schmidt,
 ) -> np.ndarray:
-    r"""Perform the modified Gram-Schmidt process twice.
+    r"""Perform the (modified) Gram-Schmidt process twice.
 
     Computes a vector :math:`v'` such that :math:`\langle v', b_i \rangle = 0` for
-    all basis vectors :math:`b_i \in B` in the orthogonal basis. This performs the modified Gram-Schmidt orthogonalization process twice, which is generally more stable than just reorthogonalizing once. [1]_ [2]_
+    all basis vectors :math:`b_i \in B` in the orthogonal basis. This performs the
+    (modified) Gram-Schmidt orthogonalization process twice, which is generally more
+    stable than just reorthogonalizing once. [1]_ [2]_
 
     Parameters
     ----------
     v
         Vector to orthogonalize.
     orthogonal_basis
         Orthogonal basis.
-    inprod
-        Inner product.
+    inner_product
+        Inner product defining orthogonality. Can be either a :class:`numpy.ndarray` or a :class:`Callable`
+        defining the inner product. Defaults to the euclidean inner product.
     normalize
         Normalize the output vector, s.t. :math:`\langle v', v' \rangle = 1`.
+    gram_schmidt_fn
+        Gram-Schmidt process to use. One of :meth:`gram_schmidt` or :meth:`modified_gram_schmidt`.
 
     Returns
     -------
-    v_orth
+    v_orth :
         Orthogonalized vector.
 
     References
@@ -157,9 +171,15 @@ def double_gram_schmidt(
     .. [2] L. Giraud, J. Langou, and M. Rozloznik, The loss of orthogonality in the
            Gram-Schmidt orthogonalization process, Comput. Math. Appl., 50 (2005)
     """
-    v_orth = modified_gram_schmidt(
-        v=v, orthogonal_basis=orthogonal_basis, inprod=inprod, normalize=normalize
+    v_orth = gram_schmidt_fn(
+        v=v,
+        orthogonal_basis=orthogonal_basis,
+        inner_product=inner_product,
+        normalize=normalize,
     )
-    return modified_gram_schmidt(
-        v=v_orth, orthogonal_basis=orthogonal_basis, inprod=inprod, normalize=normalize
+    return gram_schmidt_fn(
+        v=v_orth,
+        orthogonal_basis=orthogonal_basis,
+        inner_product=inner_product,
+        normalize=normalize,
     )