gemlib 0.9.2__py3-none-any.whl

gemlib/distributions/experimental/state_transition_marginal_model.py

@@ -0,0 +1,352 @@
+"""Describes a State Transition Model with marginalised baseline
+hazard rates.
+"""
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import (
+    dtype_util,
+    reparameterization,
+)
+
+from gemlib.distributions.discrete_markov import discrete_markov_simulation
+from gemlib.util import batch_gather, compute_state, transition_coords
+
+tla = tf.linalg
+tfd = tfp.distributions
+
+
+class StateTransitionMarginalModel(tfd.Distribution):
+    def __init__(
+        self,
+        transition_rates,
+        baseline_hazard_rate_priors,
+        stoichiometry,
+        initial_state,
+        initial_step,
+        time_delta,
+        num_steps,
+        validate_args=False,
+        allow_nan_stats=True,
+        name="StateTransitionMarginalModel",
+    ):
+        """Implements a discrete-time Markov jump process for a state transition
+           model.
+
+        :param transition_rates: a function of the form `fn(t, state)` taking
+                                 the current time `t` and state tensor `state`.
+                                 This function returns a tensor which broadcasts
+                                 to the first dimension of `stoichiometry`.
+                                 Transition rates are assumed to be risk ratios,
+                                 with the baseline hazard rate marginalised out
+                                 from the model.
+        :param baseline_hazard_rate_priors: a dictionary of `concentration` and
+                                            `rate` hyperparameters for implicit
+                                            Gamma priors on (marginalised)
+                                            baseline hazard rates. Both
+                                            `concentration` and `rate` should
+                                            broadcast with the number of rows in
+                                            `stoichiometry`.
+        :param stoichiometry: the stochiometry matrix for the state transition
+                              model with rows representing transitions and
+                              columns representing states.
+        :param initial_state: an initial state tensor with inner dimension equal
+                              to the first dimension of `stoichiometry`.
+        :param initial_step: an offset giving the time `t` of the first timestep
+                             in the model.
+        :param time_delta: the size of the time step to be used.
+        :param num_steps: the number of time steps across which the model runs.
+        """
+        parameters = dict(locals())
+        with tf.name_scope(name) as name:
+            self._transition_rates = transition_rates
+            self._stoichiometry = tf.convert_to_tensor(
+                stoichiometry,
+                dtype=initial_state.dtype,
+            )
+            self._initial_state = initial_state
+            self._initial_step = initial_step
+            self._time_delta = time_delta
+            self._num_steps = num_steps
+            self._baseline_hazard_rate_priors = baseline_hazard_rate_priors
+
+            super().__init__(
+                dtype=initial_state.dtype,
+                reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
+                validate_args=validate_args,
+                allow_nan_stats=allow_nan_stats,
+                parameters=parameters,
+                name=name,
+            )
+
+        self.dtype = initial_state.dtype
+
+    @property
+    def transition_rates(self):
+        return self._transition_rates
+
+    @property
+    def baseline_hazard_rate_priors(self):
+        return self._baseline_hazard_rate_priors
+
+    @property
+    def stoichiometry(self):
+        return self._stoichiometry
+
+    @property
+    def initial_state(self):
+        return self._initial_state
+
+    @property
+    def initial_step(self):
+        return self._initial_step
+
+    @property
+    def time_delta(self):
+        return self._time_delta
+
+    @property
+    def num_steps(self):
+        return self._num_steps
+
+    def _batch_shape(self):
+        return tf.TensorShape([])
+
+    def _event_shape(self):
+        shape = tf.TensorShape(
+            [
+                self.initial_state.shape[0],
+                tf.get_static_value(self._num_steps),
+                self._stoichiometry.shape[0],
+            ]
+        )
+        return shape
+
+    def _sample_n(self, n, seed=None):
+        """Runs a simulation from the epidemic model
+
+        :param param: a dictionary of model parameters
+        :param state_init: the initial state
+        :returns: a tuple of times and simulated states.
+        """
+        with tf.name_scope("DiscreteTimeStateTransitionModel.log_prob"):
+
+            def hazard_fn(t, state):
+                return self.transition_rates(t, state)
+
+            t, sim = discrete_markov_simulation(
+                hazard_fn=hazard_fn,
+                state=self.initial_state,
+                start=self.initial_step,
+                end=self.initial_step + self.num_steps * self.time_delta,
+                time_step=self.time_delta,
+                stoichiometry=self.stoichiometry,
+                seed=seed,
+            )
+            indices = transition_coords(self.stoichiometry)
+            sim = batch_gather(sim, indices)
+            sim = tf.transpose(sim, perm=(1, 0, 2))
+            return tf.expand_dims(sim, 0)
+
+    def _log_prob(self, y, **kwargs):
+        """Calculates the log probability of observing epidemic events y
+        :param y: a list of tensors.  The first is of shape [n_times] containing
+                  times, the second is of shape [n_times, n_states, n_states]
+                  containing event matrices.
+        :param param: a list of parameters
+        :returns: a scalar giving the log probability of the epidemic
+        """
+        dtype = dtype_util.common_dtype(
+            [y, self.initial_state], dtype_hint=self.dtype
+        )
+        events = tf.convert_to_tensor(y, dtype)
+        with tf.name_scope("StateTransitionMarginalModel.log_prob"):
+            state_timeseries = compute_state(
+                initial_state=self.initial_state,
+                events=events,
+                stoichiometry=self.stoichiometry,
+                closed=True,
+            )
+
+            tms_timeseries = tf.transpose(state_timeseries, perm=(1, 0, 2))
+            tmr_events = tf.transpose(events, perm=(1, 0, 2))
+
+            def fn(elems):
+                return tf.stack(self.transition_rates(*elems), axis=-1)
+
+            rates = tf.vectorized_map(
+                fn=fn,
+                elems=(
+                    self._initial_step + tf.range(tms_timeseries.shape[0]),
+                    tms_timeseries,
+                ),
+            )
+
+            def integrated_rate_fn():
+                """Use mid-point integration to estimate the constant rate
+                over time.
+                """
+                integrated_rates = tms_timeseries[..., :-1] * rates
+                return (
+                    integrated_rates[:-1, ...] + integrated_rates[1:, ...]
+                ) / 2.0
+
+            integrated_rates = integrated_rate_fn()
+
+            log_norm_constant = tf.reduce_sum(
+                tf.math.multiply_no_nan(
+                    tf.math.log(integrated_rates), tmr_events
+                )
+                - tf.math.lgamma(tmr_events + 1.0),
+                axis=(0, 1),
+            )
+            pi_concentration = (
+                tf.reduce_sum(tmr_events, axis=(0, 1))
+                + self.baseline_hazard_rate_priors["concentration"]
+            )
+            pi_rate = (
+                tf.reduce_sum(integrated_rates * self.time_delta, axis=(0, 1))
+                + self.baseline_hazard_rate_priors["rate"]
+            )
+
+            log_prob = (
+                log_norm_constant
+                + tf.math.lgamma(pi_concentration)
+                - (pi_concentration) * tf.math.log(pi_rate)
+            )
+
+            return tf.reduce_sum(log_prob)
+
+
+class BaselineHazardRateMarginal(tfd.Distribution):
+    def __init__(
+        self,
+        events,
+        transition_rate_fn,
+        baseline_hazard_rate_priors,
+        stoichiometry,
+        initial_state,
+        initial_step,
+        time_delta,
+        num_steps,
+        validate_args=False,
+        allow_nan_stats=True,
+        name="StateTransitionMarginalModel",
+    ):
+        """Implements a discrete-time Markov jump process for a state transition
+           model.
+
+        :param events: a [M, T, R] event tensor
+        :param transition_rates: a function of the form `fn(t, state)` taking
+                                 the current time `t` and state tensor `state`.
+                                 This function returns a tensor which broadcasts
+                                 to the first dimension of `stoichiometry`.
+                                 Transition rates are assumed to be risk ratios,
+                                 with the baseline hazard rate marginalised out
+                                 from the model.
+        :param baseline_hazard_rate_priors: a dictionary of `concentration` and
+                                            `rate` hyperparameters for implicit
+                                            Gamma priors on (marginalised)
+                                            baseline hazard rates. Both
+                                            `concentration` and `rate` should
+                                            broadcast with the number of rows in
+                                            `stoichiometry`.
+        :param stoichiometry: the stochiometry matrix for the state transition
+                              model with rows representing transitions and
+                              columns representing states.
+        :param initial_state: an initial state tensor with inner dimension equal
+                              to the first dimension of `stoichiometry`.
+        :param initial_step: an offset giving the time `t` of the first timestep
+                             in the model.
+        :param time_delta: the size of the time step to be used.
+        :param num_steps: the number of time steps across which the model runs.
+        """
+        parameters = dict(locals())
+        with tf.name_scope(name) as name:
+            self._events = events
+            self._transition_rate_fn = transition_rate_fn
+            self._stoichiometry = tf.convert_to_tensor(
+                stoichiometry,
+                dtype=initial_state.dtype,
+            )
+            self._initial_state = initial_state
+            self._initial_step = initial_step
+            self._time_delta = time_delta
+            self._num_steps = num_steps
+            self._baseline_hazard_rate_priors = baseline_hazard_rate_priors
+
+            super().__init__(
+                dtype=initial_state.dtype,
+                reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
+                validate_args=validate_args,
+                allow_nan_stats=allow_nan_stats,
+                parameters=parameters,
+                name=name,
+            )
+
+        self.dtype = initial_state.dtype
+
+    @property
+    def baseline_hazard_rate_priors(self):
+        return self._baseline_hazard_rate_priors
+
+    def _batch_shape(self):
+        return tf.TensorShape(())
+
+    def _event_shape(self):
+        shape = tf.TensorShape(self._events.shape[-1])
+        return shape
+
+    def concentration(self):
+        """Calculates the concentration parameter"""
+        return (
+            tf.reduce_sum(self._events, axis=(0, 1))
+            + self.baseline_hazard_rate_priors["concentration"]
+        )
+
+    def rate(self):
+        """Calculates the rate parameter"""
+        state = compute_state(
+            self._initial_state, self._events, self._stoichiometry, closed=True
+        )
+        tms_state = tf.transpose(state, perm=(1, 0, 2))
+
+        def fn(elems):
+            return tf.stack(self._transition_rate_fn(*elems), axis=-1)
+
+        rates = tf.vectorized_map(
+            fn=fn,
+            elems=(
+                self._initial_step + tf.range(tms_state.shape[0]),
+                tms_state,
+            ),
+        )
+
+        def integrated_rate_fn():
+            """Use mid-point integration to estimate the constant rate
+            over time.
+            """
+            integrated_rates = tms_state[..., :-1] * rates
+            return (
+                integrated_rates[:-1, ...] + integrated_rates[1:, ...]
+            ) / 2.0
+
+        return (
+            tf.reduce_sum(integrated_rate_fn(), axis=(-3, -2))
+            + self._baseline_hazard_rate_priors["rate"]
+        )
+
+    def _sample_n(self, n, seed=None):
+        tf.print("Concentration:", self.concentration())
+        tf.print("Rate:", self.rate())
+        rv = tfd.Gamma(concentration=self.concentration(), rate=self.rate())
+        return rv.sample(n, seed=seed)
+
+    def _log_prob(self, y, **kwargs):
+        """Calculates the log prob"""
+        rv = tfd.Gamma(concentration=self.concentration(), rate=self.rate())
+        return tf.reduce_sum(rv.log_prob(y))
+
+    def _mean(self):
+        rv = tfd.Gamma(concentration=self.concentration(), rate=self.rate())
+        return rv.mean()

