gemlib 0.9.2__py3-none-any.whl

gemlib/mcmc/discrete_time_state_transition_model/__init__.py

@@ -0,0 +1,21 @@
+"""DiscreteTimeStateTransitionModel-related MCMC samplers"""
+
+from gemlib.mcmc.discrete_time_state_transition_model.left_censored_events_mh import (  # noqa: E501
+    UncalibratedLeftCensoredEventTimesUpdate,
+)
+from gemlib.mcmc.discrete_time_state_transition_model.move_events import (
+    UncalibratedEventTimesUpdate,
+)
+from gemlib.mcmc.discrete_time_state_transition_model.right_censored_events_mh import (  # noqa: E501
+    UncalibratedOccultUpdate,
+)
+from gemlib.mcmc.discrete_time_state_transition_model.util import (
+    TransitionTopology,
+)
+
+__all__ = [
+    "TransitionTopology",
+    "UncalibratedEventTimesUpdate",
+    "UncalibratedLeftCensoredEventTimesUpdate",
+    "UncalibratedOccultUpdate",
+]

gemlib/mcmc/discrete_time_state_transition_model/event_time_proposal.py

@@ -0,0 +1,275 @@
+"""Mechanism for proposing event times to move"""
+
+from warnings import warn
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.mcmc.internal import util as mcmc_util
+
+from gemlib.distributions import (
+    Categorical2,
+    UniformInteger,
+    UniformKCategorical,
+)
+
+tfd = tfp.distributions
+
+
+def _events_or_inf(events, transition_id):
+    if transition_id is None:
+        return tf.fill(
+            events.shape[:-1], tf.constant(np.inf, dtype=events.dtype)
+        )
+    return tf.gather(events, transition_id, axis=-1)
+
+
+def _abscumdiff(
+    events, initial_state, topology, t, delta_t, bound_times, int_dtype=tf.int32
+):
+    """Returns the number of free events to move in target_events
+       bounded by max([N_{target_id}(t)-N_{bound_id}(t)]_{bound_t}).
+
+    :param events: a [(M), T, X] tensor of transition events
+    :param initial_state: a [M, X] tensor of the constraining initial state
+    :param target_id: the Xth index of the target event
+    :param bound_t: the times to compute the constraints
+    :param bound_id: the Xth index of the bounding event, -1 implies no bound
+
+    :returns: a tensor of shape [M] + bound_t.shape[0] +  of max free events,
+              dtype=target_events.dtype
+    """
+    with tf.name_scope("_abscumdiff"):
+        # This line prevents negative indices.  However, we must have
+        # a contract that the output of the algorithm is invalid!
+        bound_times = tf.clip_by_value(
+            bound_times, clip_value_min=0, clip_value_max=events.shape[-2] - 1
+        )
+
+        # Maybe replace with pad to avoid unstack/stack
+        prev_events = _events_or_inf(events, topology.prev)
+        target_events = tf.gather(events, topology.target, axis=-1)
+        next_events = _events_or_inf(events, topology.next)
+        event_tensor = tf.stack(
+            [prev_events, target_events, next_events], axis=-1
+        )
+
+        # Compute the absolute cumulative difference between event times
+        diff = event_tensor[..., 1:] - event_tensor[..., :-1]  # [m, T, 2]
+        cumdiff = tf.abs(tf.cumsum(diff, axis=-2))  # cumsum along time axis
+
+        # Create indices into cumdiff [m, d_max, 2].  Last dimension selects
+        # the bound for either the previous or next event.
+        indices = tf.stack(
+            [
+                tf.repeat(
+                    tf.range(events.shape[0], dtype=int_dtype),
+                    [bound_times.shape[1]],
+                ),
+                tf.reshape(bound_times, [-1]),
+                tf.repeat(tf.where(delta_t < 0, 0, 1), [bound_times.shape[1]]),
+            ],
+            axis=-1,
+        )
+        indices = tf.reshape(
+            indices, [events.shape[-3], bound_times.shape[1], 3]
+        )
+        free_events = tf.gather_nd(cumdiff, indices)
+
+        # Add on initial state
+        indices = tf.stack(
+            [
+                tf.range(events.shape[0]),
+                tf.where(
+                    delta_t[:, 0] < 0, topology.target, topology.target + 1
+                ),
+            ],
+            axis=-1,
+        )
+        bound_init_state = tf.gather_nd(initial_state, indices)
+        free_events += bound_init_state[..., tf.newaxis]
+
+        return free_events
+
+
+class Deterministic2(tfd.Deterministic):
+    def __init__(
+        self,
+        loc,
+        atol=None,
+        rtol=None,
+        validate_args=False,
+        allow_nan_stats=True,
+        log_prob_dtype=tf.float32,
+        name="Deterministic",
+    ):
+        parameters = dict(locals())
+        super().__init__(
+            loc,
+            atol=atol,
+            rtol=rtol,
+            validate_args=validate_args,
+            allow_nan_stats=allow_nan_stats,
+            name=name,
+        )
+        self.log_prob_dtype = log_prob_dtype
+
+    def _prob(self, x):
+        return tf.constant(1, dtype=self.log_prob_dtype)
+
+
+def event_time_proposal(
+    events, initial_state, topology, d_max, n_max, dtype=tf.int32, name=None
+):
+    """Draws an event time move proposal.
+    :param events: a [M, T, K] tensor of event times (M number of
+                   metapopulations, T number of full_timepoints, K number of
+                   transitions)
+    :param initial_state: a [M, S] tensor of initial metapopulation x state
+                          counts
+    :param topology: a 3-element tuple of (previous_transition,
+                     target_transition, next_transition), eg "(s->e, e->i,
+                     i->r)" (assuming we are interested presently in e->i,
+                     `None` for boundaries)
+    :param d_max: the maximum distance over which to move (in time)
+    :param n_max: the maximum number of events to move
+    """
+    target_events = tf.gather(events, topology.target, axis=-1)
+    time_interval = tf.range(d_max, dtype=dtype)
+
+    def t():
+        with tf.name_scope("t"):
+            # Waiting for fixed tf.nn.sparse_softmax_cross_entropy_with_logits
+            x = tf.cast(target_events > 0, dtype=events.dtype)  # [M, T]
+            return Categorical2(logits=tf.math.log(x), name="event_coords")
+
+    def delta_t(t):
+        with tf.name_scope("delta_t"):
+            d_max_bcast = tf.broadcast_to(d_max, [events.shape[-3]])
+            low = -tf.clip_by_value(
+                d_max_bcast, clip_value_min=0, clip_value_max=t
+            )
+            high = tf.clip_by_value(
+                d_max_bcast,
+                clip_value_min=0,
+                clip_value_max=events.shape[-2] - t - 1,
+            )
+            return UniformInteger(
+                low=low, high=high + 1, float_dtype=events.dtype
+            )
+
+    def x_star(t, delta_t):
+        with tf.name_scope("x_star"):
+            # Compute bounds
+            # The limitations of XLA mean that we must calculate bounds for
+            # intervals [t, t+delta_t) if delta_t > 0, and [t+delta_t, t) if
+            # delta_t is < 0.
+            t = t[..., tf.newaxis]
+            delta_t = delta_t[..., tf.newaxis]
+            bound_times = tf.where(
+                delta_t < 0,
+                t - time_interval - 1,
+                t + time_interval,  # [t+delta_t, t)
+            )  # [t, t+delta_t)
+            free_events = _abscumdiff(
+                events=events,
+                initial_state=initial_state,
+                topology=topology,
+                t=t,
+                delta_t=delta_t,
+                bound_times=bound_times,
+                int_dtype=dtype,
+            )
+
+            # Mask out bits of the interval we don't need for our delta_t
+            inf_mask = tf.cumsum(
+                tf.one_hot(
+                    tf.math.abs(delta_t[:, 0]),
+                    d_max,
+                    on_value=tf.constant(np.inf, events.dtype),
+                    dtype=events.dtype,
+                )
+            )
+            free_events = tf.maximum(inf_mask, free_events)
+            free_events = tf.reduce_min(free_events, axis=-1)
+
+            indices = tf.stack(
+                [tf.range(events.shape[0], dtype=dtype), t[:, 0]], axis=-1
+            )
+            available_events = tf.gather_nd(target_events, indices)
+            max_events = tf.minimum(free_events, available_events)
+            max_events = tf.clip_by_value(
+                max_events, clip_value_min=0, clip_value_max=n_max
+            )
+            # Draw x_star
+            return UniformInteger(
+                low=1, high=max_events + 1, float_dtype=events.dtype
+            )
+
+    return tfd.JointDistributionNamed(
+        {"t": t, "delta_t": delta_t, "x_star": x_star}, name=name
+    )
+
+
+def filtered_event_time_proposal(  # pylint: disable-invalid-name
+    events,
+    initial_state,
+    topology,
+    m_max,
+    d_max,
+    n_max,
+    dtype=tf.int32,
+    name=None,
+):
+    """FilteredEventTimeProposal allows us to choose a subset of indices
+    in `range(events.shape[0])` for which to propose an update.  The
+    results are then broadcast back to `events.shape[0]`.
+
+    :param events: a [M, T, X] event tensor
+    :param initial_state: a [M, S] initial state tensor
+    :param topology: a TransitionTopology named tuple describing the ordering
+                     of events
+    :param m: the number of metapopulations to move
+    :param d_max: maximum distance in time to move
+    :param n_max: maximum number of events to move (user defined)
+    :return: an instance of a JointDistributionNamed
+    """
+    if mcmc_util.is_list_like(events):
+        warn(
+            "Batched FilteredEventTimeProposals are not yet supported",
+            stacklevel=1,
+        )
+        events = events[0]
+
+    target_events = tf.gather(events, topology.target, axis=-1)
+
+    def m():
+        with tf.name_scope("m"):
+            hot_meta = tf.math.count_nonzero(target_events, axis=1) > 0
+            X = UniformKCategorical(
+                m_max, hot_meta, float_dtype=events.dtype, name="m"
+            )
+            return X
+
+    def move(m):
+        """We select out meta-population `m` from the first
+        dimension of `events`.
+        :param m: a 1-D tensor of indices of meta-populations
+        :return: a random variable of type `EventTimeProposal`
+        """
+        with tf.name_scope("move"):
+            select_meta = tf.gather(events, m, axis=0)
+            select_init = tf.gather(initial_state, m, axis=0)
+            return event_time_proposal(
+                select_meta,
+                select_init,
+                topology,
+                d_max,
+                n_max,
+                dtype=dtype,
+                name=name,
+            )
+
+    return tfd.JointDistributionNamed(
+        {"m": m, "move": move}, name="FilteredEventTimeProposal"
+    )

