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

inference/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.13.4'
-__version_tuple__ = version_tuple = (0, 13, 4)
+__version__ = version = '0.14.1'
+__version_tuple__ = version_tuple = (0, 14, 1)

inference/likelihoods.py

@@ -3,7 +3,7 @@
 """
 
 from abc import ABC, abstractmethod
-from numpy import array, log, exp, pi, sqrt, ndarray
+from numpy import array, log, exp, pi, sqrt, ndarray, logaddexp
 
 
 class Likelihood(ABC):
@@ -254,7 +254,7 @@ class LogisticLikelihood(Likelihood):
 
     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
+        return z.sum() - 2 * logaddexp(0.0, z).sum() + self.normalisation
 
     def _log_likelihood_gradient(
         self, predictions: ndarray, predictions_jacobian: ndarray

inference/mcmc/ensemble.py

@@ -3,7 +3,7 @@ import matplotlib.pyplot as plt
 
 from numpy import array, ndarray, linspace, concatenate, savez, load
 from numpy import sqrt, var, cov, diag, isfinite, triu, exp, log, median
-from numpy.random import random, randint
+from numpy.random import default_rng
 
 from inference.mcmc.utilities import Bounds, ChainProgressPrinter
 from inference.mcmc.base import MarkovChain
@@ -52,6 +52,7 @@ class EnsembleSampler(MarkovChain):
         display_progress=True,
     ):
         self.posterior = posterior
+        self.rng = default_rng()
 
         if starting_positions is not None:
             # store core data
@@ -180,9 +181,9 @@ class EnsembleSampler(MarkovChain):
 
     def __proposal(self, i: int):
         # randomly select walker that isn't 'i'
-        j = (randint(low=1, high=self.n_walkers) + i) % self.n_walkers
+        j = (self.rng.integers(low=1, high=self.n_walkers) + i) % self.n_walkers
         # sample the stretch distance
-        z = 0.5 * (self.x_lwr + self.x_width * random()) ** 2
+        z = 0.5 * (self.x_lwr + self.x_width * self.rng.random()) ** 2
         prop = self.process_proposal(
             self.walker_positions[i, :]
             + z * (self.walker_positions[j, :] - self.walker_positions[i, :])
@@ -194,7 +195,7 @@ class EnsembleSampler(MarkovChain):
             Y, z = self.__proposal(i)
             p = self.posterior(Y)
             q = exp((self.n_parameters - 1) * log(z) + p - self.walker_probs[i])
-            if random() <= q:
+            if self.rng.random() <= q:
                 self.walker_positions[i, :] = Y
                 self.walker_probs[i] = p
                 self.total_proposals[i].append(attempts)

inference/mcmc/gibbs.py

@@ -8,7 +8,7 @@ from numpy import exp, log, mean, sqrt, argmax, diff
 from numpy import percentile
 from numpy import isfinite, savez, load
 
-from numpy.random import normal, random
+from numpy.random import default_rng
 from inference.mcmc.utilities import ChainProgressPrinter, effective_sample_size
 from inference.mcmc.base import MarkovChain
 
@@ -27,6 +27,7 @@ class Parameter:
     def __init__(self, value: float, sigma: float):
         self.samples = [value]  # list to store all samples for the parameter
         self.sigma = sigma  # the width parameter for the proposal distribution
+        self.rng = default_rng()
 
         # storage for proposal width adjustment algorithm
         self.avg = 0
@@ -91,7 +92,7 @@ class Parameter:
         if self.try_count > self.max_tries:
             self.adjust_sigma(0.25)
         # return the proposed value
-        return self.samples[-1] + self.sigma * normal()
+        return self.rng.normal(loc=self.samples[-1], scale=self.sigma)
 
     def abs_proposal(self):
         # increment the try count
@@ -100,7 +101,7 @@ class Parameter:
         if self.try_count > self.max_tries:
             self.adjust_sigma(0.25)
         # return the proposed value
-        return abs(self.samples[-1] + self.sigma * normal())
+        return abs(self.rng.normal(loc=self.samples[-1], scale=self.sigma))
 
     def boundary_proposal(self):
         # increment the try count
@@ -109,7 +110,7 @@ class Parameter:
         if self.try_count > self.max_tries:
             self.adjust_sigma(0.25)
         # generate the proposed value
-        prop = self.samples[-1] + self.sigma * normal()
+        prop = self.rng.normal(loc=self.samples[-1], scale=self.sigma)
 
         # we now pass the proposal through a 'reflecting' function where
         # proposals falling outside the boundary are reflected inside
@@ -248,6 +249,7 @@ class MetropolisChain(MarkovChain):
         display_progress: bool = True,
     ):
         self.inv_temp = 1.0 / temperature
+        self.rng = default_rng()
 
         if posterior is not None:
             self.posterior = posterior
@@ -295,9 +297,8 @@ class MetropolisChain(MarkovChain):
             if pval > self.probs[-1]:
                 break
             else:
-                test = random()
                 acceptance_prob = exp(pval - self.probs[-1])
-                if test < acceptance_prob:
+                if self.rng.random() < acceptance_prob:
                     break
 
         for p, v in zip(self.params, proposal):
@@ -643,7 +644,7 @@ class GibbsChain(MetropolisChain):
                     # else calculate the acceptance probability and perform the test
                     acceptance_prob = exp(p_new - p_old)
                     p.submit_accept_prob(acceptance_prob)
-                    if random() < acceptance_prob:
+                    if self.rng.random() < acceptance_prob:
                         break
 
             p_old = deepcopy(p_new)  # NOTE - is deepcopy needed?

inference/mcmc/{hmc.py → hmc/__init__.py}

@@ -1,13 +1,14 @@
-from copy import copy
 import matplotlib.pyplot as plt
 
 from numpy import ndarray, float64
 from numpy import array, savez, savez_compressed, load, zeros
-from numpy import sqrt, var, isfinite, exp, log, dot, mean, argmax, percentile
-from numpy.random import random, normal
+from numpy import var, isfinite, exp, mean, argmax, percentile, cov
+from numpy.random import default_rng
 
 from inference.mcmc.utilities import Bounds, ChainProgressPrinter, effective_sample_size
 from inference.mcmc.base import MarkovChain
+from inference.mcmc.hmc.epsilon import EpsilonSelector
+from inference.mcmc.hmc.mass import get_particle_mass
 
 
 class HamiltonianChain(MarkovChain):
@@ -51,10 +52,12 @@ class HamiltonianChain(MarkovChain):
         in the form ``(lower_bounds, upper_bounds)``.
 
     :param inverse_mass: \
-        A vector specifying the inverse-mass value to be used for each parameter. The
-        inverse-mass is used to transform the momentum distribution in order to make
-        the problem more isotropic. Ideally, the inverse-mass for each parameter should
-        be set to the variance of the marginal distribution of that parameter.
+        The inverse-mass can be given as either a vector or matrix, and is used to
+        transform the momentum distribution so the chain can explore the posterior
+        efficiently. If given as a vector, the inverse mass should have values which
+        approximate the variance of the marginal distributions of each parameter.
+        If given as a matrix, the inverse mass should be a valid covariance matrix
+        for a multivariate normal distribution which approximates the posterior.
 
     :param bool display_progress: \
         If set as ``True``, a message is displayed during sampling
@@ -73,12 +76,10 @@ class HamiltonianChain(MarkovChain):
         display_progress=True,
     ):
         self.posterior = posterior
+        self.rng = default_rng()
         # if no gradient function is supplied, default to finite difference
         self.grad = self.finite_diff if grad is None else grad
 
-        # set the inverse mass to 1 if none supplied
-        self.inv_mass = 1.0 if inverse_mass is None else inverse_mass
-        self.sqrt_mass = 1.0 / sqrt(self.inv_mass)
         self.temperature = temperature
         self.inv_temp = 1.0 / temperature
 
@@ -90,7 +91,12 @@ class HamiltonianChain(MarkovChain):
             self.theta = [start]
             self.probs = [self.posterior(start) * self.inv_temp]
             self.leapfrog_steps = [0]
-            self.n_parameters = len(start)
+            self.n_parameters = start.size
+            self.mass = get_particle_mass(
+                inverse_mass=inverse_mass if inverse_mass is not None else 1.0,
+                n_parameters=self.n_parameters,
+            )
+
         self.chain_length = 1
 
         # set either the bounded or unbounded leapfrog update
@@ -124,23 +130,23 @@ class HamiltonianChain(MarkovChain):
         """
         steps_taken = 0
         for attempt in range(self.max_attempts):
