gemlib 0.9.2__py3-none-any.whl

gemlib/mcmc/experimental/random_walk_metropolis_test.py

@@ -0,0 +1,214 @@
+"""Test the random walk metropolis kernel"""
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+from .composable_kernel import Step
+from .mcmc_sampler import mcmc
+from .random_walk_metropolis import RwmhInfo, rwmh
+
+NUM_SAMPLES = 100000
+
+
+def split_seed(seed, n):
+    n = tf.convert_to_tensor(n)
+    return tfp.random.split_seed(seed, n=n)
+
+
+def tree_map(fn, *args):
+    return tf.nest.map_structure(fn, *args)
+
+
+def tree_flatten(tree):
+    return tf.nest.flatten(tree)
+
+
+def get_seed():
+    # jax.random.PRNGKey(42)
+    return [0, 0]
+
+
+@tfp.distributions.JointDistributionCoroutine
+def simple_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")
+
+
+def test_rwmh_1kernel():
+    seed = get_seed()
+
+    initial_position = simple_model.sample(seed=seed)
+
+    kernel = rwmh(scale=0.3)
+
+    state = kernel.init(simple_model.log_prob, initial_position)
+    new_state, results = kernel.step(simple_model.log_prob, state, seed)
+
+    assert tree_map(lambda x, y: None, new_state, state)
+
+    expected_results = (RwmhInfo(is_accepted=True),)
+    assert tree_map(lambda x, y: x == y, results, expected_results)
+
+
+def test_rwmh_2kernel():
+    seed = get_seed()
+
+    initial_position = simple_model.sample(seed=seed)
+
+    kernel = Step(rwmh(scale=0.3), ["foo"]) >> Step(
+        rwmh(scale=0.1), ["bar", "baz"]
+    )
+
+    state = kernel.init(simple_model.log_prob, initial_position)
+    new_state, results = kernel.step(simple_model.log_prob, state, seed)
+
+    assert tree_map(lambda x, y: None, new_state, state)
+
+    expected_results = (RwmhInfo(is_accepted=True), RwmhInfo(is_accepted=True))
+    assert tree_map(lambda x, y: x == y, results, expected_results)
+
+
+def test_rwmh_3kernel():
+    seed = get_seed()
+
+    initial_position = simple_model.sample(seed=seed)
+
+    kernel = (
+        Step(rwmh(scale=0.3), ["foo"])
+        >> Step(rwmh(scale=0.1), ["bar"])
+        >> Step(rwmh(scale=0.2), ["baz"])
+    )
+
+    state = kernel.init(simple_model.log_prob, initial_position)
+    new_state, results = kernel.step(simple_model.log_prob, state, seed)
+
+    assert tree_map(lambda x, y: None, new_state, state)
+
+    expected_results = (
+        RwmhInfo(is_accepted=True),
+        RwmhInfo(is_accepted=True),
+        RwmhInfo(is_accepted=True),
+    )
+    assert tree_map(lambda x, y: x == y, results, expected_results)
+
+
+def test_rwmh_1kernel_mcmc():
+    seed = get_seed()
+
+    initial_position = simple_model.sample(seed=seed)
+
+    kernel = rwmh(scale=1.8)
+
+    posterior, info = tf.function(
+        lambda: mcmc(
+            NUM_SAMPLES,
+            sampling_algorithm=kernel,
+            target_density_fn=simple_model.log_prob,
+            initial_position=initial_position,
+            seed=get_seed(),
+        ),
+        jit_compile=True,
+    )()
+
+    # Test results
+    np.testing.assert_approx_equal(
+        np.mean(info[0].is_accepted), 0.23, significant=1
+    )
+    np.testing.assert_allclose(
+        tree_map(lambda x: np.mean(x), posterior),
+        [0.0, 1.0, 2.0],
+        rtol=0.01,
+        atol=0.05,
+    )
+    np.testing.assert_allclose(
+        tree_map(lambda x: np.var(x), posterior),
+        [1.0, 1.0, 1.0],
+        rtol=0.01,
+        atol=0.05,
+    )
+
+
+def test_rwmh_2kernel_mcmc():
+    seed = get_seed()
+
+    initial_position = simple_model.sample(seed=seed)
+
+    kernel = Step(rwmh(scale=2.3), ["foo"]) >> Step(
+        rwmh(scale=1.8), ["bar", "baz"]
+    )
+
+    posterior, info = tf.function(
+        lambda: mcmc(
+            NUM_SAMPLES,
+            sampling_algorithm=kernel,
+            target_density_fn=simple_model.log_prob,
+            initial_position=initial_position,
+            seed=get_seed(),
+        ),
+        jit_compile=True,
+    )()
+
+    # Test results
+    np.testing.assert_allclose(
+        tree_flatten(tree_map(lambda x: np.mean(x), info)),
+        [0.45, 0.33],
+        atol=0.01,
+        rtol=0.05,
+    )
+    np.testing.assert_allclose(
+        tree_map(lambda x: np.mean(x), posterior),
+        [0.0, 1.0, 2.0],
+        rtol=0.01,
+        atol=0.05,
+    )
+    np.testing.assert_allclose(
+        tree_map(lambda x: np.var(x), posterior),
+        [1.0, 1.0, 1.0],
+        rtol=0.01,
+        atol=0.05,
+    )
+
+
+def test_rwmh_3kernel_mcmc():
+    seed = get_seed()
+
+    initial_position = simple_model.sample(seed=seed)
+
+    kernel = (
+        Step(rwmh(scale=2.3), ["foo"])
+        >> Step(rwmh(scale=2.3), ["bar"])
+        >> Step(rwmh(scale=2.3), ["baz"])
+    )
+
+    posterior, info = tf.function(
+        lambda: mcmc(
+            NUM_SAMPLES,
+            sampling_algorithm=kernel,
+            target_density_fn=simple_model.log_prob,
+            initial_position=initial_position,
+            seed=get_seed(),
+        ),
+        jit_compile=True,
+    )()
+
+    # Test results
+    np.testing.assert_allclose(
+        tree_flatten(tree_map(lambda x: np.mean(x), info))[0],
+        [0.45, 0.45, 0.45],
+        atol=0.01,
+        rtol=0.05,
+    )
+    np.testing.assert_allclose(
+        tree_map(lambda x: np.mean(x), posterior),
+        [0.0, 1.0, 2.0],
+        rtol=0.01,
+        atol=0.05,
+    )
+    np.testing.assert_allclose(
+        tree_map(lambda x: np.var(x), posterior),
+        [1.0, 1.0, 1.0],
+        rtol=0.01,
+        atol=0.05,
+    )