gemlib/mcmc/discrete_time_state_transition_model/fixtures.py

@@ -0,0 +1,56 @@
+"""Test fixtures"""
+
+import numpy as np
+import pytest
+
+
+@pytest.fixture(scope="module")
+def sir_metapop_example():
+    """Outcome of a simulation from a 3-metapopulation model
+    with mixing, implemented in https://colab.research.google.com/drive/1Q1PUcOnYlvCGHzRUBUAp4CxYhZ8RJzg8?usp=sharing
+    """
+
+    incidence_matrix = np.array([[-1, 0], [1, -1], [0, 1]], dtype=np.float32)
+
+    initial_conditions = np.array(
+        [[999, 50, 0], [500, 20, 0], [250, 10, 0]], dtype=np.float32
+    )
+
+    events = np.array(
+        [
+            [
+                [11.0, 11.0],
+                [11.0, 6.0],
+                [2.0, 6.0],
+                [10.0, 6.0],
+                [12.0, 7.0],
+                [11.0, 5.0],
+                [13.0, 10.0],
+            ],
+            [
+                [5.0, 4.0],
+                [5.0, 2.0],
+                [11.0, 1.0],
+                [8.0, 4.0],
+                [7.0, 4.0],
+                [5.0, 5.0],
+                [12.0, 6.0],
+            ],
+            [
+                [2.0, 2.0],
+                [4.0, 2.0],
+                [1.0, 0.0],
+                [2.0, 2.0],
+                [6.0, 1.0],
+                [4.0, 1.0],
+                [6.0, 1.0],
+            ],
+        ],
+        dtype=np.float32,
+    )
+
+    return {
+        "initial_conditions": initial_conditions,
+        "events": events,
+        "incidence_matrix": incidence_matrix,
+    }

