gemlib 0.9.2__tar.gz

gemlib-0.9.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 The GEM Authors.  All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

gemlib-0.9.2/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.1
+Name: gemlib
+Version: 0.9.2
+Summary: GEMlib scientific compute library for epidemic modelling
+Home-page: http://fhm-chicas-code.lancs.ac.uk/GEM/gemlib
+Author: Chris Jewell
+Author-email: c.jewell@lancaster.ac.uk
+Requires-Python: >=3.9.0,<3.12.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: scipy (>=1.12.0,<2.0.0)
+Requires-Dist: sphinx-book-theme (>=1.1.3,<2.0.0)
+Requires-Dist: tensorflow (>=2.15.0,<2.16.0) ; sys_platform == "linux"
+Requires-Dist: tensorflow-cpu (>=2.15.0,<2.16.0) ; sys_platform == "darwin"
+Requires-Dist: tensorflow-intel (>=2.15.0,<2.16.0) ; sys_platform == "win32"
+Requires-Dist: tensorflow-probability (>=0.23.0,<0.24.0)
+Project-URL: Repository, http://fhm-chicas-code.lancs.ac.uk/GEM/gemlib

gemlib-0.9.2/gemlib/__init__.py

@@ -0,0 +1,9 @@
+"""gemlib scientific compute library for epidemics."""
+
+import importlib.metadata
+
+__version__ = importlib.metadata.version(__package__)
+
+from gemlib import distributions, mcmc, util
+
+__all__ = ["mcmc", "distributions", "util"]

gemlib-0.9.2/gemlib/distributions/__init__.py

@@ -0,0 +1,26 @@
+"""Distribution addons for Tensorflow Probability"""
+
+from gemlib.distributions.brownian import BrownianBridge, BrownianMotion
+from gemlib.distributions.categorical2 import Categorical2
+from gemlib.distributions.continuous_time_state_transition_model import (
+    ContinuousTimeStateTransitionModel,
+)
+from gemlib.distributions.discrete_time_state_transition_model import (
+    DiscreteTimeStateTransitionModel,
+)
+from gemlib.distributions.hypergeometric import Hypergeometric
+from gemlib.distributions.kcategorical import UniformKCategorical
+from gemlib.distributions.uniform_integer import UniformInteger
+
+__all__ = [
+    "BrownianBridge",
+    "BrownianMotion",
+    "Categorical2",
+    "DiscreteApproxContStateTransitionModel",
+    "DiscreteTimeStateTransitionModel",
+    "ContinuousTimeStateTransitionModel",
+    "StateTransitionMarginalModel",
+    "UniformKCategorical",
+    "UniformInteger",
+    "Hypergeometric",
+]

gemlib-0.9.2/gemlib/distributions/brownian.py

