inference-tools 0.13.2__py3-none-any.whl → 0.13.4__py3-none-any.whl

inference/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.13.2'
-__version_tuple__ = version_tuple = (0, 13, 2)
+__version__ = version = '0.13.4'
+__version_tuple__ = version_tuple = (0, 13, 4)

inference/approx/conditional.py

@@ -12,7 +12,7 @@ class Conditional:
         self.theta = theta
         self.variable_index = variable_index
 
-    def __call__(self, x):
+    def __call__(self, x: ndarray):
         t = self.theta.copy()
         t[self.variable_index] = x
         return self.posterior(t)
@@ -58,12 +58,12 @@ def binary_search(
     return x_new
 
 
-def trapezium_full(x, dh):
+def trapezium_full(x: ndarray, dh: ndarray) -> ndarray:
     b = dh - 1
     return (b + sqrt(b**2 + 4 * x * dh)) / (2 * dh)
 
 
-def trapezium_near_zero(x, dh):
+def trapezium_near_zero(x: ndarray, dh: ndarray) -> ndarray:
     return x + (1 - x) * x * dh
 
 

inference/gp/covariance.py

@@ -11,23 +11,23 @@ class CovarianceFunction(ABC):
     """
 
     @abstractmethod
-    def pass_spatial_data(self, x):
+    def pass_spatial_data(self, x: ndarray):
         pass
 
     @abstractmethod
-    def estimate_hyperpar_bounds(self, y):
+    def estimate_hyperpar_bounds(self, y: ndarray):
         pass
 
     @abstractmethod
-    def __call__(self, u, v, theta):
+    def __call__(self, u: ndarray, v: ndarray, theta: ndarray) -> ndarray:
         pass
 
     @abstractmethod
-    def build_covariance(self, theta):
+    def build_covariance(self, theta: ndarray) -> ndarray:
         pass
 
     @abstractmethod
-    def covariance_and_gradients(self, theta):
+    def covariance_and_gradients(self, theta: ndarray):
         pass
 
     def __add__(self, other):
@@ -202,7 +202,7 @@ class SquaredExponential(CovarianceFunction):
         length-2 tuples giving the lower/upper bounds for each parameter.
     """
 
-    def __init__(self, hyperpar_bounds=None):
+    def __init__(self, hyperpar_bounds: list[tuple] = None):
         self.bounds = hyperpar_bounds
         self.n_params: int
         self.dx: ndarray
@@ -237,14 +237,14 @@ class SquaredExponential(CovarianceFunction):
             upr = log(self.dx[:, :, i].max()) + 2
             self.bounds.append((lwr, upr))
 
-    def __call__(self, u, v, theta):
+    def __call__(self, u: ndarray, v: ndarray, theta: ndarray) -> ndarray:
         a = exp(theta[0])
         L = exp(theta[1:])
         D = -0.5 * (u[:, None, :] - v[None, :, :]) ** 2
         C = exp((D / L[None, None, :] ** 2).sum(axis=2))
         return (a**2) * C
 
-    def build_covariance(self, theta):
+    def build_covariance(self, theta: ndarray) -> ndarray:
         """
         Optimized version of self.matrix() specifically for the data
         covariance matrix where the vectors v1 & v2 are both self.x.
@@ -254,7 +254,7 @@ class SquaredExponential(CovarianceFunction):
         C = exp((self.distances / L[None, None, :] ** 2).sum(axis=2)) + self.epsilon
         return (a**2) * C
 
-    def gradient_terms(self, v, x, theta):
+    def gradient_terms(self, v: ndarray, x: ndarray, theta: ndarray):
         """
         Calculates the covariance-function specific parts of
         the expression for the predictive mean and covariance
@@ -265,7 +265,7 @@ class SquaredExponential(CovarianceFunction):
         A = (x - v[None, :]) / L[None, :] ** 2
         return A.T, (a / L) ** 2
 
-    def covariance_and_gradients(self, theta):
+    def covariance_and_gradients(self, theta: ndarray):
         a = exp(theta[0])
         L = exp(theta[1:])
         C = exp((self.distances / L[None, None, :] ** 2).sum(axis=2)) + self.epsilon
@@ -303,7 +303,7 @@ class RationalQuadratic(CovarianceFunction):
         length-2 tuples giving the lower/upper bounds for each parameter.
     """
 
-    def __init__(self, hyperpar_bounds=None):
+    def __init__(self, hyperpar_bounds: list[tuple] = None):
         self.bounds = hyperpar_bounds
 
     def pass_spatial_data(self, x: ndarray):
@@ -332,7 +332,7 @@ class RationalQuadratic(CovarianceFunction):
             upr = log(self.dx[:, :, i].max()) + 2
             self.bounds.append((lwr, upr))
 
-    def __call__(self, u, v, theta):
+    def __call__(self, u: ndarray, v: ndarray, theta: ndarray) -> ndarray:
         a = exp(theta[0])
         k = exp(theta[1])
         L = exp(theta[2:])
@@ -340,14 +340,14 @@ class RationalQuadratic(CovarianceFunction):
         Z = (D / L[None, None, :] ** 2).sum(axis=2)
         return (a**2) * (1 + Z / k) ** (-k)
 
-    def build_covariance(self, theta):
+    def build_covariance(self, theta: ndarray) -> ndarray:
         a = exp(theta[0])
         k = exp(theta[1])
         L = exp(theta[2:])
         Z = (self.distances / L[None, None, :] ** 2).sum(axis=2)
         return (a**2) * ((1 + Z / k) ** (-k) + self.epsilon)
 
-    def covariance_and_gradients(self, theta):
+    def covariance_and_gradients(self, theta: ndarray):
         a = exp(theta[0])
         q = exp(theta[1])
         L = exp(theta[2:])
@@ -437,11 +437,11 @@ class ChangePoint(CovarianceFunction):
         for K in self.cov:
             if not isinstance(K, CovarianceFunction):
                 raise TypeError(
-                    """
-                    [ ChangePoint error ]
-                    >> Each of the specified covariance kernels must be an instance of
-                    >> a class which inherits from the 'CovarianceFunction' abstract
-                    >> base-class.
+                    """\n
+                    \r[ ChangePoint error ]
+                    \r>> Each of the specified covariance kernels must be an instance of
+                    \r>> a class which inherits from the 'CovarianceFunction' abstract
+                    \r>> base-class.
                     """
                 )
 