-            r0 = normal(size=self.n_parameters, scale=self.sqrt_mass)
+            r0 = self.mass.sample_momentum(self.rng)
             t0 = self.theta[-1]
-            H0 = 0.5 * dot(r0, r0 * self.inv_mass) - self.probs[-1]
+            H0 = self.kinetic_energy(r0) - self.probs[-1]
 
-            n_steps = int(self.steps * (1 + (random() - 0.5) * 0.2))
+            n_steps = int(self.steps * (1 + (self.rng.random() - 0.5) * 0.2))
             t, r = self.run_leapfrog(t0.copy(), r0.copy(), n_steps)
 
             steps_taken += n_steps
             p = self.posterior(t) * self.inv_temp
-            H = 0.5 * dot(r, r * self.inv_mass) - p
+            H = self.kinetic_energy(r) - p
             accept_prob = exp(H0 - H)
 
             self.ES.add_probability(
                 min(accept_prob, 1) if isfinite(accept_prob) else 0.0
             )
 
-            if (accept_prob >= 1) or (random() <= accept_prob):
+            if (accept_prob >= 1) or (self.rng.random() <= accept_prob):
                 break
         else:
             raise ValueError(
@@ -158,38 +164,49 @@ class HamiltonianChain(MarkovChain):
     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)
+
         for _ in range(n_steps - 1):
-            t += t_step * r
+            t += self.ES.epsilon * self.mass.get_velocity(r)
             r += r_step * self.grad(t)
-        t += t_step * r
+
+        t += self.ES.epsilon * self.mass.get_velocity(r)
         r += (0.5 * r_step) * self.grad(t)
         return t, r
 
     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)
+
         for _ in range(n_steps - 1):
-            t += t_step * r
+            t += self.ES.epsilon * self.mass.get_velocity(r)
             t, reflections = self.bounds.reflect_momenta(t)
             r *= reflections
             r += r_step * self.grad(t)
