gemlib 0.9.2__py3-none-any.whl

gemlib/mcmc/experimental/hmc.py

@@ -0,0 +1,111 @@
+"""Base Hamiltonian Monte Carlo"""
+
+from typing import Iterable, NamedTuple, Optional
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+import tensorflow_probability.python.experimental.mcmc.preconditioning_utils as pu  # noqa: E501
+
+from .mcmc_base import ChainState, SamplingAlgorithm
+
+tfd = tfp.distributions
+tfde = tfp.experimental.distributions
+
+
+class HmcKernelState(NamedTuple):
+    step_size: float
+    num_leapfrog_steps: float
+    mass_matrix: Iterable
+
+
+def hmc(
+    step_size: float = 0.1,
+    num_leapfrog_steps: int = 16,
+    mass_matrix: Optional[Iterable] = None,
+):
+    """Hamiltonian Monte Carlo
+
+    Args
+    ----
+    step_size: the step size to take
+    num_leapfrog_steps: number of leapfrog steps to take
+    mass_matrix: a mass matrix (defaults to diag(1) if None)
+
+    Returns
+    -------
+    A SamplingAlgorithm
+    """
+
+    step_size = tf.convert_to_tensor(step_size)
+    mass_matrix = (
+        tf.convert_to_tensor(mass_matrix) if mass_matrix is not None else None
+    )
+
+    def _make_momentum_distribution(position):
+        # return None
+        if mass_matrix is None:
+            return pu.make_momentum_distribution(
+                position, tf.constant([], dtype=tf.int32)
+            )
+
+        else:
+            return tfde.MultivariateNormalPrecisionFactorLinearOperator(
+                precision_factor=tf.linalg.LinearOperatorFullMatrix(
+                    mass_matrix
+                ),
+                precision=tf.linalg.LinearOperatorFullMatrix(mass_matrix),
+            )
+
+    def _build_kernel(target_log_prob_fn, momentum_distribution):
+        return tfp.experimental.mcmc.PreconditionedHamiltonianMonteCarlo(
+            target_log_prob_fn=target_log_prob_fn,
+            step_size=step_size,
+            num_leapfrog_steps=num_leapfrog_steps,
+            momentum_distribution=momentum_distribution,
+            store_parameters_in_results=True,
+        )
+
+    def init_fn(target_log_prob_fn, initial_position):
+        kernel = _build_kernel(
+            target_log_prob_fn, _make_momentum_distribution(initial_position)
+        )
+        results = kernel.bootstrap_results(initial_position)
+
+        # Repack the results data structure into our own
+        chain_state = ChainState(
+            position=initial_position,
+            log_density=results.accepted_results.target_log_prob,
+            log_density_grad=results.accepted_results.grads_target_log_prob,
+        )
+        kernel_state = results
+
+        return chain_state, kernel_state
+
+    def step_fn(target_log_prob_fn, chain_and_kernel_state, seed):
+        chain_state, kernel_state = chain_and_kernel_state
+
+        # Pack kernel state into results here
+        seed = tfp.random.sanitize_seed(seed)
+
+        kernel = _build_kernel(
+            target_log_prob_fn,
+            _make_momentum_distribution(chain_state.position),
+        )
+
+        new_position, results = kernel.one_step(
+            chain_state.position,
+            kernel.bootstrap_results(chain_state.position),
+            seed,
+        )
+
+        info = results
+
+        chain_state = ChainState(
+            position=new_position,
+            log_density=results.accepted_results.target_log_prob,
+            log_density_grad=results.accepted_results.grads_target_log_prob,
+        )
+
+        return (chain_state, results), info
+
+    return SamplingAlgorithm(init_fn, step_fn)

gemlib/mcmc/experimental/hmc_test.py