gemlib/mcmc/experimental/test_util.py

@@ -0,0 +1,49 @@
+"""Simple counting kernel for testing"""
+
+from typing import NamedTuple
+
+import tensorflow as tf
+
+from .mcmc_base import ChainState, SamplingAlgorithm
+
+__all__ = ["CountingKernelInfo", "CountingKernelState", "counting_kernel"]
+
+
+class CountingKernelState(NamedTuple):
+    invocation: int
+
+
+class CountingKernelInfo(NamedTuple):
+    is_accepted: bool
+
+
+def counting_kernel():
+    def init_fn(target_log_prob_fn, position):
+        chain_state = ChainState(
+            position=position,
+            log_density=target_log_prob_fn(**position._asdict()),
+            log_density_grad=(),
+        )
+        kernel_state = CountingKernelState(tf.constant(0))
+
+        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
+
+        new_position = chain_state.position.__class__(
+            **{k: v + 1.0 for k, v in chain_state.position._asdict().items()}
+        )
+
+        new_chain_state = ChainState(
+            position=new_position,
+            log_density=target_log_prob_fn(**new_position._asdict()),
+            log_density_grad=(),
+        )
+        new_kernel_state = CountingKernelState(kernel_state.invocation + 1)
+
+        return (new_chain_state, new_kernel_state), CountingKernelInfo(
+            tf.constant(True)
+        )
+
+    return SamplingAlgorithm(init_fn, step_fn)