@@ -0,0 +1,141 @@
+"""Brownian motion as a distribution"""
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.internal import (
+    distribution_util as dist_util,
+)
+from tensorflow_probability.python.internal import (
+    dtype_util,
+    reparameterization,
+)
+from tensorflow_probability.python.internal.tensor_util import (
+    convert_nonref_to_tensor,
+)
+
+tfd = tfp.distributions
+
+
+class BrownianMotion(tfd.Distribution):
+    def __init__(
+        self,
+        index_points,
+        x0=0.0,
+        scale=1.0,
+        validate_args=False,
+        allow_nan_stats=True,
+        name="BrownianMotion",
+    ):
+        parameters = dict(locals())
+        dtype = dtype_util.common_dtype([x0, index_points, scale])
+        self._x0 = convert_nonref_to_tensor(x0, dtype_hint=dtype)
+
+        self._index_points = convert_nonref_to_tensor(
+            index_points, dtype_hint=dtype
+        )
+        self._scale = tf.convert_to_tensor(scale, dtype_hint=dtype)
+
+        self._increments = tfd.MultivariateNormalDiag(
+            loc=tf.zeros_like(self._index_points[..., 1:]),
+            scale_diag=tf.math.sqrt(
+                self._index_points[..., 1:] - self._index_points[..., :-1]
+            )
+            * self._scale,
+            validate_args=validate_args,
+            allow_nan_stats=allow_nan_stats,
+            name="bm_increments",
+        )  # iid increments
+
+        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._x0.shape)
+
+    def _event_shape(self):
+        return tf.TensorShape(self._index_points.shape[-1] - 1)
+
+    def _sample_n(self, n, seed=None):
+        return self._x0 + tf.math.cumsum(
+            self._increments.sample(n, seed=seed), axis=-1
+        )
+
+    def _log_prob(self, x):
+        path = dist_util.pad(x, axis=-1, front=True, value=self._x0)
+        diff = path[..., 1:] - path[..., :-1]
+        return self._increments.log_prob(diff)
+
+
+class BrownianBridge(tfd.Distribution):
+    def __init__(
+        self,
+        index_points,
+        x0=0.0,
+        x1=0.0,
+        scale=1.0,
+        validate_args=False,
+        allow_nan_stats=True,
+        name="BrownianBridge",
+    ):
+        parameters = dict(locals())
+        dtype = dtype_util.common_dtype([index_points, x0, x1, scale])
+        self._index_points = convert_nonref_to_tensor(
+            index_points, dtype_hint=dtype
+        )
+        self._x0 = convert_nonref_to_tensor(x0, dtype_hint=dtype)
+        self._x1 = convert_nonref_to_tensor(x1, dtype_hint=dtype)
+        self._scale = convert_nonref_to_tensor(scale, dtype_hint=dtype)
+
+        self._increments = tfd.MultivariateNormalDiag(
+            loc=0.0,
+            scale_diag=tf.math.sqrt(
+                self._index_points[..., 1:] - self._index_points[..., :-1]
+            )
+            * self._scale,
+            name="bb_increments",
+        )
+
+        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,
+                name=name,
+            )
+
+    def _batch_shape(self):
+        return tf.TensorShape(self._x0.shape)
+
+    def _event_shape(self):
+        return tf.TensorShape(self._index_points.shape[-1] - 2)
+
+    def _sample_n(self, n, seed=None):
+        """Sampling based on re-leveling pure
+        Brownian motion
+        """
+        z = self._increments.sample(n)
+        z = tf.cumsum(z, axis=-1)
+
+        y_ref_0 = tf.stack([tf.zeros_like(z[..., 0]), z[..., -1]], axis=-1)
+        y_ref_1 = tf.stack([self._x0, self._x1], axis=-1)
+        line = tfp.math.interp_regular_1d_grid(
+            x=self._index_points[..., 1:-1],
+            x_ref_min=self._index_points[..., 0],
+            x_ref_max=self._index_points[..., -1],
+            y_ref=y_ref_1 - y_ref_0,
+        )
+        return z[..., :-1] + line
+
+    def _log_prob(self, x):
+        path = dist_util.pad(x, -1, front=True, value=self._x0)
+        path = dist_util.pad(path, -1, back=True, value=self._x1)
+        diff = path[..., 1:] - path[..., :-1]
+        return self._increments.log_prob(diff)

gemlib-0.9.2/gemlib/distributions/categorical2.py

@@ -0,0 +1,33 @@
+"""Categorical2 corrects a bug in the tfd.Categorical.log_prob"""
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+from tensorflow_probability.python.distributions.categorical import (
+    _broadcast_cat_event_and_params,
+)
+from tensorflow_probability.python.internal import distribution_util, dtype_util
+
+tfd = tfp.distributions
+
+
+# Todo remove this class when https://github.com/tensorflow/tensorflow/issues/40606
+#   is fixed
+class Categorical2(tfd.Categorical):
+    """Done to override the faulty log_prob in tfd.Categorical due to
+    https://github.com/tensorflow/tensorflow/issues/40606
+    """
+
+    def _log_prob(self, k):
+        with tf.name_scope("Cat2log_prob"):
+            logits = self.logits_parameter()
+            if self.validate_args:
+                k = distribution_util.embed_check_integer_casting_closed(
+                    k, target_dtype=self.dtype
+                )
+            k, logits = _broadcast_cat_event_and_params(
+                k, logits, base_dtype=dtype_util.base_dtype(self.dtype)
+            )
+            logits_normalised = tf.math.log(tf.math.softmax(logits))
+            return tf.cast(
+                tf.gather(logits_normalised, k, batch_dims=1), logits.dtype
+            )

gemlib-0.9.2/gemlib/distributions/continuous_markov.py