@@ -450,9 +450,9 @@ class ChangePoint(CovarianceFunction):
         if location_bounds is not None:
             if len(location_bounds) != self.n_kernels - 1:
                 raise ValueError(
-                    """
-                    [ ChangePoint error ]
-                    >> The length of 'location_bounds' must be one less than the number of kernels
+                    """\n
+                    \r[ ChangePoint error ]
+                    \r>> The length of 'location_bounds' must be one less than the number of kernels
                     """
                 )
             self.location_bounds = [check_bounds(lb) for lb in location_bounds]
@@ -462,9 +462,9 @@ class ChangePoint(CovarianceFunction):
         if width_bounds is not None:
             if len(width_bounds) != self.n_kernels - 1:
                 raise ValueError(
-                    """
-                    [ ChangePoint error ]
-                    >> The length of 'width_bounds' must be one less than the number of kernels
+                    """\n
+                    \r[ ChangePoint error ]
+                    \r>> The length of 'width_bounds' must be one less than the number of kernels
                     """
                 )
             self.width_bounds = [check_bounds(wb) for wb in width_bounds]
@@ -526,7 +526,7 @@ class ChangePoint(CovarianceFunction):
         # check for consistency of length of bounds
         assert self.n_params == len(self.bounds)
 
-    def __call__(self, u: ndarray, v: ndarray, theta):
+    def __call__(self, u: ndarray, v: ndarray, theta: ndarray) -> ndarray:
         kernel_coeffs = [1.0]
         for slc in self.cp_slc:
             w_u = self.logistic(u[:, self.axis], theta[slc])
@@ -543,7 +543,7 @@ class ChangePoint(CovarianceFunction):
             for i in range(self.n_kernels)
         )
 