-        t += t_step * r
+
+        t += self.ES.epsilon * self.mass.get_velocity(r)
         t, reflections = self.bounds.reflect_momenta(t)
         r *= reflections
         r += (0.5 * r_step) * self.grad(t)
         return t, r
 
     def hamiltonian(self, t: ndarray, r: ndarray) -> float:
-        return 0.5 * dot(r, r * self.inv_mass) - self.posterior(t) * self.inv_temp
+        return 0.5 * (r @ self.mass.get_velocity(r)) - self.posterior(t) * self.inv_temp
+
+    def kinetic_energy(self, r: ndarray) -> float:
+        return 0.5 * (r @ self.mass.get_velocity(r))
 
-    def estimate_mass(self, burn=1, thin=1):
-        self.inv_mass = var(array(self.theta[burn::thin]), axis=0)
+    def estimate_mass(self, burn=1, thin=1, diagonal=True):
+        if diagonal:
+            inverse_mass = var(array(self.theta[burn::thin]), axis=0)
+        else:
+            inverse_mass = cov(array(self.theta[burn::thin]).T)
+        self.mass = get_particle_mass(
+            inverse_mass=inverse_mass, n_parameters=self.n_parameters
+        )
 
     def finite_diff(self, t: ndarray) -> ndarray:
         p = self.posterior(t) * self.inv_temp
