# Colocalisation: Pearson's and Manders' coefficients, and comparing conditions with the right n.
# Python 3.13, NumPy, pandas, SciPy.
import numpy as np, pandas as pd
from scipy import stats

im = pd.read_csv("images.csv")                       # columns: image, x, y, ch1, ch2; one row per pixel
print("Coefficients for the three example images (threshold = background mean + 3 SD = 16 in both channels):")
for name, g in im.groupby("image"):
    a, b = g.ch1.to_numpy(), g.ch2.to_numpy(); t = 16.0
    r = np.corrcoef(a, b)[0, 1]
    m1 = a[(a > t) & (b > t)].sum() / a[a > t].sum()  # thresholded Manders M1: share of channel-1 signal where channel 2 is present
    m2 = b[(a > t) & (b > t)].sum() / b[b > t].sum()  # thresholded Manders M2: share of channel-2 signal where channel 1 is present
    print(f"  image {name}: Pearson {r:.3f}; Manders M1 {m1:.3f}, M2 {m2:.3f}")

d = pd.read_csv("cells.csv")                         # columns: experiment, condition, cell, pearson; one row per cell
d["z"] = np.arctanh(d.pearson)                       # Fisher z, so that coefficients can be averaged and compared
t = stats.ttest_ind(d.z[d.condition == "treated"], d.z[d.condition == "control"], equal_var=False)
print(f"Per-cell Welch's t-test (150 cells treated as independent): t = {t.statistic:.2f}, p = {t.pvalue:.4f}")
m = d.groupby(["experiment", "condition"]).z.mean().unstack()
diff = m.treated - m.control
t2 = stats.ttest_rel(m.treated, m.control)
print("Mean Pearson per experiment (back-transformed): " + "; ".join(
    f"experiment {e}: control {np.tanh(m.loc[e, 'control']):.3f}, treated {np.tanh(m.loc[e, 'treated']):.3f}" for e in m.index))
hw = stats.t.ppf(0.975, len(diff) - 1) * diff.std(ddof=1) / np.sqrt(len(diff))
print(f"Paired t-test on experiment means (n = 3 experiments): difference in z {diff.mean():.3f} "
      f"(95% CI {diff.mean() - hw:.3f} to {diff.mean() + hw:.3f}), t = {t2.statistic:.2f}, p = {t2.pvalue:.4f}")