-    def build_covariance(self, theta):
+    def build_covariance(self, theta: ndarray) -> ndarray:
         kernel_coeffs = [1.0]
         for slc in self.cp_slc:
             w = self.logistic(self.x_cp, theta[slc])
@@ -558,7 +558,7 @@ class ChangePoint(CovarianceFunction):
             for i in range(self.n_kernels)
         )
 
-    def covariance_and_gradients(self, theta):
+    def covariance_and_gradients(self, theta: ndarray):
         K_vals = []
         K_grads = []
         for i in range(self.n_kernels):
@@ -593,12 +593,12 @@ class ChangePoint(CovarianceFunction):
         return covar, gradients
 
     @staticmethod
-    def logistic(x, theta):
+    def logistic(x, theta: ndarray):
         z = (x - theta[0]) / theta[1]
         return 1.0 / (1.0 + exp(-z))
 
     @staticmethod
-    def logistic_and_gradient(x, theta):
+    def logistic_and_gradient(x, theta: ndarray):
         z = (x - theta[0]) / theta[1]
         f = 1.0 / (1.0 + exp(-z))
         dfdc = -f * (1 - f) / theta[1]
@@ -643,7 +643,7 @@ class HeteroscedasticNoise(CovarianceFunction):
         sequence of length-2 tuples giving the lower/upper bounds.
     """
 
-    def __init__(self, hyperpar_bounds=None):
+    def __init__(self, hyperpar_bounds: list[tuple] = None):
         self.bounds = hyperpar_bounds
 
     def pass_spatial_data(self, x: ndarray):
@@ -668,10 +668,10 @@ class HeteroscedasticNoise(CovarianceFunction):
         s = log(ptp(y))
         self.bounds = [(s - 8, s + 2) for _ in range(self.n_params)]
 
-    def __call__(self, u, v, theta):
+    def __call__(self, u: ndarray, v: ndarray, theta: ndarray) -> ndarray:
         return zeros([u.size, v.size])
 
-    def build_covariance(self, theta):
+    def build_covariance(self, theta: ndarray) -> ndarray:
         """
         Optimized version of self.matrix() specifically for the data
         covariance matrix where the vectors v1 & v2 are both self.x.
@@ -679,7 +679,7 @@ class HeteroscedasticNoise(CovarianceFunction):
         sigma_sq = exp(2 * theta)
         return diag(sigma_sq)
 
-    def covariance_and_gradients(self, theta):
+    def covariance_and_gradients(self, theta: ndarray):
         sigma_sq = exp(2 * theta)
         K = diag(sigma_sq)
         grads = [s * dk for s, dk in zip(sigma_sq, self.dK)]

inference/gp/mean.py

