All tests run in single precision now

Author: pnkraemer

probdiffeq/backend/numpy.py

@@ -52,8 +52,8 @@ def flip(arr, /, axis=None):
     return jnp.flip(arr, axis=axis)
 
 
-def asarray(x, /):
-    return jnp.asarray(x)
+def asarray(x, /, dtype=None):
+    return jnp.asarray(x, dtype=dtype)
 
 
 def squeeze(arr, /):

probdiffeq/backend/testing.py

@@ -67,9 +67,11 @@ def _allclose(a, b, /, *, atol: float | None, rtol: float | None):
     a = jnp.asarray(1.0 * a)
     b = jnp.asarray(1.0 * b)
 
+    # numpy.allclose uses defaults atol=1e-8 and rtol=1e-5;
+    # we mirror this as atol=sqrt(tol) and rtol slightly larger.
     tol = jnp.sqrt(jnp.finfo(b.dtype).eps)
     if atol is None:
         atol = tol
     if rtol is None:
-        rtol = tol
+        rtol = 10 * tol
     return jnp.allclose(a, b, atol=atol, rtol=rtol)

probdiffeq/impl/_stats.py

@@ -63,16 +63,16 @@ def mahalanobis_norm_relative(self, u, /, rv):
         return np.reshape(np.abs(mahalanobis) / np.sqrt(rv.mean.size), ())
 
     def logpdf(self, u, /, rv):
-        # The cholesky factor is triangular, so we compute a cheap slogdet.
-        diagonal = linalg.diagonal_along_axis(rv.cholesky, axis1=-1, axis2=-2)
+        cholesky = linalg.qr_r(rv.cholesky.T).T
+        diagonal = linalg.diagonal_along_axis(cholesky, axis1=-1, axis2=-2)
         slogdet = np.sum(np.log(np.abs(diagonal)))
 
         dx = u - rv.mean
-        residual_white = linalg.solve_triangular(rv.cholesky.T, dx, trans="T")
-        x1 = linalg.vector_dot(residual_white, residual_white)
-        x2 = 2.0 * slogdet
-        x3 = u.size * np.log(np.pi() * 2)
-        return -0.5 * (x1 + x2 + x3)
+        residual_white = linalg.solve_triangular(cholesky, dx, lower=True, trans=0)
+        sqrnorm = linalg.vector_dot(residual_white, residual_white)
+
+        const = np.log(np.pi() * 2)
+        return -1 / 2 * sqrnorm - u.size / 2 * const - slogdet
 
     def mean(self, rv):
         return rv.mean
@@ -128,12 +128,14 @@ def logpdf(self, u, /, rv):
             u = u[None, :]
 
         def logpdf_scalar(x, r):
+            cholesky = linalg.qr_r(r.cholesky.T).T
+
             dx = x - r.mean
-            w = linalg.solve_triangular(r.cholesky.T, dx, trans="T")
+            w = linalg.solve_triangular(cholesky.T, dx, trans="T")
 
             maha_term = linalg.vector_dot(w, w)
 
-            diagonal = linalg.diagonal_along_axis(r.cholesky, axis1=-1, axis2=-2)
+            diagonal = linalg.diagonal_along_axis(cholesky, axis1=-1, axis2=-2)
             slogdet = np.sum(np.log(np.abs(diagonal)))
             logdet_term = 2.0 * slogdet
             return -0.5 * (logdet_term + maha_term + x.size * np.log(np.pi() * 2))
@@ -195,12 +197,14 @@ def mahalanobis_norm_relative(self, u, /, rv):
 
     def logpdf(self, u, /, rv):
         def logpdf_scalar(x, r):
+            cholesky = linalg.qr_r(r.cholesky.T).T
+
             dx = x - r.mean
-            w = linalg.solve_triangular(r.cholesky.T, dx, trans="T")
+            w = linalg.solve_triangular(cholesky.T, dx, trans="T")
 
             maha_term = linalg.vector_dot(w, w)
 
-            diagonal = linalg.diagonal_along_axis(r.cholesky, axis1=-1, axis2=-2)
+            diagonal = linalg.diagonal_along_axis(cholesky, axis1=-1, axis2=-2)
             slogdet = np.sum(np.log(np.abs(diagonal)))
             logdet_term = 2.0 * slogdet
             return -0.5 * (logdet_term + maha_term + x.size * np.log(np.pi() * 2))

tests/conftest.py

