import sys, tempfile, shutil
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from IPython.display import display

DOCS = Path('/Users/sarahurbut/aladynoulli2/docs/reparam')

# Import the same plot functions that produced the published Fig S27 PDFs.
sys.path.insert(0, str(Path('/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures')))
from generate_prs_signature_plots import (
    plot_top_associations_bar,
    plot_significant_heatmap,
    plot_full_heatmap,
)

# Those functions call plt.close(fig) after savefig. Disable that so we can
# also show inline below.
_orig_close = plt.close
plt.close = lambda *a, **k: None

C = pd.read_csv(DOCS / 'prs_signature_centered_gamma_associations.csv')
N = pd.read_csv(DOCS / 'prs_signature_400k_gamma_associations.csv')

print(f'Centered (published) rows: {len(C)}, FDR<0.05: {C.significant_fdr.sum()}')
print(f'Reparam 400K-init    rows: {len(N)}, FDR<0.05: {N.significant_fdr.sum()}')

m = (C.rename(columns=lambda c: c + '_c' if c not in ('prs', 'signature') else c)
     .merge(N.rename(columns=lambda c: c + '_n' if c not in ('prs', 'signature') else c),
            on=['prs', 'signature']))
print(f'Merged rows: {len(m)} (expected 756)')

Centered (published) rows: 756, FDR<0.05: 116
Reparam 400K-init    rows: 756, FDR<0.05: 416
Merged rows: 756 (expected 756)

named = [('CAD','Sig 5','Cardiovascular'),
         ('LDL_SF','Sig 5','Cardiovascular'),
         ('T2D','Sig 15','Metabolic'),
         ('BMI','Sig 15','Metabolic'),
         ('AF','Sig 0','Cardiovascular'),
         ('HT','Sig 5','Cardiovascular'),
         ('CVD','Sig 5','Cardiovascular')]
rows = []
for p, s, cat in named:
    sub = m[(m['prs']==p) & (m['signature']==s)]
    if not len(sub):
        continue
    r = sub.iloc[0]
    rows.append({
        'PRS': p, 'Signature': s, 'Category': cat,
        'centered_gamma': f'{r.effect_c:+.4f}', 'centered_Z': f'{r.z_score_c:+.1f}',
        'centered_FDR': bool(r.significant_fdr_c),
        '400k_gamma':    f'{r.effect_n:+.4f}', '400k_Z':    f'{r.z_score_n:+.1f}',
        '400k_FDR':    bool(r.significant_fdr_n),
        'sign_agree': int(np.sign(r.effect_c) == np.sign(r.effect_n)),
    })
display(pd.DataFrame(rows))

def render_pool(pool_tag, df, label, save_dir=DOCS):
    """Run the 3 supplement plot functions, show each inline, and copy the
    saved PDFs to save_dir with the pool prefix."""
    print(f'\n=== {label} ===')
    with tempfile.TemporaryDirectory() as td:
        out = Path(td)
        plot_top_associations_bar(df.copy(), out)
        plt.show()
        plot_significant_heatmap(df.copy(), out)
        plt.show()
        plot_full_heatmap(df.copy(), out)
        plt.show()
        for stem in ['top_prs_associations.pdf',
                     'significant_prs_heatmap.pdf',
                     'complete_prs_heatmap.pdf']:
            s = out / stem
            if s.exists():
                shutil.copy2(s, save_dir / f'prs_signature_{pool_tag}_{stem}')

render_pool('centered', C, 'Centered (published Fig S27)')

=== Centered (published Fig S27) ===

✓ Saved to: /var/folders/fl/ng5crz0x0fnb6c6x8dk7tfth0000gn/T/tmpvvlmm66c/top_prs_associations.pdf

  Filtering by FDR-corrected p-value < 0.05
    Found 116 significant associations

✓ Saved to: /var/folders/fl/ng5crz0x0fnb6c6x8dk7tfth0000gn/T/tmpvvlmm66c/significant_prs_heatmap.pdf

✓ Saved to: /var/folders/fl/ng5crz0x0fnb6c6x8dk7tfth0000gn/T/tmpvvlmm66c/complete_prs_heatmap.pdf

render_pool('400k', N, 'Reparam 400K-init (v1.1)')

=== Reparam 400K-init (v1.1) ===
✓ Saved to: /var/folders/fl/ng5crz0x0fnb6c6x8dk7tfth0000gn/T/tmp5bkcdf2x/top_prs_associations.pdf

  Filtering by FDR-corrected p-value < 0.05
    Found 416 significant associations