@@ -8,23 +8,23 @@ class MeanFunction(ABC):
     """
 
     @abstractmethod
-    def pass_spatial_data(self, x):
+    def pass_spatial_data(self, x: ndarray):
         pass
 
     @abstractmethod
-    def estimate_hyperpar_bounds(self, y):
+    def estimate_hyperpar_bounds(self, y: ndarray):
         pass
 
     @abstractmethod
-    def __call__(self, q, theta):
+    def __call__(self, q, theta: ndarray):
         pass
 
     @abstractmethod
-    def build_mean(self, theta):
+    def build_mean(self, theta: ndarray):
         pass
 
     @abstractmethod
-    def mean_and_gradients(self, theta):
+    def mean_and_gradients(self, theta: ndarray):
         pass
 
 

inference/gp/optimisation.py

@@ -7,7 +7,7 @@ import matplotlib.pyplot as plt
 
 from inference.gp.regression import GpRegressor
 from inference.gp.covariance import CovarianceFunction, SquaredExponential
-from inference.gp.acquisition import ExpectedImprovement
+from inference.gp.acquisition import AcquisitionFunction, ExpectedImprovement
 from inference.gp.mean import MeanFunction, ConstantMean
 
 
@@ -92,7 +92,7 @@ class GpOptimiser:
         kernel: CovarianceFunction = SquaredExponential,
         mean: MeanFunction = ConstantMean,
         cross_val: bool = False,
-        acquisition=ExpectedImprovement,
+        acquisition: AcquisitionFunction = ExpectedImprovement,
         optimizer: str = "bfgs",
         n_processes: int = 1,
     ):
@@ -165,7 +165,12 @@ class GpOptimiser:
                 self.y_err = append(self.y_err, new_y_err)
             else:
                 raise ValueError(
-                    "y_err must be specified for new evaluations if y_err was specified during __init__"
+                    """\n
+                    \r[ GpOptimiser error ]
+                    \r>> 'new_y_err' argument of the 'add_evaluation' method must be
+                    \r>> specified if the 'y_err' argument was specified when the
+                    \r>> instance of GpOptimiser was initialised.
+                    """
                 )
 
         # re-train the GP
@@ -243,7 +248,7 @@ class GpOptimiser:
             proposed_ev = proposed_ev[0]
         return proposed_ev
 
-    def plot_results(self, filename=None, show_plot=True):
+    def plot_results(self, filename: str = None, show_plot=True):
         fig = plt.figure(figsize=(10, 4))
         ax1 = fig.add_subplot(121)
         maxvals = maximum.accumulate(self.y)

inference/gp/regression.py

@@ -1,4 +1,4 @@
-from numpy import diagonal, diag, dot, sqrt, log
+from numpy import diagonal, diag, sqrt, log
 from numpy import array, eye, ndarray, zeros
 from numpy.random import random
 from numpy.linalg import cholesky, LinAlgError
@@ -98,8 +98,8 @@ class GpRegressor:
         if self.y.ndim != 1:
             raise ValueError(
                 f"""\n
-                [ GpRegressor error ]
-                >> 'y' argument must be a 1D array, but instead has shape {self.y.shape}
+                \r[ GpRegressor error ]
+                \r>> 'y' argument must be a 1D array, but instead has shape {self.y.shape}
                 """
             )
 
@@ -113,19 +113,19 @@ class GpRegressor:
         else:
             raise ValueError(
                 f"""\n
-                [ GpRegressor Error ]
-                >> 'x' argument must be a 2D array, but instead has
-                >> {self.x.ndim} dimensions and shape {self.x.shape}.
+                \r[ GpRegressor Error ]
+                \r>> 'x' argument must be a 2D array, but instead has
+                \r>> {self.x.ndim} dimensions and shape {self.x.shape}.
                 """
             )
 
         if self.x.shape[0] != self.n_points:
             raise ValueError(
                 f"""\n
-                [ GpRegressor Error ]
-                >> The first dimension of the 'x' array must be equal in size
-                >> to the 'y' array.
-                >> 'x' has shape {self.x.shape}, but 'y' has size {self.y.size}.
+                \r[ GpRegressor Error ]
+                \r>> The first dimension of the 'x' array must be equal in size
+                \r>> to the 'y' array.
+                \r>> 'x' has shape {self.x.shape}, but 'y' has size {self.y.size}.
                 """
             )
 
@@ -209,13 +209,13 @@ class GpRegressor:
             K_qx = self.cov(q, self.x, self.cov_hyperpars)
             K_qq = self.cov(q, q, self.cov_hyperpars)
 
-            mu_q.append(dot(K_qx, self.alpha)[0] + self.mean(q, self.mean_hyperpars))
+            mu_q.append((K_qx @ self.alpha)[0] + self.mean(q, self.mean_hyperpars))
             v = solve_triangular(self.L, K_qx.T, lower=True)
             errs.append(K_qq[0, 0] - (v**2).sum())
 
         return array(mu_q), sqrt(abs(array(errs)))
 
-    def set_hyperparameters(self, hyperpars):
+    def set_hyperparameters(self, hyperpars: ndarray):
         """
         Update the hyper-parameter values of the model.
 
@@ -243,7 +243,7 @@ class GpRegressor:
             self.L.T, solve_triangular(self.L, self.y - self.mu, lower=True)
         )
 
