# Clonogenic survival, sensitiser vs control: linear-quadratic fit per experiment and condition, SF2 compared across
# experiments (paired), versus one Poisson or quasi-Poisson model of all dishes pooled. Python 3.13, NumPy, pandas, SciPy.
import numpy as np, pandas as pd
from scipy import stats

d = pd.read_csv("colonies.csv")                      # one row per dish: experiment, condition, dose_gy, dish, seeded, colonies
def lq(x):                                           # per experiment and condition: PE, then SF and a least-squares LQ fit
    m = x.groupby("dose_gy").apply(lambda g: (g.colonies / g.seeded).mean(), include_groups=False)
    pe = m.loc[0.0]; sf = m.drop(0.0) / pe; D = sf.index.to_numpy()
    A = np.column_stack([D, D ** 2]); (a, b), *_ = np.linalg.lstsq(A, -np.log(sf.to_numpy()), rcond=None)
    return pe, a, b, np.exp(-(2 * a + 4 * b))
res = {(e, c): lq(g) for (e, c), g in d.groupby(["experiment", "condition"])}
for e in sorted(d.experiment.unique()):
    pc, ac, bc, s2c = res[(e, "control")]; ps, as_, bs, s2s = res[(e, "sensitiser")]
    print(f"Experiment {e}: control PE {pc:.2f}, alpha {ac:.3f}, beta {bc:.3f}, SF2 {s2c:.3f}; sensitiser PE {ps:.2f}, SF2 {s2s:.3f}")
r = np.array([np.log(res[(e, "sensitiser")][3] / res[(e, "control")][3]) for e in sorted(d.experiment.unique())])
t = stats.ttest_1samp(r, 0); ci = t.confidence_interval()
print(f"SF2, sensitiser / control, one ratio per experiment (paired, n = {len(r)}): {np.exp(r.mean()):.2f} (95% CI {np.exp(ci.low):.2f} to {np.exp(ci.high):.2f}), p = {t.pvalue:.2g}")
# Dose enhancement ratio at SF = 0.1: dose for control / dose for sensitiser, from each experiment's LQ fits (>1 = sensitised)
d10 = lambda a, b: (-a + np.sqrt(a ** 2 - 4 * b * np.log(0.1))) / (2 * b)
q = np.array([np.log(d10(*res[(e, "control")][1:3]) / d10(*res[(e, "sensitiser")][1:3])) for e in sorted(d.experiment.unique())])
t = stats.ttest_1samp(q, 0); ci = t.confidence_interval()
print(f"Dose enhancement ratio at SF = 0.1, one per experiment (paired, n = {len(q)}): {np.exp(q.mean()):.2f} (95% CI {np.exp(ci.low):.2f} to {np.exp(ci.high):.2f}), p = {t.pvalue:.2g}")
# Pooled Poisson model of all dishes: log mean = log seeded + a plating efficiency for each experiment and condition
# + dose and dose^2, plus condition x (dose, dose^2) in the larger model. Dishes are treated as independent counts.
y = d.colonies.to_numpy(float); off = np.log(d.seeded.to_numpy(float)); D = d.dose_gy.to_numpy(); s = (d.condition == "sensitiser").to_numpy(float)
P = pd.get_dummies(d.experiment.astype(str) + d.condition).to_numpy(float)
def poisson_fit(X):                                  # Fisher scoring; returns residual deviance and Pearson chi-squared
    b = np.linalg.lstsq(X, np.log(y + 0.5) - off, rcond=None)[0]
    for _ in range(100):
        mu = np.exp(off + X @ b); step = np.linalg.solve(X.T @ (X * mu[:, None]), X.T @ (y - mu)); b = b + step
        if np.max(np.abs(step)) < 1e-10: break
    mu = np.exp(off + X @ b); return 2 * np.sum(np.where(y > 0, y * np.log(np.where(y > 0, y, 1) / mu), 0) - (y - mu)), np.sum((y - mu) ** 2 / mu)
X0 = np.column_stack([P, D, D ** 2]); X1 = np.column_stack([X0, s * D, s * D ** 2])
(dev0, _), (dev1, pear1) = poisson_fit(X0), poisson_fit(X1); lr = dev0 - dev1; df_res = len(y) - X1.shape[1]
phi = pear1 / df_res; Fq = (lr / 2) / phi            # Pearson dispersion, as glm(family = quasipoisson)
print(f"All {len(y)} dishes pooled, Poisson model (plating efficiency per experiment and condition): likelihood-ratio chi2(2) = {lr:.1f}, "
      f"p = {stats.chi2.sf(lr, 2):.2g}; Pearson dispersion {phi:.1f}")
print(f"Same model with quasi-Poisson scaling: F(2, {df_res}) = {Fq:.2f}, p = {stats.f.sf(Fq, 2, df_res):.2g}")