Fix normal arithmetic and cholesky updates

Author: marvinpfoertner

src/probnum/randvars/_arithmetic.py

@@ -124,7 +124,7 @@ def _rv_binary_op(rv1: _RandomVariable, rv2: _RandomVariable) -> _RandomVariable
 
 def _make_rv_binary_op_result_shape_dtype_sample_fn(op_fn, rv1, rv2):
     def sample(seed, sample_shape):
-        seed1, seed2, _ = backend.random.split(seed, 3)
+        seed1, seed2 = backend.random.split(seed, 2)
 
         return op_fn(
             rv1.sample(seed=seed1, sample_shape=sample_shape),
@@ -294,9 +294,17 @@ def _matmul_normal_constant(norm_rv: _Normal, constant_rv: _Constant) -> _Normal
         mean = norm_rv.mean @ constant_rv.support
         cov = constant_rv.support.T @ (norm_rv.cov @ constant_rv.support)
 
-        if cov.shape == () and mean.shape == (1,):
+        if mean.shape == ():
+            cov = cov.reshape(())
+
+            if cov_cholesky is not None:
+                cov_cholesky = cov_cholesky.reshape(())
+        elif mean.shape == (1,):
             cov = cov.reshape((1, 1))
 
+            if cov_cholesky is not None:
+                cov_cholesky = cov_cholesky.reshape((1, 1))
+
         return _Normal(mean=mean, cov=cov, cov_cholesky=cov_cholesky)
 
     # This part does not do the Cholesky update,
@@ -335,11 +343,22 @@ def _matmul_constant_normal(constant_rv: _Constant, norm_rv: _Normal) -> _Normal
             )
         else:
             cov_cholesky = None
-        return _Normal(
-            mean=constant_rv.support @ norm_rv.mean,
-            cov=constant_rv.support @ (norm_rv.cov @ constant_rv.support.T),
-            cov_cholesky=cov_cholesky,
-        )
+
+        mean = constant_rv.support @ norm_rv.mean
+        cov = constant_rv.support @ (norm_rv.cov @ constant_rv.support.T)
+
+        if mean.shape == ():
+            cov = cov.reshape(())
+
+            if cov_cholesky is not None:
+                cov_cholesky = cov_cholesky.reshape(())
+        elif mean.shape == (1,):
+            cov = cov.reshape((1, 1))
+
+            if cov_cholesky is not None:
+                cov_cholesky = cov_cholesky.reshape((1, 1))
+
+        return _Normal(mean=mean, cov=cov, cov_cholesky=cov_cholesky)
 
     # This part does not do the Cholesky update,
     # because of performance configurations: currently, there is no way of switching

src/probnum/randvars/_normal.py