-    def check_error_data(self, y_err, y_cov):
+    def check_error_data(self, y_err, y_cov) -> ndarray:
         if y_cov is not None:
             # if y_cov is given as a list or tuple, attempt conversion to an array
             if any([type(y_cov) is t for t in [list, tuple]]):
@@ -321,7 +321,7 @@ class GpRegressor:
         else:
             return zeros([self.n_points, self.n_points])
 
-    def process_points(self, points):
+    def process_points(self, points: ndarray) -> ndarray:
         x = points if isinstance(points, ndarray) else array(points)
 
         if x.ndim <= 1 and self.n_dimensions == 1:
@@ -373,12 +373,11 @@ class GpRegressor:
             K_qx = self.cov(pnt, self.x, self.cov_hyperpars)
             A, R = self.cov.gradient_terms(pnt[0, :], self.x, self.cov_hyperpars)
 
-            B = (K_qx * self.alpha).T
             Q = solve_triangular(self.L, (A * K_qx).T, lower=True)
 
             # calculate the mean and covariance
-            mean = dot(A, B)
-            covariance = R - Q.T.dot(Q)
+            mean = A @ (K_qx * self.alpha).T
+            covariance = R - (Q.T @ Q)
 
             # store the results for the current point
             mu_q.append(mean)
@@ -408,19 +407,18 @@ class GpRegressor:
         for pnt in p[:, None, :]:
             K_qx = self.cov(pnt, self.x, self.cov_hyperpars)
             A, _ = self.cov.gradient_terms(pnt[0, :], self.x, self.cov_hyperpars)
-            B = (K_qx * self.alpha).T
             Q = solve_triangular(self.L.T, solve_triangular(self.L, K_qx.T, lower=True))
 
             # calculate the mean and covariance
-            dmu_dx = dot(A, B)
-            dV_dx = -2 * (A * K_qx[None, :]).dot(Q)
+            dmu_dx = A @ (K_qx * self.alpha).T
+            dV_dx = -2 * (A * K_qx[None, :]) @ Q
 
             # store the results for the current point
             mu_gradients.append(dmu_dx)
             var_gradients.append(dV_dx)
         return array(mu_gradients).squeeze(), array(var_gradients).squeeze()
 
-    def build_posterior(self, points: ndarray):
+    def build_posterior(self, points: ndarray, mean_only=False):
         """
         Generates the full mean vector and covariance matrix for the Gaussian-process
         posterior distribution at a set of specified points.
@@ -431,22 +429,26 @@ class GpRegressor:
             number of dimensions). Alternatively, a list of array-like objects can be
             given, which will be converted to a ``ndarray`` internally.
 
+        :param mean_only: \
+            If set to ``True``, only the mean vector of the posterior is calculated
+            and returned, instead of both the mean and covariance.
+
         :return: \
             The mean vector as a 1D array, followed by the covariance matrix as a 2D array.
         """
         v = self.process_points(points)
         K_qx = self.cov(v, self.x, self.cov_hyperpars)
         K_qq = self.cov(v, v, self.cov_hyperpars)
-        mu = dot(K_qx, self.alpha) + array(
-            [self.mean(p, self.mean_hyperpars) for p in v]
-        )
-        sigma = K_qq - dot(
-            K_qx,
-            solve_triangular(self.L.T, solve_triangular(self.L, K_qx.T, lower=True)),
-        )
-        return mu, sigma
+        mu = (K_qx @ self.alpha) + array([self.mean(p, self.mean_hyperpars) for p in v])
+
+        if mean_only:
+            return mu
+        else:
+            Q = solve_triangular(self.L, K_qx.T, lower=True)
+            sigma = K_qq - (Q.T @ Q)
+            return mu, sigma
 
-    def loo_predictions(self):
+    def loo_predictions(self) -> tuple[ndarray, ndarray]:
         """
         Calculates the 'leave-one out' (LOO) predictions for the data, where each data
         point is removed from the training set and then has its value predicted using
@@ -463,7 +465,7 @@ class GpRegressor:
         sigma = sqrt(var)
         return mu, sigma
 
-    def loo_likelihood(self, theta: ndarray):
+    def loo_likelihood(self, theta: ndarray) -> float:
         """
         Calculates the 'leave-one out' (LOO) log-likelihood.
 
