inference-tools 0.13.3__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.3'
-__version_tuple__ = version_tuple = (0, 13, 3)
+__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

@@ -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}.
                 """
             )
 
@@ -215,7 +215,7 @@ class GpRegressor:
 
         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:
@@ -448,7 +448,7 @@ class GpRegressor:
             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
@@ -465,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.
 
@@ -525,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.
 
@@ -566,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

inference/mcmc/gibbs.py

@@ -120,7 +120,7 @@ class Parameter:
         else:
             return self.upper - d % self.width
 
-    def submit_accept_prob(self, p):
+    def submit_accept_prob(self, p: float):
         self.num += 1
         self.avg += p
         self.var += p * (1 - p)
@@ -375,23 +375,26 @@ class MetropolisChain(MarkovChain):
         ind = argmax(self.probs)
         return array([p.samples[ind] for p in self.params])
 
-    def set_non_negative(self, parameter, flag=True):
+    def set_non_negative(self, parameter: int, flag=True):
         """
         Constrain a particular parameter to have non-negative values.
 
-        :param int parameter: Index of the parameter which is to be set \
-                              as non-negative.
+        :param int parameter: \
+            Index of the parameter which is to be set as non-negative.
         """
         self.params[parameter].non_negative = flag
 
-    def set_boundaries(self, parameter, boundaries, remove=False):
+    def set_boundaries(
+        self, parameter: int, boundaries: tuple[float, float], remove=False
+    ):
         """
         Constrain the value of a particular parameter to specified boundaries.
 
-        :param int parameter: Index of the parameter for which boundaries \
-                              are to be set.
+        :param int parameter: \
+            Index of the parameter for which boundaries are to be set.
 
-        :param boundaries: Tuple of boundaries in the format (lower_limit, upper_limit)
+        :param boundaries: \
+            Tuple of boundaries in the format (lower_limit, upper_limit)
         """
         if remove:
             self.params[parameter].remove_boundaries()
@@ -402,7 +405,7 @@ class MetropolisChain(MarkovChain):
         """
         Plot diagnostic traces that give information on how the chain is progressing.
 
-        Currently this method plots:
+        Currently, this method plots:
 
         - The posterior log-probability as a function of step number, which is useful
           for checking if the chain has reached a maximum. Any early parts of the chain

inference/mcmc/hmc.py

@@ -145,8 +145,8 @@ class HamiltonianChain(MarkovChain):
         else:
             raise ValueError(
                 f"""\n
-                [ HamiltonianChain error ]
-                >> Failed to take step within maximum allowed attempts of {self.max_attempts}
+                \r[ HamiltonianChain error ]
+                \r>> Failed to take step within maximum allowed attempts of {self.max_attempts}
                 """
             )
 
@@ -155,7 +155,9 @@ class HamiltonianChain(MarkovChain):
         self.leapfrog_steps.append(steps_taken)
         self.chain_length += 1
 
-    def standard_leapfrog(self, t: ndarray, r: ndarray, n_steps: int):
+    def standard_leapfrog(
+        self, t: ndarray, r: ndarray, n_steps: int
+    ) -> tuple[ndarray, ndarray]:
         t_step = self.inv_mass * self.ES.epsilon
         r_step = self.inv_temp * self.ES.epsilon
         r += (0.5 * r_step) * self.grad(t)
@@ -166,7 +168,9 @@ class HamiltonianChain(MarkovChain):
         r += (0.5 * r_step) * self.grad(t)
         return t, r
 
-    def bounded_leapfrog(self, t: ndarray, r: ndarray, n_steps: int):
+    def bounded_leapfrog(
+        self, t: ndarray, r: ndarray, n_steps: int
+    ) -> tuple[ndarray, ndarray]:
         t_step = self.inv_mass * self.ES.epsilon
         r_step = self.inv_temp * self.ES.epsilon
         r += (0.5 * r_step) * self.grad(t)

inference/mcmc/pca.py

@@ -33,13 +33,13 @@ class PcaChain(MetropolisChain):
         and returns the posterior log-probability.
 
     :param start: \
-        Vector of model parameters which correspond to the parameter-space coordinates
-        at which the chain will start.
+        Values of the model parameters as a ``numpy.ndarray`` which correspond to the
+        parameter-space coordinates at which the chain will start.
 
     :param widths: \
-        Vector of standard deviations which serve as initial guesses for the widths of
-        the proposal distribution for each model parameter. If not specified, the starting
-        widths will be approximated as 5% of the values in 'start'.
+        A ``numpy.ndarray`` of standard deviations which serve as initial guesses for
+        the widths of the proposal distribution for each model parameter. If not
+        specified, the starting widths will be approximated as 5% of the values in 'start'.
 
     :param bounds: \
         An instance of the ``inference.mcmc.Bounds`` class, or a sequence of two