gemlib/distributions/hypergeometric.py

@@ -0,0 +1,144 @@
+"""Hypergeometric random variable"""
+
+# ruff: noqa: N803, N802
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import (
+    dtype_util,
+    parameter_properties,
+    reparameterization,
+)
+
+from gemlib.distributions.hypergeometric_sampler import sample_hypergeometric
+
+tfd = tfp.distributions
+
+__all__ = ["Hypergeometric"]
+
+
+def _log_factorial(x):
+    """Computes x!"""
+    return tf.math.lgamma(x + 1.0)
+
+
+def _log_choose(n, k):
+    """Computes nCk"""
+    return _log_factorial(n) - _log_factorial(k) - _log_factorial(n - k)
+
+
+class Hypergeometric(tfd.Distribution):
+    def __init__(
+        self,
+        N,
+        K,
+        n,
+        validate_args=False,
+        allow_nan_stats=True,
+        name="Hypergeometric",
+    ):
+        """Hypergeometric distribution
+
+        Args:
+        ----
+          N: Population size
+          K: number of units of interest in the population
+          n: size of sample drawn from the population
+          validate_args: should arguments be validated for correctness
+          allow_nan_stats: allow NaN to be returned for mode, mean, variance...
+
+        """
+        parameters = dict(locals())
+        with tf.name_scope(name) as name:
+            dtype = dtype_util.common_dtype([N, K, n], tf.float32)
+            self._N = tf.cast(N, dtype=dtype)
+            self._K = tf.cast(K, dtype=dtype)
+            self._n = tf.cast(n, dtype=dtype)
+            self._n_positive_mask = N > 0.0
+            super().__init__(
+                dtype=dtype,
+                reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
+                validate_args=validate_args,
+                allow_nan_stats=allow_nan_stats,
+                parameters=parameters,
+                name=name,
+            )
+            if validate_args is True:
+                tf.debugging.assert_non_negative(
+                    N, message="N must be non-negative"
+                )
+                tf.debugging.assert_less_equal(K, N, message="K must be <= N")
+                tf.debugging.assert_less_equal(n, N, message="n must be <= N")
+
+    @classmethod
+    def _parameter_properties(cls, dtype, num_classes=None):
+        return {
+            "N": parameter_properties.ParameterProperties(
+                default_constraining_bijector_fn=parameter_properties.BIJECTOR_NOT_IMPLEMENTED
+            ),
+            "K": parameter_properties.ParameterProperties(),
+            "n": parameter_properties.ParameterProperties(),
+        }
+
+    @staticmethod
+    def _param_shapes(sample_shape):
+        return dict(
+            zip(
+                ("N", "K", "n"),
+                ([tf.convert_to_tensor(sample_shape, dtype=tf.int32)] * 3),
+            )
+        )
+
+    @classmethod
+    def _params_event_ndims(cls):
+        return {"N": 0, "K": 0, "n": 0}
+
+    @property
+    def N(self):
+        """Population size"""
+        return self._parameters["N"]
+
+    @property
+    def K(self):
+        """Number of units of interest in population"""
+        return self._parameters["K"]
+
+    @property
+    def n(self):
+        """Sample size"""
+        return self._parameters["n"]
+
+    def _default_event_space_bijector(self):
+        return
+
+    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(self.name + "/sample_n"):
+            sample = sample_hypergeometric(n, self.N, self.K, self.n, seed=seed)
+            sample = tf.where(self._n_positive_mask, sample, 0.0)
+            return sample
+
+    def _log_prob(self, x):
+        numerator = _log_choose(self._K, x) + _log_choose(
+            self._N - self._K, self._n - x
+        )
+        denominator = _log_choose(self._N, self._n)
+        return numerator - denominator
+
+    def _mode(self):
+        return tf.math.floor((self._n + 1) * (self._K + 1) / (self._N + 2))
+
+    def _mean(self):
+        return self._n * self._K / self._N
+
+    def _variance(self):
+        n = self._n
+        N = self._N
+        K = self._K
+
+        return n * (K / N) * (N - K) / N * (N - n) / (N - 1)

