gemlib 0.9.2__py3-none-any.whl

gemlib/distributions/continuous_time_state_transition_model.py

@@ -0,0 +1,185 @@
+"""Continuous time state transition model"""
+
+from typing import Callable, Optional
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import reparameterization
+
+from gemlib.distributions.continuous_markov import (
+    EpidemicEvent,
+    compute_state,
+    continuous_markov_simulation,
+    continuous_time_log_likelihood,
+)
+
+# aliasing for convenience
+tfd = tfp.distributions
+Tensor = tf.Tensor
+DTYPE = tf.float32
+
+
+class ContinuousTimeStateTransitionModel(tfd.Distribution):
+    """Continuous time state transition model."""
+
+    def __init__(
+        self,
+        transition_rate_fn: Callable[[Tensor], Tensor],
+        incidence_matrix: Tensor,
+        initial_state: Tensor,
+        num_events: int,
+        initial_time: Optional[float] = 0.0,
+        validate_args: Optional[bool] = False,
+        allow_nan_stats: Optional[bool] = True,
+        name: Optional[str] = "ContinuousTimeStateTransitionModel",
+    ) -> None:
+        """
+        Initializes a ContinuousTimeStateTransitionModel object.
+
+        Args:
+            transition_rate_fn: Python callable of the form `fn(t, state)`
+                                taking the current time `t: float` and state
+                                tensor `state`, and returning a tuple of tensors
+                                containing transition rates between states.
+            incidence_matrix: Matrix representing the incidence of transitions
+                              between states.
+            initial_state: A `[N, S]` tensor containing the initial state of the
+                           population of `N` individuals in `S` epidemiological
+                           classes.
+            num_events: the number of events to simulate
+            initial_time: Initial time of the model. Defaults to 0.0.
+            name: Name of the model. Defaults to
+                  "ContinuousTimeStateTransitionModel".
+        """
+        parameters = dict(locals())
+
+        self._incidence_matrix = tf.convert_to_tensor(incidence_matrix)
+        self._initial_state = tf.convert_to_tensor(initial_state)
+        self._initial_time = tf.convert_to_tensor(initial_time)
+
+        super().__init__(
+            dtype=self._incidence_matrix.dtype,
+            reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
+            validate_args=validate_args,
+            allow_nan_stats=allow_nan_stats,
+            parameters=parameters,
+            name=name,
+        )
+
+    @property
+    def transition_rate_fn(self):
+        """Transition rate function for the model."""
+        return self._parameters["transition_rate_fn"]
+
+    @property
+    def incidence_matrix(self):
+        """Incidence matrix for the model."""
+        return self._parameters["incidence_matrix"]
+
+    @property
+    def initial_state(self):
+        """Initial state of the model."""
+        return self._parameters["initial_state"]
+
+    @property
+    def num_events(self):
+        """Number of events to simulate."""
+        return self._parameters["num_events"]
+
+    @property
+    def initial_time(self):
+        """Initial wall clock for the model. Sets the time scale."""
+        return self._parameters["initial_time"]
+
+    def compute_state(
+        self, event_list: EpidemicEvent, include_final_state: bool = False
+    ) -> Tensor:
+        """Given an event list `event_list`, compute a timeseries
+           of state given the model.
+
+        Args
+        ----
+            event_list: the event list, assumed to be sorted by time.
+            include_final_state: should the final state be included in the
+                returned timeseries?  If `True`, then the time dimension of
+                the returned tensor will be 1 greater than the length of the
+                event list.  If `False` (default) these will be equal.
+
+        Return
+        ------
+        A `[T, N, S]` tensor where `T` is the number of events, `N` is the
+        number of individuals, and `S` is the number of states.
+        """
+        return compute_state(
+            self.incidence_matrix,
+            self.initial_state,
+            event_list,
+            include_final_state,
+        )
+
+    # Bypass the reshaping that tfd.Distribution._call_sample_n does
+    def _call_sample_n(self, sample_shape, seed) -> EpidemicEvent:
+        return self._sample_n(sample_shape, seed)
+
+    def _sample_n(self, sample_shape: int, seed=None) -> EpidemicEvent:
+        """
+        Samples n outcomes from the continuous time state transition model.
+
+        Args:
+            n (int): The number of realisations of the Markov process to sample
+                     (currently ignored).
+            seed (int, optional): The seed value for random number generation.
+                                  Defaults to None.
+
+        Returns:
+            EpidemicEvent: A list of n outcomes sampled from the continuous time
+                           state transition model.
+        """
+
+        outcome = continuous_markov_simulation(
+            transition_rate_fn=self.transition_rate_fn,
+            incidence_matrix=self._incidence_matrix,
+            initial_state=self._initial_state,
+            initial_time=self._initial_time,
+            num_markov_jumps=self.num_events,
+            seed=seed,
+        )
+
+        return outcome
+
+    def log_prob(self, value: EpidemicEvent) -> float:
+        return self._log_prob(value)
+
+    def _log_prob(self, value: EpidemicEvent) -> float:
+        """
+        Computes the log probability of the given outcomes.
+
+        Args:
+            value (EpidemicEvent): an EpidemicEvent object representing the
+                                   outcomes.
+
+        Returns:
+            float: The log probability of the given outcomes.
+        """
+        log_lik = continuous_time_log_likelihood(
+            transition_rate_fn=self.transition_rate_fn,
+            incidence_matrix=self.incidence_matrix,
+            initial_state=self.initial_state,
+            initial_time=self.initial_time,
+            num_jumps=self.num_events,
+            event_list=value,
+        )
+
+        return log_lik
+
+    def _event_shape_tensor(self) -> Tensor:
+        return tf.constant([self.num_events], dtype=tf.int32)
+
+    def _event_shape(self) -> tf.TensorShape:
+        return tf.TensorShape([self.num_events])
+
+    def _batch_shape_tensor(self) -> Tensor:
+        return tf.constant([])
+
+    def _batch_shape(self) -> tf.TensorShape:
+        return tf.TensorShape([])