@@ -0,0 +1,371 @@
+"""Function for continuous time simulation"""
+
+from typing import Callable, List, NamedTuple
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+from gemlib.util import batch_gather, transition_coords
+
+# aliasing for convenience
+tfd = tfp.distributions
+Tensor = tf.Tensor
+DTYPE = tf.float32
+
+
+class EpidemicEvent(NamedTuple):
+    """Tracker of an event in an epidemic simulation
+
+    Attributes:
+        time (float): The time at which the event occurred.
+        transition (int): The type of transition that occurred.
+        individual (int): The individual involved in the event.
+    """
+
+    time: float
+    transition: int
+    individual: int
+
+
+def _one_hot_expand_state(condensed_state: tf.Tensor) -> tf.Tensor:
+    """Expand the state of the epidemic to a one-hot representation
+    Args:
+        epidemic_state: The state of the epidemic
+    Returns:
+        The one-hot representation of the epidemic state
+    """
+    # Create one-hot encoded vectors for each state
+    one_hot_states = tf.one_hot(
+        tf.range(len(condensed_state)),
+        depth=len(condensed_state),
+        dtype=tf.float32,
+    )
+    # Repeat each one-hot state based on its corresponding count
+    repeated_states = tf.repeat(one_hot_states, condensed_state, axis=0)
+
+    # Reshape and transpose to get state per row representation
+    return repeated_states
+
+
+def _total_flux(transition_rates, state, incidence_matrix):
+    """Multiplies `transition_rates` by source `state`s to return
+       the total flux along transitions given `state`.
+
+    Args
+    ----
+        transition_rates: a `[R,N]` tensor of per-individual transition rates
+                          for `R` transitions and `N` aggregation units.
+        state: a `[N, S]` tensor of `N` aggregation units and `S` states.
+        incidence_matrix: a `[S, R]` matrix describing the change in `S` for
+                          each transition `R`.
+
+    Returns
+    -------
+    A [R,N] tensor of total flux along each transition, taking into account the
+    availability of individuals in the source state.
+    """
+    source_state_idx = transition_coords(incidence_matrix)[:, 0]
+    source_states = batch_gather(state, indices=source_state_idx[:, tf.newaxis])
+    transition_rates = tf.stack(transition_rates, axis=-1)
+
+    return tf.einsum("...nr,...nr->...rn", transition_rates, source_states)
+
+
+def compute_state(
+    incidence_matrix: Tensor,
+    initial_state: Tensor,
+    event_list: EpidemicEvent,
+    include_final_state: bool = False,
+):
+    """Given an event list `event_list`, compute a timeseries
+       of state given the model.
+
+    Args
+    ----
+        incidence_matrix: a `[S,R]` graph incidence matrix for `S`
+            compartments and `R` transitions.
+        initial_state: a `[N,S]` representing the initial state of `N`
+            units by `S` compartments.
+        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.
+    """
+    event_list = event_list.__class__(
+        *[tf.convert_to_tensor(x) for x in event_list]
+    )
+
+    initial_state = tf.convert_to_tensor(initial_state)
+    incidence_matrix = tf.convert_to_tensor(incidence_matrix)
+
+    # Compute one-hot encoding of event timeseries
+    event_tensor_shape = (
+        event_list.time.shape[0],  # T
+        initial_state.shape[-2],  # N
+        incidence_matrix.shape[-1] + 1,  # R + pad for ghost events
+    )
+    hot_indices = tf.stack(
+        [
+            tf.range(event_list.time.shape[0]),  # T
+            event_list.individual,  # N
+            event_list.transition,  # R
+        ],
+        axis=-1,
+    )
+    event_tensor = tf.scatter_nd(
+        indices=hot_indices,
+        updates=tf.ones([event_list.individual.shape[-1]], initial_state.dtype),
+        shape=event_tensor_shape,
+    )[..., :-1]  # Clip last dimension to remove ghost events
+
+    # Compute deltas and cumsum over the state
+    delta = tf.matmul(event_tensor, incidence_matrix, transpose_b=True)
+
+    if include_final_state is False:
+        delta = delta[:-1]
+
+    state = tf.cumsum(
+        tf.concat([initial_state[tf.newaxis, ...], delta], axis=-3)
+    )
+
+    return state
+
+
+def exponential_propogate(
+    transition_rate_fn: Callable, incidence_matrix: Tensor
+) -> EpidemicEvent:
+    """Generates a function for propogating an epidemic forward in time
+
+    Closure over the transition rate function and the incidence matrix
+    which outline the epidemic dynamics and model structure. The returned
+    function can be used to simulate the epidemic forward in time one step.
+
+    Args:
+        transition_rate_fn (Callable): a function that takes the current
+            state of the epidemic and returns the transition rates for each
+            individual/meta-population.
+        incidence_matrix (tensor): A `[R, S]` matrix that describes the graph
+            structure of the state transition mode. The rows correspond to the
+            `R` transitions and the columns correspond to the `S` states.
+
+    Returns:
+        EpidemicEvent: A NamedTuple that describes the next event in the
+            epidemic.
+    """
+    tr_incidence_matrix = tf.transpose(incidence_matrix)
+
+    def propogate_fn(time: float, state: Tensor, seed: int) -> List:
+        """Propogates the state of the epidemic forward in time
+
+        Args:
+            time (float): Wall clock of the epidemic - can easily recover
+            the time delta
+            state (tensor): `[N,S]` representing the current state.
+
+        Returns:
+            EpidemicEvent: The next event in the epidemic.
+        """
+        seed_exp, seed_cat = tfp.random.split_seed(seed, n=2)
+        num_units = state.shape[-2]
+
+        # compute event rates for all possible events
+        transition_rates = _total_flux(
+            transition_rate_fn(time, state), state, incidence_matrix
+        )
+
+        # simulate next time
+        t_next = tfd.Exponential(rate=tf.reduce_sum(transition_rates)).sample(
+            seed=seed_exp
+        )
+
+        # use categorical distribution to get event type and indiviudal id
+        event_id = tfd.Categorical(
+            probs=tf.reshape(transition_rates, shape=(-1,)),
+            dtype=tf.int32,
+        ).sample(seed=seed_cat)
+
+        unit_idx = tf.math.floormod(event_id, num_units)
+        transition_idx = tf.math.floordiv(event_id, num_units)
+
+        # update the state
+        new_state = tf.tensor_scatter_nd_add(
+            state,
+            [[unit_idx]],
+            [tr_incidence_matrix[transition_idx]],
+        )
+
+        return (
+            time + t_next,
+            new_state,
+            EpidemicEvent(time + t_next, transition_idx, unit_idx),
+        )
+
+    return propogate_fn
+
+
+def continuous_markov_simulation(
+    transition_rate_fn: Callable,
+    initial_state: Tensor,
+    incidence_matrix: Tensor,
+    num_markov_jumps: int,
+    initial_time: float = 0.0,
+    seed=None,
+) -> EpidemicEvent:
+    """
+    Simulates a continuous-time Markov process
+
+    Args:
+        transition_rate_fn (Callable): A function that computes the transition
+            rates given the current state and incidence matrix.
+        initial_state (Tensor): A [N, S] tensor, respresenting a population of N
+                        units and S states.
+        num_markov_jumps (int): The number of iterations to simulate.
+        incidence_matrix (Tensor): The `[S,R]` incidence matrix representing the
+            state transition model with S states and R transitions.
+        seed (Optional[List(int,int)): The random seed.
+    Returns:
+        EpidemicEvent: An object containing the simulated epidemic events.
+
+    """
+    initial_state = tf.convert_to_tensor(initial_state)
+    incidence_matrix = tf.convert_to_tensor(incidence_matrix)
+    dtype = initial_state.dtype
+    seed = tfp.random.sanitize_seed(seed, salt="continuous_markov_simulation")
+
+    propagate_fn = exponential_propogate(transition_rate_fn, incidence_matrix)
+
+    accum = EpidemicEvent(
+        time=tf.TensorArray(dtype, size=num_markov_jumps, dynamic_size=False),
+        transition=tf.TensorArray(
+            tf.int32, size=num_markov_jumps, dynamic_size=False
+        ),
+        individual=tf.TensorArray(
+            tf.int32, size=num_markov_jumps, dynamic_size=False
+        ),
+    )
+
+    def cond(i, time, state, *_):
+        transition_rates = _total_flux(
+            transition_rate_fn(time, state), state, incidence_matrix
+        )
+        cont = (i < num_markov_jumps) & (tf.reduce_sum(transition_rates) > 0.0)
+        return cont
+
+    def body(i, time, state, seed, accum):
+        next_seed, this_seed = tfp.random.split_seed(seed, salt="body")
+        next_time, next_state, event = propagate_fn(time, state, this_seed)
+        accum = EpidemicEvent(*[x.write(i, y) for x, y in zip(accum, event)])
+        return i + 1, next_time, next_state, next_seed, accum
+
+    actual_markov_jumps, _, _, _, accum = tf.while_loop(
+        cond, body, loop_vars=(0, initial_time, initial_state, seed, accum)
+    )
+
+    # Pad unused parts of the output TensorArrays if the
+    # loop terminates before num_markov_jumps
+    indices = tf.range(actual_markov_jumps, num_markov_jumps)
+    fills = tf.fill([num_markov_jumps - actual_markov_jumps], np.inf)
+
+    output = EpidemicEvent(
+        time=accum.time.scatter(
+            indices,
+            fills,
+        ),
+        individual=accum.individual.scatter(indices, tf.zeros_like(indices)),
+        transition=accum.transition.scatter(
+            indices,
+            tf.fill(
+                [num_markov_jumps - actual_markov_jumps],
+                incidence_matrix.shape[1],
+            ),
+        ),
+    )
+
+    output = tf.nest.map_structure(lambda x: x.stack(), output)
+
+    return output
+
+
+def continuous_time_log_likelihood(
+    transition_rate_fn: Callable,
+    incidence_matrix: Tensor,
+    initial_state: Tensor,
+    initial_time: float,
+    num_jumps: int,
+    event_list: EpidemicEvent,
+) -> float:
+    """
+    Computes the log-likelihood of a continuous-time Markov process
+    given the transition rate function,
+    incidence matrix, initial state, number of jumps, and event data.
+
+    Args:
+        transition_rate_fn (Callable): A function that computes the
+        transition rate given the current state and time.
+        incidence_matrix: The incidence matrix representing
+        the connections between states in `[S,R]` format.
+        initial_state: The initial state of the process as a `[N,R]`.
+        num_jumps (int): The number of jumps to simulate.
+        event (EpidemicEvent): The event data containing the times
+        and states.
+
+    Returns:
+        Tensor: The log-likelihood of the continuous-time Markov process.
+    """
+    # construct the epidemic states [T, N, S]
+    states = compute_state(
+        incidence_matrix=incidence_matrix,
+        initial_state=initial_state,
+        event_list=event_list,
+    )
+
+    # compute the transition rates for each of the states in the event
+    time = tf.concat([[initial_time], event_list.time], axis=-1)
+    rates = tf.vectorized_map(
+        fn=lambda x: transition_rate_fn(*x),
+        elems=(time[:-1], states),
+    )  # R-tuple of [T,N] tensors
+
+    # zero out impossible events - this is a hacky way to do it (element wise
+    # matrix multiplication)
+    total_flux = _total_flux(rates, states, incidence_matrix)
+
+    indices = tf.stack(
+        [
+            tf.range(event_list.time.shape[0]),
+            tf.clip_by_value(  # Clip due to ghost events (values zeroed later)
+                event_list.transition,
+                clip_value_min=0,
+                clip_value_max=incidence_matrix.shape[-1] - 1,
+            ),
+            event_list.individual,
+        ],
+        axis=-1,
+    )
+
+    # compute event specific rate - get indices of the event that happened
+    event_rate = tf.gather_nd(total_flux, indices)
+
+    # compute total rate per timestep
+    total_rate = tf.reduce_sum(total_flux, axis=(-2, -1))
+    # total_rate = tf.einsum("tns -> t", total_flux)
+
+    # compute time deltas
+    time_delta = time[1:] - time[:-1]
+
+    # compute the log-likelihood
+    loglik_t = -total_rate * time_delta + tf.math.log(event_rate)
+
+    # Zero out for any inf times (i.e. possible padding of event_list chunk)
+    loglik_t = tf.where(
+        tf.math.is_finite(time_delta), loglik_t, tf.zeros_like(loglik_t)
+    )
+
+    return tf.reduce_sum(loglik_t)