gemlib 0.9.2__tar.gz → 0.9.4__tar.gz

{gemlib-0.9.2 → gemlib-0.9.4}/PKG-INFO

@@ -1,10 +1,13 @@
 Metadata-Version: 2.1
 Name: gemlib
-Version: 0.9.2
+Version: 0.9.4
 Summary: GEMlib scientific compute library for epidemic modelling
-Home-page: http://fhm-chicas-code.lancs.ac.uk/GEM/gemlib
+Home-page: https://gem-epidemics.gitlab.io/gemlib
+Keywords: epidemic,Bayesian,inference,infectious disease model,probabilistic programming
 Author: Chris Jewell
 Author-email: c.jewell@lancaster.ac.uk
+Maintainer: Jessica Bridgen
+Maintainer-email: j.bridgen@lancaster.ac.uk
 Requires-Python: >=3.9.0,<3.12.0
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
@@ -16,4 +19,4 @@ 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
+Project-URL: Repository, https://gitlab.com/gem-epidemics/gemlib

{gemlib-0.9.2 → gemlib-0.9.4}/gemlib/distributions/continuous_markov.py

@@ -14,7 +14,7 @@ Tensor = tf.Tensor
 DTYPE = tf.float32
 
 
-class EpidemicEvent(NamedTuple):
+class EventList(NamedTuple):
     """Tracker of an event in an epidemic simulation
 
     Attributes:
@@ -65,6 +65,7 @@ def _total_flux(transition_rates, state, incidence_matrix):
     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)
@@ -75,7 +76,7 @@ def _total_flux(transition_rates, state, incidence_matrix):
 def compute_state(
     incidence_matrix: Tensor,
     initial_state: Tensor,
-    event_list: EpidemicEvent,
+    event_list: EventList,
     include_final_state: bool = False,
 ):
     """Given an event list `event_list`, compute a timeseries
@@ -140,7 +141,7 @@ def compute_state(
 
 def exponential_propogate(
     transition_rate_fn: Callable, incidence_matrix: Tensor
-) -> EpidemicEvent:
+) -> EventList:
     """Generates a function for propogating an epidemic forward in time
 
     Closure over the transition rate function and the incidence matrix
@@ -156,7 +157,7 @@ def exponential_propogate(
             `R` transitions and the columns correspond to the `S` states.
 
     Returns:
-        EpidemicEvent: A NamedTuple that describes the next event in the
+        EventList: A NamedTuple that describes the next event in the
             epidemic.
     """
     tr_incidence_matrix = tf.transpose(incidence_matrix)
