Skip to content
GitLab
Commit ef41975d authored by Chris Jewell's avatar Chris Jewell
Browse files

First refactor. Doesn't work with CUDA.

parent 047e9bd0
Hide whitespace changes
Inline Side-by-side
covid/model.py
View file @ ef41975d
......@@ -76,113 +76,33 @@ def load_data(paths, settings, dtype=DTYPE):
}
class CovidUK:
class CovidUKStochastic:
def __init__(
self,
W: np.float64,
C: np.float64,
N: np.float64,
xi_freq: int,
params: dict,
initial_state: np.float64,
initial_time: np.float64,
time_step: np.int64,
num_steps: np.int64,
transition_rates,
stoichiometry,
initial_state,
initial_step,
time_delta,
num_steps,
):
"""Represents a CovidUK ODE model
:param W: Commuting volume
:param C: a n_ladsxn_lads matrix of inter-LAD commuting
:param N: a vector of population sizes in each LAD
:param date_range: a time range [start, end)
:param beta_freq: the frequency at which beta changes
:param time_step: a time step to use in the discrete time simulation
"""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`.
: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.
"""
dtype = dtype_util.common_dtype([W, C, N, initial_state], dtype_hint=DTYPE)
self.initial_state = tf.convert_to_tensor(initial_state, dtype=dtype)
self.initial_time = initial_time
self.n_lads = C.shape[0]
C = tf.convert_to_tensor(C, dtype=dtype)
self.C = C + tf.transpose(C)
self.C = tf.linalg.set_diag(self.C, tf.zeros(self.C.shape[0], dtype=dtype))
self.W = tf.constant(W, dtype=dtype)
self.N = tf.constant(N, dtype=dtype)
self.time_step = time_step
self.transition_rates = transition_rates
self.stoichiometry = stoichiometry
self.initial_state = initial_state
self.initial_step = initial_step
self.time_delta = time_delta
self.num_steps = num_steps
self.params = {
k: tf.convert_to_tensor(v, dtype=dtype) for k, v in params.items()
}
self.xi_freq = np.int32(xi_freq)
self.xi_select = np.arange(self.num_steps, dtype=np.int32) // self.xi_freq
self.max_t = self.xi_select.shape[0] - 1
def create_initial_state(self, init_matrix=None):
I = tf.convert_to_tensor(init_matrix, dtype=DTYPE)
S = self.N - I
E = tf.zeros(self.N.shape, dtype=DTYPE)
R = tf.zeros(self.N.shape, dtype=DTYPE)
return tf.stack([S, E, I, R], axis=-1)
class CovidUKStochastic(CovidUK):
stoichiometry = tf.constant(
[[-1, 1, 0, 0], [0, -1, 1, 0], [0, 0, -1, 1]], dtype=DTYPE
)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def make_h(self, param=None):
"""Constructs a function that takes `state` and outputs a
transition rate matrix (with 0 diagonal).
"""
if param is None:
param = self.params
def h(t, state):
"""Computes a transition rate matrix
:param state: a tensor of shape [M, S] for S states and M population strata. States
are S, E, I, R. We arrange the state like this because the state vectors are then arranged
contiguously in memory for fast calculation below.
:return a tensor of shape [M, M, S] containing transition matric for each i=0,...,(c-1)
"""
w_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, self.W.shape[0])
commute_volume = tf.gather(self.W, w_idx)
xi_idx = tf.cast(
tf.clip_by_value(t // self.xi_freq, 0, self.params["xi"].shape[0] - 1),
dtype=tf.int64,
)
xi = tf.gather(self.params["xi"], xi_idx)
beta = self.params["beta1"] * tf.math.exp(xi)
infec_rate = beta * (
state[..., 2]
+ self.params["beta2"]
* commute_volume
* tf.linalg.matvec(self.C, state[..., 2] / self.N)
)
infec_rate = infec_rate / self.N + 0.000000001 # Vector of length nc
ei = tf.broadcast_to(
[self.params["nu"]], shape=[state.shape[0]]
) # Vector of length nc
ir = tf.broadcast_to(
[self.params["gamma"]], shape=[state.shape[0]]
) # Vector of length nc
return [infec_rate, ei, ir]
return h
def ngm(self, t, state, param):
"""Computes a next generation matrix -- pressure from i to j is G_{ij}
:param t: the time step
......@@ -215,13 +135,12 @@ class CovidUKStochastic(CovidUK):
:param state_init: the initial state
:returns: a tuple of times and simulated states.
"""
hazard = self.make_h()
t, sim = discrete_markov_simulation(
hazard_fn=hazard,
hazard_fn=self.transition_rates,
state=self.initial_state,
start=self.initial_time,
end=self.initial_time + self.num_steps * self.time_step,
time_step=self.time_step,
start=self.initial_step,
end=self.initial_step + self.num_steps * self.time_delta,
time_step=self.time_delta,
seed=seed,
)
return t, sim
......@@ -238,7 +157,7 @@ class CovidUKStochastic(CovidUK):
)
y = tf.convert_to_tensor(y, dtype)
with tf.name_scope("CovidUKStochastic.log_prob"):
hazard = self.make_h()
hazard = self.transition_rates
return discrete_markov_log_prob(
y, self.initial_state, hazard, self.time_step, self.stoichiometry
y, self.initial_state, hazard, self.time_delta, self.stoichiometry
)
covid/pydata.py
View file @ ef41975d
"""Python-based data munging"""
import os
import re
from warnings import warn
import geopandas as gp
import numpy as np
import pandas as pd
import pyreadr as pyr
......@@ -87,7 +86,19 @@ def linelist2timeseries(date, region_code, date_range=None):
def phe_case_data(linelisting_file, date_range=None, pillar=None):
ll = pd.read_excel(linelisting_file)
read_file = dict(csv=pd.read_csv, xlsx=pd.read_excel)
match_extension = re.match(r"(.*)\.(.*)$", linelisting_file)
if match_extension is None:
raise ValueError(
f"Linelisting filename '{linelisting_file}' is not in name.extension format"
)
filetype = match_extension.group(2)
try:
ll = read_file[filetype](linelisting_file)
except KeyError:
raise ValueError(f"No handler implemented for file type '{filetype}'")
if pillar is not None:
ll = ll.loc[ll["pillar"] == pillar]
date = ll["specimen_date"]
......
mcmc.py
View file @ ef41975d
......@@ -29,6 +29,7 @@ tfd = tfp.distributions
tfb = tfp.bijectors
DTYPE = config.floatX
STOICHIOMETRY = tf.constant([[-1, 1, 0, 0], [0, -1, 1, 0], [0, 0, -1, 1]], dtype=DTYPE)
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
# os.environ["XLA_FLAGS"] = '--xla_dump_to=xla_dump --xla_dump_hlo_pass_re=".*"'
......@@ -78,7 +79,7 @@ state = compute_state(
[covar_data["pop"][:, tf.newaxis], tf.zeros_like(events[:, 0, :])], axis=-1
),
events=events,
stoichiometry=CovidUKStochastic.stoichiometry,
stoichiometry=STOICHIOMETRY,
)
start_time = state.shape[1] - cases.shape[1]
initial_state = state[:, start_time, :]
......@@ -88,17 +89,44 @@ num_xi = events.shape[1] // xi_freq
num_metapop = covar_data["pop"].shape[0]
# Create the model
# Create the epidemic model given parameters
def build_epidemic(param):
def transition_rates(t, state):
C = tf.convert_to_tensor(covar_data["C"], dtype=DTYPE)
C = tf.linalg.set_diag(C + tf.transpose(C), tf.zeros(C.shape[0], dtype=DTYPE))
W = tf.constant(covar_data["W"], dtype=DTYPE)
N = tf.constant(covar_data["pop"], dtype=DTYPE)
w_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, W.shape[0] - 1)
commute_volume = tf.gather(W, w_idx)
xi_idx = tf.cast(
tf.clip_by_value(t // 14, 0, param["xi"].shape[0] - 1), dtype=tf.int64,
)
xi = tf.gather(param["xi"], xi_idx)
beta = param["beta1"] * tf.math.exp(xi)
infec_rate = beta * (
state[..., 2]
+ param["beta2"] * commute_volume * tf.linalg.matvec(C, state[..., 2] / N)
)
infec_rate = infec_rate / N + 0.000000001 # Vector of length nc
ei = tf.broadcast_to(
[param["nu"]], shape=[state.shape[0]]
) # Vector of length nc
ir = tf.broadcast_to(
[param["gamma"]], shape=[state.shape[0]]
) # Vector of length nc
return [infec_rate, ei, ir]
return CovidUKStochastic(
C=covar_data["C"],
N=covar_data["pop"],
W=covar_data["W"],
xi_freq=xi_freq,
params=param,
transition_rates=transition_rates,
stoichiometry=STOICHIOMETRY,
initial_state=initial_state,
initial_time=0.0,
time_step=1.0,
initial_step=0,
time_delta=1.0,
num_steps=events.shape[1],
)
......@@ -237,7 +265,7 @@ def trace_results_fn(_, results):
@tf.function(autograph=False, experimental_compile=True)
def sample(n_samples, init_state, sigma_theta, sigma_xi, num_event_updates):
def sample(n_samples, init_state, sigma_theta, sigma_xi):
with tf.name_scope("main_mcmc_sample_loop"):
init_state = init_state.copy()
......@@ -359,7 +387,6 @@ for i in tqdm.tqdm(range(NUM_BURSTS), unit_scale=NUM_BURST_SAMPLES):
init_state=current_state,
sigma_theta=theta_scale,
sigma_xi=xi_scale,
num_event_updates=tf.constant(NUM_EVENT_TIME_UPDATES, tf.int32),
)
current_state = [s[-1] for s in samples]
s = slice(i * THIN_BURST_SAMPLES, i * THIN_BURST_SAMPLES + THIN_BURST_SAMPLES)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment