gemlib 0.9.2__py3-none-any.whl

gemlib/distributions/kcategorical.py

@@ -0,0 +1,113 @@
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import (
+    reparameterization,
+    tensorshape_util,
+)
+from tensorflow_probability.python.internal.tensor_util import (
+    convert_nonref_to_tensor,
+)
+
+tfd = tfp.distributions
+
+
+def _log_choose(N, k):  # noqa: N803
+    return (
+        tf.math.lgamma(N + 1.0)
+        - tf.math.lgamma(k + 1.0)
+        - tf.math.lgamma(N - k + 1.0)
+    )
+
+
+class UniformKCategorical(tfd.Distribution):
+    def __init__(
+        self,
+        k,
+        mask,
+        float_dtype=tf.float32,
+        validate_args=False,
+        allow_nan_stats=True,
+        name="UniformKCategorical",
+    ):
+        """Uniform K-Categorical distribution.
+
+        Given a set of items indexed $1,...,n$ and a boolean mask of the same
+            shape sample $k$ indices without replacement.
+
+        :param k: the number of indices to sample
+        :param mask: a boolean mask with `True` where an element is valid,
+                     otherwise `False`
+        :param validate_args: Whether to validate args
+        :param allow_nan_stats: allow nan stats
+        :param name: name of the distribution
+
+        Example 1: Generate 4 samples of size k given a mask
+            import numpy as np
+            import tensorflow as tf
+            import tensorflow_probability as tfp
+            from gemlib.distributions.kcategorical import UniformKCategorical
+
+            # Mask determines which indices are valid and returned by sample().
+            # Below combinations of the indices 0, 3, 4, and 6 will be realised
+            # when sampling.
+            mask = [True, False, False, True, True, False, True]
+            X = UniformKCategorical(k=3, mask=mask)
+            x = X.sample(4)
+            tf.print(x)
+
+        Example 2: Probability of a given sample
+            import numpy as np
+            import tensorflow as tf
+            import tensorflow_probability as tfp
+            from gemlib.distributions.kcategorical import UniformKCategorical
+
+            # Probability of drawing an unordered sample of
+            # size k from N items where N equals the number
+            # of True states in the mask determines N.
+            mask = [True, False, False, True, True, False, True]
+            sample = tf.convert_to_tensor([4, 3, 6])
+            X = UniformKCategorical(k=sample.shape[-1], mask=mask)
+            lp = X.log_prob(sample)
+            tf.print('prob:', tf.exp(lp))
+
+        """
+        parameters = dict(locals())
+        self._mask = convert_nonref_to_tensor(mask, dtype_hint=tf.bool)
+        self._k = convert_nonref_to_tensor(k, dtype_hint=tf.int32)
+        self._float_dtype = float_dtype
+        dtype = self._k.dtype
+
+        with tf.name_scope(name) as name:
+            super().__init__(
+                dtype=dtype,
+                reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
+                validate_args=validate_args,
+                allow_nan_stats=allow_nan_stats,
+                parameters=parameters,
+                name=name,
+            )
+
+    def _batch_shape(self):
+        return tf.TensorShape(self._mask.shape[:-1])
+
+    def _event_shape(self):
+        return tensorshape_util.constant_value_as_shape(
+            tf.expand_dims(self._k, axis=0)
+        )
+
+    def _sample_n(self, n, seed=None):
+        seed = tfp.random.sanitize_seed(seed, salt="KCategorical._sample_n")
+        u = tfd.Uniform(
+            low=tf.zeros(self._mask.shape, dtype=tf.float32),
+            high=tf.ones(self._mask.shape, dtype=tf.float32),
+        ).sample(n, seed=seed)
+        u = u * tf.cast(self._mask, u.dtype)
+        _, x = tf.math.top_k(u, k=self._k, sorted=True)
+        return x
+
+    def _log_prob(self, x):
+        N = tf.math.count_nonzero(self._mask, axis=-1)
+        return -_log_choose(
+            tf.cast(N, dtype=self._float_dtype),
+            tf.cast(self._k, dtype=self._float_dtype),
+        )

gemlib/distributions/kcategorical_test.py