@@ -120,9 +120,6 @@ def __init__(
 
         if mean.ndim == 0:
             # Scalar Gaussian
-            if self._cov_cholesky is None:
-                self._cov_cholesky = backend.sqrt(cov)
-
             self.__cov_op_cholesky = None
 
             super().__init__(
@@ -255,6 +252,8 @@ def compute_cov_cholesky(
                     lower=True,
                 ),
             )
+        elif self.ndim == 0:
+            self._cov_cholesky = backend.sqrt(self.cov)
         else:
             self.__cov_op_cholesky = linops.aslinop(
                 backend.to_numpy(
@@ -275,10 +274,7 @@ def cov_cholesky_is_precomputed(self) -> bool:
         initialization or if (ii) the property `self.cov_cholesky` has
         been called before.
         """
-        if self.__cov_op_cholesky is None:
-            return False
-
-        return True
+        return self._cov_cholesky is not None or self.__cov_op_cholesky is not None
 
     def __getitem__(self, key: ArrayLikeGetitemArgType) -> "Normal":
         """Marginalization in multi- and matrixvariate normal random variables,

tests/test_randvars/test_arithmetic/conftest.py

@@ -2,58 +2,67 @@
 import numpy as np
 import pytest
 
-from probnum import linops, randvars
+from probnum import backend, linops, randvars
 from probnum.problems.zoo.linalg import random_spd_matrix
 from probnum.typing import ShapeArgType
+from tests.testing import seed_from_args
 
 
 @pytest.fixture
-def rng() -> np.random.Generator:
-    return np.random.default_rng(42)
+def constant(shape_const: ShapeArgType) -> randvars.Constant:
+    seed = seed_from_args(shape_const, 19836)
 
-
-@pytest.fixture
-def constant(shape_const: ShapeArgType, rng: np.random.Generator) -> randvars.Constant:
-    return randvars.Constant(support=rng.normal(size=shape_const))
+    return randvars.Constant(
+        support=backend.random.standard_normal(seed, shape=shape_const)
+    )
 
 
 @pytest.fixture
 def multivariate_normal(
-    shape: ShapeArgType, precompute_cov_cholesky: bool, rng: np.random.Generator
+    shape: ShapeArgType, precompute_cov_cholesky: bool
 ) -> randvars.Normal:
+    seed = seed_from_args(shape, precompute_cov_cholesky, 1908)
+    seed_mean, seed_cov = backend.random.split(seed)
+
     rv = randvars.Normal(
-        mean=rng.normal(size=shape),
-        cov=random_spd_matrix(rng=rng, dim=shape[0]),
+        mean=backend.random.standard_normal(seed_mean, shape=shape),
+        cov=random_spd_matrix(seed_cov, dim=shape[0]),
     )
     if precompute_cov_cholesky:
-        rv.precompute_cov_cholesky()
+        rv.compute_cov_cholesky()
     return rv
 
 
 @pytest.fixture
 def matrixvariate_normal(
-    shape: ShapeArgType, precompute_cov_cholesky: bool, rng: np.random.Generator
+    shape: ShapeArgType, precompute_cov_cholesky: bool
 ) -> randvars.Normal:
+    seed = seed_from_args(shape, precompute_cov_cholesky, 354)
+    seed_mean, seed_cov_A, seed_cov_B = backend.random.split(seed, num=3)
+
     rv = randvars.Normal(
-        mean=rng.normal(size=shape),
+        mean=backend.random.standard_normal(seed_mean, shape=shape),
         cov=linops.Kronecker(
-            A=random_spd_matrix(dim=shape[0], rng=rng),
-            B=random_spd_matrix(dim=shape[1], rng=rng),
+            A=random_spd_matrix(seed_cov_A, dim=shape[0]),
+            B=random_spd_matrix(seed_cov_B, dim=shape[1]),
         ),
     )
     if precompute_cov_cholesky:
-        rv.precompute_cov_cholesky()
+        rv.compute_cov_cholesky()
     return rv
 
 
 @pytest.fixture
 def symmetric_matrixvariate_normal(
-    shape: ShapeArgType, precompute_cov_cholesky: bool, rng: np.random.Generator
+    shape: ShapeArgType, precompute_cov_cholesky: bool
 ) -> randvars.Normal:
+    seed = seed_from_args(shape, precompute_cov_cholesky, 246)
+    seed_mean, seed_cov = backend.random.split(seed)
+
     rv = randvars.Normal(
-        mean=random_spd_matrix(dim=shape[0], rng=rng),
-        cov=linops.SymmetricKronecker(A=random_spd_matrix(dim=shape[0], rng=rng)),
+        mean=random_spd_matrix(seed_mean, dim=shape[0]),
+        cov=linops.SymmetricKronecker(A=random_spd_matrix(seed_cov, dim=shape[0])),
     )
     if precompute_cov_cholesky:
-        rv.precompute_cov_cholesky()
+        rv.compute_cov_cholesky()
     return rv

tests/test_randvars/test_arithmetic/test_generic.py

@@ -11,7 +11,9 @@
 @pytest.mark.parametrize("shape,dtype", [((5,), np.single), ((2, 3), np.double)])
 def test_generic_randvar_dtype_shape_inference(shape: ShapeArgType, dtype: DTypeLike):
     x = randvars.RandomVariable(
-        shape=shape, dtype=dtype, sample=lambda size, rng: np.zeros(size + shape)
+        shape=shape,
+        dtype=dtype,
+        sample=lambda seed, sample_shape: np.zeros(sample_shape + shape),
     )
     y = np.array(5.0)
     z = x + y

tests/test_utils/test_linalg/test_cholesky_updates.py

@@ -2,6 +2,7 @@
 import pytest
 
 import probnum.utils.linalg as utlin
+from probnum import backend
 from probnum.problems.zoo.linalg import random_spd_matrix
 
 
@@ -13,20 +14,28 @@ def even_ndim():
 
 
 @pytest.fixture
-def rng():
-    return np.random.default_rng(seed=123)
+def spdmats(even_ndim):
+    seed = backend.random.seed(abs(hash(even_ndim)))
+    seed1, seed2 = backend.random.split(seed, num=2)
+
+    spdmat1 = random_spd_matrix(seed1, dim=even_ndim)
+    spdmat2 = random_spd_matrix(seed2, dim=even_ndim)
+
+    return spdmat1, spdmat2
 
 
 @pytest.fixture
-def spdmat1(even_ndim, rng):
-    return random_spd_matrix(rng, dim=even_ndim)
+def spdmat1(spdmats):
+    return spdmats[0]
 
 
 @pytest.fixture
-def spdmat2(even_ndim, rng):
-    return random_spd_matrix(rng, dim=even_ndim)
+def spdmat2(spdmats):
+    return spdmats[1]
 
 
+@pytest.mark.skipif_backend(backend.Backend.JAX)
+@pytest.mark.skipif_backend(backend.Backend.TORCH)
 def test_cholesky_update(spdmat1, spdmat2):
     expected = np.linalg.cholesky(spdmat1 + spdmat2)
 
@@ -36,6 +45,8 @@ def test_cholesky_update(spdmat1, spdmat2):
     np.testing.assert_allclose(expected, received)
 
 
+@pytest.mark.skipif_backend(backend.Backend.JAX)
+@pytest.mark.skipif_backend(backend.Backend.TORCH)
 def test_cholesky_optional(spdmat1, even_ndim):
     """Assert that cholesky_update() transforms a non-square matrix square-root into a
     correct Cholesky factor."""
@@ -46,6 +57,8 @@ def test_cholesky_optional(spdmat1, even_ndim):
     np.testing.assert_allclose(expected, received)
 
 
+@pytest.mark.skipif_backend(backend.Backend.JAX)
+@pytest.mark.skipif_backend(backend.Backend.TORCH)
 def test_tril_to_positive_tril():
 
     # Make a random tril matrix

tests/testing/__init__.py

@@ -1,2 +1,3 @@
 from .assertions import *
+from .random import seed_from_args
 from .statistics import *

tests/testing/random.py

@@ -0,0 +1,8 @@
+from collections.abc import Hashable
+
+from probnum import backend
+from probnum.typing import SeedType
+
+
+def seed_from_args(*args: Hashable) -> SeedType:
+    return backend.random.seed(abs(sum(map(hash, args))))