@@ -477,7 +479,7 @@ class GpRegressor:
             L = cholesky(K_xx)
             iK = solve_triangular(L, eye(L.shape[0]), lower=True)
             iK = iK.T @ iK
-            alpha = iK.dot(self.y - mu)
+            alpha = iK @ (self.y - mu)
             var = 1.0 / diag(iK)
             return -0.5 * (var * alpha**2 + log(var)).sum()
         except LinAlgError:
@@ -499,7 +501,7 @@ class GpRegressor:
         L = cholesky(K_xx)
         iK = solve_triangular(L, eye(L.shape[0]), lower=True)
         iK = iK.T @ iK
-        alpha = iK.dot(self.y - mu)
+        alpha = iK @ (self.y - mu)
         var = 1.0 / diag(iK)
         LOO = -0.5 * (var * alpha**2 + log(var)).sum()
 
@@ -507,13 +509,13 @@ class GpRegressor:
         c1 = alpha * var
         c2 = 0.5 * var * (1 + var * alpha**2)
         for dK in grad_K:
-            Z = iK.dot(dK)
-            g = (c1 * Z.dot(alpha) - c2 * diag(Z.dot(iK))).sum()
+            Z = iK @ dK
+            g = (c1 * (Z @ alpha) - c2 * diag(Z @ iK)).sum()
             cov_gradients.append(g)
 
         mean_gradients = []
         for dmu in grad_mu:
-            Z = iK.dot(dmu)
+            Z = iK @ dmu
             g = (c1 * Z).sum()
             mean_gradients.append(g)
 
@@ -523,7 +525,7 @@ class GpRegressor:
 
         return LOO, grad
 
-    def marginal_likelihood(self, theta: ndarray):
+    def marginal_likelihood(self, theta: ndarray) -> float:
         """
         returns the log-marginal likelihood for the supplied hyper-parameter values.
 
@@ -554,8 +556,8 @@ class GpRegressor:
         iK = solve_triangular(L, eye(L.shape[0]), lower=True)
         iK = iK.T @ iK
         # calculate the log-marginal likelihood
-        alpha = iK.dot(self.y - mu)
-        LML = -0.5 * dot((self.y - mu).T, alpha) - log(diagonal(L)).sum()
+        alpha = iK @ (self.y - mu)
+        LML = -0.5 * ((self.y - mu).T @ alpha) - log(diagonal(L)).sum()
         # calculate the mean parameter gradients
         grad = zeros(self.n_hyperpars)
         grad[self.mean_slice] = array([(alpha * dmu).sum() for dmu in grad_mu])
@@ -564,7 +566,7 @@ class GpRegressor:
         grad[self.cov_slice] = array([0.5 * (Q * dK.T).sum() for dK in grad_K])
         return LML, grad
 
-    def differential_evo(self):
+    def differential_evo(self) -> ndarray:
         # optimise the hyper-parameters
         opt_result = differential_evolution(
             func=lambda x: -self.model_selector(x), bounds=self.hp_bounds

inference/likelihoods.py

@@ -73,14 +73,16 @@ class Likelihood(ABC):
             )
 
     @abstractmethod
-    def _log_likelihood(self, predictions):
+    def _log_likelihood(self, predictions: ndarray) -> float:
         pass
 
     @abstractmethod
-    def _log_likelihood_gradient(self, predictions, predictions_jacobian):
+    def _log_likelihood_gradient(
+        self, predictions: ndarray, predictions_jacobian: ndarray
+    ) -> ndarray:
         pass
 
-    def __call__(self, theta):
+    def __call__(self, theta: ndarray) -> float:
         """
         Returns the log-likelihood value for the given set of model parameters.
 
@@ -92,7 +94,7 @@ class Likelihood(ABC):
         """
         return self._log_likelihood(predictions=self.model(theta))
 
-    def gradient(self, theta):
+    def gradient(self, theta: ndarray) -> ndarray:
         """
         Returns the gradient of the log-likelihood with respect to model parameters.
 