@@ -0,0 +1,61 @@
+"""Tests for HMC sampler"""
+
+import numpy as np
+import pytest
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+from .hmc import hmc
+from .mcmc_sampler import mcmc
+
+tfd = tfp.distributions
+
+NUM_SAMPLES = 100000
+
+
+@pytest.fixture
+def simple_model():
+    @tfp.distributions.JointDistributionCoroutine
+    def model():
+        yield tfp.distributions.Normal(loc=0.0, scale=1.0, name="foo")
+        yield tfp.distributions.Normal(loc=1.0, scale=1.0, name="bar")
+        yield tfp.distributions.Normal(loc=2.0, scale=1.0, name="baz")
+
+    return model
+
+
+def test_hmc(simple_model):
+    mcmc_init_state = simple_model.sample(seed=[0, 0])
+    algorithm = hmc(step_size=0.1, num_leapfrog_steps=16)
+
+    state = algorithm.init(simple_model.log_prob, mcmc_init_state)
+    new_state, info = algorithm.step(simple_model.log_prob, state, [0, 0])
+
+    tf.nest.assert_same_structure(state, new_state)
+
+
+def test_many_hmc(simple_model):
+    mcmc_init_state = simple_model.sample(seed=[0, 0])
+    algorithm = hmc(step_size=1.2, num_leapfrog_steps=16)
+
+    samples, info = tf.function(
+        lambda: mcmc(
+            NUM_SAMPLES,
+            sampling_algorithm=algorithm,
+            target_density_fn=simple_model.log_prob,
+            initial_position=mcmc_init_state,
+            seed=[0, 0],
+        ),
+        jit_compile=True,
+    )()
+
+    np.testing.assert_allclose(
+        np.array([np.mean(x) for x in samples]),
+        np.array([0.0, 1.0, 2.0]),
+        atol=1e-2,
+    )
+    np.testing.assert_allclose(
+        np.array([np.var(x) for x in samples]),
+        np.array([1.0, 1.0, 1.0]),
+        atol=1e-2,
+    )

gemlib/mcmc/experimental/mcmc_base.py

@@ -0,0 +1,33 @@
+"""Base MCMC datatypes"""
+
+from typing import Callable, NamedTuple, Optional, Tuple
+
+Position = NamedTuple
+KernelInfo = NamedTuple
+
+
+class ChainState(NamedTuple):
+    """Represent the state of an MCMC probability space"""
+
+    position: Position
+    log_density: float
+    log_density_grad: Optional[float] = None
+
+
+class KernelState(NamedTuple):
+    """Represent the state of a stateful MCMC kernel"""
+
+    pass
+
+
+class SamplingAlgorithm(NamedTuple):
+    """Represent a sampling algorithm"""
+
+    init: Callable[[NamedTuple], Tuple[ChainState, KernelState]]
+    step: Callable[
+        [
+            ChainState,
+            Callable[[NamedTuple], float],
+        ],
+        Callable[[ChainState], Tuple[ChainState, KernelInfo]],
+    ]

gemlib/mcmc/experimental/mcmc_sampler.py

@@ -0,0 +1,94 @@
+"""Higher-order functions to run MCMC"""
+
+from functools import partial
+from typing import Any, Callable, Iterable, Tuple
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+from .mcmc_base import SamplingAlgorithm
+
+
+def split_seed(seed, n):
+    n = tf.convert_to_tensor(n)
+    return tfp.random.split_seed(seed, n=n)
+
+
+def _tensor_array_from_element(elem, size):
+    return tf.TensorArray(
+        elem.dtype,
+        size=size,
+        element_shape=elem.shape,
+    )
+
+
+def scan(fn, init, xs):
+    """Scan
+
+    This function is equivalent to
+
+    ```
+    scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])
+    ```
+    """
+    # Set up results accumulator
+    _, initial_trace = fn(init, xs[0])
+
+    flat_initial_trace = tf.nest.flatten(initial_trace, expand_composites=True)
+    trace_arrays = []
+    for trace_elem in flat_initial_trace:
+        trace_arrays.append(
+            _tensor_array_from_element(trace_elem, size=xs.shape[0])
+        )
+
+    def trace_one_step(i, trace_arrays, trace):
+        return [
+            ta.write(i, x)
+            for ta, x in zip(
+                trace_arrays, tf.nest.flatten(trace, expand_composites=True)
+            )
+        ]
+
+    def cond(i, carry, accum):
+        return i < xs.shape[0]
+
+    def body(i, carry, accum):
+        new_carry, result = fn(carry, xs[i])
+        new_accum = trace_one_step(i, accum, result)
+        return i + 1, new_carry, new_accum
+
+    _, final_state, trace_arrays = tf.while_loop(
+        cond=cond, body=body, loop_vars=(0, init, trace_arrays)
+    )
+
+    stacked_trace = tf.nest.pack_sequence_as(
+        initial_trace,
+        [ta.stack() for ta in trace_arrays],
+        expand_composites=True,
+    )
+
+    return final_state, stacked_trace
+
+
+def mcmc(
+    num_samples: int,
+    sampling_algorithm: SamplingAlgorithm,
+    target_density_fn: Callable[[Any, ...], float],
+    initial_position: Iterable,
+    seed: Tuple[int, int],
+):
+    initial_position = tf.nest.map_structure(
+        lambda x: tf.convert_to_tensor(x), initial_position
+    )
+    initial_state = sampling_algorithm.init(target_density_fn, initial_position)
+    kernel_step_fn = partial(sampling_algorithm.step, target_density_fn)
+
+    def one_step(state, rng_key):
+        new_state, info = kernel_step_fn(state, rng_key)
+        return new_state, (new_state[0].position, info)
+
+    keys = split_seed(seed, num_samples)
+
+    _, trace = scan(one_step, initial_state, keys)
+
+    return trace