@@ -392,7 +409,7 @@ class HamiltonianChain(MarkovChain):
 
     def save(self, filename, compressed=False):
         items = {
-            "inv_mass": self.inv_mass,
+            "inv_mass": self.mass.inv_mass,
             "inv_temp": self.inv_temp,
             "theta": self.theta,
             "probs": self.probs,
@@ -450,69 +467,3 @@ class HamiltonianChain(MarkovChain):
         # build the epsilon selector
         chain.ES.load_items(D)
         return chain
-
-
-class EpsilonSelector:
-    def __init__(self, epsilon: float):
-        # storage
-        self.epsilon = epsilon
-        self.epsilon_values = [copy(epsilon)]  # sigma values after each assessment
-        self.epsilon_checks = [0.0]  # chain locations at which sigma was assessed
-
-        # tracking variables
-        self.avg = 0
-        self.var = 0
-        self.num = 0
-
-        # settings for epsilon adjustment algorithm
-        self.accept_rate = 0.65
-        self.chk_int = 15  # interval of steps at which proposal widths are adjusted
-        self.growth_factor = 1.4  # growth factor for self.chk_int
-
-    def add_probability(self, p):
-        self.num += 1
-        self.avg += p
-        self.var += max(p * (1 - p), 0.03)
-
-        if self.num >= self.chk_int:
-            self.update_epsilon()
-
-    def update_epsilon(self):
-        """
-        looks at average tries over recent steps, and adjusts proposal
-        widths self.sigma to bring the average towards self.target_tries.
-        """
-        # normal approximation of poisson binomial distribution
-        mu = self.avg / self.num
-        std = sqrt(self.var) / self.num
-
-        # now check if the desired success rate is within 2-sigma
-        if ~(mu - 2 * std < self.accept_rate < mu + 2 * std):
-            adj = (log(self.accept_rate) / log(mu)) ** 0.15
-            adj = min(adj, 2.0)
-            adj = max(adj, 0.5)
-            self.adjust_epsilon(adj)
-        else:  # increase the check interval
-            self.chk_int = int((self.growth_factor * self.chk_int) * 0.1) * 10
-
-    def adjust_epsilon(self, ratio: float):
-        self.epsilon *= ratio
-        self.epsilon_values.append(copy(self.epsilon))
-        self.epsilon_checks.append(self.epsilon_checks[-1] + self.num)
-        self.avg = 0
-        self.var = 0
-        self.num = 0
-
-    def get_items(self):
-        return self.__dict__
-
-    def load_items(self, dictionary: dict):
-        self.epsilon = float(dictionary["epsilon"])
-        self.epsilon_values = list(dictionary["epsilon_values"])
-        self.epsilon_checks = list(dictionary["epsilon_checks"])
-        self.avg = float(dictionary["avg"])
-        self.var = float(dictionary["var"])
-        self.num = float(dictionary["num"])
-        self.accept_rate = float(dictionary["accept_rate"])
-        self.chk_int = int(dictionary["chk_int"])
-        self.growth_factor = float(dictionary["growth_factor"])

inference/mcmc/hmc/epsilon.py

@@ -0,0 +1,68 @@
+from copy import copy
+from numpy import sqrt, log
+
+
+class EpsilonSelector:
+    def __init__(self, epsilon: float):
+        # storage
+        self.epsilon = epsilon
+        self.epsilon_values = [copy(epsilon)]  # sigma values after each assessment
+        self.epsilon_checks = [0.0]  # chain locations at which sigma was assessed
+
+        # tracking variables
+        self.avg = 0
+        self.var = 0
+        self.num = 0
+
+        # settings for epsilon adjustment algorithm
+        self.accept_rate = 0.65
+        self.chk_int = 15  # interval of steps at which proposal widths are adjusted
+        self.growth_factor = 1.4  # growth factor for self.chk_int
+
+    def add_probability(self, p: float):
+        self.num += 1
+        self.avg += p
+        self.var += max(p * (1 - p), 0.03)
+
+        if self.num >= self.chk_int:
+            self.update_epsilon()
+
+    def update_epsilon(self):
+        """
+        looks at the acceptance rate of proposed steps and adjusts the epsilon
+        value to bring the acceptance rate toward its target value.
+        """
+        # normal approximation of poisson binomial distribution
+        mu = self.avg / self.num
+        std = sqrt(self.var) / self.num
+
+        # now check if the desired success rate is within 2-sigma
+        if ~(mu - 2 * std < self.accept_rate < mu + 2 * std):
+            adj = (log(self.accept_rate) / log(mu)) ** 0.15
+            adj = min(adj, 2.0)
+            adj = max(adj, 0.5)
+            self.adjust_epsilon(adj)
+        else:  # increase the check interval
+            self.chk_int = int((self.growth_factor * self.chk_int) * 0.1) * 10
+
+    def adjust_epsilon(self, ratio: float):
+        self.epsilon *= ratio
+        self.epsilon_values.append(copy(self.epsilon))
+        self.epsilon_checks.append(self.epsilon_checks[-1] + self.num)
+        self.avg = 0
+        self.var = 0
+        self.num = 0
+
+    def get_items(self):
+        return self.__dict__
+
+    def load_items(self, dictionary: dict):
+        self.epsilon = float(dictionary["epsilon"])
+        self.epsilon_values = list(dictionary["epsilon_values"])
+        self.epsilon_checks = list(dictionary["epsilon_checks"])
+        self.avg = float(dictionary["avg"])
+        self.var = float(dictionary["var"])
+        self.num = float(dictionary["num"])
+        self.accept_rate = float(dictionary["accept_rate"])
+        self.chk_int = int(dictionary["chk_int"])
+        self.growth_factor = float(dictionary["growth_factor"])

inference/mcmc/hmc/mass.py

@@ -0,0 +1,80 @@
+from abc import ABC, abstractmethod
+from typing import Union
+from numpy import ndarray, sqrt, eye, isscalar
+from numpy.random import Generator
+from numpy.linalg import cholesky
+from scipy.linalg import solve_triangular
+
+
+class ParticleMass(ABC):
+    inv_mass: Union[float, ndarray]
+
+    @abstractmethod
+    def get_velocity(self, r: ndarray) -> ndarray:
+        pass
+
+    @abstractmethod
+    def sample_momentum(self, rng: Generator) -> ndarray:
+        pass
+
+
+class ScalarMass(ParticleMass):
+    def __init__(self, inv_mass: float, n_parameters: int):
+        self.inv_mass = inv_mass
+        self.sqrt_mass = 1 / sqrt(self.inv_mass)
+        self.n_parameters = n_parameters
+
+    def get_velocity(self, r: ndarray) -> ndarray:
+        return r * self.inv_mass
+
+    def sample_momentum(self, rng: Generator) -> ndarray:
+        return rng.normal(size=self.n_parameters, scale=self.sqrt_mass)
+
+
+class VectorMass(ScalarMass):
+    def __init__(self, inv_mass: ndarray, n_parameters: int):
+        super().__init__(inv_mass, n_parameters)
+        assert inv_mass.ndim == 1
+        assert inv_mass.size == n_parameters
+
+
+class MatrixMass(ParticleMass):
+    def __init__(self, inv_mass: ndarray, n_parameters: int):
+        assert inv_mass.ndim == 2
+        assert inv_mass.shape[0] == inv_mass.shape[1] == n_parameters
+        assert (inv_mass == inv_mass.T).all()
+
+        self.inv_mass = inv_mass
+        self.n_parameters = n_parameters
+        # find the cholesky decomp of the mass matrix
+        iL = cholesky(inv_mass)
+        self.L = solve_triangular(iL, eye(self.n_parameters), lower=True).T
+
+    def get_velocity(self, r: ndarray) -> ndarray:
+        return self.inv_mass @ r
+
+    def sample_momentum(self, rng: Generator) -> ndarray:
+        return self.L @ rng.normal(size=self.n_parameters)
+
+
+def get_particle_mass(
+    inverse_mass: Union[float, ndarray], n_parameters: int
+) -> ParticleMass:
+    if isscalar(inverse_mass):
+        return ScalarMass(inverse_mass, n_parameters)
+
+    if not isinstance(inverse_mass, ndarray):
+        raise TypeError(
+            f"""\n
+            \r[ HamiltonianChain error ]
+            \r>> The value given to the 'inverse_mass' keyword argument must be either
+            \r>> a scalar type (e.g. int or float), or a numpy.ndarray.
+            \r>> Instead, the given value has type:
+            \r>> {type(inverse_mass)}
+            """
+        )
+
+    if inverse_mass.ndim == 1:
+        return VectorMass(inverse_mass, n_parameters)
+    else:
+        return MatrixMass(inverse_mass, n_parameters)

inference/mcmc/parallel.py

@@ -7,7 +7,7 @@ from random import choice
 
 import matplotlib.pyplot as plt
 from numpy import arange, exp, identity, zeros
-from numpy.random import random, shuffle, seed, randint
+from numpy.random import default_rng
 from inference.plotting import transition_matrix_plot
 from inference.mcmc.base import MarkovChain
 
@@ -18,23 +18,19 @@ class ChainPool:
         self.pool_size = len(self.chains)
         self.pool = Pool(self.pool_size)
 
-    def advance(self, n):
+    def advance(self, n: int):
         self.chains = self.pool.map(
             self.adv_func, [(n, chain) for chain in self.chains]
         )
 
     @staticmethod
-    def adv_func(arg):
+    def adv_func(arg: tuple[int, MarkovChain]) -> MarkovChain:
         n, chain = arg
         chain.advance(n)
         return chain
 
 
-def tempering_process(
-    chain: MarkovChain, connection: Connection, end: Event, proc_seed: int
-):
-    # used to ensure each process has a different random seed
-    seed(proc_seed)
+def tempering_process(chain: MarkovChain, connection: Connection, end: Event):
     # main loop
     while not end.is_set():
         # poll the pipe until there is something to read
@@ -108,6 +104,7 @@ class ParallelTempering:
     """
 
     def __init__(self, chains: list[MarkovChain]):
+        self.rng = default_rng()
         self.shutdown_evt = Event()
         self.connections = []
         self.processes = []
@@ -132,7 +129,7 @@ class ParallelTempering:
             self.connections.append(parent_ctn)
             p = Process(
                 target=tempering_process,
-                args=(chn, child_ctn, self.shutdown_evt, randint(30000)),
+                args=(chn, child_ctn, self.shutdown_evt),
             )
             self.processes.append(p)
 
@@ -159,7 +156,7 @@ class ParallelTempering:
         Randomly pair up each chain, with uniform sampling across all possible pairings
         """
         proposed_swaps = arange(self.N_chains)
-        shuffle(proposed_swaps)
+        self.rng.shuffle(proposed_swaps)
         return [p for p in zip(proposed_swaps[::2], proposed_swaps[1::2])]
 
     def tight_pairs(self):
@@ -181,7 +178,7 @@ class ParallelTempering:
             leftovers = [
                 i for i in range(self.N_chains) if not any(i in p for p in sample)
             ]
-            shuffle(leftovers)
+            self.rng.shuffle(leftovers)
             sample.extend(
                 [
                     p if p[0] < p[1] else (p[1], p[0])
@@ -216,7 +213,7 @@ class ParallelTempering:
             pj = probabilities[j] / self.inv_temps[j]
             dp = pi - pj
 
-            if random() <= exp(-dt * dp):  # check if the swap is successful
+            if self.rng.random() <= exp(-dt * dp):  # check if the swap is successful
                 Di = {
                     "task": "update_position",
                     "position": positions[i],
@@ -233,7 +230,7 @@ class ParallelTempering:
                 self.connections[j].send(Di)
                 self.successful_swaps[i, j] += 1
 
-    def advance(self, n, swap_interval=10):
+    def advance(self, n: int, swap_interval=10):
         """
         Advances each chain by a total of *n* steps, performing swap attempts
         at intervals set by the *swap_interval* keyword.

inference/mcmc/pca.py

@@ -4,7 +4,6 @@ import matplotlib.pyplot as plt
 
 from numpy import array, savez, load, zeros
 from numpy import sqrt, exp, dot, cov
-from numpy.random import random, normal
 from scipy.linalg import eigh
 
 from inference.mcmc.gibbs import MetropolisChain, Parameter
@@ -157,7 +156,7 @@ class PcaChain(MetropolisChain):
         # loop over each eigenvector and take a step along each
         for v, p in zip(self.directions, self.params):
             while True:
-                prop = theta0 + v * p.sigma * normal()
+                prop = theta0 + v * p.sigma * self.rng.normal()
                 prop = self.process_proposal(prop)
                 p_new = self.posterior(prop) * self.inv_temp
 
@@ -165,10 +164,9 @@ class PcaChain(MetropolisChain):
                     p.submit_accept_prob(1.0)
                     break
                 else:
-                    test = random()
                     acceptance_prob = exp(p_new - p_old)
                     p.submit_accept_prob(acceptance_prob)
-                    if test < acceptance_prob:
+                    if self.rng.random() < acceptance_prob:
                         break
 
             theta0 = copy(prop)
@@ -184,7 +182,7 @@ class PcaChain(MetropolisChain):
         if self.chain_length == self.next_update:
             self.update_directions()
 
-    def save(self, filename):
+    def save(self, filename: str):
         """
         Save the entire state of the chain object as an .npz file.
 
@@ -221,7 +219,7 @@ class PcaChain(MetropolisChain):
         savez(filename, **items)
 
     @classmethod
-    def load(cls, filename, posterior=None):
+    def load(cls, filename: str, posterior=None):
         """
         Load a chain object which has been previously saved using the save() method.
 

inference/pdf/hdi.py

@@ -1,11 +1,9 @@
 from _warnings import warn
 from typing import Sequence
+from numpy import ndarray, array, sort, zeros, take_along_axis, expand_dims
 
-from numpy import ndarray, array, sort
-from scipy.optimize import differential_evolution
 
-
-def sample_hdi(sample: ndarray, fraction: float, allow_double=False):
+def sample_hdi(sample: ndarray, fraction: float) -> ndarray:
     """
     Estimate the highest-density interval(s) for a given sample.
 
@@ -13,26 +11,25 @@ def sample_hdi(sample: ndarray, fraction: float, allow_double=False):
     fraction of the elements in the given sample.
 
     :param sample: \
-        A sample for which the interval will be determined.
+        A sample for which the interval will be determined. If the sample is given
+        as a 2D numpy array, the interval calculation will be distributed over the
+        second dimension of the array, i.e. given a sample array of shape ``(m, n)``
+        the highest-density intervals are returned as an array of shape ``(2, n)``.
 
     :param float fraction: \
         The fraction of the total probability to be contained by the interval.
 
-    :param bool allow_double: \
-        When set to True, a double-interval is returned instead if one exists whose
-        total length is meaningfully shorter than the optimal single interval.
-
     :return: \
-        Tuple(s) specifying the lower and upper bounds of the highest-density interval(s).
+        The lower and upper bounds of the highest-density interval(s) as a numpy array.
     """
 
     # verify inputs are valid
     if not 0.0 < fraction < 1.0:
         raise ValueError(
             f"""\n
-            [ sample_hdi error ]
-            >> The 'fraction' argument must be a float between 0 and 1,
-            >> but the value given was {fraction}.
+            \r[ sample_hdi error ]
+            \r>> The 'fraction' argument must be a float between 0 and 1,
+            \r>> but the value given was {fraction}.
             """
         )
 
@@ -43,66 +40,72 @@ def sample_hdi(sample: ndarray, fraction: float, allow_double=False):
     else:
         raise ValueError(
             f"""\n
-            [ sample_hdi error ]
-            >> The 'sample' argument should be a numpy.ndarray or a
-            >> Sequence which can be converted to an array, but
-            >> instead has type {type(sample)}.
+            \r[ sample_hdi error ]
+            \r>> The 'sample' argument should be a numpy.ndarray or a
+            \r>> Sequence which can be converted to an array, but
+            \r>> instead has type {type(sample)}.
             """
         )
 
-    if s.size < 2:
+    if s.ndim > 2 or s.ndim == 0:
         raise ValueError(
             f"""\n
-            [ sample_hdi error ]
-            >> The given 'sample' array must contain at least 2 values.
+            \r[ sample_hdi error ]
+            \r>> The 'sample' argument should be a numpy.ndarray
+            \r>> with either one or two dimensions, but the given
+            \r>> array has dimensionality {s.ndim}.
             """
         )
 
-    if s.ndim > 1:
-        s = s.flatten()
-    s.sort()
-    n = s.size
-    L = int(fraction * n)
+    if s.ndim == 1:
+        s.resize([s.size, 1])
+
+    n_samples, n_intervals = s.shape
+    L = int(fraction * n_samples)
+
+    if n_samples < 2:
+        raise ValueError(
+            f"""\n
+            \r[ sample_hdi error ]
+            \r>> The first dimension of the given 'sample' array must 
+            \r>> have have a length of at least 2.
+            """
+        )
 
     # check that we have enough samples to estimate the HDI for the chosen fraction
-    if n <= L:
+    if n_samples <= L:
         warn(
             f"""\n
-            [ sample_hdi warning ]
-            >> The given number of samples is insufficient to estimate the interval
-            >> for the given fraction.
+            \r[ sample_hdi warning ]
+            \r>> The given number of samples is insufficient to estimate the interval
+            \r>> for the given fraction.
             """
         )
-        return s[0], s[-1]
-    elif n - L < 20:
+
+    elif n_samples - L < 20:
         warn(
             f"""\n
-            [ sample_hdi warning ]
-            >> len(sample)*(1 - fraction) is small - calculated interval may be inaccurate.
+            \r[ sample_hdi warning ]
+            \r>> n_samples * (1 - fraction) is small - calculated interval may be inaccurate.
             """
         )
 
-    # find the optimal single HDI
-    widths = s[L:] - s[: n - L]
-    i = widths.argmin()
-    r1, w1 = (s[i], s[i + L]), s[i + L] - s[i]
-
-    if allow_double:
-        # now get the best 2-interval solution
-        minfunc = dbl_interval_length(sample, fraction)
-        bounds = minfunc.get_bounds()
-        de_result = differential_evolution(minfunc, bounds)
-        I1, I2 = minfunc.return_intervals(de_result.x)
-        w2 = (I2[1] - I2[0]) + (I1[1] - I1[0])
-
-    # return the split interval if the width reduction is non-trivial:
-    if allow_double and w2 < w1 * 0.99:
-        return I1, I2
+    # check that we have enough samples to estimate the HDI for the chosen fraction
+    s.sort(axis=0)
+    hdi = zeros([2, n_intervals])
+    if n_samples > L:
+        # find the optimal single HDI
+        widths = s[L:, :] - s[: n_samples - L, :]
+        i = expand_dims(widths.argmin(axis=0), axis=0)
+        hdi[0, :] = take_along_axis(s, i, 0).squeeze()
+        hdi[1, :] = take_along_axis(s, i + L, 0).squeeze()
     else:
-        return r1
+        hdi[0, :] = s[0, :]
+        hdi[1, :] = s[-1, :]
+    return hdi.squeeze()
 
 
-class dbl_interval_length:
+class DoubleIntervalLength:
     def __init__(self, sample, fraction):
         self.sample = sort(sample)
         self.f = fraction

inference/plotting.py

@@ -445,29 +445,11 @@ def hdi_plot(
     if axis is None:
         _, axis = plt.subplots()
 
-    from numpy import take_along_axis, expand_dims
-
     # iterate over the intervals and plot each
     for frac, col in zip(intervals, colors):
-        L = int(frac * n)
-
-        # check that we have enough samples to estimate the HDI for the chosen fraction
-        if n > L:
-            # find the optimal single HDI
-            widths = s[L:, :] - s[: n - L, :]
-            i = expand_dims(widths.argmin(axis=0), axis=0)
-            lwr = take_along_axis(s, i, 0).squeeze()
-            upr = take_along_axis(s, i + L, 0).squeeze()
-        else:
-            lwr = s[0, :]
-            upr = s[-1, :]
-
-        if label_intervals:
-            axis.fill_between(
-                x, lwr, upr, color=col, label="{}% HDI".format(int(100 * frac))
-            )
-        else:
-            axis.fill_between(x, lwr, upr, color=col)
+        lwr, upr = sample_hdi(s, fraction=frac)
+        lab = f"{int(100 * frac)}% HDI" if label_intervals else None
+        axis.fill_between(x, lwr, upr, color=col, label=lab)
 
     return axis
 

inference/priors.py

@@ -4,9 +4,11 @@
 
 from abc import ABC, abstractmethod
 from typing import Union, Iterable
-from numpy import atleast_1d, log, pi, zeros, concatenate, where, ndarray, isfinite
-from numpy.random import normal, exponential, uniform
 from itertools import chain
+from numpy import atleast_1d, log, pi, zeros, concatenate, where, ndarray, isfinite
+from numpy.random import default_rng
+
+rng = default_rng()
 
 
 class BasePrior(ABC):
@@ -294,7 +296,7 @@ class GaussianPrior(BasePrior):
         :returns: \
             A single sample from the prior distribution as a 1D ``numpy.ndarray``.
         """
-        return normal(loc=self.mean, scale=self.sigma)
+        return rng.normal(loc=self.mean, scale=self.sigma)
 
     @classmethod
     def combine(cls, priors):
@@ -377,7 +379,7 @@ class ExponentialPrior(BasePrior):
         :returns: \
             A single sample from the prior distribution as a 1D ``numpy.ndarray``.
         """
-        return exponential(scale=self.beta)
+        return rng.exponential(scale=self.beta)
 
     @classmethod
     def combine(cls, priors: list[BasePrior]):
@@ -470,7 +472,7 @@ class UniformPrior(BasePrior):
         :returns: \
             A single sample from the prior distribution as a 1D ``numpy.ndarray``.
         """
-        return uniform(low=self.lower, high=self.upper)
+        return rng.uniform(low=self.lower, high=self.upper)
 
     @classmethod
     def combine(cls, priors):

{inference_tools-0.13.4.dist-info → inference_tools-0.14.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: inference-tools
-Version: 0.13.4
+Version: 0.14.1
 Summary: A collection of python tools for Bayesian data analysis
 Author-email: Chris Bowman <chris.bowman.physics@gmail.com>
 License: MIT License
@@ -42,6 +42,7 @@ Requires-Dist: pytest-cov>=3.0.0; extra == "tests"
 Requires-Dist: pyqt5>=5.15; extra == "tests"
 Requires-Dist: hypothesis>=6.24; extra == "tests"
 Requires-Dist: freezegun>=1.1.0; extra == "tests"
+Dynamic: license-file
 
 # inference-tools
 

{inference_tools-0.13.4.dist-info → inference_tools-0.14.1.dist-info}/RECORD

@@ -1,9 +1,9 @@
 inference/__init__.py,sha256=Wheq9bSUF5Y_jAc_w_Avi4WW2kphDK0qHGM6FsIKSxY,275
-inference/_version.py,sha256=dhPsd2j9Al5Z6JN-zXmB6-Uti2Ily2tvtEtKPTNHaCQ,413
-inference/likelihoods.py,sha256=fS_k3mRr7bv6kgDt29u_OB6emU-ARVZktf7j-eXA-2U,10008
-inference/plotting.py,sha256=U1M_F5I-UMtfHiaN1YihcxYq5gg_2MNyPm7MxF1LecY,19747
+inference/_version.py,sha256=9-B5HerO_wiKUcm3zqJZazE8kjqwU6_WcIM1m-vWIoQ,513
+inference/likelihoods.py,sha256=0mRn9S7CaX6hNv1fKVeaAFYk50bALvVbyX7E2aH3Bn8,10021
+inference/plotting.py,sha256=vMpRGiZMMlVgAcVaKC2wtvjzVlBmOkC2BM90A3wSwJ8,19194
 inference/posterior.py,sha256=ptPZgzT--ehbpu57nW9GmFuyovFOSmw56HWfuC-8GGA,3584
-inference/priors.py,sha256=67cgKw7jDurda9UByFJ7jOoEJH1FyZDOHC9-nvr0nWY,19352
+inference/priors.py,sha256=zDuIgJTZrqEqkp8rE-aBRlAuqBacR9aC_QNm8jNIYl8,19368
 inference/approx/__init__.py,sha256=b8xCdshVeGHyao6-P0038QB71WOMLrcYXCOYiYjK7Tk,132
 inference/approx/conditional.py,sha256=IeUismbo25qa1BUIqsZ2noum9_mLXNaORsg57abxBec,9515
 inference/gp/__init__.py,sha256=R4iPgf8TdunkOv_VLwue7Fz3AjGWDTBop58nCmbmMQ0,801
@@ -15,19 +15,21 @@ inference/gp/optimisation.py,sha256=sPhakklWIgg1yEUhUzA-m5vl0kVPvHdcgnQ0OAGT8qs,
 inference/gp/regression.py,sha256=10TzqVeUzUkuw8-Cbe4LbxevByTi5iE5QDdRClN7Nhk,25677
 inference/mcmc/__init__.py,sha256=IsEhVSIpZCDNIqgSq_21M6DH6x8F1jJbYWM0e3S3QG4,445
 inference/mcmc/base.py,sha256=cEh1LPmKd6JMop8EcuH3dvAeJYei88pcPTw1xe7tGKY,10496
-inference/mcmc/ensemble.py,sha256=s9Xspq5r360_XmpRHCplN5cscD60UoYXlYqx3yVEhsM,15528
-inference/mcmc/gibbs.py,sha256=9US0VqLEI_f70vrHg0sFZQneJMyjm8BF_l_0bD-ZqKI,24190
-inference/mcmc/hmc.py,sha256=rfTqvD3aZqqHXcM17_Yj8U_2mt2eTQ_BI6hOeFqycoo,19420
-inference/mcmc/parallel.py,sha256=HRK1Ka02iO5Q6m3282lqZeAlCZPXHIglC8RAlDE6Xd4,14082
-inference/mcmc/pca.py,sha256=-XVs25hH8FRA6XY4xWEK1cUZ8oDDllW7t_vlK6FU7Gs,10739
+inference/mcmc/ensemble.py,sha256=JRXu7SBYXN4Y9RzgA6kGUHpZNw4q4A9wf0KOAQdlz0E,15585
+inference/mcmc/gibbs.py,sha256=f-eccDBILfaZercZii3vuJ29V505VUsCHoxhD9gZ7xA,24288
+inference/mcmc/parallel.py,sha256=SKLzMP4aqIj1xsxKuByA1lr1GdgIu5pPzVw7hlfXZEQ,14053
+inference/mcmc/pca.py,sha256=NxC81NghGlBQslFVOk2HzpsnCjlEdDnv_w8es4Qe7PU,10695
 inference/mcmc/utilities.py,sha256=YjpK3FvV0Q98jLusrZrvGck-bjm6uZZ1U7HHH3aly8g,6048
+inference/mcmc/hmc/__init__.py,sha256=R2ZjKca1CjWwzAHRKetZOAbiJKo2YZVmX4jrz2EAyL4,17661
+inference/mcmc/hmc/epsilon.py,sha256=t2kNi10MSVFXjmAx5zRUARDuPu_yWbwoK2McMuaaAUs,2467
+inference/mcmc/hmc/mass.py,sha256=qnxsbkogZFeqGbssZ2w4tsaUGytEXL-I0Gqs4UZzcAg,2545
 inference/pdf/__init__.py,sha256=gVmQ1HLTab6_oWMQN26A1r7PkqbApaJmBK-c7TIFxjY,270
 inference/pdf/base.py,sha256=Zj5mfFmDqTe5cFz0biBxcvEaxdOUC-SsOUjebUEX7HM,5442
-inference/pdf/hdi.py,sha256=j_W4kv70weXR7C2ltTHR6OUNkAK-kLQhnrnpPrjiLxQ,4282
+inference/pdf/hdi.py,sha256=soFw3fKQdzxbGNhU9BvFHdt0uGKfhus3E3vM6L47yhY,4638
 inference/pdf/kde.py,sha256=KSl8y---602MlxoSVH8VknNQYZ2KAOTky50QU3jRw28,12999
 inference/pdf/unimodal.py,sha256=9S05c0hq_rF-MLoDJgUmaJKRdcP8F9_Idj7Ncb6m9q0,6218
-inference_tools-0.13.4.dist-info/LICENSE,sha256=Y0-EfO5pdxf6d0J6Er13ZSWiPZ2o6kHvM37eRgnJdww,1069
-inference_tools-0.13.4.dist-info/METADATA,sha256=l2x2GqQSfSrgrLeZLLqQ-LX00bxB1CKjj3AxHvTJ7F8,5378
-inference_tools-0.13.4.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
-inference_tools-0.13.4.dist-info/top_level.txt,sha256=I7bsb71rLtH3yvVH_HSLXUosY2AwCxEG3vctNsEhbEM,10
-inference_tools-0.13.4.dist-info/RECORD,,
+inference_tools-0.14.1.dist-info/licenses/LICENSE,sha256=Y0-EfO5pdxf6d0J6Er13ZSWiPZ2o6kHvM37eRgnJdww,1069
+inference_tools-0.14.1.dist-info/METADATA,sha256=flIdPaL3VFTzogPS-k2CwjIu_kbww--4libVTB9zegY,5400
+inference_tools-0.14.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+inference_tools-0.14.1.dist-info/top_level.txt,sha256=I7bsb71rLtH3yvVH_HSLXUosY2AwCxEG3vctNsEhbEM,10
+inference_tools-0.14.1.dist-info/RECORD,,

{inference_tools-0.13.4.dist-info → inference_tools-0.14.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (73.0.1)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any