Spatial features of measles transmission in England and Wales during the pre-vaccine era¶

class Importation:
    def __init__(self, model, infdurdist, infdurmin: int =1, period: int = 30, count: int = 3, end_tick: int = 10*365):
        self.model = model
        self.infdurdist = infdurdist
        self.infdurmin = infdurmin
        self.period = period
        self.count = count
        self.end_tick = end_tick if end_tick is not None else model.params.nticks

        self.model.nodes.add_vector_property("imports", model.params.nticks + 1, dtype=np.uint32, default=0)

        return
    
    def step(self, tick: int) -> None:
        if tick > 0 and tick % self.period == 0 and tick < self.end_tick:
            i_susceptible = np.nonzero(self.model.people.state == SEIR.State.SUSCEPTIBLE.value)[0]
            if len(i_susceptible) > 0:
                count = min(self.count, len(i_susceptible))
                i_infect = np.random.choice(i_susceptible, size=count, replace=False)
                self.model.people.state[i_infect] = SEIR.State.INFECTIOUS.value
                samples = dists.sample_floats(self.infdurdist, np.zeros(count, np.float32))
                samples = np.round(samples)
                samples = np.maximum(samples, self.infdurmin).astype(self.model.people.itimer.dtype)
                self.model.people.itimer[i_infect] = samples
                inf_by_node = np.bincount(self.model.people.nodeid[i_infect], minlength=len(self.model.nodes)).astype(self.model.nodes.S.dtype)
                self.model.nodes.S[tick + 1] -= inf_by_node
                self.model.nodes.I[tick + 1] += inf_by_node
                self.model.nodes.imports[tick] = inf_by_node
            # else:
            #     print(f"No susceptibles to infect at tick {tick}")

        return

class SeasonalTransmission(SEIR.Transmission):
        
    def step(self, tick: int) -> None:
        ft = self.model.nodes.forces[tick]

        N = self.model.nodes.S[tick] + self.model.nodes.E[tick] + (I := self.model.nodes.I[tick])  # noqa: E741
        # Might have R
        if hasattr(self.model.nodes, "R"):
            N += self.model.nodes.R[tick]

        ft[:] = self.model.params.beta * I / N * self.model.params.beta_season[tick % 365]
        transfer = ft[:, None] * self.model.network
        ft += transfer.sum(axis=0)
        ft -= transfer.sum(axis=1)
        ft = -np.expm1(-ft)  # Convert to probability of infection

        newly_infected_by_node = np.zeros((nb.get_num_threads(), self.model.nodes.count), dtype=np.int32)
        self.nb_transmission_step(
            self.model.people.state,
            self.model.people.nodeid,
            ft,
            newly_infected_by_node,
            self.model.people.etimer,
            self.expdurdist,
            self.expdurmin,
            tick,
        )
        newly_infected_by_node = newly_infected_by_node.sum(axis=0).astype(self.model.nodes.S.dtype)  # Sum over threads

        # state(t+1) = state(t) + ∆state(t)
        self.model.nodes.S[tick + 1] -= newly_infected_by_node
        self.model.nodes.E[tick + 1] += newly_infected_by_node
        # Record today's ∆
        self.model.nodes.newly_infected[tick] = newly_infected_by_node

        return

# See ../../../data/readme.md for a description of the England and Wales measles data contained there.  
# We provide pre-computed pairwise distances between the population centers in EnglandWalesDistances.npy.
# Here, we compute a population-wide average crude birth rate (CBR) for England and Wales from the data.

average_pop = np.array([np.mean(place.population) for place in engwal.places.values()])
average_birthrate = np.array([np.mean(place.births/place.population) for place in engwal.places.values()])
average_cbr = 1000*np.sum(average_pop*average_birthrate)/np.sum(average_pop)
birthrates = np.array([place.births/place.population *1000 for place in engwal.places.values()])