gemlib/distributions/hypergeometric_sampler.py

@@ -0,0 +1,103 @@
+"""Hypergeometric sampling algorithm"""
+
+# ruff: noqa: N803
+
+import numpy as np
+import tensorflow as tf
+from tensorflow_probability.python.internal import (
+    batched_rejection_sampler as brs,
+)
+from tensorflow_probability.python.internal import (
+    dtype_util,
+    samplers,
+    tensor_util,
+)
+from tensorflow_probability.python.internal import prefer_static as ps
+
+
+def sample_hypergeometric(num_samples, N, K, n, seed=None):
+    dtype = dtype_util.common_dtype([N, K, n], tf.float32)
+    N = tensor_util.convert_nonref_to_tensor(N, dtype, name="N")
+    K = tensor_util.convert_nonref_to_tensor(K, dtype, name="K")
+    n = tensor_util.convert_nonref_to_tensor(n, dtype, name="n")
+    good_params_mask = (N >= 1.0) & (N >= K) & (n <= N)
+    N = tf.where(good_params_mask, N, 100.0)
+    K = tf.where(good_params_mask, K, 50.0)
+    n = tf.where(good_params_mask, n, 50.0)
+    sample_shape = ps.concat(
+        [
+            [num_samples],
+            ps.broadcast_shape(
+                ps.broadcast_shape(ps.shape(N), ps.shape(K)), ps.shape(n)
+            ),
+        ],
+        axis=0,
+    )
+
+    # First Transform N, K, n such that
+    # N / 2 >= K, N / 2 >= n
+    is_k_small = 0.5 * N >= K
+    is_n_small = n <= 0.5 * N
+    previous_K = K
+    previous_n = n
+    K = tf.where(is_k_small, K, N - K)
+    n = tf.where(is_n_small, n, N - n)
+
+    # TODO: Can we write this in a more numerically stable way?
+    def _log_hypergeometric_coeff(x):
+        return (
+            tf.math.lgamma(x + 1.0)
+            + tf.math.lgamma(K - x + 1.0)
+            + tf.math.lgamma(n - x + 1.0)
+            + tf.math.lgamma(N - K - n + x + 1.0)
+        )
+
+    p = K / N
+    q = 1 - p
+    a = n * p + 0.5
+    c = tf.math.sqrt(2.0 * a * q * (1.0 - n / N))
+    k = tf.math.floor((n + 1) * (K + 1) / (N + 2))
+    g = _log_hypergeometric_coeff(k)
+    diff = tf.math.floor(a - c)
+    x = (a - diff - 1) / (a - diff)
+    diff = tf.where(
+        (n - diff) * (p - diff / N) * tf.math.square(x)
+        > (diff + 1.0) * (q - (n - diff - 1) / N),
+        diff + 1.0,
+        diff,
+    )
+    # TODO: Can we write this difference of lgammas more numerically stably?
+    h = (a - diff) * tf.math.exp(
+        0.5 * (g - _log_hypergeometric_coeff(diff)) + np.log(2.0)
+    )
+    b = tf.math.minimum(tf.math.minimum(n, K) + 1, tf.math.floor(a + 5 * c))
+
+    def generate_and_test_samples(seed):
+        v_seed, u_seed = samplers.split_seed(seed)
+        U = samplers.uniform(sample_shape, dtype=dtype, seed=u_seed)
+        V = samplers.uniform(sample_shape, dtype=dtype, seed=v_seed)
+        # Guard against 0.
+
+        X = a + h * (V - 0.5) / (1.0 - U)
+        samples = tf.math.floor(X)
+        good_sample_mask = (samples >= 0.0) & (samples < b)
+        T = g - _log_hypergeometric_coeff(samples)
+        # Uses slow pass since we are trying to do this in a vectorized way.
+        good_sample_mask = good_sample_mask & (2 * tf.math.log1p(-U) <= T)
+        return samples, good_sample_mask
+
+    samples = brs.batched_las_vegas_algorithm(
+        generate_and_test_samples, seed=seed
+    )[0]
+    samples = tf.where(good_params_mask, samples, np.nan)
+    # Now transform the samples depending on if we constrained N and / or k
+    samples = tf.where(
+        ~is_k_small & ~is_n_small,
+        samples + previous_K + previous_n - N,
+        tf.where(
+            ~is_k_small,
+            previous_n - samples,
+            tf.where(~is_n_small, previous_K - samples, samples),
+        ),
+    )
+    return samples