gemlib/distributions/continuous_time_state_transition_model_test.py

@@ -0,0 +1,289 @@
+"""Test ContinuousTimeStateTransitionModel"""
+
+import numpy as np
+import pytest
+import tensorflow as tf
+from scipy.optimize import minimize
+
+from gemlib.distributions.continuous_time_state_transition_model import (
+    ContinuousTimeStateTransitionModel,
+    EpidemicEvent,
+    compute_state,
+)
+
+NUM_EVENTS = 1999
+
+
+@pytest.fixture
+def example_ilm():
+    """A simple event list with 4 individuals, SIR model"""
+    return {
+        "incidence_matrix": np.array(
+            [[-1, 0], [1, -1], [0, 1]], dtype=np.float32
+        ),
+        "event_list": EpidemicEvent(
+            time=np.array(
+                [0.4, 1.3, 1.5, 1.9, 2.3, np.inf, np.inf], dtype=np.float32
+            ),
+            transition=np.array([0, 0, 1, 1, 1, 2, 2], dtype=np.int32),
+            individual=np.array([1, 2, 0, 2, 1, 0, 0], dtype=np.int32),
+        ),
+        "initial_conditions": np.array(
+            [[0, 1, 0], [1, 0, 0], [1, 0, 0], [1, 0, 0]], dtype=np.float32
+        ),
+    }
+
+
+@pytest.fixture
+def simple_sir_model():
+    def rate_fn(t, state):
+        si_rate = 0.25 * state[:, 1] / tf.reduce_sum(state, axis=-1)
+        ir_rate = tf.broadcast_to([0.14], si_rate.shape)
+
+        return si_rate, ir_rate
+
+    # [3 species, 2 reactions]
+    incidence_matrix = np.array([[-1, 0], [1, -1], [0, 1]], dtype=np.float32)
+
+    initial_state = np.array(
+        [[999, 1, 0]], dtype=np.float32
+    )  # [1 unit, 3 classes]
+
+    return ContinuousTimeStateTransitionModel(
+        transition_rate_fn=rate_fn,
+        incidence_matrix=incidence_matrix,
+        initial_state=initial_state,
+        num_events=NUM_EVENTS,
+        initial_time=0.0,
+    )
+
+
+def test_simple_sir_shapes(simple_sir_model):
+    """Test expected output shape"""
+    tf.debugging.assert_equal(
+        simple_sir_model.event_shape_tensor(), tf.constant(NUM_EVENTS)
+    )
+    tf.debugging.assert_equal(
+        simple_sir_model.event_shape, tf.TensorShape([NUM_EVENTS])
+    )
+    tf.debugging.assert_equal(
+        simple_sir_model.batch_shape_tensor(), tf.constant([], tf.int32)
+    )
+    tf.debugging.assert_equal(simple_sir_model.batch_shape, tf.TensorShape([]))
+
+
+def test_simple_sir_eager(simple_sir_model):
+    """Test a simple SIR model"""
+
+    sample = simple_sir_model.sample(seed=[0, 0])
+
+    assert isinstance(sample, EpidemicEvent)
+
+    state = simple_sir_model.compute_state(sample)
+    tf.debugging.assert_non_negative(state)
+
+
+def test_simple_sir_graph(simple_sir_model):
+    """Test a simple SIR model"""
+
+    @tf.function
+    def fn():
+        return simple_sir_model.sample(seed=[0, 0])
+
+    sample = fn()
+
+    assert isinstance(sample, EpidemicEvent)
+
+    state = simple_sir_model.compute_state(sample)
+    tf.debugging.assert_non_negative(state)
+
+
+def test_compute_state(example_ilm):
+    expected_state_eager = compute_state(
+        example_ilm["incidence_matrix"],
+        example_ilm["initial_conditions"],
+        example_ilm["event_list"],
+        include_final_state=True,
+    )
+
+    @tf.function
+    def compute_state_graph(*args):
+        return compute_state(*args)
+
+    expected_state_graph = compute_state_graph(
+        example_ilm["incidence_matrix"],
+        example_ilm["initial_conditions"],
+        example_ilm["event_list"],
+        True,
+    )
+
+    actual_state = np.array(
+        [
+            [
+                [0, 1, 0],  # T=0
+                [1, 0, 0],
+                [1, 0, 0],
+                [1, 0, 0],
+            ],
+            [
+                [0, 1, 0],  # T=1
+                [0, 1, 0],
+                [1, 0, 0],
+                [1, 0, 0],
+            ],
+            [
+                [0, 1, 0],  # T=2
+                [0, 1, 0],
+                [0, 1, 0],
+                [1, 0, 0],
+            ],
+            [
+                [0, 0, 1],  # T=3
+                [0, 1, 0],
+                [0, 1, 0],
+                [1, 0, 0],
+            ],
+            [
+                [0, 0, 1],  # T=4
+                [0, 1, 0],
+                [0, 0, 1],
+                [1, 0, 0],
+            ],
+            [
+                [0, 0, 1],  # T=5
+                [0, 0, 1],
+                [0, 0, 1],
+                [1, 0, 0],
+            ],
+            [
+                [0, 0, 1],  # T=6
+                [0, 0, 1],
+                [0, 0, 1],
+                [1, 0, 0],
+            ],
+            [
+                [0, 0, 1],  # T=7
+                [0, 0, 1],
+                [0, 0, 1],
+                [1, 0, 0],
+            ],
+        ],
+        dtype=np.float32,
+    )
+
+    np.testing.assert_array_equal(expected_state_eager, actual_state)
+    np.testing.assert_array_equal(expected_state_graph, actual_state)
+
+
+def test_simple_sir_loglik(example_ilm):
+    """Test loglikelihood function"""
+    # epi constants
+    incidence_matrix = example_ilm["incidence_matrix"]
+    initial_population = example_ilm["initial_conditions"]
+
+    def rate_fn(t, state):
+        si_rate = tf.broadcast_to([0.5], [state.shape[0]])
+        ir_rate = tf.broadcast_to([0.7], si_rate.shape)
+
+        return si_rate, ir_rate
+
+    # create an instance of the model
+    epi_model = ContinuousTimeStateTransitionModel(
+        transition_rate_fn=rate_fn,
+        incidence_matrix=incidence_matrix,
+        initial_state=initial_population,
+        num_events=NUM_EVENTS,
+        initial_time=0.0,
+    )
+
+    log_lik = epi_model.log_prob(example_ilm["event_list"])
+    # hand calculated log likelihood
+    actual_loglik = -7.256319192936088
+
+    np.testing.assert_almost_equal(log_lik, desired=actual_loglik, decimal=5)
+
+
+def test_simple_sir_loglik_graph_mode(example_ilm):
+    """Test loglikelihood function"""
+    # epi constants
+    incidence_matrix = example_ilm["incidence_matrix"]
+    initial_population = example_ilm["initial_conditions"]
+
+    def rate_fn(t, state):
+        si_rate = tf.broadcast_to([0.5], [state.shape[0]])
+        ir_rate = tf.broadcast_to([0.7], si_rate.shape)
+
+        return si_rate, ir_rate
+
+    # create an instance of the model
+    epi_model = ContinuousTimeStateTransitionModel(
+        transition_rate_fn=rate_fn,
+        incidence_matrix=incidence_matrix,
+        initial_state=initial_population,
+        num_events=NUM_EVENTS,
+        initial_time=0.0,
+    )
+
+    @tf.function
+    def fn():
+        log_lik = epi_model.log_prob(example_ilm["event_list"])
+        return log_lik
+
+    log_lik = fn()
+    # hand calculated log likelihood
+    actual_loglik = -7.256319192936088
+
+    np.testing.assert_almost_equal(log_lik, desired=actual_loglik, decimal=5)
+
+
+def test_simple_sir_workflow(simple_sir_model):
+    """Using an instance of the ContinuousTimeStateTransitionModel"""
+
+    # sample from the model
+    sample_epi_path = simple_sir_model.sample(seed=[20240820, 1347])
+
+    # maximize the likelihood to estimate the parameters of the model
+    def make_rate_fn(rate_parameters):
+        SI_rate = rate_parameters[0]
+        IR_rate = rate_parameters[1]
+
+        def rate_fn(t, state):
+            si_rate = SI_rate * state[:, 1] / tf.reduce_sum(state, axis=-1)
+            ir_rate = tf.broadcast_to([IR_rate], si_rate.shape)
+
+            return si_rate, ir_rate
+
+        return rate_fn
+
+    def mle_fn(log_rate_parameters):
+        rate_parameters = tf.math.exp(log_rate_parameters)
+        rate_fn = make_rate_fn(rate_parameters)
+
+        model = ContinuousTimeStateTransitionModel(
+            transition_rate_fn=rate_fn,
+            incidence_matrix=simple_sir_model.incidence_matrix,
+            initial_state=simple_sir_model.initial_state,
+            num_events=NUM_EVENTS,
+            initial_time=0.0,
+        )
+
+        log_lik = model.log_prob(sample_epi_path)
+        return -log_lik
+
+    initial_parameters = np.array([-0.1, -0.1], dtype=np.float32)
+    opt = minimize(
+        mle_fn,
+        initial_parameters,
+        method="L-BFGS-B",
+        jac="3-point",
+        options={"finite_diff_rel_step": None},
+    )
+
+    std_errors = np.sqrt(np.diagonal(opt.hess_inv.todense()))
+    lower_ci = np.exp(opt.x - 1.96 * std_errors)  # 95% CI
+    upper_ci = np.exp(opt.x + 1.96 * std_errors)
+
+    actuals = np.array([0.25, 0.14])
+
+    assert opt.success
+    assert np.all((lower_ci < actuals) & (actuals < upper_ci))