@@ -0,0 +1,52 @@
+# Dependency imports
+import numpy as np
+import tensorflow as tf
+from tensorflow_probability.python.internal import test_util
+
+from gemlib.distributions.kcategorical import UniformKCategorical
+
+
+@test_util.test_all_tf_execution_regimes
+class TestUniformInteger(test_util.TestCase):
+    def setUp(self):
+        self.seed = 10402302
+        self.mask = [True, False, False, True, True, False, True]
+
+    def test_sample(self):
+        """Sample draws one sample with shape (1,3) ."""
+        tf.random.set_seed(self.seed)
+        target = tf.convert_to_tensor([[0, 6, 3]])
+        X = UniformKCategorical(k=target.shape[-1], mask=self.mask)
+        x = X.sample(1, seed=1234)
+        x_ = self.evaluate(x)
+        target_ = self.evaluate(target)
+        self.assertAllEqual(target_, x_)
+        self.assertDTypeEqual(x_, np.int32)
+
+    def test_log_prob_float32(self):
+        """Log probability of 1 realisations using float32."""
+        target = tf.convert_to_tensor([[0, 6, 3]])
+        X = UniformKCategorical(
+            k=target.shape[-1], mask=self.mask, float_dtype=tf.float32
+        )
+        lp = X.log_prob(target)
+        lp_ = self.evaluate(lp)
+        print("32", lp)
+        self.assertAlmostEqual(lp_, -1.3862944, places=5)
+        self.assertDTypeEqual(lp_, np.float32)
+
+    def test_log_prob_float64(self):
+        """Log probability of 1 realisations using float32."""
+        target = tf.convert_to_tensor([[0, 6, 3]])
+        X = UniformKCategorical(
+            k=target.shape[-1], mask=self.mask, float_dtype=tf.float64
+        )
+        lp = X.log_prob(target)
+        lp_ = self.evaluate(lp)
+        print("64", lp)
+        self.assertAlmostEqual(lp_, -1.3862944, places=5)
+        self.assertDTypeEqual(lp_, np.float64)
+
+
+if __name__ == "__main__":
+    tf.test.main()

gemlib/distributions/uniform_integer.py

@@ -0,0 +1,169 @@
+"""The UniformInteger distribution class"""
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import (
+    reparameterization,
+    samplers,
+)
+
+tfd = tfp.distributions
+
+
+class UniformInteger(tfd.Distribution):
+    def __init__(
+        self,
+        low=0,
+        high=1,
+        validate_args=False,
+        allow_nan_stats=True,
+        dtype=tf.int32,
+        float_dtype=tf.float32,
+        name="UniformInteger",
+    ):
+        """Integer uniform distribution.
+
+        Args:
+        ----
+          low: Integer tensor, lower boundary of the output interval. Must have
+            `low <= high`.
+          high: Integer tensor, _inclusive_ upper boundary of the output
+            interval.  Must have `low <= high`.
+          validate_args: Python `bool`, default `False`. When `True`
+            distribution parameters are checked for validity despite possibly
+            degrading runtime performance. When `False` invalid inputs may
+            silently render incorrect outputs.
+          allow_nan_stats: Python `bool`, default `True`. When `True`,
+           statistics (e.g., mean, mode, variance) use the value "`NaN`" to
+           indicate the result is undefined. When `False`, an exception is
+           raised if one or more of the statistic's batch members are undefined.
+          dtype: returned integer dtype when sampling.
+          float_dtype: returned float dtype of log probability.
+          name: Python `str` name prefixed to Ops created by this class.
+
+        Example 1: sampling
+        ```python
+        import tensorflow as tf
+        from gemlib.distributions.uniform_integer import UniformInteger
+
+        tf.random.set_seed(10402302)
+        X = UniformInteger(0, 10, dtype=tf.int32)
+        x = X.sample([3, 3], seed=1)
+        tf.print("samples:", x, "=", [[8, 4, 8], [2, 7, 9], [6, 0, 9]])
+        ```
+
+        Example 2: log probability
+        ```python
+        import tensorflow as tf
+        from gemlib.distributions.uniform_integer import UniformInteger
+
+        X = UniformInteger(0, 10, float_dtype=tf.float32)
+        lp = X.log_prob([[8, 4, 8], [2, 7, 9], [6, 0, 9]])
+        total_lp = tf.math.round(tf.math.reduce_sum(lp) * 1e5) / 1e5
+        tf.print("total lp:", total_lp, "= -20.72327")
+        ```
+
+        Raises:
+        ------
+          InvalidArgument if `low > high` and `validate_args=False`.
+
+        """
+        parameters = dict(locals())
+        with tf.name_scope(name) as name:
+            self._low = tf.cast(low, name="low", dtype=dtype)
+            self._high = tf.cast(high, name="high", dtype=dtype)
+            super().__init__(
+                dtype=dtype,
+                reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
+                validate_args=validate_args,
+                allow_nan_stats=allow_nan_stats,
+                parameters=parameters,
+                name=name,
+            )
+        self.float_dtype = float_dtype
+        if validate_args is True:
+            tf.assert_greater(
+                self._high, self._low, "Condition low < high failed"
+            )
+
+    @staticmethod
+    def _param_shapes(sample_shape):
+        return dict(
+            zip(
+                ("low", "high"),
+                ([tf.convert_to_tensor(sample_shape, dtype=tf.int32)] * 2),
+            )
+        )
+
+    @classmethod
+    def _params_event_ndims(cls):
+        return {"low": 0, "high": 0}
+
+    @property
+    def low(self):
+        """Lower boundary of the output interval."""
+        return self._low
+
+    @property
+    def high(self):
+        """Upper boundary of the output interval."""
+        return self._high
+
+    def range(self, name="range"):
+        """`high - low`."""
+        with self._name_and_control_scope(name):
+            return self._range()
+
+    def _range(self, low=None, high=None):
+        low = self.low if low is None else low
+        high = self.high if high is None else high
+        return high - low
+
+    def _batch_shape_tensor(self, low=None, high=None):
+        return tf.broadcast_dynamic_shape(
+            tf.shape(self.low if low is None else low),
+            tf.shape(self.high if high is None else high),
+        )
+
+    def _batch_shape(self):
+        return tf.broadcast_static_shape(self.low.shape, self.high.shape)
+
+    def _event_shape_tensor(self):
+        return tf.constant([], dtype=tf.int32)
+
+    def _event_shape(self):
+        return tf.TensorShape([])
+
+    def _sample_n(self, n, seed=None):
+        with tf.name_scope("sample_n"):
+            low = tf.convert_to_tensor(self.low)
+            high = tf.convert_to_tensor(self.high)
+            shape = tf.concat(
+                [[n], self._batch_shape_tensor(low=low, high=high)], 0
+            )
+            samples = samplers.uniform(shape=shape, dtype=tf.float32, seed=seed)
+            return low + tf.cast(
+                tf.cast(self._range(low=low, high=high), tf.float32) * samples,
+                self.dtype,
+            )
+
+    def _prob(self, x):
+        with tf.name_scope("prob"):
+            low = tf.cast(self.low, self.float_dtype)
+            high = tf.cast(self.high, self.float_dtype)
+            x = tf.cast(x, dtype=self.float_dtype)
+
+            return tf.where(
+                tf.math.is_nan(x),
+                x,
+                tf.where(
+                    (x < low) | (x >= high),
+                    tf.zeros_like(x),
+                    tf.ones_like(x) / self._range(low=low, high=high),
+                ),
+            )
+
+    def _log_prob(self, x):
+        with tf.name_scope("log_prob"):
+            res = tf.math.log(self._prob(x))
+            return res