gemlib/distributions/hypergeometric_test.py

@@ -0,0 +1,46 @@
+"""Test the Hypergeometric random vaiable"""
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import test_util
+
+from gemlib.distributions.hypergeometric import Hypergeometric
+
+tfd = tfp.distributions
+
+
+@test_util.test_all_tf_execution_regimes
+class TestHypergeometric(test_util.TestCase):
+    def setUp(self):
+        self._rng = np.random.RandomState(5)
+        super().setUp()
+
+    def test_neg_args(self):
+        """Test for invalid arguments"""
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, "N must be non-negative"
+        ):
+            self.evaluate(Hypergeometric(N=-3, K=1, n=2, validate_args=True))
+
+    def test_sample_n_float32(self):
+        """Sample returning float32 args"""
+
+        X = Hypergeometric(345.0, 35.0, 100.0)
+        x = X.sample([1000, 1000], seed=1)
+
+        x = self.evaluate(x)
+        self.assertDTypeEqual(x, np.float32)
+        self.assertAllClose(tf.reduce_mean(x), X.mean(), atol=1e-3, rtol=1e-3)
+
+    def test_sample_n_float64(self):
+        """Sample returning float32 args"""
+
+        X = Hypergeometric(
+            np.float64(345.0), np.float64(35.0), np.float64(100.0)
+        )
+        x = X.sample([1000, 1000], seed=1)
+
+        x = self.evaluate(x)
+        self.assertDTypeEqual(x, np.float64)
+        self.assertAllClose(tf.reduce_mean(x), X.mean(), atol=1e-3, rtol=1e-3)