gemlib 0.9.2__py3-none-any.whl

gemlib/mcmc/adaptive_random_walk_metropolis_test.py

@@ -0,0 +1,144 @@
+# Dependency imports
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python import distributions as tfd
+from tensorflow_probability.python.internal import test_util
+
+from gemlib.mcmc.adaptive_random_walk_metropolis import (
+    AdaptiveRandomWalkMetropolis,
+)
+
+
+@test_util.test_all_tf_execution_regimes
+class TestAdaptiveRandomWalkMetropolis(test_util.TestCase):
+    def test_1d_normal(self):
+        """Sample from Standard Normal Distribution."""
+        dtype = np.float32
+
+        target = tfd.Normal(loc=dtype(0), scale=dtype(1))
+
+        kernel = AdaptiveRandomWalkMetropolis(
+            target_log_prob_fn=target.log_prob,
+            target_accept_ratio=0.44,
+            initial_covariance=dtype(0.001),
+        )
+        samples = tfp.mcmc.sample_chain(
+            num_results=2000,
+            current_state=dtype([0.1]),
+            kernel=kernel,
+            num_burnin_steps=500,
+            trace_fn=None,
+        )
+
+        sample_mean = tf.math.reduce_mean(samples, axis=0)
+        sample_std = tf.math.reduce_std(samples, axis=0)
+        [sample_mean_, sample_std_] = self.evaluate([sample_mean, sample_std])
+
+        self.assertAllClose([0.0], sample_mean_, atol=0.17, rtol=0.0)
+        self.assertAllClose([1.0], sample_std_, atol=0.2, rtol=0.0)
+
+    def test_3d_mvn(self):
+        """Sample from 3-variate Gaussian Distribution."""
+        dtype = np.float32
+
+        true_mean = dtype([1.0, 2.0, 3.0])
+        true_cov = dtype(
+            [[0.36, 0.12, 0.06], [0.12, 0.29, -0.13], [0.06, -0.13, 0.26]]
+        )
+        target = tfd.MultivariateNormalTriL(
+            loc=true_mean, scale_tril=tf.linalg.cholesky(true_cov)
+        )
+        kernel = AdaptiveRandomWalkMetropolis(
+            target_log_prob_fn=target.log_prob,
+            initial_covariance=dtype(0.001) * np.eye(3, dtype=dtype),
+        )
+        samples = tfp.mcmc.sample_chain(
+            num_results=2000,
+            current_state=dtype([0.1, 0.1, 0.1]),
+            kernel=kernel,
+            num_burnin_steps=500,
+            trace_fn=None,
+        )
+
+        sample_mean = tf.math.reduce_mean(samples, axis=0)
+        [sample_mean_] = self.evaluate([sample_mean])
+        self.assertAllClose(sample_mean_, true_mean, atol=0.1, rtol=0.1)
+
+        sample_cov = tfp.stats.covariance(samples)
+        sample_cov_ = self.evaluate(sample_cov)
+        self.assertAllClose(sample_cov_, true_cov, atol=0.1, rtol=0.1)
+
+    def test_float64(self):
+        """Sample with dtype float64."""
+        dtype = np.float64
+
+        target = tfd.Normal(loc=dtype(0), scale=dtype(1))
+
+        kernel = AdaptiveRandomWalkMetropolis(
+            target_log_prob_fn=target.log_prob,
+            target_accept_ratio=0.44,
+            initial_covariance=dtype(0.001),
+        )
+        tfp.mcmc.sample_chain(
+            num_results=20,
+            current_state=dtype([0.1]),
+            kernel=kernel,
+            trace_fn=None,
+        )
+
+    def test_bijector(self):
+        """Employ bijector when sampling."""
+        dtype = np.float32
+
+        target = tfd.Normal(loc=dtype(0), scale=dtype(1))
+
+        kernel = AdaptiveRandomWalkMetropolis(
+            target_log_prob_fn=target.log_prob,
+            target_accept_ratio=0.44,
+            initial_covariance=dtype(0.001),
+        )
+        kernel = tfp.mcmc.TransformedTransitionKernel(
+            inner_kernel=kernel, bijector=tfp.bijectors.Exp()
+        )
+        tfp.mcmc.sample_chain(
+            num_results=20,
+            current_state=dtype([0.1]),
+            kernel=kernel,
+            trace_fn=None,
+        )
+
+    def test_is_calibrated(self):
+        dtype = np.float32
+        kernel = AdaptiveRandomWalkMetropolis(
+            target_log_prob_fn=lambda x: -tf.square(x) / 2.0,
+            initial_covariance=dtype(0.001),
+        )
+        self.assertTrue(kernel.is_calibrated)
+
+
+# *** RUN TESTS ***
+
+# ** Method 1 **
+# When using tf.test.main() from within jupyter notebook use
+#   import sys
+#   sys.argv = sys.argv[:1]
+# to eliminate the ERROR "unknown command line flag 'f"
+# if __name__ == '__main__':
+#    try:
+#        tf.test.main()
+#    except SystemExit as inst:
+#        if inst.args[0] is True: # raised by sys.exit(True) if tests fail
+#            raise
+
+# ** Method 2 **
+# Althernatively use unittest.main() instead of tf.test.main()
+# when running tests within jupyter notebook:
+#   import unittest
+#   unittest.main(argv=['first-arg-is-ignored'], exit=False)
+
+# ** Method 3 **
+# If not using jupyter notebook (or similar environment) the
+# following should work:
+if __name__ == "__main__":
+    tf.test.main()