@@ -3,14 +3,15 @@
 Necessary because the implementation has been faulty in the past. Never again.
 """
 
-from probdiffeq.backend import functools, stats, testing
+from probdiffeq.backend import functools, random, stats, testing, tree_util
 from probdiffeq.backend import numpy as np
 from probdiffeq.impl import impl
 
 
 @testing.parametrize("fact", ["dense", "isotropic", "blockdiag"])
 def test_logpdf(fact):
     rv, ssm = random_variable(fact=fact)
+
     u = np.ones_like(ssm.stats.mean(rv))
 
     (mean_dense, cov_dense) = ssm.stats.to_multivariate_normal(rv)
@@ -37,4 +38,9 @@ def random_variable(fact):
     output_scale = np.ones_like(ssm.prototypes.output_scale())
     discretize = ssm.conditional.ibm_transitions(output_scale)
     rv = discretize(0.1, output_scale)
-    return rv.noise, ssm
+
+    key = random.prng_key(seed=1)
+    noise_flat, unravel = tree_util.ravel_pytree(rv.noise)
+    noise_flat = random.normal(key, shape=noise_flat.shape)
+    noise = unravel(noise_flat)
+    return noise, ssm

tests/test_impl/test_logpdfs.py

@@ -104,5 +104,6 @@ def test_fixedpoint_smoother_equivalent_different_grid(solver_setup, solution_sm
     # Compare QOI and marginals
     marginals_allclose_func = functools.partial(testing.marginals_allclose, ssm=ssm)
     marginals_allclose_func = functools.vmap(marginals_allclose_func)
+
     assert testing.allclose(u_fixedpoint, u_interp)
     assert np.all(marginals_allclose_func(marginals_fixedpoint, marginals_interp))

tests/test_ivpsolvers/test_strategy_smoother_vs_fixedpoint.py

@@ -5,18 +5,6 @@
 from probdiffeq.backend import numpy as np
 
 
-def set_environment():
-    """Set the environment (e.g., 64-bit precision).
-
-    The setup used to precompute references should match that of the other tests.
-    """
-    # Test on CPU.
-    config.update("platform_name", "cpu")
-
-    # Double precision
-    config.update("enable_x64", True)
-
-
 def three_body_first(num_derivatives_max=6):
     vf, (u0,), (t0, _) = ode.ivp_three_body_1st()
     vf = functools.partial(vf, t=t0)
@@ -30,8 +18,8 @@ def van_der_pol_second(num_derivatives_max=6):
 
 
 if __name__ == "__main__":
-    # 64-bit precision and the like
-    set_environment()
+    # Double precision
+    config.update("enable_x64", True)
 
     solution1 = three_body_first()
     np.save("./tests/test_taylor/data/three_body_first_solution.npy", solution1)

tests/test_taylor/data/generate_reference_solutions.py

@@ -30,7 +30,7 @@ def fixture_pb_with_solution():
 
 
 @testing.parametrize_with_cases("taylor_fun", cases=".", prefix="case_")
-@testing.parametrize("num", [1, 6])
+@testing.parametrize("num", [1, 4])
 def test_approximation_identical_to_reference(pb_with_solution, taylor_fun, num):
     (f, init), solution = pb_with_solution
 

tests/test_taylor/test_exact_first_order.py

@@ -5,7 +5,7 @@
 from probdiffeq.backend import numpy as np
 
 
-@testing.parametrize("num", [0, 1, 4])
+@testing.parametrize("num", [0, 1, 3])
 @testing.parametrize("fact", ["isotropic", "dense", "blockdiag"])
 def test_initialised_correct_shape_and_values(num, fact):
     vf, (u0,), (t0, _) = ode.ivp_lotka_volterra()

tests/test_taylor/test_inexact_first_order.py

@@ -3,9 +3,7 @@
 These are so crucial and annoying to debug that they need their own test set.
 """
 
-from math import prod
-
-from probdiffeq.backend import functools, linalg, testing, tree_util
+from probdiffeq.backend import functools, linalg, random, testing, tree_util
 from probdiffeq.backend import numpy as np
 from probdiffeq.util import cholesky_util
 
@@ -34,14 +32,16 @@ def cov(x):
     assert testing.allclose(cov(bw_noise), C1)
 
 
-@testing.parametrize("Cshape, HCshape", ([(3, 3), (2, 3)],))
-def test_revert_kernel_noisefree(Cshape, HCshape):
+@testing.parametrize("Cshape, Hshape", ([(3, 3), (2, 3)],))
+def test_revert_kernel_noisefree(Cshape, Hshape):
     C = _some_array(Cshape) + 1.0
-    HC = _some_array(HCshape) + 2.0
+    H = _some_array(Hshape) + 2.0
+    HC = H @ C
 
     S = HC @ HC.T
     K = C @ HC.T @ linalg.inv(S)
-    C1 = C @ C.T - K @ S @ K.T
+
+    C1 = (np.eye(Cshape[0]) - K @ H) @ C @ C.T @ (np.eye(Cshape[0]) - K @ H).T
 
     extra, (bw_noise, g) = cholesky_util.revert_conditional_noisefree(
         R_X_F=HC.T, R_X=C.T
@@ -56,7 +56,8 @@ def cov(x):
 
 
 def _some_array(shape):
-    return np.arange(1.0, 1.0 + prod(shape)).reshape(shape)
+    key = random.prng_key(seed=1)
+    return random.normal(key, shape=shape)
 
 
 def test_sqrt_sum_square_scalar():