# Define the geodataframe
cells = []
for placename, place in engwal.places.items():
    lat, long = place.latitude, place.longitude
    nodeid = engwal.placenames.index(placename)
    name = placename
    # As long as we are here, record the proportion of weeks with zero cases.  
    # Wavelet analysis is more complex and we will come back to that later.
    prop_zero = np.mean(np.array(place.cases) == 0)

    cells.append({"nodeid": nodeid, "name": name, "population": place.population[0], "geometry": Point(long, lat), "prop_zero": prop_zero})

scenario = gpd.GeoDataFrame(cells, columns=["nodeid", "name", "population", "geometry", "prop_zero"], crs="EPSG:4326")

run_sims = False
start_from_last = False

output_file = Path("outputs") / "EW_outputs.npy"
params_file = Path("outputs") / "EW_parameter_samples.csv"
nsims = 150
nticks = 40*365
npatches = len(scenario)
if run_sims:
    istart = 0
    if start_from_last and (output_file.exists() 
                            and params_file.exists()):

        print("Resuming from last saved simulation...")

        outputs =np.load(output_file)
        while istart < outputs.shape[0] and np.any(outputs[istart]):
            istart += 1
        params_df = pd.read_csv(params_file)
        beta_samples = params_df["beta"].values
        amplitude_samples = params_df["amplitude"].values
        k_samples = params_df["k"].values
        b_samples = params_df["b"].values
        c_samples = params_df["c"].values

    else:

        print("Starting new set of simulations...")

        beta_samples = np.random.uniform(3, 4.5, nsims) #Ro from 24 to 36
        amplitude_samples = np.random.uniform(0.5, 1.5, nsims) #1 will be exactly the seasonal forcing profile from reference above
        k_samples = 10 ** np.random.uniform(-3.5, -1, nsims)
        #a_samples = np.random.uniform(-0.5, 0.5, nsims)
        b_samples = np.random.uniform(0.25, 1.0, nsims)
        c_samples = np.random.uniform(1, 2, nsims)
        params_df = pd.DataFrame({
            "beta": beta_samples,
            "amplitude": amplitude_samples,
            "k": k_samples,
            "b": b_samples,
            "c": c_samples
        })
        Path("outputs").mkdir(exist_ok=True)
        params_df.to_csv(params_file, index=False)
        outputs = np.zeros((nsims, (nticks+1) // 7, npatches), dtype=np.uint32)
    
    log_betas = np.array([0.155, 0.571, 0.46, 0.34, 0.30, 0.34, 0.24, 0.15, 0.31, 0.40, 0.323, 0.238, 0.202, 0.203, 0.074,
        -0.095, -0.218, -0.031, 0.433, 0.531, 0.479, 0.397, 0.444, 0.411, 0.291, 0.509])
    beta_season = np.repeat(log_betas, np.floor(365 / len(log_betas)).astype(int))
    beta_season = np.append(beta_season, beta_season[-1])
    beta_season = np.exp( (beta_season - np.mean(beta_season)) )

    for i in range(istart, nsims):
        print(f"Running simulation {i+1} of {nsims}...")
        parameters = PropertySet(
            {"seed": 4, "nticks": nticks,
            "exp_shape": 40, "exp_scale": 0.25,
            "verbose": True, "beta": beta_samples[i], "inf_mean": 8, "inf_sigma": 2, "cbr": average_cbr,
            "beta_season": beta_season**amplitude_samples[i],
            "importation_period": 30, "importation_count": 3, "importation_start": 1, "importation_end": 10*365,
            "gravity_k": k_samples[i], "gravity_b": b_samples[i], "gravity_c": c_samples[i],
            "skip_capacity": False, "capacity_safety_factor": 3.0}
        )

        scenario["E"] = 0
        scenario["I"] = 3
        scenario["R"] = np.round(0.95 * scenario.population).astype(np.uint32)
        scenario["S"] = scenario.population - scenario["E"] - scenario["I"] - scenario["R"]

        expdist = dists.gamma(shape=parameters.exp_shape, scale=parameters.exp_scale)
        infdist = dists.normal(loc=parameters.inf_mean, scale=parameters.inf_sigma)

        birthrate_map = ValuesMap.from_scalar(parameters.cbr, nticks=parameters.nticks, nnodes=len(scenario))
        rate_const = 365 * ((1 + parameters.cbr / 1000) ** (1/365) - 1) # approximate annual to daily
        stable_age_dist = np.array(1000*np.exp(-rate_const*np.arange(89)))
        pyramid = AliasedDistribution(stable_age_dist)
        survival = KaplanMeierEstimator(stable_age_dist.cumsum())
        model = Model(scenario, parameters, birthrates=birthrate_map.values)

        model.components = [
            SEIR.Susceptible(model),
            SEIR.Exposed(model, expdist, infdist),
            SEIR.Infectious(model, infdist),
            SEIR.Recovered(model),
            Importation(model, infdist),
            SeasonalTransmission(model, infdist),
            BirthsByCBR(model, birthrates=birthrate_map.values, pyramid=pyramid),
            MortalityByEstimator(model, estimator=survival),
        ]

        model.network = gravity(np.array(scenario.population), distances, 1, 0, model.params.gravity_b, model.params.gravity_c)
        # The scale of the parameter k is actually unit-ful, and depends on the values of b and c.  
        # To make k more interpretable, and avoid having to sample large swaths of useless parameter space,
        # we rescale the network here so that the average export fraction is 1, then multiply by k.
        average_export_frac = np.mean(model.network.sum(axis=1))
        model.network = model.network / average_export_frac * model.params.gravity_k 
        model.network = row_normalizer(model.network, 0.2) #Don't let any individual node send more than 20% of its infectivity elsewhere.

        model.run(f"Sim {i+1}/{nsims}")
        incidence = model.nodes.newly_infected
        num_weeks = incidence.shape[0] // 7
        weekly_incidence = incidence[:num_weeks*7, :].reshape(num_weeks, 7, incidence.shape[1]).sum(axis=1)        
        outputs[i, :, :] = weekly_incidence
        np.save(output_file, outputs)

outputs_file = Path("outputs") / "EW_outputs.npy"
params_file = Path("outputs") / "EW_parameter_samples.csv"

if outputs_file.exists() and params_file.exists():

    weekly_incidence = np.load(outputs_file)
    params_df = pd.read_csv(params_file)
    results_df = params_df.copy()

    prop_zero = np.mean(weekly_incidence[:, 1040:, :]==0, axis=1)
    prop_low = np.mean(weekly_incidence[:, 1040:, :]<=2, axis=1)

    results_df["prop_zero"] = list(prop_zero)
    results_df["prop_low"] = list(prop_low)

else:

    print("No simulation outputs found.  Please run the simulations first.")
    results_df = None

from scipy.optimize import curve_fit
from mpl_toolkits.mplot3d import Axes3D

# 1. Define fitting functions for mean and variance
def logistic(x, x0, k):
    # Logistic function bounded between 0 and 1, transitions from 1 to 0 as x increases
    return 1 / (1 + np.exp(k * (x - x0)))

def fit_mean_var(x, y):
    # Fit a logistic to the mean
    # Initial guess: midpoint at median(x), width=1
    p0 = [np.median(x), 1.0]
    bounds = ([-np.inf, 0.01], [np.inf, 10.0])
    popt_mean, _ = curve_fit(logistic, x, y, p0=p0, bounds=bounds, maxfev=10000)
    return popt_mean

# 3. Define similarity metric (sum of squared differences of mean and variance fits)
def similarity_metric(mean_data, mean_sim):
    # Evaluate both curves on a common grid
    x_grid = np.linspace(2.5, 6.5, 200)
    data_curve = logistic(x_grid, *mean_data)
    sim_curve = logistic(x_grid, *mean_sim)
    # Similarity: sum of squared differences between the curves
    mean_diff = np.sum((sim_curve - data_curve) ** 2)
    return mean_diff

similarity = []

if results_df is not None:

    for i in results_df.index:
        # Plot observed and simulated data
        logpop_obs = np.log10(scenario.population)
        prop_zero_obs = scenario.prop_zero
        logpop_sim = np.log10(scenario.population)
        prop_zero_sim = results_df.loc[i]['prop_zero']

        # Fit to observed
        popt_obs = fit_mean_var(logpop_obs, prop_zero_obs)
        popt_sim = fit_mean_var(logpop_sim, prop_zero_sim)
        # Similarity metric
        similarity.append(similarity_metric(popt_obs, popt_sim))

    results_df["similarity_CCS"] = similarity

else:

    print("No results dataframe found.  Please run the simulations first.")

def pad_data(x):
    """
    Pad data to the next power of 2
    """
    nx = len(x) # number of samples
    nx2 = (2**np.ceil(np.log(nx)/np.log(2))).astype(int) # next power of 2
    x2 = np.zeros(nx2, dtype=x.dtype) # pad to next power of 2
    offset = (nx2-nx)//2 # offset
    x2[offset:(offset+nx)] = x # copy
    return x2

def log_transform(x, debug=1):
    """
    Log transform for case data
    """ 
    # add one and take log
    x = np.log(x+1)
    # set mean=0 and std=1
    m = np.mean(x)
    s = np.std(x)
    x = (x - m)/s
    return x


def calc_Ws(cases):
    # transform case data
    log_cases = pad_data(log_transform(cases))

    # setup and execute wavelet transform
    # https://pywavelets.readthedocs.io/en/latest/ref/cwt.html#morlet-wavelet
    wavelet = pywt.ContinuousWavelet('cmor2-1')

    dt = 1 # 2 weeks
    widths = np.logspace(np.log10(1), np.log10(7*52), int(7*52))
    [cwt, frequencies] = pywt.cwt(log_cases, widths, wavelet, dt)

    # Number of time steps in padded time series
    nt = len(cases)
    # trim matrix
    offset = (cwt.shape[1] - nt) // 2
    cwt = cwt[:, offset:offset + nt]

    return cwt, frequencies

def wavelet_phase_diff(data, distances, ref_cwt):

    if isinstance(data, pd.DataFrame):
        x = np.zeros(len(data))
        y = np.zeros(len(data))
        y2 = np.zeros(len(data))
        for i, row in data.iterrows():
            if distances[i] > 30:
                continue
            cwt, frequencies = calc_Ws(row["cases"].flatten())
            diff = ref_cwt*np.conj(cwt)
            ind = np.where(np.logical_and(frequencies < 1/(1.5 * 52), frequencies > 1 / (3 * 52))) #Frequencies around 2 year periods.
            diff1 = diff[ind[0], :]
            x[i] = distances[i]
            y[i] = np.angle(np.mean(diff1))
            ind2 = np.where(np.logical_and(frequencies < 1/(0.75 * 52), frequencies > 1 / (1.25 * 52))) #Frequencies around 1 year periods
            diff2 = diff[ind2[0], :]
            y2[i] = np.angle(np.mean(diff2))
            
    elif isinstance(data, np.ndarray):
        # initialize output arrays
        x = np.zeros(data.shape[1])
        y = np.zeros(data.shape[1])
        y2 = np.zeros(data.shape[1])
        for i in range(data.shape[1]):
            if distances[i] > 30:
                continue
            cwt, frequencies = calc_Ws(data[:, i].flatten())
            diff = ref_cwt*np.conj(cwt)
            ind = np.where(np.logical_and(frequencies < 1/(1.5 * 52), frequencies > 1 / (3 * 52)))
            diff1 = diff[ind[0], :]
            x[i] = distances[i]
            y[i] = np.angle(np.mean(diff1))
            ind2 = np.where(np.logical_and(frequencies < 1/(0.75 * 52), frequencies > 1 / (1.25 * 52)))
            diff2 = diff[ind2[0], :]
            y2[i] = np.angle(np.mean(diff2))

    return x, y, y2