Rewrite asserts with exceptions

geometric_kernels/_logging.py

@@ -1,4 +1,4 @@
-""" Setup logging """
+"""Setup logging"""
 
 import logging
 

geometric_kernels/feature_maps/deterministic.py

@@ -43,15 +43,28 @@ def __init__(
         self.num_levels = num_levels
 
         if repeated_eigenvalues_laplacian is None:
-            assert eigenfunctions is None
+            if eigenfunctions is not None:
+                raise ValueError(
+                    "You must either provide both `repeated_eigenvalues_laplacian` and `eigenfunctions` or none of the two."
+                )
             repeated_eigenvalues_laplacian = self.space.get_repeated_eigenvalues(
                 self.num_levels
             )
             eigenfunctions = self.space.get_eigenfunctions(self.num_levels)
         else:
-            assert eigenfunctions is not None
-            assert repeated_eigenvalues_laplacian.shape == (num_levels, 1)
-            assert eigenfunctions.num_levels == num_levels
+            if eigenfunctions is None:
+                raise ValueError(
+                    "You must either provide both `repeated_eigenvalues_laplacian` and `eigenfunctions` or none of the two."
+                )
+            if repeated_eigenvalues_laplacian.shape != (num_levels, 1):
+                raise ValueError(
+                    f"Expected `repeated_eigenvalues_laplacian` to have shape [num_levels={num_levels}, 1] but got {B.shape(repeated_eigenvalues_laplacian)}."
+                )
+            if eigenfunctions.num_levels != num_levels:
+                raise ValueError(
+                    f"`num_levels` must coincide with `num_levels` in the provided `eigenfunctions`,"
+                    f"but `num_levels`={num_levels} and `eigenfunctions.num_levels`={eigenfunctions.num_levels}"
+                )
 
         self._repeated_eigenvalues = repeated_eigenvalues_laplacian
         self._eigenfunctions = eigenfunctions

geometric_kernels/feature_maps/probability_densities.py

@@ -22,7 +22,13 @@
     eigvalsh,
     from_numpy,
 )