@@ -170,7 +171,7 @@ def exponential_propogate(
             state (tensor): `[N,S]` representing the current state.
 
         Returns:
-            EpidemicEvent: The next event in the epidemic.
+            EventList: The next event in the epidemic.
         """
         seed_exp, seed_cat = tfp.random.split_seed(seed, n=2)
         num_units = state.shape[-2]
@@ -204,7 +205,7 @@ def exponential_propogate(
         return (
             time + t_next,
             new_state,
-            EpidemicEvent(time + t_next, transition_idx, unit_idx),
+            EventList(time + t_next, transition_idx, unit_idx),
         )
 
     return propogate_fn
@@ -217,7 +218,7 @@ def continuous_markov_simulation(
     num_markov_jumps: int,
     initial_time: float = 0.0,
     seed=None,
-) -> EpidemicEvent:
+) -> EventList:
     """
     Simulates a continuous-time Markov process
 
@@ -231,17 +232,17 @@ def continuous_markov_simulation(
             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.
+        EventList: 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
+    dtype = tf.float32
     seed = tfp.random.sanitize_seed(seed, salt="continuous_markov_simulation")
 
     propagate_fn = exponential_propogate(transition_rate_fn, incidence_matrix)
 
-    accum = EpidemicEvent(
+    accum = EventList(
         time=tf.TensorArray(dtype, size=num_markov_jumps, dynamic_size=False),
         transition=tf.TensorArray(
             tf.int32, size=num_markov_jumps, dynamic_size=False
@@ -261,7 +262,7 @@ def continuous_markov_simulation(
     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)])
+        accum = EventList(*[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(
@@ -273,7 +274,7 @@ def continuous_markov_simulation(
     indices = tf.range(actual_markov_jumps, num_markov_jumps)
     fills = tf.fill([num_markov_jumps - actual_markov_jumps], np.inf)
 
-    output = EpidemicEvent(
+    output = EventList(
         time=accum.time.scatter(
             indices,
             fills,
@@ -299,7 +300,7 @@ def continuous_time_log_likelihood(
     initial_state: Tensor,
     initial_time: float,
     num_jumps: int,
-    event_list: EpidemicEvent,
+    event_list: EventList,
 ) -> float:
     """
     Computes the log-likelihood of a continuous-time Markov process
@@ -313,7 +314,7 @@ def continuous_time_log_likelihood(
         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
+        event (EventList): The event data containing the times
         and states.
 
     Returns:

{gemlib-0.9.2 → gemlib-0.9.4}/gemlib/distributions/continuous_time_state_transition_model.py

@@ -7,7 +7,7 @@ import tensorflow_probability as tfp
 from tensorflow_probability.python.internal import reparameterization
 
 from gemlib.distributions.continuous_markov import (
-    EpidemicEvent,
+    EventList,
     compute_state,
     continuous_markov_simulation,
     continuous_time_log_likelihood,
@@ -55,10 +55,18 @@ class ContinuousTimeStateTransitionModel(tfd.Distribution):
 
         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)
+        self._initial_time = tf.convert_to_tensor(
+            initial_time, dtype=self._initial_state.dtype
+        )
+
+        dtype = EventList(
+            time=self._initial_time.dtype,
+            transition=tf.int32,
+            individual=tf.int32,
+        )
 
         super().__init__(
-            dtype=self._incidence_matrix.dtype,
+            dtype=dtype,
             reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
             validate_args=validate_args,
             allow_nan_stats=allow_nan_stats,
@@ -92,7 +100,7 @@ class ContinuousTimeStateTransitionModel(tfd.Distribution):
         return self._parameters["initial_time"]
 
     def compute_state(
-        self, event_list: EpidemicEvent, include_final_state: bool = False
+        self, event_list: EventList, include_final_state: bool = False
     ) -> Tensor:
         """Given an event list `event_list`, compute a timeseries
            of state given the model.
@@ -118,10 +126,10 @@ class ContinuousTimeStateTransitionModel(tfd.Distribution):
         )
 
     # Bypass the reshaping that tfd.Distribution._call_sample_n does
-    def _call_sample_n(self, sample_shape, seed) -> EpidemicEvent:
+    def _call_sample_n(self, sample_shape, seed) -> EventList:
         return self._sample_n(sample_shape, seed)
 
-    def _sample_n(self, sample_shape: int, seed=None) -> EpidemicEvent:
+    def _sample_n(self, sample_shape: int, seed=None) -> EventList:
         """
         Samples n outcomes from the continuous time state transition model.
 
@@ -132,7 +140,7 @@ class ContinuousTimeStateTransitionModel(tfd.Distribution):
                                   Defaults to None.
 
         Returns:
-            EpidemicEvent: A list of n outcomes sampled from the continuous time
+            EventList: A list of n outcomes sampled from the continuous time
                            state transition model.
         """
 
@@ -147,15 +155,12 @@ class ContinuousTimeStateTransitionModel(tfd.Distribution):
 
         return outcome
 
-    def log_prob(self, value: EpidemicEvent) -> float:
-        return self._log_prob(value)
-
-    def _log_prob(self, value: EpidemicEvent) -> float:
+    def _log_prob(self, value: EventList) -> float:
         """
         Computes the log probability of the given outcomes.
 
         Args:
-            value (EpidemicEvent): an EpidemicEvent object representing the
+            value (EventList): an EventList object representing the
                                    outcomes.
 
         Returns:
@@ -172,14 +177,30 @@ class ContinuousTimeStateTransitionModel(tfd.Distribution):
 
         return log_lik
 
-    def _event_shape_tensor(self) -> Tensor:
-        return tf.constant([self.num_events], dtype=tf.int32)
+    def _event_shape_tensor(self) -> EventList:
+        return EventList(
+            time=tf.constant([self.num_events], dtype=tf.int32),
+            transition=tf.constant([self.num_events], dtype=tf.int32),
+            individual=tf.constant([self.num_events], dtype=tf.int32),
+        )
 
-    def _event_shape(self) -> tf.TensorShape:
-        return tf.TensorShape([self.num_events])
+    def _event_shape(self) -> EventList:
+        return EventList(
+            time=tf.TensorShape([self.num_events]),
+            transition=tf.TensorShape([self.num_events]),
+            individual=tf.TensorShape([self.num_events]),
+        )
 
-    def _batch_shape_tensor(self) -> Tensor:
-        return tf.constant([])
+    def _batch_shape_tensor(self) -> EventList:
+        return EventList(
+            time=tf.constant([]),
+            transition=tf.constant([]),
+            individual=tf.constant([]),
+        )
 
-    def _batch_shape(self) -> tf.TensorShape:
-        return tf.TensorShape([])
+    def _batch_shape(self) -> EventList:
+        return EventList(
+            time=tf.TensorShape([]),
+            transition=tf.TensorShape([]),
+            individual=tf.TensorShape([]),
+        )

{gemlib-0.9.2 → gemlib-0.9.4}/gemlib/distributions/continuous_time_state_transition_model_test.py

@@ -1,17 +1,21 @@
 """Test ContinuousTimeStateTransitionModel"""
 
+from collections import namedtuple
+
 import numpy as np
 import pytest
 import tensorflow as tf
+import tensorflow_probability as tfp
 from scipy.optimize import minimize
 
 from gemlib.distributions.continuous_time_state_transition_model import (
     ContinuousTimeStateTransitionModel,
-    EpidemicEvent,
+    EventList,
     compute_state,
 )
 
 NUM_EVENTS = 1999
+tfd = tfp.distributions
 
 
 @pytest.fixture
@@ -21,7 +25,7 @@ def example_ilm():
         "incidence_matrix": np.array(
             [[-1, 0], [1, -1], [0, 1]], dtype=np.float32
         ),
-        "event_list": EpidemicEvent(
+        "event_list": EventList(
             time=np.array(
                 [0.4, 1.3, 1.5, 1.9, 2.3, np.inf, np.inf], dtype=np.float32
             ),
@@ -58,18 +62,163 @@ def simple_sir_model():
     )
 
 
+@pytest.fixture
+def bayesian_sir_model():
+    DTYPE = np.float32
+
+    @tfd.JointDistributionCoroutine
+    def model():
+        # Priors
+        beta = yield tfd.Gamma(
+            concentration=DTYPE(0.1),
+            rate=DTYPE(0.1),
+            name="beta",
+        )
+        gamma = yield tfd.Gamma(
+            concentration=DTYPE(2.0), rate=DTYPE(8.0), name="gamma"
+        )
+
+        # Epidemic model
+        incidence_matrix = np.array(
+            [  #  SI  IR
+                [-1, 0],  # S
+                [1, -1],  # I
+                [0, 1],  # R
+            ],
+            dtype=DTYPE,
+        )
+
+        initial_state = np.array([[99, 1, 0]]).astype(DTYPE)
+
+        def transition_rates(t, state):
+            si_rate = beta * state[:, 1] / tf.reduce_sum(state, axis=-1)
+            ir_rate = tf.fill((state.shape[0],), gamma)
+            return si_rate, ir_rate
+
+        NUM_EVENTS = 200
+
+        yield ContinuousTimeStateTransitionModel(
+            transition_rate_fn=transition_rates,
+            incidence_matrix=incidence_matrix,
+            initial_state=initial_state,
+            num_events=NUM_EVENTS,
+            name="sir",
+        )
+
+    ModelType = namedtuple("StructTuple", ["beta", "gamma", "sir"])
+
+    example = ModelType(
+        beta=0.5,
+        gamma=0.14,
+        sir=EventList(
+            time=np.array(
+                [ 3.8722684,  4.0912385,  4.5234103,  4.722512 ,  4.9462056,
+                  5.7793403,  5.7914076,  6.028009 ,  6.1961355,  6.8580866,
+                  7.6919003,  7.8862953,  8.106527 ,  8.489728 ,  8.565479 ,
+                  8.710205 ,  8.7175045,  8.807663 ,  8.824789 ,  8.863979 ,
+                  9.035543 ,  9.428677 ,  9.470956 ,  9.492353 ,  9.529175 ,
+                  9.57752  ,  9.618181 ,  9.834693 ,  9.88009  ,  9.963752 ,
+                  10.158042 , 10.233447 , 10.283622 , 10.464923 , 10.472176 ,
+                  10.509909 , 10.713008 , 10.932794 , 10.937864 , 11.025746 ,
+                  11.3029   , 11.452712 , 11.45593  , 11.801975 , 11.944026 ,
+                  12.169918 , 12.2015085, 12.289305 , 12.379549 , 12.658843 ,
+                  12.673543 , 12.675543 , 12.758116 , 12.839079 , 12.887654 ,
+                  13.002403 , 13.063832 , 13.06984  , 13.204261 , 13.279876 ,
+                  13.368993 , 13.465008 , 13.545585 , 13.6250305, 13.652922 ,
+                  13.6884165, 13.722584 , 13.7678585, 13.773789 , 13.814655 ,
+                  13.861033 , 14.0029125, 14.039422 , 14.068574 , 14.082346 ,
+                  14.352536 , 14.35859  , 14.422744 , 14.575989 , 14.657471 ,
+                  14.69332  , 14.700054 , 14.848383 , 14.882924 , 15.0288925,
+                  15.093655 , 15.108234 , 15.236706 , 15.322691 , 15.328567 ,
+                  15.404174 , 15.421979 , 15.657191 , 15.922981 , 15.9314375,
+                  16.213312 , 16.311094 , 16.331137 , 16.50571  , 16.542715 ,
+                  16.634466 , 16.743319 , 16.77621  , 16.864067 , 17.02836  ,
+                  17.149601 , 17.375605 , 17.418259 , 17.424673 , 17.46273  ,
+                  17.503548 , 17.648726 , 17.659864 , 17.879429 , 18.013472 ,
+                  18.117981 , 18.781914 , 18.822414 , 18.923925 , 18.97823  ,
+                  19.034103 , 19.124441 , 19.150362 , 19.70805  , 19.848843 ,
+                  19.925968 , 19.967867 , 20.104042 , 20.159836 , 20.608583 ,
+                  21.046917 , 21.156258 , 21.233568 , 21.242342 , 21.574625 ,
+                  21.72764  , 21.865955 , 22.06923  , 22.242723 , 23.101812 ,
+                  23.451588 , 23.622063 , 23.676891 , 23.758053 , 24.117605 ,
+                  24.329521 , 25.265572 , 25.2762   , 25.319914 , 25.386583 ,
+                  25.476671 , 25.518967 , 25.865902 , 26.031075 , 26.163708 ,
+                  26.169014 , 26.238436 , 26.445032 , 26.973305 , 27.094973 ,
+                  27.237108 , 27.26869  , 27.58751  , 27.857409 , 28.816969 ,
+                  29.008768 , 29.6066   , 30.18358  , 31.610064 , 32.21432  ,
+                  32.693485 , 33.329193 , 33.356907 , 33.576965 , 33.786686 ,
+                  34.548702 , 35.99837  , 36.461132 , 36.63711  , 36.97347  ,
+                  37.104862 , 38.354027 , 38.914967 , 39.609924 , 39.66717  ,
+                  42.54664  , 42.610905 , 42.873867 , 43.515656 , 43.913628 ,
+                  45.195133 , 45.952778 , 46.862324 , 47.602505 , 50.113663 ,
+                  54.11174  , 56.933784 , np.inf, np.inf, np.inf,
+                 ],
+                dtype=np.float32,
+            ),
+            transition=np.array(
+              [0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0,
+               1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 1,
+               1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1,
+               0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1,
+               0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 1, 0, 1, 1,
+               1, 0, 1, 1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1,
+               0, 1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0,
+               0, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1,
+               1, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 2,
+               2, 2,
+               ],
+               dtype=np.int32,
+            ),
+            individual=np.array(
+              [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+               0, 0,
+              ],
+              dtype=np.int32,
+            ),
+        )
+    )  # fmt: skip
+
+    return {"model": model, "example": example}
+
+
 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)
+        simple_sir_model.event_shape_tensor(),
+        EventList(
+            tf.constant(NUM_EVENTS),
+            tf.constant(NUM_EVENTS),
+            tf.constant(NUM_EVENTS),
+        ),
     )
     tf.debugging.assert_equal(
-        simple_sir_model.event_shape, tf.TensorShape([NUM_EVENTS])
+        simple_sir_model.event_shape,
+        EventList(
+            tf.TensorShape([NUM_EVENTS]),
+            tf.TensorShape([NUM_EVENTS]),
+            tf.TensorShape([NUM_EVENTS]),
+        ),
     )
     tf.debugging.assert_equal(
-        simple_sir_model.batch_shape_tensor(), tf.constant([], tf.int32)
+        simple_sir_model.batch_shape_tensor(),
+        EventList(
+            tf.constant([], tf.int32),
+            tf.constant([], tf.int32),
+            tf.constant([], tf.int32),
+        ),
+    )
+    tf.debugging.assert_equal(
+        simple_sir_model.batch_shape,
+        EventList(tf.TensorShape([]), tf.TensorShape([]), tf.TensorShape([])),
     )
-    tf.debugging.assert_equal(simple_sir_model.batch_shape, tf.TensorShape([]))
 
 
 def test_simple_sir_eager(simple_sir_model):
@@ -77,7 +226,7 @@ def test_simple_sir_eager(simple_sir_model):
 
     sample = simple_sir_model.sample(seed=[0, 0])
 
-    assert isinstance(sample, EpidemicEvent)
+    assert isinstance(sample, EventList)
 
     state = simple_sir_model.compute_state(sample)
     tf.debugging.assert_non_negative(state)
@@ -92,7 +241,7 @@ def test_simple_sir_graph(simple_sir_model):
 
     sample = fn()
 
-    assert isinstance(sample, EpidemicEvent)
+    assert isinstance(sample, EventList)
 
     state = simple_sir_model.compute_state(sample)
     tf.debugging.assert_non_negative(state)
@@ -287,3 +436,15 @@ def test_simple_sir_workflow(simple_sir_model):
 
     assert opt.success
     assert np.all((lower_ci < actuals) & (actuals < upper_ci))
+
+
+def test_tfp_jd_integration(bayesian_sir_model):
+    model = bayesian_sir_model["model"]
+    example = bayesian_sir_model["example"]
+
+    model.sample(seed=[20240714, 1139])
+
+    conditioned_model = model.experimental_pin(sir=example.sir)
+    lp = conditioned_model.log_prob(beta=0.5, gamma=0.14)
+
+    np.testing.assert_approx_equal(lp, 10.0128927)

{gemlib-0.9.2 → gemlib-0.9.4}/gemlib/distributions/discrete_time_state_transition_model.py

@@ -35,7 +35,7 @@ class DiscreteTimeStateTransitionModel(tfd.Distribution):
 
     def __init__(
         self,
-        transition_rates,
+        transition_rate_fn,
         incidence_matrix,
         initial_state,
         initial_step,
@@ -49,7 +49,7 @@ class DiscreteTimeStateTransitionModel(tfd.Distribution):
 
         Args:
         ----
-          transition_rates: Python callable of the form `fn(t, state)` taking
+          transition_rate_fn: Python callable of the form `fn(t, state)` taking
             the current time `t` (Python float) and state tensor `state`. This
             function returns a tensor which broadcasts to the first dimension of
             `incidence_matrix`.
@@ -117,7 +117,7 @@ class DiscreteTimeStateTransitionModel(tfd.Distribution):
 
             # Instantiate model
             sir = DiscreteTimeStateTransitionModel(
-                    transition_rates=txrates,
+                    transition_rate_fn=txrates,
                     incidence_matrix=incidence_matrix,
                     initial_state=initial_state,
                     initial_step=initial_step,
@@ -161,7 +161,7 @@ class DiscreteTimeStateTransitionModel(tfd.Distribution):
         """
         parameters = dict(locals())
         with tf.name_scope(name) as name:
-            self._transition_rates = transition_rates
+            self._transition_rate_fn = transition_rate_fn
             self._incidence_matrix = tf.convert_to_tensor(
                 incidence_matrix, dtype=initial_state.dtype
             )
@@ -183,8 +183,8 @@ class DiscreteTimeStateTransitionModel(tfd.Distribution):
         self.dtype = initial_state.dtype
 
     @property
-    def transition_rates(self):
-        return self._transition_rates
+    def transition_rate_fn(self):
+        return self._transition_rate_fn
 
     @property
     def incidence_matrix(self):
@@ -261,7 +261,7 @@ class DiscreteTimeStateTransitionModel(tfd.Distribution):
             seed, salt="DiscreteTimeStateTransitionModel"
         )
         t, sim = discrete_markov_simulation(
-            hazard_fn=self.transition_rates,
+            hazard_fn=self.transition_rate_fn,
             state=self.initial_state,
             start=self.initial_step,
             end=self.initial_step + self.num_steps * self.time_delta,
@@ -288,7 +288,7 @@ class DiscreteTimeStateTransitionModel(tfd.Distribution):
         )
         y = tf.convert_to_tensor(y, dtype)
 
-        hazard = self.transition_rates
+        hazard = self.transition_rate_fn
         return discrete_markov_log_prob(
             events=y,
             init_state=self.initial_state,

{gemlib-0.9.2 → gemlib-0.9.4}/gemlib/distributions/discrete_time_state_transition_model_examples.py

@@ -56,7 +56,7 @@ def txrates(t, state):
 
 # Instantiate model
 sir = DiscreteTimeStateTransitionModel(
-    transition_rates=txrates,
+    transition_rate_fn=txrates,
     stoichiometry=stoichiometry,
     initial_state=initial_state,
     initial_step=initial_step,
@@ -166,7 +166,7 @@ def txrates(t, state):
 
 # Instantiate model
 sirs = DiscreteTimeStateTransitionModel(
-    transition_rates=txrates,
+    transition_rate_fn=txrates,
     stoichiometry=stoichiometry,
     initial_state=initial_state,
     initial_step=initial_step,
@@ -252,7 +252,7 @@ def txrates(t, state):
 
 # Instantiate model
 seir = DiscreteTimeStateTransitionModel(
-    transition_rates=txrates,
+    transition_rate_fn=txrates,
     stoichiometry=stoichiometry,
     initial_state=initial_state,
     initial_step=initial_step,
@@ -336,7 +336,7 @@ def txrates(t, state):
 initial_step, time_delta, num_steps = 0.0, 1.0, 100
 
 sirc = DiscreteTimeStateTransitionModel(
-    transition_rates=txrates,
+    transition_rate_fn=txrates,
     stoichiometry=stoichiometry,
     initial_state=initial_state,
     initial_step=initial_step,
@@ -409,7 +409,7 @@ def txrates(t, state):
 
 # Instantiate model
 sivr = DiscreteTimeStateTransitionModel(
-    transition_rates=txrates,
+    transition_rate_fn=txrates,
     stoichiometry=stoichiometry,
     initial_state=initial_state,
     initial_step=initial_step,

{gemlib-0.9.2 → gemlib-0.9.4}/gemlib/distributions/discrete_time_state_transition_model_test.py

@@ -38,7 +38,7 @@ class TestDiscreteTimeStateTransitionModel(test_util.TestCase):
             return [si, ir]
 
         return DiscreteTimeStateTransitionModel(
-            transition_rates=txrates,
+            transition_rate_fn=txrates,
             incidence_matrix=incidence_matrix,
             initial_state=initial_state,
             initial_step=initial_step,
@@ -296,7 +296,7 @@ class TestDiscreteTimeStateTransitionModelLogProbMaxima(test_util.TestCase):
 
         beta, gamma = tf.unstack(params)
         return DiscreteTimeStateTransitionModel(
-            transition_rates=txrates,
+            transition_rate_fn=txrates,
             incidence_matrix=incidence_matrix,
             initial_state=initial_state,
             initial_step=0.0,

{gemlib-0.9.2 → gemlib-0.9.4}/gemlib/mcmc/chain_binomial_rippler.py

@@ -209,7 +209,7 @@ def default_initial_ripple(model, current_events, current_state, seed):
 
     # Choose new infection events at time t
     proposed_transition_rates = tf.stack(
-        model.transition_rates(
+        model.transition_rate_fn(
             proposed_time_idx, tf.transpose(current_state_t)
         ),
         axis=0,
@@ -267,7 +267,7 @@ def chain_binomial_rippler(model, current_events, initial_ripple_fn, seed=None):
             # Calculate transition rates for current and new states
             def transition_probs(time, state):
                 rates = tf.stack(
-                    model.transition_rates(time, tf.transpose(state)), axis=-2
+                    model.transition_rate_fn(time, tf.transpose(state)), axis=-2
                 )
                 return 1.0 - tf.math.exp(-rates * model.time_delta)
 

{gemlib-0.9.2 → gemlib-0.9.4}/gemlib/mcmc/chain_binomial_rippler_test.py

@@ -33,7 +33,7 @@ def _make_model():
         return [si, ir]
 
     model = DiscreteTimeStateTransitionModel(
-        transition_rates=hazard_fn,
+        transition_rate_fn=hazard_fn,
         initial_state=init_state,
         initial_step=0,
         time_delta=1.0,
@@ -87,7 +87,7 @@ class CBRSIRTest(test_util.TestCase):
             return [si, ir]
 
         model = DiscreteTimeStateTransitionModel(
-            transition_rates=hazard_fn,
+            transition_rate_fn=hazard_fn,
             initial_state=init_state,
             initial_step=0,
             time_delta=1.0,

{gemlib-0.9.2 → gemlib-0.9.4}/pyproject.toml

@@ -1,10 +1,14 @@
 [tool.poetry]
 name = "gemlib"
-version = "0.9.2"
+version = "0.9.4"
 description = "GEMlib scientific compute library for epidemic modelling"
 authors = ["Chris Jewell <c.jewell@lancaster.ac.uk>", 
            "Alison Hale <haleac@lancaster.ac.uk>"]
-repository = "http://fhm-chicas-code.lancs.ac.uk/GEM/gemlib"
+maintainers = ["Jessica Bridgen <j.bridgen@lancaster.ac.uk>",
+	       "Alin Morariu <a.morariu@lancaster.ac.uk"]
+repository = "https://gitlab.com/gem-epidemics/gemlib"
+homepage = "https://gem-epidemics.gitlab.io/gemlib"
+keywords = ["epidemic", "Bayesian", "inference", "infectious disease model", "probabilistic programming"]
 
 [tool.poetry.dependencies]
 python = ">=3.9.0, <3.12.0"