✓ Saved to: /var/folders/fl/ng5crz0x0fnb6c6x8dk7tfth0000gn/T/tmp5bkcdf2x/significant_prs_heatmap.pdf

✓ Saved to: /var/folders/fl/ng5crz0x0fnb6c6x8dk7tfth0000gn/T/tmp5bkcdf2x/complete_prs_heatmap.pdf

mask = m.z_score_c.notna() & m.z_score_n.notna()
x = m.loc[mask, 'z_score_c'].values
y = m.loc[mask, 'z_score_n'].values

fig, ax = plt.subplots(figsize=(7, 7))
ax.scatter(x, y, s=12, alpha=0.55, c='#1a3a5c')
lim = max(np.abs(x).max(), np.abs(y).max()) * 1.05
ax.plot([-lim, lim], [-lim, lim], '--', color='red', alpha=0.5, lw=1)
ax.axhline(0, color='gray', lw=0.5); ax.axvline(0, color='gray', lw=0.5)
ax.set_xlabel('Z (Centered / published)')
ax.set_ylabel('Z (Reparam 400K-init)')
r = np.corrcoef(x, y)[0, 1]
strong = (np.abs(x) > 2) | (np.abs(y) > 2)
r_strong = np.corrcoef(x[strong], y[strong])[0, 1] if strong.sum() > 1 else np.nan
ax.set_title(f'r(all) = {r:.3f}   r(|Z|>2) = {r_strong:.3f}    (red dashed = y=x)', fontsize=11)
ax.set_xlim(-lim, lim); ax.set_ylim(-lim, lim)
plt.tight_layout()
plt.savefig(DOCS / 'prs_signature_centered_vs_400k_scatter.pdf', bbox_inches='tight', dpi=150)
plt.savefig(DOCS / 'prs_signature_centered_vs_400k_scatter.png', bbox_inches='tight', dpi=150)
plt.show()

any_sig = m.significant_fdr_c | m.significant_fdr_n
sub = m[any_sig]
print(f'\nCells with any-pool FDR<0.05: {any_sig.sum()}')
print(f'Sign agreement on any-significant cells: {(np.sign(sub.effect_c)==np.sign(sub.effect_n)).mean()*100:.1f}%')
strong_m = (m.z_score_c.abs() > 5) | (m.z_score_n.abs() > 5)
ss = m[strong_m]
print(f'Sign agreement on |Z|>5 cells ({len(ss)}): {(np.sign(ss.effect_c)==np.sign(ss.effect_n)).mean()*100:.1f}%')

Cells with any-pool FDR<0.05: 433
Sign agreement on any-significant cells: 75.8%
Sign agreement on |Z|>5 cells (223): 80.3%

plt.close = _orig_close  # restore

Pool	Source	γ definition / init	FDR<0.05
Centered (published)	`paper_figs/fig4/gamma_adventures/prs_signatures_nolr/gamma_associations.csv` ≡ `SuppDataFiles-3/Annotation/gamma_associations.csv`	λ ~ GP(μ + Gγ, K), λ free; γ trained only by GP prior	116
Reparam 400K-init	40 batches in `censor_e_batchrun_vectorized_REPARAM_v3_nokappa_400k_init/`, all warm-started from `gamma_level_pooled_400k.pt`	λ = μ + Gγ + δ, δ ~ GP(0, K); γ trained by full NLL	416

	PRS	Signature	Category	centered_gamma	centered_Z	centered_FDR	400k_gamma	400k_Z	400k_FDR	sign_agree
0	CAD	Sig 5	Cardiovascular	+0.1534	+27.2	True	+0.3345	+35.8	True	1
1	LDL_SF	Sig 5	Cardiovascular	+0.0713	+22.7	True	+0.1380	+13.7	True	1
2	T2D	Sig 15	Metabolic	+0.1539	+58.3	True	+1.1458	+79.1	True	1
3	BMI	Sig 15	Metabolic	+0.0220	+13.7	True	+0.2309	+18.0	True	1
4	AF	Sig 0	Cardiovascular	+0.0379	+22.1	True	+0.3006	+35.1	True	1
5	HT	Sig 5	Cardiovascular	+0.0475	+12.3	True	+0.2088	+26.4	True	1
6	CVD	Sig 5	Cardiovascular	+0.0770	+14.2	True	+0.1687	+22.0	True	1

PRS-Signature Associations: Published (centered) vs Reparam-400K-init (v1.1)¶

Headline named hits¶

Centered (published Fig S27) — three plots¶

Reparam 400K-init — three plots¶

Z-score consistency: centered vs 400K-init¶

Interpretation¶

Provenance¶