gemlib/distributions/uniform_integer_test.py

@@ -0,0 +1,55 @@
+# Dependency imports
+import numpy as np
+import tensorflow as tf
+from tensorflow_probability.python.internal import test_util
+
+from gemlib.distributions.uniform_integer import UniformInteger
+
+
+@test_util.test_all_tf_execution_regimes
+class TestUniformInteger(test_util.TestCase):
+    def setUp(self):
+        self.seed = 10402302
+        self.fixture = [[8, 4, 8], [2, 7, 9], [6, 0, 9]]
+
+    def test_sample_n_int32(self):
+        """Sample returning dtype int32."""
+        tf.random.set_seed(self.seed)
+        X = UniformInteger(0, 10)
+        x = X.sample([3, 3], seed=1)
+        x_ = self.evaluate(x)
+        self.fixture_ = self.evaluate(tf.convert_to_tensor(self.fixture))
+        self.assertAllEqual(self.fixture, x_)
+        self.assertDTypeEqual(x_, np.int32)
+
+    def test_sample_n_int64(self):
+        """Sample returning int64."""
+        tf.random.set_seed(self.seed)
+        X = UniformInteger(0, 10, dtype=tf.int64)
+        x = X.sample([3, 3], seed=1)
+        x_ = self.evaluate(x)
+        self.fixture_ = self.evaluate(tf.convert_to_tensor(self.fixture))
+        self.assertAllEqual(self.fixture_, x_)
+        self.assertDTypeEqual(x_, np.int64)
+
+    def test_log_prob_float32(self):
+        """log_prob returning float32."""
+        X = UniformInteger(0, 10)
+        lp = X.log_prob(self.fixture)
+        self.assertSequenceEqual(lp.shape, [3, 3])
+        lp_ = self.evaluate(lp)
+        self.assertAlmostEqual(np.sum(lp_), -20.723265, places=5)
+        self.assertDTypeEqual(lp_, np.float32)
+
+    def test_log_prob_float64(self):
+        """log_prob returning float64."""
+        X = UniformInteger(0, 10, float_dtype=tf.float64)
+        lp = X.log_prob(self.fixture)
+        self.assertSequenceEqual(lp.shape, [3, 3])
+        lp_ = self.evaluate(lp)
+        self.assertAlmostEqual(np.sum(lp_), -20.723265, places=5)
+        self.assertDTypeEqual(lp_, np.float64)
+
+
+if __name__ == "__main__":
+    tf.test.main()

gemlib/mcmc/__init__.py

@@ -0,0 +1,23 @@
+"""MCMC kernel addons"""
+
+import gemlib.mcmc.discrete_time_state_transition_model as discrete_time
+from gemlib.mcmc.adaptive_random_walk_metropolis import (
+    AdaptiveRandomWalkMetropolis,
+)
+from gemlib.mcmc.chain_binomial_rippler import CBRKernel
+from gemlib.mcmc.compound_kernel import CompoundKernel
+from gemlib.mcmc.damped_chain_binomial_rippler import DampedCBRKernel
+from gemlib.mcmc.gibbs_kernel import GibbsKernel
+from gemlib.mcmc.h5_posterior import Posterior
+from gemlib.mcmc.multi_scan_kernel import MultiScanKernel
+
+__all__ = [
+    "AdaptiveRandomWalkMetropolis",
+    "CBRKernel",
+    "CompoundKernel",
+    "DampedCBRKernel",
+    "GibbsKernel",
+    "MultiScanKernel",
+    "Posterior",
+    "discrete_time",
+]