-from geometric_kernels.utils.utils import ordered_pairwise_differences
+from geometric_kernels.utils.utils import (
+    _check_1_vector,
+    _check_field_in_params,
+    _check_matrix,
+    _check_rank_1_array,
+    ordered_pairwise_differences,
+)
 
 
 def student_t_sample(
@@ -74,13 +80,17 @@ def student_t_sample(
         samples of type `dtype`, and `key` is the updated random key for `jax`,
         or the similar random state (generator) for any other backend.
     """
-    assert B.shape(df) == (1,), "df must be a 1-vector."
+    _check_1_vector(df, "df")
 
-    n = int(B.length(loc))
+    _check_rank_1_array(loc, "loc")
+    _check_matrix(shape, "shape")
 
-    assert B.shape(loc) == (n,), "loc must be a 1-dim vector"
-    assert B.shape(shape) == (n, n), "shape must be a matrix"
+    n = B.shape(loc)[0]
 
+    if tuple(B.shape(shape)) != (n, n):
+        raise ValueError(
+            f"`Expected `shape` matrix to have shape [{n}, {n}], but got {B.shape(shape)}."
+        )
     shape_sqrt = B.chol(shape)
     dtype = dtype or dtype_double(key)
     key, z = B.randn(key, dtype, *size, n)
@@ -140,10 +150,11 @@ def base_density_sample(
         of samples, and `key` is the updated random key for `jax`, or the
         similar random state (generator) for any other backend.
     """
-    assert "lengthscale" in params
-    assert params["lengthscale"].shape == (1,)
-    assert "nu" in params
-    assert params["nu"].shape == (1,)
+    _check_field_in_params(params, "lengthscale")
+    _check_1_vector(params["lengthscale"], 'params["lengthscale"]')
+
+    _check_field_in_params(params, "nu")
+    _check_1_vector(params["nu"], 'params["nu"]')
 
     nu = params["nu"]
     L = params["lengthscale"]
@@ -232,7 +243,8 @@ def _alphas(n: int) -> B.Numeric:
     .. todo::
         Update proposition numbers when the paper gets published.
     """
-    assert n >= 2
+    if n < 2:
+        raise ValueError("Dimension of the hyperbolic space `n` must be >= 2.")
     x, j = symbols("x, j")
     if (n % 2) == 0:
         m = n // 2
@@ -269,9 +281,10 @@ def _sample_mixture_heat(
     .. todo::
         Update proposition numbers when the paper gets published.
     """
-    assert B.rank(alpha) == 1
+    _check_rank_1_array(alpha, "alpha")
     m = B.shape(alpha)[0] - 1
-    assert m >= 0
+    if m < 0:
+        raise ValueError("The mixture must contain at least 1 component.")
     dtype = B.dtype(lengthscale)
     js = B.range(dtype, 0, m + 1)
 
@@ -332,9 +345,10 @@ def _sample_mixture_matern(
     .. todo::
         Update proposition numbers when the paper gets published.
     """
-    assert B.rank(alpha) == 1
+    _check_rank_1_array(alpha, "alpha")
     m = B.shape(alpha)[0] - 1
-    assert m >= 0
+    if m < 0:
+        raise ValueError("The mixture must contain at least 1 component.")
     dtype = B.dtype(lengthscale)
     js = B.range(dtype, 0, m + 1)
     if shifted_laplacian:
@@ -397,10 +411,11 @@ def hyperbolic_density_sample(
         samples, and `key` is the updated random key for `jax`, or the similar
         random state (generator) for any other backend.
     """
-    assert "lengthscale" in params
-    assert params["lengthscale"].shape == (1,)
-    assert "nu" in params
-    assert params["nu"].shape == (1,)
+    _check_field_in_params(params, "lengthscale")
+    _check_1_vector(params["lengthscale"], 'params["lengthscale"]')
+
+    _check_field_in_params(params, "nu")
+    _check_1_vector(params["nu"], 'params["nu"]')
 
     nu = params["nu"]
     L = params["lengthscale"]
@@ -477,10 +492,11 @@ def spd_density_sample(
         samples, and `key` is the updated random key for `jax`, or the similar
         random state (generator) for any other backend.
     """
-    assert "lengthscale" in params
-    assert params["lengthscale"].shape == (1,)
-    assert "nu" in params
-    assert params["nu"].shape == (1,)
+    _check_field_in_params(params, "nu")
+    _check_1_vector(params["nu"], 'params["nu"]')
+
+    _check_field_in_params(params, "lengthscale")
+    _check_1_vector(params["lengthscale"], 'params["lengthscale"]')
 
     nu = params["nu"]
     L = params["lengthscale"]
@@ -514,7 +530,6 @@ def spd_density_sample(
         diffp = B.pi * B.abs(diffp)
         logprod = B.sum(B.log(B.tanh(diffp)), axis=-1)
         prod = B.exp(logprod)
-        assert B.all(prod > 0)
 
         # accept with probability `prod`
         key, u = B.rand(key, B.dtype(L), 1)

geometric_kernels/kernels/feature_map.py

@@ -11,7 +11,11 @@
 from geometric_kernels.feature_maps import FeatureMap
 from geometric_kernels.kernels.base import BaseGeometricKernel
 from geometric_kernels.spaces.base import Space
-from geometric_kernels.utils.utils import make_deterministic
+from geometric_kernels.utils.utils import (
+    _check_1_vector,
+    _check_field_in_params,
+    make_deterministic,
+)
 
 
 class MaternFeatureMapKernel(BaseGeometricKernel):
@@ -108,10 +112,11 @@ def K(
         X2: Optional[B.Numeric] = None,
         **kwargs,
     ):
-        assert "lengthscale" in params
-        assert params["lengthscale"].shape == (1,)
-        assert "nu" in params
-        assert params["nu"].shape == (1,)
+        _check_field_in_params(params, "lengthscale")
+        _check_1_vector(params["lengthscale"], 'params["lengthscale"]')
+
+        _check_field_in_params(params, "nu")
+        _check_1_vector(params["nu"], 'params["nu"]')
 
         _, features_X = self.feature_map(
             X, params, normalize=self.normalize, **kwargs
@@ -127,10 +132,11 @@ def K(
         return feature_product
 
     def K_diag(self, params: Dict[str, B.Numeric], X: B.Numeric, **kwargs):
-        assert "lengthscale" in params
-        assert params["lengthscale"].shape == (1,)
-        assert "nu" in params
-        assert params["nu"].shape == (1,)
+        _check_field_in_params(params, "lengthscale")
+        _check_1_vector(params["lengthscale"], 'params["lengthscale"]')
+
+        _check_field_in_params(params, "nu")
+        _check_1_vector(params["nu"], 'params["nu"]')
 
         _, features_X = self.feature_map(
             X, params, normalize=self.normalize, **kwargs

geometric_kernels/kernels/hodge_compositional.py

@@ -11,6 +11,7 @@
 from geometric_kernels.kernels.base import BaseGeometricKernel
 from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel
 from geometric_kernels.spaces import HodgeDiscreteSpectrumSpace
+from geometric_kernels.utils.utils import _check_1_vector, _check_field_in_params
 
 
 class MaternHodgeCompositionalKernel(BaseGeometricKernel):
@@ -129,13 +130,9 @@ def K(
         inputs, or batches of matrices of inputs, depending on the space.
         """
 
-        assert all(
-            key in params for key in ["harmonic", "gradient", "curl"]
-        ), "MaternHodgeCompositionalKernel's parameters must contain keys 'harmonic', 'gradient', 'curl'."
-        assert all(
-            B.shape(params[key]["logit"]) == (1,)
-            for key in ["harmonic", "gradient", "curl"]
-        ), "The 'logit' parameters of MaternHodgeCompositionalKernel must have shape (1,)."
+        for key in ("harmonic", "gradient", "curl"):
+            _check_field_in_params(params, key)
+            _check_1_vector(params[key]["logit"], f'params["{key}"]["logit"]')
 
         # Copy the parameters to avoid modifying the original dict.
         params = {key: params[key].copy() for key in ["harmonic", "gradient", "curl"]}
@@ -162,13 +159,9 @@ def K_diag(
         diagonal.
         """
 
-        assert all(
-            key in params for key in ["harmonic", "gradient", "curl"]
-        ), "MaternHodgeCompositionalKernel's parameters must contain keys 'harmonic', 'gradient', 'curl'."
-        assert all(
-            B.shape(params[key]["logit"]) == (1,)
-            for key in ["harmonic", "gradient", "curl"]
-        ), "The 'logit' parameters of MaternHodgeCompositionalKernel must have shape (1,)."
+        for key in ("harmonic", "gradient", "curl"):
+            _check_field_in_params(params, key)
+            _check_1_vector(params[key]["logit"], f'params["{key}"]["logit"]')
 
         # Copy the parameters to avoid modifying the original dict.
         params = {key: params[key].copy() for key in ["harmonic", "gradient", "curl"]}

geometric_kernels/kernels/karhunen_loeve.py

@@ -11,6 +11,7 @@
 from geometric_kernels.lab_extras import from_numpy, is_complex
 from geometric_kernels.spaces import DiscreteSpectrumSpace
 from geometric_kernels.spaces.eigenfunctions import Eigenfunctions
+from geometric_kernels.utils.utils import _check_1_vector, _check_field_in_params
 
 
 class MaternKarhunenLoeveKernel(BaseGeometricKernel):
@@ -73,13 +74,26 @@ def __init__(
         self.num_levels = num_levels  # in code referred to as `L`.
 
         if eigenvalues_laplacian is None:
-            assert eigenfunctions is None
+            if eigenfunctions is not None:
+                raise ValueError(
+                    "You must either provide both `eigenfunctions` and `eigenvalues_laplacian`, or none of the two."
+                )
             eigenvalues_laplacian = self.space.get_eigenvalues(self.num_levels)
             eigenfunctions = self.space.get_eigenfunctions(self.num_levels)
         else:
-            assert eigenfunctions is not None
-            assert eigenvalues_laplacian.shape == (num_levels, 1)
-            assert eigenfunctions.num_levels == num_levels
+            if eigenfunctions is None:
+                raise ValueError(
+                    "You must either provide both `eigenfunctions` and `eigenvalues_laplacian`, or none of the two."
+                )
+            if eigenvalues_laplacian.shape != (num_levels, 1):
+                raise ValueError(
+                    f"Expected `eigenvalues_laplacian` to have shape [num_levels={num_levels}, 1] but got {eigenvalues_laplacian.shape}"
+                )
+            if eigenfunctions.num_levels != num_levels:
+                raise ValueError(
+                    f"`num_levels` must coincide with `num_levels` in the provided `eigenfunctions`,"
+                    f"but `num_levels`={num_levels} and `eigenfunctions.num_levels`={eigenfunctions.num_levels}"
+                )
 
         self._eigenvalues_laplacian = eigenvalues_laplacian
         self._eigenfunctions = eigenfunctions
@@ -134,8 +148,8 @@ def spectrum(
         :return:
             The spectrum of the Matérn kernel.
         """
-        assert lengthscale.shape == (1,)
-        assert nu.shape == (1,)
+        _check_1_vector(lengthscale, "lengthscale")
+        _check_1_vector(nu, "nu")
 
         # Note: 1.0 in safe_nu can be replaced by any finite positive value
         safe_nu = B.where(nu == np.inf, B.cast(B.dtype(lengthscale), np.r_[1.0]), nu)
@@ -180,10 +194,11 @@ def eigenvalues(self, params: Dict[str, B.Numeric]) -> B.Numeric:
         :return:
             An [L, 1]-shaped array.
         """
-        assert "lengthscale" in params
-        assert params["lengthscale"].shape == (1,)
-        assert "nu" in params
-        assert params["nu"].shape == (1,)
+        _check_field_in_params(params, "lengthscale")
+        _check_1_vector(params["lengthscale"], 'params["lengthscale"]')
+
+        _check_field_in_params(params, "nu")
+        _check_1_vector(params["nu"], 'params["nu"]')
 
         spectral_values = self.spectrum(
             self.eigenvalues_laplacian,
@@ -210,10 +225,11 @@ def eigenvalues(self, params: Dict[str, B.Numeric]) -> B.Numeric:
     def K(
         self, params: Dict[str, B.Numeric], X: B.Numeric, X2: Optional[B.Numeric] = None, **kwargs  # type: ignore
     ) -> B.Numeric:
-        assert "lengthscale" in params
-        assert params["lengthscale"].shape == (1,)
-        assert "nu" in params
-        assert params["nu"].shape == (1,)
+        _check_field_in_params(params, "lengthscale")
+        _check_1_vector(params["lengthscale"], 'params["lengthscale"]')
+
+        _check_field_in_params(params, "nu")
+        _check_1_vector(params["nu"], 'params["nu"]')
 
         weights = B.cast(B.dtype(params["nu"]), self.eigenvalues(params))  # [L, 1]
         Phi = self.eigenfunctions
@@ -224,10 +240,11 @@ def K(
             return K
 
     def K_diag(self, params: Dict[str, B.Numeric], X: B.Numeric, **kwargs) -> B.Numeric:
-        assert "lengthscale" in params
-        assert params["lengthscale"].shape == (1,)
-        assert "nu" in params
-        assert params["nu"].shape == (1,)
+        _check_field_in_params(params, "lengthscale")
+        _check_1_vector(params["lengthscale"], 'params["lengthscale"]')
+
+        _check_field_in_params(params, "nu")
+        _check_1_vector(params["nu"], 'params["nu"]')
 
         weights = B.cast(B.dtype(params["nu"]), self.eigenvalues(params))  # [L, 1]
         Phi = self.eigenfunctions