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

inference/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.14.0'
-__version_tuple__ = version_tuple = (0, 14, 0)
+__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/{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 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
@@ -77,9 +80,6 @@ class HamiltonianChain(MarkovChain):
         # 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
 
@@ -91,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
@@ -125,16 +130,16 @@ class HamiltonianChain(MarkovChain):
         """
         steps_taken = 0
         for attempt in range(self.max_attempts):
-            r0 = self.rng.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 + (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(
@@ -159,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
@@ -393,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,
@@ -451,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: 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 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/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_tools-0.14.0.dist-info → inference_tools-0.14.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: inference-tools
-Version: 0.14.0
+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.14.0.dist-info → inference_tools-0.14.1.dist-info}/RECORD

@@ -1,7 +1,7 @@
 inference/__init__.py,sha256=Wheq9bSUF5Y_jAc_w_Avi4WW2kphDK0qHGM6FsIKSxY,275
-inference/_version.py,sha256=CNbGkYuFNWjd3DIUYNizEoC18xKEHqri0oYO6s7QgJ8,513
-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=zDuIgJTZrqEqkp8rE-aBRlAuqBacR9aC_QNm8jNIYl8,19368
 inference/approx/__init__.py,sha256=b8xCdshVeGHyao6-P0038QB71WOMLrcYXCOYiYjK7Tk,132
@@ -17,17 +17,19 @@ inference/mcmc/__init__.py,sha256=IsEhVSIpZCDNIqgSq_21M6DH6x8F1jJbYWM0e3S3QG4,44
 inference/mcmc/base.py,sha256=cEh1LPmKd6JMop8EcuH3dvAeJYei88pcPTw1xe7tGKY,10496
 inference/mcmc/ensemble.py,sha256=JRXu7SBYXN4Y9RzgA6kGUHpZNw4q4A9wf0KOAQdlz0E,15585
 inference/mcmc/gibbs.py,sha256=f-eccDBILfaZercZii3vuJ29V505VUsCHoxhD9gZ7xA,24288
-inference/mcmc/hmc.py,sha256=7SDjiwzVCqme1g8v65XldQpW5dnt7O3p1IG2AYGBb4o,19484
 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=soFw3fKQdzxbGNhU9BvFHdt0uGKfhus3E3vM6L47yhY,4638
 inference/pdf/kde.py,sha256=KSl8y---602MlxoSVH8VknNQYZ2KAOTky50QU3jRw28,12999
 inference/pdf/unimodal.py,sha256=9S05c0hq_rF-MLoDJgUmaJKRdcP8F9_Idj7Ncb6m9q0,6218
-inference_tools-0.14.0.dist-info/LICENSE,sha256=Y0-EfO5pdxf6d0J6Er13ZSWiPZ2o6kHvM37eRgnJdww,1069
-inference_tools-0.14.0.dist-info/METADATA,sha256=vlamppwmRyKBi3LuZP-IVRMsDy2UhdzoEiabRWw0aoE,5378
-inference_tools-0.14.0.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-inference_tools-0.14.0.dist-info/top_level.txt,sha256=I7bsb71rLtH3yvVH_HSLXUosY2AwCxEG3vctNsEhbEM,10
-inference_tools-0.14.0.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.14.0.dist-info → inference_tools-0.14.1.dist-info}/WHEEL

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