gemlib/mcmc/experimental/mcmc_sampler_test.py

@@ -0,0 +1,44 @@
+"""Test mcmc_sampler"""
+
+from typing import NamedTuple
+
+import tensorflow as tf
+
+from .mcmc_sampler import mcmc
+from .test_util import CountingKernelInfo, counting_kernel
+
+
+class TestPosition(NamedTuple):
+    x: float
+    y: float
+
+
+NUM_SAMPLES = 100
+
+
+def test_mcmc():
+    sampling_algorithm = counting_kernel()
+
+    initial_position = TestPosition(0.0, -100.0)
+
+    def tlp(x, y):
+        return tf.constant(0.0)
+
+    samples, info = mcmc(
+        num_samples=NUM_SAMPLES,
+        sampling_algorithm=sampling_algorithm,
+        target_density_fn=tlp,
+        initial_position=initial_position,
+        seed=[0, 0],
+    )
+    print(samples)
+    tf.debugging.assert_equal(
+        samples,
+        TestPosition(
+            x=tf.range(1.0, 1.0 + NUM_SAMPLES, delta=1.0),
+            y=tf.range(-99.0, -99.0 + NUM_SAMPLES, delta=1.0),
+        ),
+    )
+    tf.debugging.assert_equal(
+        info, CountingKernelInfo(tf.fill(NUM_SAMPLES, True))
+    )

gemlib/mcmc/experimental/multi_scan.py

@@ -0,0 +1,53 @@
+"""MultiScanKernel calls one_step a number of times on an inner kernel"""
+
+from functools import partial
+
+import tensorflow as tf
+from tensorflow_probability.python.internal import samplers
+
+from .mcmc_base import SamplingAlgorithm
+
+__all__ = ["multi_scan"]
+
+
+def multi_scan(
+    num_updates: int, sampling_algorithm: SamplingAlgorithm
+) -> SamplingAlgorithm:
+    """Performs multiple applications of a kernel
+
+    `sampling_algorithm` is invoked `num_updates` times
+    returning the state and info after the last step.
+
+    Args
+    ----
+    num_updates: integer giving the number of updates
+    sampling_algorithm: an instance of `SamplingAlgorithm`
+
+    Return
+    ------
+    An instance of `SamplingAlgorithm`
+    """
+
+    num_updates_ = tf.convert_to_tensor(num_updates)
+    init_fn = sampling_algorithm.init
+
+    def step_fn(target_log_prob_fn, current_state, seed=None):
+        seed = samplers.sanitize_seed(seed, salt="multi_scan_kernel")
+        seeds = samplers.split_seed(seed, n=num_updates)
+        kernel = partial(sampling_algorithm.step, target_log_prob_fn)
+
+        def body(i, state, _):
+            state, info = kernel(state, seeds[i])
+            return i + 1, state, info
+
+        def cond(i, *_):
+            return i < num_updates_
+
+        init_state, init_info = kernel(current_state, seed)  # unrolled first it
+
+        _, last_state, last_info = tf.while_loop(
+            cond, body, loop_vars=(1, init_state, init_info)
+        )
+        return last_state, last_info
+
+    return SamplingAlgorithm(init_fn, step_fn)

gemlib/mcmc/experimental/multi_scan_test.py