gemlib/mcmc/brownian_bridge_kernel.py

@@ -0,0 +1,291 @@
+"""A Brownian Bridge kernel is intended to operate on
+a timeseries
+"""
+
+# ruff: noqa: B023
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import samplers
+from tensorflow_probability.python.mcmc.internal import util as mcmc_util
+from tensorflow_probability.python.mcmc.random_walk_metropolis import (
+    UncalibratedRandomWalkResults,
+)
+
+from gemlib.distributions import BrownianBridge, BrownianMotion, UniformInteger
+
+tfd = tfp.distributions
+mcmc = tfp.mcmc
+
+MIN_SPAN = 3
+
+
+def _slide_left(x, shift):
+    x_right = x[..., -1:]
+    y = tf.roll(x, -shift, axis=-1)
+    mask = tf.range(x.shape[-1]) >= (x.shape[-1] - shift)
+    pad = x_right * tf.cast(mask, x_right.dtype)
+    return y * tf.cast(~mask, y.dtype) + pad
+
+
+def _slide_right(x, shift):
+    x_left = x[..., 0:]
+    y = tf.roll(x, shift, axis=-1)
+    mask = tf.range(x.shape[-1]) < shift
+    pad = x_left * tf.cast(mask, x_left.dtype)
+    return y * tf.cast(~mask, x_left.dtype) + pad
+
+
+class UncalibratedBrownianBridgeKernel(mcmc.TransitionKernel):
+    def __init__(
+        self,
+        target_log_prob_fn,
+        index_points,
+        span=3,
+        scale=0.1,
+        left=True,
+        right=True,
+        name="UncalibratedBrownianBridgeKernel",
+    ):
+        with tf.name_scope(
+            mcmc_util.make_name(
+                name, "UncalibratedBrownianBridgeKernel", "__init__"
+            )
+        ) as name:
+            if span < MIN_SPAN:
+                raise ValueError(
+                    f"`span` must be at least {MIN_SPAN} timepoints"
+                )
+            if scale <= 0.0:
+                raise ValueError("`scale` must be positive")
+
+            span_parts = list(span) if mcmc_util.is_list_like(span) else [span]
+            self.span_parts = [
+                tf.convert_to_tensor(s, name="span") for s in span_parts
+            ]
+
+            if mcmc_util.is_list_like(scale):
+                scale_parts = list(scale)
+            else:
+                scale_parts = [scale]
+            self.scale_parts = [
+                tf.convert_to_tensor(s, name="scale") for s in scale_parts
+            ]
+
+            self._index_points = tf.convert_to_tensor(index_points)
+            self._left = tf.cast(left, tf.int32)
+            self._right = tf.cast(right, tf.int32)
+
+            self.dtype = self.scale_parts[0].dtype
+            cls_name = mcmc_util.make_name(
+                name, "UncalibratedBrownianBridgeKernel", ""
+            )
+
+            self._parameters = {
+                "target_log_prob_fn": target_log_prob_fn,
+                "index_points": index_points,
+                "span": span,
+                "scale": scale,
+                "left": left,
+                "right": right,
+                "name": cls_name,
+            }
+
+    @property
+    def is_calibrated(self):
+        return False
+
+    @property
+    def name(self):
+        return self._parameters["name"]
+
+    @property
+    def span(self):
+        return self._parameters["span"]
+
+    @property
+    def scale(self):
+        return self._parameters["scale"]
+
+    @property
+    def jitter(self):
+        return self._parameters["jitter"]
+
+    @property
+    def target_log_prob_fn(self):
+        return self._parameters["target_log_prob_fn"]
+
+    def one_step(self, current_state, previous_results, seed=None):
+        with tf.name_scope(mcmc_util.make_name(self.name, "bbmh", "one_step")):
+            with tf.name_scope("initialize"):
+                if mcmc_util.is_list_like(current_state):
+                    current_state_parts = list(current_state)
+                else:
+                    current_state_parts = [current_state]
+                current_state_parts = [
+                    tf.convert_to_tensor(s, name="current_state")
+                    for s in current_state_parts
+                ]
+            seed = samplers.sanitize_seed(
+                seed, salt="UncalibratedBrownianBridgeKernel"
+            )
+
+            new_state_parts = []
+            log_acceptance_correction_parts = []
+            for current_state_part, span_part, scale_part in zip(
+                current_state_parts,
+                self.span_parts,
+                self.scale_parts,
+            ):
+                t_low_seed, bridge_seed = samplers.split_seed(seed)
+
+                # Evaluate bridge limits
+                t_low = UniformInteger(
+                    0 - self._left,
+                    current_state_part.shape[-1] - span_part + self._right,
+                ).sample(seed=seed)
+
+                # We have 3 cases:
+                #  0. If t_low > 0 and t_high < (current_state.shape[-1]-1):
+                #     Brownian Bridge
+                #  1. If t_low == 0: reverse Brownian motion
+                #  2. If t_high >= (current_state.shape[-1]-1): Brownian motion
+                def brownian_bridge_proposal():
+                    with tf.name_scope("brownian_bridge_proposal"):
+                        indices = t_low + tf.range(span_part)
+                        current_bridge = tf.gather(
+                            current_state_part,
+                            indices=indices,
+                        )
+                        bridge = BrownianBridge(
+                            index_points=tf.gather(self._index_points, indices),
+                            x0=current_bridge[..., 0],
+                            x1=current_bridge[..., -1],
+                            scale=scale_part,
+                        )
+                        new_bridge = bridge.sample(seed=bridge_seed)
+                        log_acceptance_correction = bridge.log_prob(
+                            current_bridge[1:-1]
+                        ) - bridge.log_prob(new_bridge)
+
+                        new_state = tf.tensor_scatter_nd_update(
+                            current_state_part,
+                            indices=indices[1:-1][:, tf.newaxis],
+                            updates=new_bridge,
+                            name="update_new_state",
+                        )
+                        return new_state, log_acceptance_correction
+
+                def brownian_motion_right_proposal():
+                    with tf.name_scope("brownian_motion_proposal"):
+                        indices = tf.range(
+                            current_state_part.shape[0] - span_part,
+                            current_state_part.shape[0],
+                        )  # Index into current_state_part
+                        current_bridge = tf.gather(
+                            current_state_part,
+                            indices=indices,
+                            name="current_state_slice",
+                        )
+                        bridge = BrownianMotion(
+                            index_points=tf.gather(self._index_points, indices),
+                            x0=current_bridge[..., 0],
+                            scale=scale_part,
+                        )
+                        new_bridge = bridge.sample(seed=bridge_seed)
+                        log_acceptance_correction = bridge.log_prob(
+                            current_bridge[1:]
+                        ) - bridge.log_prob(new_bridge)
+                        new_state = tf.tensor_scatter_nd_update(
+                            current_state_part,
+                            indices=tf.expand_dims(indices[1:], -1),
+                            updates=new_bridge,
+                            name="update_new_state",
+                        )
+                        return new_state, log_acceptance_correction
+
+                def brownian_motion_left_proposal():
+                    with tf.name_scope("brownian_motion_left_proposal"):
+                        indices = tf.range(span_part)
+                        current_bridge = tf.gather(
+                            current_state_part,
+                            indices=indices,
+                            name="current_state_slice",
+                        )
+                        bridge = BrownianMotion(
+                            index_points=tf.gather(self._index_points, indices),
+                            x0=current_bridge[..., -1],
+                            scale=scale_part,
+                        )
+                        new_bridge = bridge.sample(seed=bridge_seed)
+                        log_acceptance_correction = bridge.log_prob(
+                            tf.reverse(
+                                current_bridge[:-1],
+                                axis=[-1],
+                                name="reverse_current_bridge",
+                            ),
+                        ) - bridge.log_prob(new_bridge)
+                        new_state = tf.tensor_scatter_nd_update(
+                            current_state_part,
+                            indices=tf.expand_dims(indices[:-1], -1),
+                            updates=tf.reverse(
+                                new_bridge, axis=[-1], name="reverse_new_bridge"
+                            ),
+                            name="update_new_state",
+                        )
+                        return new_state, log_acceptance_correction
+
+                case_enum = (  # 0=bridge, 1=right, 2=left
+                    tf.cast(
+                        t_low == (current_state_part.shape[0] - span_part),
+                        tf.int32,
+                    )
+                    + tf.cast(t_low == -1, tf.int32) * 2
+                )
+
+                (
+                    new_state,
+                    log_acceptance_correction_part,
+                ) = tf.switch_case(
+                    case_enum,
+                    [
+                        brownian_bridge_proposal,
+                        brownian_motion_right_proposal,
+                        brownian_motion_left_proposal,
+                    ],
+                )
+
+                new_state_parts.append(new_state)
+                log_acceptance_correction_parts.append(
+                    log_acceptance_correction_part
+                )
+
+            target_log_prob = self.target_log_prob_fn(*new_state_parts)
+
+            def maybe_flatten(x):
+                return x if mcmc_util.is_list_like(current_state) else x[0]
+
+            return [
+                maybe_flatten(new_state_parts),
+                UncalibratedRandomWalkResults(
+                    log_acceptance_correction=maybe_flatten(
+                        log_acceptance_correction_parts
+                    ),
+                    target_log_prob=target_log_prob,
+                    seed=seed,
+                ),
+            ]
+
+    def bootstrap_results(self, current_state):
+        if mcmc_util.is_list_like(current_state):
+            current_state_parts = list(current_state)
+        else:
+            current_state_parts = [current_state]
+
+        init_target_log_prob = self.target_log_prob_fn(*current_state_parts)
+
+        return UncalibratedRandomWalkResults(
+            log_acceptance_correction=tf.zeros_like(init_target_log_prob),
+            target_log_prob=init_target_log_prob,
+            seed=samplers.zeros_seed(),
+        )