gemlib/mcmc/discrete_time_state_transition_model/left_censored_events_mh.py

@@ -0,0 +1,258 @@
+"""Metropolis Hastings implementation for left-censored events
+in a discrete-time metapopulation epidemic model
+"""
+
+from typing import NamedTuple, Tuple
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import prefer_static as ps
+from tensorflow_probability.python.internal import samplers
+from tensorflow_probability.python.mcmc.internal import util as mcmc_util
+
+from .left_censored_events_proposal import (
+    left_censored_event_time_proposal,
+)
+
+tfd = tfp.distributions
+
+__all__ = ["UncalibratedLeftCensoredEventTimesUpdate"]
+
+LEFT_CENSORED_EVENT_TIMES_UPDATE_TUPLE_LEN = 2
+
+
+def _update_state(update, current_state, transition_idx, incidence_matrix):
+    current_initial_conditions, current_events = current_state
+    update = {k: tf.convert_to_tensor(v) for k, v in update.items()}
+
+    # -1 is moving events forward in time
+    sign = tf.gather(
+        [-1, 1],
+        update["direction"],
+    )
+    events_delta = update["num_events"] * sign
+
+    # Update initial conditions
+    indices = tf.stack(
+        [
+            tf.broadcast_to(
+                update["unit"], [current_initial_conditions.shape[-1]]
+            ),
+            ps.range(current_initial_conditions.shape[-1]),
+        ],
+        axis=-1,
+    )
+    new_initial_conditions = tf.tensor_scatter_nd_add(
+        current_initial_conditions,
+        indices=indices,
+        updates=tf.cast(
+            tf.cast(
+                events_delta, incidence_matrix.dtype
+            )  # TODO sort this out: casts are usually a code smell!
+            * ps.gather(incidence_matrix, transition_idx, axis=-1),
+            current_initial_conditions.dtype,
+        ),
+    )
+
+    # Update events
+    indices = tf.stack(
+        [
+            update["unit"],
+            update["timepoint"],
+            ps.broadcast_to(transition_idx, update["unit"].shape),
+        ],
+        axis=-1,
+    )
+    new_events = tf.tensor_scatter_nd_sub(
+        current_events,
+        indices=indices,
+        updates=tf.cast(events_delta, dtype=current_events.dtype),
+    )
+
+    return new_initial_conditions, new_events
+
+
+def _reverse_update(update):
+    direction = (update["direction"] + 1) % 2
+    return {
+        "unit": update["unit"],
+        "timepoint": update["timepoint"],
+        "direction": direction,
+        "num_events": update["num_events"],
+    }
+
+
+class LeftCensoredEventTimeResults(NamedTuple):
+    log_acceptance_correction: float
+    target_log_prob: float
+    unit: int
+    timepoint: int
+    direction: int
+    num_events: int
+    seed: Tuple[int, int]
+
+
+class UncalibratedLeftCensoredEventTimesUpdate(tfp.mcmc.TransitionKernel):
+    """UncalibratedLeftCensoredEventTimesUpdate"""
+
+    def __init__(
+        self,
+        target_log_prob_fn,
+        transition_index,
+        incidence_matrix,
+        max_timepoint,
+        max_events,
+        name=None,
+    ):
+        """An uncalibrated random walk for initial conditions.
+        :param target_log_prob_fn: the log density of the target distribution
+        :param transition_index: the index of the transition to adjust
+        :param incidence_matrix: the `[S,R]` incidence matrix
+        :param max_timepoint: max timepoint up to which to move events
+        :param max_events: max number of events per unit/timepoint to move
+        """
+        self._name = name
+        self._parameters = {
+            "target_log_prob_fn": target_log_prob_fn,
+            "transition_index": transition_index,
+            "incidence_matrix": incidence_matrix,
+            "max_timepoint": max_timepoint,
+            "max_events": max_events,
+            "name": name,
+        }
+
+    @property
+    def target_log_prob_fn(self):
+        return self._parameters["target_log_prob_fn"]
+
+    @property
+    def transition_index(self):
+        return self._parameters["transition_index"]
+
+    @property
+    def incidence_matrix(self):
+        return self._parameters["incidence_matrix"]
+
+    @property
+    def max_timepoint(self):
+        return self._parameters["max_timepoint"]
+
+    @property
+    def max_events(self):
+        return self._parameters["max_events"]
+
+    @property
+    def name(self):
+        return self._parameters["name"]
+
+    @property
+    def parameters(self):
+        return self._parameters
+
+    @property
+    def is_calibrated(self):
+        return False
+
+    def one_step(self, current_state, previous_kernel_results, seed=None):
+        """Update the initial conditions
+
+        :param current_state: a tuple of `(current_initial_conditions,
+                              current_events)`
+        :param previous_kernel_results: previous kernel results tuple
+        :param seed: optional seed tuple `(int32, int32)`
+        :returns: new state tuple `(next_initial_conditions, next_events)`
+        """
+        with tf.name_scope("uncalibrated_left_censored_events_mh/one_step"):
+            seed = samplers.sanitize_seed(
+                seed, salt="uncalibrated_left_censored_events_mh"
+            )
+
+            if (not mcmc_util.is_list_like(current_state)) and (
+                len(current_state) == LEFT_CENSORED_EVENT_TIMES_UPDATE_TUPLE_LEN
+            ):
+                raise ValueError(
+                    f"State for LeftCensoredEventTimesUpdate must be a\
+ list/tuple of length {LEFT_CENSORED_EVENT_TIMES_UPDATE_TUPLE_LEN}"
+                )
+
+            proposal = left_censored_event_time_proposal(
+                events=current_state[1],
+                initial_state=current_state[0],
+                transition=self.transition_index,
+                incidence_matrix=self.incidence_matrix,
+                num_units=1,
+                max_timepoint=self.max_timepoint,
+                max_events=self.max_events,
+                name=f"{self.name}/fwd_proposal",
+            )
+            fwd_update = proposal.sample(seed=seed)
+            fwd_proposal_log_prob = proposal.log_prob(
+                fwd_update, name="fwd_proposal_log_prob"
+            )
+
+            next_state = _update_state(
+                fwd_update,
+                current_state,
+                self.transition_index,
+                self.incidence_matrix,
+            )
+            next_target_log_prob = self.target_log_prob_fn(*next_state)
+
+            rev_update = _reverse_update(fwd_update)
+            rev_proposal = left_censored_event_time_proposal(
+                events=next_state[1],
+                initial_state=next_state[0],
+                transition=self.transition_index,
+                incidence_matrix=self.incidence_matrix,
+                num_units=1,
+                max_timepoint=self.max_timepoint,
+                max_events=self.max_events,
+                name=f"{self.name}/rev_proposal",
+            )
+            rev_proposal_log_prob = rev_proposal.log_prob(rev_update)
+            log_acceptance_correction = tf.reduce_sum(
+                rev_proposal_log_prob - fwd_proposal_log_prob
+            )
+            results = (
+                next_state,
+                LeftCensoredEventTimeResults(
+                    log_acceptance_correction=log_acceptance_correction,
+                    target_log_prob=next_target_log_prob,
+                    unit=fwd_update["unit"],
+                    timepoint=fwd_update["timepoint"],
+                    direction=fwd_update["direction"],
+                    num_events=fwd_update["num_events"],
+                    seed=seed,
+                ),
+            )
+
+            return results
+
+    def bootstrap_results(self, init_state):
+        with tf.name_scope(
+            "uncalibrated_left_censored_events_mh/boostrap_results"
+        ):
+            if (not mcmc_util.is_list_like(init_state)) and (
+                len(init_state) == 2  # noqa: PLR2004
+            ):
+                raise ValueError(
+                    f"State for LeftCensoredEventTimesUpdate must be a \
+list/tuple of length {LEFT_CENSORED_EVENT_TIMES_UPDATE_TUPLE_LEN}"
+                )
+
+            initial_conditions = tf.convert_to_tensor(init_state[0])
+            events = tf.convert_to_tensor(init_state[1])
+            init_target_log_prob = self.target_log_prob_fn(
+                initial_conditions, events
+            )
+            return LeftCensoredEventTimeResults(
+                log_acceptance_correction=tf.constant(
+                    0.0, dtype=init_target_log_prob.dtype
+                ),
+                target_log_prob=init_target_log_prob,
+                unit=tf.zeros([1], dtype=tf.int32),
+                timepoint=tf.zeros([1], dtype=tf.int32),
+                direction=tf.constant(0, dtype=tf.int32),
+                num_events=tf.ones([1], dtype=tf.int32),
+                seed=samplers.zeros_seed(),
+            )