@@ -0,0 +1,71 @@
+"""Test modules for multiscan kernel"""
+
+from typing import NamedTuple
+
+import numpy as np
+import tensorflow as tf
+
+from .mcmc_sampler import mcmc
+from .multi_scan import multi_scan
+from .test_util import CountingKernelInfo, counting_kernel
+
+
+class TestPosition(NamedTuple):
+    x: float
+
+
+def test_one_multi_scan():
+    multi_scan_iterations = 100
+
+    sampler = multi_scan(multi_scan_iterations, counting_kernel())
+
+    def tlp(x):
+        return x
+
+    initial_position = TestPosition(0.0)
+
+    state = sampler.init(tlp, initial_position)
+    (chain_state, kernel_state), info = sampler.step(
+        target_log_prob_fn=tlp, current_state=state, seed=[0, 0]
+    )
+
+    tf.debugging.assert_equal(
+        chain_state.position, np.float32(multi_scan_iterations)
+    )
+    tf.debugging.assert_equal(
+        chain_state.log_density, np.float32(multi_scan_iterations)
+    )
+    tf.debugging.assert_equal(kernel_state.invocation, multi_scan_iterations)
+    tf.debugging.assert_equal(info, CountingKernelInfo(True))
+
+
+def test_many_multi_scan():
+    num_samples = 5
+    multi_scan_iterations = 100
+
+    sampler = multi_scan(multi_scan_iterations, counting_kernel())
+
+    def tlp(x):
+        return x
+
+    initial_position = TestPosition(0.0)
+
+    samples, info = mcmc(
+        num_samples=num_samples,
+        sampling_algorithm=sampler,
+        target_density_fn=tlp,
+        initial_position=initial_position,
+        seed=[0, 0],
+    )
+
+    tf.debugging.assert_equal(
+        samples,
+        tf.range(
+            100.0,
+            100.0 + (num_samples * multi_scan_iterations),
+            delta=multi_scan_iterations,
+        ),
+    )
+    tf.debugging.assert_equal(
+        info, CountingKernelInfo(tf.fill(num_samples, True))
+    )

gemlib/mcmc/experimental/random_walk_metropolis.py

@@ -0,0 +1,107 @@
+"""Implementation of the Random Walk Metropolis algorithm"""
+
+from typing import Callable, NamedTuple, Tuple
+
+import tensorflow_probability as tfp
+
+from .mcmc_base import ChainState, SamplingAlgorithm
+
+
+class RwmhInfo(NamedTuple):
+    """Represents the information about a random walk Metropolis-Hastings (RWMH)
+      step.
+    This can be expanded to include more information in the future (as needed
+      for a specific kernel).
+
+    Attributes
+    ----------
+        is_accepted (bool): Indicates whether the proposal was accepted or not.
+
+    """
+
+    is_accepted: bool
+
+
+class RwmhKernelState(NamedTuple):
+    """Represents the kernel state of an arbitrary MCMC kernel.
+
+    Attributes
+    ----------
+        scale (float): The scale parameter for the RWMH kernel.
+
+    """
+
+    scale: float
+
+
+def rwmh(scale=1.0):
+    def _build_kernel(log_prob_fn):
+        """Partial"""
+        return tfp.mcmc.RandomWalkMetropolis(
+            target_log_prob_fn=log_prob_fn,
+            new_state_fn=tfp.mcmc.random_walk_normal_fn(scale=scale),
+        )
+
+    def init_fn(target_log_prob_fn, target_state):
+        kernel = _build_kernel(target_log_prob_fn)
+        results = kernel.bootstrap_results(target_state)
+
+        chain_state = ChainState(
+            position=target_state,
+            log_density=results.accepted_results.target_log_prob,
+            log_density_grad=(),
+        )
+        kernel_state = RwmhKernelState(scale=scale)
+
+        return chain_state, kernel_state
+
+    def step_fn(
+        target_log_prob_fn: Callable[[NamedTuple], float],
+        target_and_kernel_state: Tuple[ChainState, RwmhKernelState],
+        seed,
+    ) -> Callable[[ChainState], Tuple[ChainState, RwmhInfo]]:
+        """Computation that calls a kernel.
+
+        Args:
+        ----
+        target_log_prob_fn: the conditional log target density/mass function
+        conditioned_state: Parts of the global state that are not updated by
+                           the kernel, but may be needed to instantiate it.
+        chain_and_kernel_state: a tuple containing a ChainState object and
+                        kernel-specific state.
+        target_state: the sub-state on which the kernel operates.
+
+        Returns:
+        -------
+        a tuple of the new target sub-state and information about the sampler
+
+        """
+        # This could be replaced with BlackJAX easily
+        kernel = _build_kernel(target_log_prob_fn)
+
+        target_chain_state, kernel_state = target_and_kernel_state
+
+        new_target_position, results = kernel.one_step(
+            target_chain_state.position,
+            kernel.bootstrap_results(target_chain_state.position),
+            seed=seed,
+        )
+
+        new_chain_and_kernel_state = (
+            ChainState(
+                position=new_target_position,
+                log_density=results.accepted_results.target_log_prob,
+                log_density_grad=(),
+            ),
+            kernel_state,
+        )
+
+        info = (
+            RwmhInfo(
+                results.is_accepted,
+            ),
+        )
+
+        return new_chain_and_kernel_state, info
+
+    return SamplingAlgorithm(init_fn, step_fn)