gemlib/mcmc/brownian_bridge_kernel_test.py

@@ -0,0 +1,164 @@
+"""Tests Brownian Bridge kernel"""
+
+import os
+import pickle as pkl
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import test_util
+
+from gemlib.distributions import BrownianMotion
+from gemlib.mcmc.brownian_bridge_kernel import UncalibratedBrownianBridgeKernel
+
+tfd = tfp.distributions
+
+DTYPE = tf.float64
+
+
+def model_fixture():
+    """Fixture from model below"""
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+    with open(os.path.join(dir_path, "bb_fixture.pkl"), "rb") as f:
+        return pkl.load(f)
+
+
+class TestBrownianBridgeKernel(test_util.TestCase):
+    def test_simple_brownian_motion(self):
+        x = tf.range(0.0, 10.0, 0.1, dtype=DTYPE)
+        Y = BrownianMotion(x)
+        y = Y.sample()
+
+        kernel = tfp.mcmc.MetropolisHastings(
+            inner_kernel=UncalibratedBrownianBridgeKernel(
+                Y.log_prob,
+                index_points=x,
+                span=90,
+                scale=tf.constant(1.0, DTYPE),
+            )
+        )
+
+        # kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
+        #     inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
+        #         target_log_prob_fn=Y.log_prob,
+        #         num_leapfrog_steps=3,
+        #         step_size=0.1,
+        #     ),
+        #     num_adaptation_steps=500,
+        # )
+
+        samples, results = tf.function(
+            lambda: tfp.mcmc.sample_chain(
+                num_results=10000, kernel=kernel, current_state=y
+            )
+        )()
+
+        print(
+            "Acceptance rate:",
+            tf.reduce_mean(tf.cast(results.is_accepted, tf.float32)),
+        )
+
+        # fig, ax = plt.subplots(1, 2)
+        # ax[0].plot(
+        #     x[1:],
+        #     samples.numpy().T,
+        #     color="lightblue",
+        #     alpha=0.3,
+        # )
+        # ax[0].plot(
+        #     x[1:],
+        #     y,
+        #     "o",
+        #     color="black",
+        # )
+        # ax[1].plot(samples[:, 75])
+        # plt.show()
+
+        self.assertAllClose(0.0, np.mean(samples[:, 0]), atol=1.0, rtol=0.1)
+        self.assertAllClose(10.0, np.var(samples[:, -1]), atol=1.0, rtol=0.1)
+
+    def test_poisson_with_brownian_mean(self):
+        x = tf.range(0.0, 10.0, 0.1, dtype=DTYPE)
+
+        def model():
+            mu0 = tfd.Normal(
+                loc=tf.constant(1.0, DTYPE),
+                scale=tf.constant(1.0, DTYPE),
+            )
+
+            def mu(mu0):
+                return BrownianMotion(x, x0=mu0)
+
+            def y(mu):
+                rate = tf.concat([[0.0], mu], axis=-1)
+                return tfd.Independent(
+                    tfd.Poisson(rate=tf.math.exp(rate)),
+                    reinterpreted_batch_ndims=1,
+                )
+
+            return tfd.JointDistributionNamed({"mu0": mu0, "mu": mu, "y": y})
+
+        model = model()
+        trial = model_fixture()
+
+        def logp(mu):
+            return model.log_prob(
+                {"mu0": trial["mu0"], "mu": mu, "y": trial["y"]}
+            )
+
+        mcmc_kernel = tfp.mcmc.MetropolisHastings(
+            inner_kernel=UncalibratedBrownianBridgeKernel(
+                logp,
+                index_points=x,
+                span=5,
+                scale=tf.constant(1.0, DTYPE),
+                left=True,
+                right=True,
+            )
+        )
+
+        # mcmc_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
+        #     inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
+        #         target_log_prob_fn=logp, num_leapfrog_steps=3, step_size=1.0
+        #     ),
+        #     num_adaptation_steps=500,
+        # )
+
+        samples, results = tf.function(
+            lambda: tfp.mcmc.sample_chain(
+                num_results=5000,
+                kernel=mcmc_kernel,
+                current_state=tf.fill(
+                    trial["mu"].shape, tf.constant(3.0, DTYPE)
+                ),
+            )
+        )()
+        print(
+            "Acceptance rate:",
+            tf.reduce_mean(tf.cast(results.is_accepted, tf.float32)),
+        )
+
+        # fig, ax = plt.subplots(1, 2)
+        # ax[0].plot(
+        #     np.exp(samples[500:].numpy().T),
+        #     color="lightblue",
+        #     alpha=0.3,
+        # )
+        # ax[0].plot(
+        #     np.exp(trial["mu"]),
+        #     "o",
+        #     color="black",
+        # )
+        # ax[1].plot(np.exp(samples[:, 75]))
+        # plt.show()
+
+        self.assertAllClose(
+            0.0,
+            0.0,
+            rtol=1.5,
+            atol=2.0,
+        )
+
+
+if __name__ == "__main__":
+    TestBrownianBridgeKernel().test_poisson_with_brownian_mean()