code.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression

from SBI_infection.simulator import simulate

INFECTED_DF = pd.read_csv("SBI_infection/data/infected_timeseries.csv")
REWIRE_DF   = pd.read_csv("SBI_infection/data/rewiring_timeseries.csv")
DEGREE_DF   = pd.read_csv("SBI_infection/data/final_degree_histograms.csv")


def observed_summary():
    infected_mean = INFECTED_DF.groupby("time")["infected_fraction"].mean().values
    rewires_mean  = REWIRE_DF.groupby("time")["rewire_count"].mean().values

    deg = DEGREE_DF.groupby("degree")["count"].mean()

    deg_vals = deg.index.values
    deg_counts = deg.values

    peak_inf = infected_mean.max()
    peak_time = infected_mean.argmax()
    total_inf = infected_mean.sum()

    total_rewire = rewires_mean.sum()

    mean_deg = np.sum(deg_vals * deg_counts) / np.sum(deg_counts)

    var_deg = np.sum(
        ((deg_vals - mean_deg)**2) * deg_counts
    ) / np.sum(deg_counts)

    return np.array([peak_inf, peak_time, total_inf, total_rewire, mean_deg, var_deg])

s_obs = observed_summary()

def sim_summary(beta, gamma, rho):
    infected, rewires, degree_hist = simulate(
        beta=beta,
        gamma=gamma,
        rho=rho
    )

    peak_inf = infected.max()
    peak_time = infected.argmax()
    total_inf = infected.sum()

    total_rewire = rewires.sum()

    deg_vals = np.arange(len(degree_hist))
    deg_counts = degree_hist

    mean_deg = np.sum(deg_vals * deg_counts) / np.sum(deg_counts)

    var_deg = np.sum(
        ((deg_vals - mean_deg)**2) * deg_counts
    ) / np.sum(deg_counts)

    return np.array([peak_inf, peak_time, total_inf, total_rewire, mean_deg, var_deg])

def sample_prior():
    beta = np.random.uniform(0,1)
    gamma = np.random.uniform(0,1)
    rho = np.random.uniform(0,1)
    return beta, gamma, rho

summary_scale = np.maximum(np.abs(s_obs), 1e-6)
def distance(a,b):
    return np.linalg.norm((a-b)/summary_scale)

def rejection_abc(n_samples=10000, keep=0.01):
    rows = []
    for i in range(n_samples):
        beta, gamma, rho = sample_prior()
        s_sim = sim_summary(beta, gamma, rho)
        d = distance(s_obs, s_sim)
        rows.append([beta, gamma, rho, d, *s_sim])
        if i % 500 == 0:
            print("Completed:", i)

    df = pd.DataFrame(rows,
        columns=["beta","gamma","rho","dist",
                 "s1","s2","s3","s4","s5","s6"]
    )
    cutoff = df["dist"].quantile(keep)
    accepted = df[df["dist"] <= cutoff].copy()
    return accepted

def plot_posteriors(df):
    fig, ax = plt.subplots(1,3, figsize=(15,4))

    ax[0].hist(df["beta"], bins=20)
    ax[0].set_title("Posterior beta")

    ax[1].hist(df["gamma"], bins=20)
    ax[1].set_title("Posterior gamma")

    ax[2].hist(df["rho"], bins=20)
    ax[2].set_title("Posterior rho")

    plt.tight_layout()
    plt.show()

def pair_plot(df):
    plt.figure(figsize=(6,5))
    plt.scatter(df["beta"], df["gamma"], alpha=0.5)
    plt.xlabel("beta")
    plt.ylabel("gamma")
    plt.title("beta vs gamma")
    plt.show()

accepted = rejection_abc(
    n_samples=10000,
    keep=0.01
)

## Regression Adjustment ABC
def regression_adjustment(df, s_obs):
    summary_cols = ["s1","s2","s3","s4","s5","s6"]
    X = df[summary_cols].values
    X_centered = X - s_obs
    adjusted = df.copy()

    for param in ["beta","gamma","rho"]:
        y = df[param].values
        model = LinearRegression()
        model.fit(X_centered, y)
        pred_shift = model.predict(X_centered)
        adjusted[param] = y - pred_shift + model.intercept_

    return adjusted

print(accepted.describe())

plot_posteriors(accepted)
pair_plot(accepted)

accepted.to_csv("posterior_samples.csv", index=False)

adjusted = regression_adjustment(accepted, s_obs)
plot_posteriors(adjusted)
pair_plot(adjusted)

adjusted.to_csv("posterior_adjusted.csv", index=False)