@@ -110,10 +112,10 @@ class Likelihood(ABC):
             predictions_jacobian=self.model_jacobian(theta),
         )
 
-    def cost(self, theta):
+    def cost(self, theta: ndarray) -> float:
         return -self.__call__(theta)
 
-    def cost_gradient(self, theta):
+    def cost_gradient(self, theta: ndarray) -> ndarray:
         return -self.gradient(theta)
 
 
@@ -154,11 +156,13 @@ class GaussianLikelihood(Likelihood):
         self.inv_sigma_sqr = self.inv_sigma**2
         self.normalisation = -log(self.sigma).sum() - 0.5 * log(2 * pi) * self.n_data
 
-    def _log_likelihood(self, predictions):
+    def _log_likelihood(self, predictions: ndarray) -> float:
         z = (self.y - predictions) * self.inv_sigma
         return -0.5 * (z**2).sum() + self.normalisation
 
-    def _log_likelihood_gradient(self, predictions, predictions_jacobian):
+    def _log_likelihood_gradient(
+        self, predictions: ndarray, predictions_jacobian: ndarray
+    ) -> ndarray:
         dL_dF = (self.y - predictions) * self.inv_sigma_sqr
         return dL_dF @ predictions_jacobian
 
@@ -199,11 +203,13 @@ class CauchyLikelihood(Likelihood):
         self.inv_gamma = 1.0 / self.gamma
         self.normalisation = -log(pi * self.gamma).sum()
 
-    def _log_likelihood(self, predictions):
+    def _log_likelihood(self, predictions: ndarray) -> float:
         z = (self.y - predictions) * self.inv_gamma
         return -log(1 + z**2).sum() + self.normalisation
 
-    def _log_likelihood_gradient(self, predictions, predictions_jacobian):
+    def _log_likelihood_gradient(
+        self, predictions: ndarray, predictions_jacobian: ndarray
+    ) -> ndarray:
         z = (self.y - predictions) * self.inv_gamma
         dL_dF = 2 * self.inv_gamma * z / (1 + z**2)
         return dL_dF @ predictions_jacobian
@@ -246,11 +252,13 @@ class LogisticLikelihood(Likelihood):
         self.inv_scale = 1.0 / self.scale
         self.normalisation = -log(self.scale).sum()
 
-    def _log_likelihood(self, predictions):
+    def _log_likelihood(self, predictions: ndarray) -> float:
         z = (self.y - predictions) * self.inv_scale
         return z.sum() - 2 * log(1 + exp(z)).sum() + self.normalisation
 
-    def _log_likelihood_gradient(self, predictions, predictions_jacobian):
+    def _log_likelihood_gradient(
+        self, predictions: ndarray, predictions_jacobian: ndarray
+    ) -> ndarray:
         z = (self.y - predictions) * self.inv_scale
         dL_dF = (2 / (1 + exp(-z)) - 1) * self.inv_scale
         return dL_dF @ predictions_jacobian

inference/mcmc/base.py

@@ -106,7 +106,9 @@ class MarkovChain(ABC):
         else:
             return GaussianKDE(self.get_parameter(index, burn=burn, thin=thin))
 
-    def get_interval(self, interval=0.95, burn: int = 1, thin: int = 1, samples=None):
+    def get_interval(
+        self, interval: float = 0.95, burn: int = 1, thin: int = 1, samples: int = None
+    ) -> tuple[ndarray, ndarray]:
         """
         Return the samples from the chain which lie inside a chosen highest-density interval.
 
@@ -126,8 +128,8 @@ class MarkovChain(ABC):
             that specifying ``samples`` overrides the value of ``thin``.
 
         :return: \
-            List containing sample points stored as tuples, and a corresponding list of
-            log-probability values.
+            Samples from the chosen interval as a 2D ``numpy.ndarray``, followed by the
+            corresponding log-probability values as a 1D ``numpy.ndarray``.
         """
 
         # get the sorting indices for the probabilities