Figure 2: Population-Level Patterns¶
Purpose¶
Show population-level disease signature patterns and their temporal evolution.
Panels Required:¶
- Panel A: Heatmap of discovered signatures with top diseases for each
- Panel B: Patient clustering based on signature weights
- Panel C: Signature prevalence across age cohorts (young vs. old)
- Panel D: Temporal evolution of signature prevalence in population
Key Message:¶
Focus on trajectory aspects, not just clusters - emphasize temporal dynamics
In [ ]:
# Setup
import sys
import os
sys.path.append('/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/')
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
# Set style for publication-quality figures
sns.set_style("whitegrid")
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['font.size'] = 11
plt.rcParams['axes.labelsize'] = 12
plt.rcParams['axes.titlesize'] = 14
plt.rcParams['xtick.labelsize'] = 10
plt.rcParams['ytick.labelsize'] = 10
plt.rcParams['legend.fontsize'] = 9
plt.rcParams['figure.titlesize'] = 16
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = ['Arial', 'DejaVu Sans', 'Helvetica', 'Liberation Sans']
print("Setup complete")
In [ ]:
# Define plotting functions
from scipy.special import expit as sigmoid
def plot_single_signature(phi_np, clusters_np, disease_names, signature_idx,
age_offset=30, figsize=(10, 7), top_n_diseases=None,
plot_probability=False, prevalence_t_np=None, logit_prev_t_np=None):
"""
Plot a single signature with prettier styling for publication.
Parameters:
-----------
phi_np : np.ndarray, shape (K, D, T)
Phi values (log hazard ratios) for each signature, disease, and timepoint
clusters_np : np.ndarray, shape (D,)
Disease-to-signature assignments
disease_names : list
List of disease names
signature_idx : int
Signature index to plot
age_offset : int
Age offset (timepoint 0 = age_offset)
figsize : tuple
Figure size (width, height)
top_n_diseases : int or None
If specified, plot only top N signature-specific diseases based on max phi value
plot_probability : bool
If True, convert phi to probability using sigmoid. If False, plot log hazard ratio.
prevalence_t_np : np.ndarray or None, shape (D, T)
Smoothed prevalence for each disease and timepoint. If provided, will be overlaid.
logit_prev_t_np : np.ndarray or None, shape (D, T)
Logit prevalence for each disease and timepoint. If provided, will be overlaid as probability.
Returns:
--------
fig : matplotlib.figure.Figure
The figure object
"""
K, D, T = phi_np.shape
age_points = np.arange(T) + age_offset
# Convert to probability if requested
if plot_probability:
plot_values = sigmoid(phi_np)
y_label = 'Probability'
else:
plot_values = phi_np
y_label = 'Log Hazard Ratio'
# Create single figure
fig, ax = plt.subplots(figsize=figsize)
k = signature_idx
# Find diseases assigned to this signature
sig_disease_indices = np.where(clusters_np == k)[0]
# Find all other diseases (background)
background_disease_indices = np.where(clusters_np != k)[0]
# Plot background diseases in light grey
for d in background_disease_indices:
ax.plot(age_points, plot_values[k, d, :],
color='lightgray', alpha=0.15, linewidth=0.5, zorder=1)
# Plot signature-specific diseases in colored lines
if len(sig_disease_indices) > 0:
# Optionally select top N diseases
if top_n_diseases is not None and len(sig_disease_indices) > top_n_diseases:
# Rank by max value
max_values = np.max(plot_values[k, sig_disease_indices, :], axis=1)
top_indices = np.argsort(max_values)[-top_n_diseases:]
sig_disease_indices = sig_disease_indices[top_indices]
# Use a nice color palette
n_colors = len(sig_disease_indices)
if n_colors <= 10:
colors = plt.cm.tab10(np.linspace(0, 1, n_colors))
else:
colors = plt.cm.Set3(np.linspace(0, 1, n_colors))
for i, d in enumerate(sig_disease_indices):
color = colors[i]
disease_name = disease_names[d] if d < len(disease_names) else f"Disease_{d}"
# Plot model's learned phi
ax.plot(age_points, plot_values[k, d, :],
color=color, linewidth=2.5, label=disease_name, zorder=3, alpha=0.9)
# Overlay smoothed prevalence if available
if prevalence_t_np is not None and d < prevalence_t_np.shape[0]:
if isinstance(prevalence_t_np, torch.Tensor):
prev_values = prevalence_t_np[d, :].detach().cpu().numpy()
else:
prev_values = prevalence_t_np[d, :]
ax.plot(age_points, prev_values,
color=color, linewidth=1.8, linestyle='--',
alpha=0.5, label=f'{disease_name} (prev)', zorder=2)
# Styling
ax.set_title(f'Signature {k}', fontsize=16, fontweight='bold', pad=15)
ax.set_xlabel('Age (years)', fontsize=13, fontweight='bold')
ax.set_ylabel(y_label, fontsize=13, fontweight='bold')
ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5, color='gray')
ax.set_xlim(age_points[0], age_points[-1])
# Set y-axis limits based on data range
if len(sig_disease_indices) > 0:
sig_data = plot_values[k, sig_disease_indices, :]
y_min = max(0, np.min(sig_data) * 0.95) if plot_probability else np.min(sig_data) * 1.05
y_max = np.max(sig_data) * 1.05
# Also consider prevalence if available
if prevalence_t_np is not None:
if isinstance(prevalence_t_np, torch.Tensor):
prev_data = prevalence_t_np[sig_disease_indices, :].detach().cpu().numpy()
else:
prev_data = prevalence_t_np[sig_disease_indices, :]
y_min = min(y_min, max(0, np.min(prev_data) * 0.95))
y_max = max(y_max, np.max(prev_data) * 1.05)
ax.set_ylim(y_min, y_max)
# Add legend
if len(sig_disease_indices) > 0:
ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left', fontsize=9,
framealpha=0.95, edgecolor='gray', fancybox=True, frameon=True)
plt.tight_layout()
return fig
def plot_signature_phi_patterns(phi_np, clusters_np, disease_names, selected_signatures=None,
age_offset=30, figsize=(8, 5), top_n_diseases=None,
plot_probability=True, prevalence_t_np=None, logit_prev_t_np=None):
"""
Plot phi values (as probability or log hazard ratio) for signature-specific diseases (colored)
vs background diseases (in grey), arranged vertically. Optionally overlay prevalence and logit prevalence.
Parameters:
-----------
phi_np : np.ndarray, shape (K, D, T)
Phi values (log hazard ratios) for each signature, disease, and timepoint
clusters_np : np.ndarray, shape (D,)
Disease-to-signature assignments
disease_names : list
List of disease names
selected_signatures : list or None
List of signature indices to plot. If None, plot all signatures.
age_offset : int
Age offset (timepoint 0 = age_offset)
figsize : tuple
Figure size per signature subplot (width, height)
top_n_diseases : int or None
If specified, plot only top N signature-specific diseases based on max phi value
plot_probability : bool
If True, convert phi to probability using sigmoid. If False, plot log hazard ratio.
prevalence_t_np : np.ndarray or None, shape (D, T)
Smoothed prevalence for each disease and timepoint. If provided, will be overlaid.
logit_prev_t_np : np.ndarray or None, shape (D, T)
Logit prevalence for each disease and timepoint. If provided, will be overlaid as probability.
"""
K, D, T = phi_np.shape
if selected_signatures is None:
selected_signatures = list(range(K))
n_sigs = len(selected_signatures)
age_points = np.arange(T) + age_offset
# Convert to probability if requested
if plot_probability:
plot_values = sigmoid(phi_np)
y_label = 'Prob (disease | sig k, age)'
else:
plot_values = phi_np
y_label = 'Log hazard ratio'
# Create subplots in a grid (3 rows x 5 cols for 15 signatures)
n_cols = 5
n_rows = (n_sigs + n_cols - 1) // n_cols # Ceiling division
fig, axes = plt.subplots(n_rows, n_cols, figsize=(figsize[0] * n_cols, figsize[1] * n_rows),
sharex=True, sharey=False)
# Convert axes to a flat array for easier indexing
if n_sigs == 1:
axes_flat = np.array([axes]).flatten()
elif n_rows == 1:
axes_flat = np.array(axes).flatten()
else:
axes_flat = axes.flatten()
for idx, k in enumerate(selected_signatures):
ax = axes_flat[idx]
# Find diseases assigned to this signature
sig_disease_indices = np.where(clusters_np == k)[0]
# Find all other diseases (background)
background_disease_indices = np.where(clusters_np != k)[0]
# Plot background diseases in grey
for d in background_disease_indices:
ax.plot(age_points, plot_values[k, d, :],
color='grey', alpha=0.2, linewidth=0.3, zorder=1)
# Plot signature-specific diseases in colored lines
if len(sig_disease_indices) > 0:
# Optionally select top N diseases
if top_n_diseases is not None and len(sig_disease_indices) > top_n_diseases:
# Rank by max value
max_values = np.max(plot_values[k, sig_disease_indices, :], axis=1)
top_indices = np.argsort(max_values)[-top_n_diseases:]
sig_disease_indices = sig_disease_indices[top_indices]
# Use different colors for signature-specific diseases
colors = plt.cm.tab10(np.linspace(0, 1, min(len(sig_disease_indices), 10)))
for i, d in enumerate(sig_disease_indices):
color = colors[i % len(colors)] if len(sig_disease_indices) <= 10 else 'red'
disease_name = disease_names[d] if d < len(disease_names) else f"Disease_{d}"
# Plot model's learned phi (as probability)
ax.plot(age_points, plot_values[k, d, :],
color=color, linewidth=2.5, label=disease_name, zorder=2)
# Overlay smoothed prevalence if available
if prevalence_t_np is not None and d < prevalence_t_np.shape[0]:
# Convert to numpy if it's a torch tensor
if isinstance(prevalence_t_np, torch.Tensor):
prev_values = prevalence_t_np[d, :].detach().cpu().numpy()
else:
prev_values = prevalence_t_np[d, :]
ax.plot(age_points, prev_values,
color=color, linewidth=1.5, linestyle='--',
alpha=0.6, label=f'{disease_name} (prev)', zorder=3)
# Overlay logit prevalence (converted to probability) if available
if logit_prev_t_np is not None and d < logit_prev_t_np.shape[0]:
# Convert to numpy if it's a torch tensor
if isinstance(logit_prev_t_np, torch.Tensor):
logit_prev_values = logit_prev_t_np[d, :].detach().cpu().numpy()
else:
logit_prev_values = logit_prev_t_np[d, :]
logit_prev_prob = sigmoid(logit_prev_values)
ax.plot(age_points, logit_prev_prob,
color=color, linewidth=1.5, linestyle=':',
alpha=0.6, label=f'{disease_name} (logit prev)', zorder=3)
ax.set_title(f'Signature {k}', fontsize=12, fontweight='bold', pad=5)
if idx >= (n_rows - 1) * n_cols: # Only label x-axis on bottom row
ax.set_xlabel('Age (yr)', fontsize=10)
if idx % n_cols == 0: # Only label y-axis on left column
ax.set_ylabel(y_label, fontsize=10)
ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
ax.set_xlim(age_points[0], age_points[-1])
# Set y-axis limits based on data range for this signature
if len(sig_disease_indices) > 0:
sig_data = plot_values[k, sig_disease_indices, :]
y_min = max(0, np.min(sig_data) * 0.9) if plot_probability else np.min(sig_data) * 1.1
y_max = np.max(sig_data) * 1.1
# Also consider prevalence and logit prevalence if available
if prevalence_t_np is not None:
# Convert to numpy if it's a torch tensor
if isinstance(prevalence_t_np, torch.Tensor):
prev_data = prevalence_t_np[sig_disease_indices, :].detach().cpu().numpy()
else:
prev_data = prevalence_t_np[sig_disease_indices, :]
y_min = min(y_min, max(0, np.min(prev_data) * 0.9))
y_max = max(y_max, np.max(prev_data) * 1.1)
if logit_prev_t_np is not None:
# Convert to numpy if it's a torch tensor
if isinstance(logit_prev_t_np, torch.Tensor):
logit_prev_t_np_numpy = logit_prev_t_np.detach().cpu().numpy()
else:
logit_prev_t_np_numpy = logit_prev_t_np
logit_prev_prob_data = sigmoid(logit_prev_t_np_numpy[sig_disease_indices, :])
y_min = min(y_min, max(0, np.min(logit_prev_prob_data) * 0.9))
y_max = max(y_max, np.max(logit_prev_prob_data) * 1.1)
ax.set_ylim(y_min, y_max)
# Add legend if there are signature-specific diseases (only for first few to avoid clutter)
if len(sig_disease_indices) > 0 and len(sig_disease_indices) <= 10 and idx < 3: # Only show legend for first 3 signatures
ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left', fontsize=6,
framealpha=0.9, edgecolor='gray', fancybox=True)
# Hide unused subplots
for idx in range(n_sigs, n_rows * n_cols):
axes_flat[idx].set_visible(False)
plt.tight_layout(rect=[0, 0, 0.95, 1]) # Leave space for legend
return fig
print("Plotting functions defined!")
Panel A: Signature-Disease Heatmap¶
Show heatmap of signatures with top diseases for each.
Note: We can use the existing plot_disease_blocks function from utils.py which creates heatmaps showing cluster correspondence between biobanks. This shows how diseases are assigned to signatures (clusters).
In [ ]:
# Load data
from helper_py.pathway_discovery import load_full_data
# Load UKB model checkpoint to get clusters (disease-signature assignments)
import torch
# Load model checkpoints for cross-cohort comparison
# Note: plot_disease_blocks shows cluster correspondence between biobanks
# For main paper, we may want to show UKB signature-disease assignments
ukb_checkpoint_path = '/Users/sarahurbut/Dropbox-Personal/model_with_kappa_bigam.pt'
ukb_checkpoint = torch.load(ukb_checkpoint_path, map_location='cpu')
print(f"UKB checkpoint keys: {list(ukb_checkpoint.keys())}")
print(f"Number of diseases: {len(ukb_checkpoint['disease_names'])}")
print(f"Number of signatures: {ukb_checkpoint['clusters'].max() + 1 if hasattr(ukb_checkpoint['clusters'], 'max') else 'N/A'}")
# Extract clusters (disease-to-signature assignments)
clusters = ukb_checkpoint['clusters']
disease_names = ukb_checkpoint['disease_names']
# Convert to list if needed
if isinstance(disease_names, (list, tuple)):
disease_names = list(disease_names)
elif hasattr(disease_names, 'values'):
disease_names = disease_names.values.tolist()
In [ ]:
# Create signature-disease heatmap
# Option 1: Use existing plot_disease_blocks for cross-cohort comparison
# Option 2: Create new heatmap showing UKB signature-disease assignments
# For main paper Figure 2A, we want to show:
# - Each signature (column) with top diseases (rows)
# - Diseases clustered within each signature based on their psi/phi associations
# Load psi matrix (disease-signature associations)
if 'psi' in ukb_checkpoint:
psi = ukb_checkpoint['psi']
print(f"Psi shape: {psi.shape}")
# Create heatmap of psi values
# Rows: diseases, Columns: signatures
# This shows the strength of association between each disease and signature
# TODO: Create publication-ready heatmap
# - Show top diseases for each signature
# - Cluster diseases within signatures
# - Use appropriate color scale
# - Add labels and legend
else:
print("Psi matrix not found in checkpoint")
print("Available keys:", list(ukb_checkpoint.keys()))
In [ ]:
# Load master checkpoint with pooled phi from corrected E
master_checkpoint = torch.load('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/master_for_fitting_pooled_correctedE.pt', map_location='cpu')
phi_np = master_checkpoint['model_state_dict']['phi'].numpy()
# Load clusters from old checkpoint (clusters don't change)
ukb_checkpoint_path = '/Users/sarahurbut/Dropbox-Personal/model_with_kappa_bigam.pt'
ukb_checkpoint = torch.load(ukb_checkpoint_path, map_location='cpu')
clusters_np = ukb_checkpoint['clusters']
# Load corrected prevalence
prevalence_t_np = torch.load('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/prevalence_t_corrected.pt', map_location='cpu')
if torch.is_tensor(prevalence_t_np):
prevalence_t_np = prevalence_t_np.numpy()
# Get disease names
disease_names = ukb_checkpoint['disease_names']
if isinstance(disease_names, (list, tuple)):
disease_names = list(disease_names)
elif hasattr(disease_names, 'values'):
disease_names = disease_names.values.tolist()
print(f"Loaded master checkpoint:")
print(f" Phi shape: {phi_np.shape}")
print(f" Clusters shape: {clusters_np.shape if hasattr(clusters_np, 'shape') else len(clusters_np)}")
print(f" Prevalence shape: {prevalence_t_np.shape}")
selected_sigs = [5, 6, 11, 14]
fig = plot_signature_phi_patterns(phi_np, clusters_np, disease_names,
selected_signatures=selected_sigs,
top_n_diseases=7, # Show top 7 signature-specific diseases
plot_probability=False, # Plot log hazard ratio (not probability)
prevalence_t_np=prevalence_t_np) # Overlay prevalence
plt.show()
In [ ]:
%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures/generate_fig2_signatures.py
In [ ]:
%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/pool_aou_phi_from_batches.py
In [ ]:
## Panel E: Cross-Cohort Comparison - Cardiovascular and Malignancy Signatures
# Load all three model checkpoints
# UKB: master checkpoint with pooled phi
ukb_master = torch.load('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/master_for_fitting_pooled_correctedE.pt', map_location='cpu')
ukb_phi = ukb_master['model_state_dict']['phi'].numpy()
# AOU: initialized model
aou_checkpoint = torch.load('/Users/sarahurbut/aladynoulli2/aou_model_master_correctedE.pt', map_location='cpu')
aou_phi = aou_checkpoint['model_state_dict']['phi'].numpy()
# MGB: initialized model
mgb_checkpoint = torch.load('/Users/sarahurbut/aladynoulli2/mgb_model_initialized.pt', map_location='cpu')
mgb_phi = mgb_checkpoint['model_state_dict']['phi'].numpy()
mgb_checkpoint2 = torch.load('/Users/sarahurbut/Dropbox-Personal/model_with_kappa_bigam_MGB.pt')
# Get clusters and disease names from UKB (for reference)
ukb_checkpoint_ref = torch.load('/Users/sarahurbut/Dropbox-Personal/model_with_kappa_bigam.pt', map_location='cpu')
ukb_clusters = ukb_checkpoint_ref['clusters']
ukb_disease_names = ukb_checkpoint_ref['disease_names']
if isinstance(ukb_disease_names, (list, tuple)):
ukb_disease_names = list(ukb_disease_names)
elif hasattr(ukb_disease_names, 'values'):
ukb_disease_names = ukb_disease_names.values.tolist()
# Get AOU and MGB disease names
aou_disease_names = aou_checkpoint['disease_names']
if isinstance(aou_disease_names, (list, tuple)):
aou_disease_names = list(aou_disease_names)
elif hasattr(aou_disease_names, 'values'):
aou_disease_names = aou_disease_names.values.tolist()
mgb_disease_names = mgb_checkpoint['disease_names']
if isinstance(mgb_disease_names, (list, tuple)):
mgb_disease_names = list(mgb_disease_names)
elif hasattr(mgb_disease_names, 'values'):
mgb_disease_names = mgb_disease_names.values.tolist()
# Define signature mappings
# Cardiovascular: UKB 5, AOU 16, MGB 5
# Malignancy: UKB 6, AOU 11, MGB 11
cv_signatures = {'ukb': 5, 'aou': 16, 'mgb': 5}
malignancy_signatures = {'ukb': 6, 'aou': 11, 'mgb': 11}
# Define diseases to plot for each signature type
# Cardiovascular diseases (from UKB signature 5)
cv_diseases_ukb = [
('Coronary atherosclerosis', 111), # Need to find actual indices
('Myocardial infarction', 112),
('Other chronic ischemic heart disease, unspecified', 115),
('Unstable angina (intermediate coronary syndrome)', 113)
]
# Malignancy diseases (from UKB signature 6)
malignancy_diseases_ukb = [
('Malignant neoplasm, other', 18),
('Secondary malignancy of bone', 20),
('Secondary malignancy of lymph nodes', 21),
('Secondary malignancy of respiratory organs', 22),
('Secondary malignant neoplasm', 19),
('Secondary malignant neoplasm of digestive systems', 23),
('Secondary malignant neoplasm of liver', 24)
]
# Helper function to find disease index by name
def find_disease_index(disease_name, disease_names_list):
"""Find disease index by name (fuzzy matching)"""
for i, name in enumerate(disease_names_list):
if disease_name.lower() in str(name).lower() or str(name).lower() in disease_name.lower():
return i
return None
# Create the 3x2 plot
fig, axes = plt.subplots(3, 2, figsize=(14, 12))
# Get actual time dimension from phi (should be 51 or 52)
T = ukb_phi.shape[2]
age_points = np.arange(30, 30 + T) # Ages starting at 30
print(f"Time dimension: {T}, Age range: {age_points[0]}-{age_points[-1]}")
# Colors for diseases
cv_colors = ['green', 'orange', 'lightblue', 'purple']
malignancy_colors = ['lightblue', 'orange', 'green', 'brown', 'yellow', 'pink', 'peachpuff'] # Fixed: peachpuff instead of lightorange
# First, identify specific cardiovascular diseases in UKB signature 5 by name
# These are the diseases we want to plot: Coronary atherosclerosis, Myocardial infarction, etc.
ukb_cv_sig = 5
if isinstance(ukb_clusters, torch.Tensor):
ukb_clusters_np = ukb_clusters.numpy()
else:
ukb_clusters_np = ukb_clusters
ukb_cv_indices = np.where(ukb_clusters_np == ukb_cv_sig)[0]
ukb_cv_names = [ukb_disease_names[i] for i in ukb_cv_indices if i < len(ukb_disease_names)]
# Target cardiovascular disease names to match
target_cv_diseases = [
'Coronary atherosclerosis',
'Myocardial infarction',
'Other chronic ischemic heart disease',
'Unstable angina'
]
# Find these diseases in UKB signature 5
ukb_cv_target_names = []
ukb_cv_target_indices = []
for target_name in target_cv_diseases:
for i, ukb_name in enumerate(ukb_cv_names):
if (target_name.lower() in str(ukb_name).lower() or
str(ukb_name).lower() in target_name.lower()):
if ukb_cv_indices[i] not in ukb_cv_target_indices:
ukb_cv_target_names.append(ukb_name)
ukb_cv_target_indices.append(ukb_cv_indices[i])
break
print(f"\nUKB cardiovascular signature {ukb_cv_sig} target diseases ({len(ukb_cv_target_names)}):")
for i, (idx, name) in enumerate(zip(ukb_cv_target_indices, ukb_cv_target_names)):
print(f" {i}: {name} (idx {idx})")
# Plot Cardiovascular signatures (left column) - MATCH BY NAME
for row_idx, (cohort, sig_idx) in enumerate(cv_signatures.items()):
ax = axes[row_idx, 0]
if cohort == 'ukb':
phi_data = ukb_phi
clusters = ukb_clusters
disease_names_list = ukb_disease_names
elif cohort == 'aou':
phi_data = aou_phi
clusters = aou_checkpoint.get('clusters', None)
disease_names_list = aou_disease_names
else: # mgb
phi_data = mgb_phi
clusters = mgb_checkpoint.get('clusters', None)
disease_names_list = mgb_disease_names
# Debug: Print info
print(f"\n{cohort.upper()} - Cardiovascular signature {sig_idx}:")
print(f" Phi shape: {phi_data.shape}")
print(f" Clusters type: {type(clusters)}")
# Get diseases in this signature
if clusters is None:
print(f" WARNING: No clusters found for {cohort}!")
sig_disease_indices = []
matched_disease_indices = []
else:
if isinstance(clusters, torch.Tensor):
clusters_np = clusters.numpy()
else:
clusters_np = clusters
print(f" Clusters shape: {clusters_np.shape if hasattr(clusters_np, 'shape') else len(clusters_np)}")
print(f" Unique signatures: {np.unique(clusters_np) if hasattr(clusters_np, 'shape') else 'N/A'}")
sig_disease_indices = np.where(clusters_np == sig_idx)[0]
print(f" Diseases in signature {sig_idx}: {len(sig_disease_indices)}")
# Match diseases by name with UKB cardiovascular target diseases
matched_disease_indices = []
matched_names = []
for ukb_name in ukb_cv_target_names:
# Find this disease in the current cohort's signature
for d_idx in sig_disease_indices:
if d_idx < len(disease_names_list):
cohort_disease_name = str(disease_names_list[d_idx])
# Fuzzy matching: check if names are similar
if (ukb_name.lower() in cohort_disease_name.lower() or
cohort_disease_name.lower() in ukb_name.lower() or
ukb_name.lower().replace(' ', '') == cohort_disease_name.lower().replace(' ', '')):
if d_idx not in matched_disease_indices:
matched_disease_indices.append(d_idx)
matched_names.append(cohort_disease_name)
break
print(f" Matched {len(matched_disease_indices)} diseases with UKB:")
for i, (idx, name) in enumerate(zip(matched_disease_indices[:10], matched_names[:10])):
print(f" {i}: {name} (idx {idx})")
# Plot matched diseases (or all if no matches found)
diseases_to_plot = matched_disease_indices if len(matched_disease_indices) > 0 else sig_disease_indices[:4]
if len(diseases_to_plot) > 0:
# Get disease names and plot
for i, d_idx in enumerate(diseases_to_plot):
if d_idx < phi_data.shape[1]:
phi_values = phi_data[sig_idx, d_idx, :]
# Ensure age_points matches phi_values length
if len(phi_values) != len(age_points):
age_points_adj = np.arange(30, 30 + len(phi_values))
else:
age_points_adj = age_points
disease_name = disease_names_list[d_idx] if d_idx < len(disease_names_list) else f"Disease {d_idx}"
ax.plot(age_points_adj, phi_values, label=disease_name,
color=cv_colors[i % len(cv_colors)], linewidth=2)
else:
print(f" WARNING: No diseases found in {cohort} signature {sig_idx}!")
ax.text(0.5, 0.5, f'No diseases in signature {sig_idx}',
transform=ax.transAxes, ha='center', va='center', fontsize=12)
ax.set_xlabel('Age (yr)', fontsize=10)
ax.set_ylabel('Log hazard ratio', fontsize=10)
ax.set_title(f'{cohort.upper()}', fontsize=12, fontweight='bold')
ax.grid(True, alpha=0.3)
ax.set_ylim(-12, -3)
ax.set_xlim(30, 80)
if row_idx == 0:
ax.set_title('Cardiovascular signatures', fontsize=12, fontweight='bold')
if row_idx == 2 and len(sig_disease_indices) > 0:
ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=8)
# First, identify diseases in UKB signature 6 (malignancy) by name
# Then match those diseases across AOU and MGB
ukb_malignancy_sig = 6
if isinstance(ukb_clusters, torch.Tensor):
ukb_clusters_np = ukb_clusters.numpy()
else:
ukb_clusters_np = ukb_clusters
ukb_malignancy_indices = np.where(ukb_clusters_np == ukb_malignancy_sig)[0]
ukb_malignancy_names = [ukb_disease_names[i] for i in ukb_malignancy_indices if i < len(ukb_disease_names)]
print(f"\nUKB malignancy signature {ukb_malignancy_sig} diseases ({len(ukb_malignancy_names)}):")
for i, name in enumerate(ukb_malignancy_names[:10]): # Print first 10
print(f" {i}: {name}")
# Plot Malignancy signatures (right column) - MATCH BY NAME
for row_idx, (cohort, sig_idx) in enumerate(malignancy_signatures.items()):
ax = axes[row_idx, 1]
if cohort == 'ukb':
phi_data = ukb_phi
clusters = ukb_clusters
disease_names_list = ukb_disease_names
elif cohort == 'aou':
phi_data = aou_phi
clusters = aou_checkpoint.get('clusters', None)
disease_names_list = aou_disease_names
else: # mgb
phi_data = mgb_phi
clusters = mgb_checkpoint.get('clusters', None)
disease_names_list = mgb_disease_names
# Debug: Print info
print(f"\n{cohort.upper()} - Malignancy signature {sig_idx}:")
print(f" Phi shape: {phi_data.shape}")
print(f" Clusters type: {type(clusters)}")
# Get diseases in this signature
if clusters is None:
print(f" WARNING: No clusters found for {cohort}!")
sig_disease_indices = []
matched_disease_indices = []
else:
if isinstance(clusters, torch.Tensor):
clusters_np = clusters.numpy()
else:
clusters_np = clusters
print(f" Clusters shape: {clusters_np.shape if hasattr(clusters_np, 'shape') else len(clusters_np)}")
print(f" Unique signatures: {np.unique(clusters_np) if hasattr(clusters_np, 'shape') else 'N/A'}")
sig_disease_indices = np.where(clusters_np == sig_idx)[0]
print(f" Diseases in signature {sig_idx}: {len(sig_disease_indices)}")
# Match diseases by name with UKB malignancy diseases
matched_disease_indices = []
matched_names = []
for ukb_name in ukb_malignancy_names:
# Find this disease in the current cohort's signature
for d_idx in sig_disease_indices:
if d_idx < len(disease_names_list):
cohort_disease_name = str(disease_names_list[d_idx])
# Fuzzy matching: check if names are similar
if (ukb_name.lower() in cohort_disease_name.lower() or
cohort_disease_name.lower() in ukb_name.lower() or
ukb_name.lower().replace(' ', '') == cohort_disease_name.lower().replace(' ', '')):
if d_idx not in matched_disease_indices:
matched_disease_indices.append(d_idx)
matched_names.append(cohort_disease_name)
break
print(f" Matched {len(matched_disease_indices)} diseases with UKB:")
for i, (idx, name) in enumerate(zip(matched_disease_indices[:10], matched_names[:10])):
print(f" {i}: {name} (idx {idx})")
# Plot matched diseases (or all if no matches found)
diseases_to_plot = matched_disease_indices if len(matched_disease_indices) > 0 else sig_disease_indices[:7]
if len(diseases_to_plot) > 0:
# Get disease names and plot
for i, d_idx in enumerate(diseases_to_plot):
if d_idx < phi_data.shape[1]:
phi_values = phi_data[sig_idx, d_idx, :]
# Ensure age_points matches phi_values length
if len(phi_values) != len(age_points):
age_points_adj = np.arange(30, 30 + len(phi_values))
else:
age_points_adj = age_points
disease_name = disease_names_list[d_idx] if d_idx < len(disease_names_list) else f"Disease {d_idx}"
ax.plot(age_points_adj, phi_values, label=disease_name,
color=malignancy_colors[i % len(malignancy_colors)], linewidth=2)
else:
print(f" WARNING: No diseases found in {cohort} signature {sig_idx}!")
ax.text(0.5, 0.5, f'No diseases in signature {sig_idx}',
transform=ax.transAxes, ha='center', va='center', fontsize=12)
ax.set_xlabel('Age (yr)', fontsize=10)
ax.set_ylabel('Log hazard ratio', fontsize=10)
ax.set_title(f'{cohort.upper()}', fontsize=12, fontweight='bold')
ax.grid(True, alpha=0.3)
ax.set_ylim(-16, -4)
ax.set_xlim(30, 80)
if row_idx == 0:
ax.set_title('Malignancy signatures', fontsize=12, fontweight='bold')
if row_idx == 2 and len(sig_disease_indices) > 0:
ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=8)
plt.suptitle('Panel E: Cross-Cohort Comparison', fontsize=14, fontweight='bold', y=0.995)
plt.tight_layout()
plt.show()
Calculate Empirical Hazards from MGB/AOU ICD10 Data¶
Calculate empirical hazards from raw ICD10 data to compare with model's learned phi values. This addresses the competing risk concern by showing the model captures empirical patterns.
Plot Signature Patterns for All Three Cohorts (UKB, AOU, MGB)¶
Create signature plots showing phi values for all signatures (0-14) for each cohort. Each plot shows diseases in the signature (colored) vs background diseases (grey).
In [ ]:
# Define plotting functions
from scipy.special import expit as sigmoid
import sys
sys.path.append('/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures')
from plot_signature_patterns import plot_signature_phi_patterns
print("Plotting functions loaded!")
In [ ]:
%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/plot_ukb_sigs.py
%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/plot_aou_all_signatures.py
In [ ]:
import torch
mgb_checkpoint = torch.load('/Users/sarahurbut/aladynoulli2/mgb_model_initialized.pt', map_location='cpu')
if 'model_state_dict' in mgb_checkpoint and 'psi' in mgb_checkpoint['model_state_dict']:
psi = mgb_checkpoint['model_state_dict']['psi']
print(f"Psi range: [{psi.min():.3f}, {psi.max():.3f}]")
print(f"Psi mean: {psi.mean():.3f}")
# Check if values are in initialized range (-2 to +1) or trained range
if psi.max() > 2.0:
print("✓ Looks like TRAINED psi (values > 2)")
else:
print("⚠ Looks like INITIALIZED psi (values in -2 to +1 range)")
In [ ]:
%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/plot_mgb_all_signatures.py
Panel B & C: PSI Heatmap and Disease Prevalence by Age (Corrected E)¶
Create two plots:
- Panel B: PSI heatmap showing disease-signature associations (log odds ratios)
- Panel C: Disease prevalence by age, clustered by signature (average pi over people)
These plots demonstrate that we no longer have censoring bias with the corrected E approach.
In [ ]:
import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.special import expit as sigmoid
import glob
from pathlib import Path
# Load PSI from one batch (or average from 40 batches)
batch_pattern = '/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/enrollment_model_W0.0001_batch_*_*.pt'
batch_files = sorted(glob.glob(batch_pattern))
print(f"Found {len(batch_files)} batch files")
# Option 1: Use first batch
if len(batch_files) > 0:
batch_checkpoint = torch.load(batch_files[0], map_location='cpu', weights_only=False)
psi = batch_checkpoint['model_state_dict']['psi']
if torch.is_tensor(psi):
psi = psi.detach().cpu().numpy()
print(f"Loaded PSI from first batch: {batch_files[0]}")
print(f"PSI shape: {psi.shape}")
# Option 2: Average across all batches (uncomment to use)
all_psis = []
for batch_file in batch_files:
try:
checkpoint = torch.load(batch_file, map_location='cpu', weights_only=False)
psi_batch = checkpoint['model_state_dict']['psi']
if torch.is_tensor(psi_batch):
psi_batch = psi_batch.detach().cpu().numpy()
all_psis.append(psi_batch)
except:
continue
psi = np.mean(np.stack(all_psis, axis=0), axis=0)
print(f"Averaged PSI across {len(all_psis)} batches")
else:
raise FileNotFoundError("No batch files found!")
# Load clusters and disease names
ukb_checkpoint_ref = torch.load('/Users/sarahurbut/Dropbox-Personal/model_with_kappa_bigam.pt', map_location='cpu')
clusters = ukb_checkpoint_ref['clusters']
if torch.is_tensor(clusters):
clusters = clusters.detach().cpu().numpy()
disease_names = ukb_checkpoint_ref['disease_names']
if isinstance(disease_names, (list, tuple)):
disease_names = list(disease_names)
elif hasattr(disease_names, 'values'):
disease_names = disease_names.values.tolist()
print(f"Clusters shape: {clusters.shape}")
print(f"Number of diseases: {len(disease_names)}")
print(f"Number of signatures: {psi.shape[0]}")
In [ ]:
# Load pi tensor and average over people
pi_path = '/Users/sarahurbut/Library/CloudStorage/Dropbox/enrollment_predictions_fixedphi_correctedE_vectorized/pi_enroll_fixedphi_sex_FULL.pt'
print(f"Loading pi tensor from: {pi_path}")
pi_full = torch.load(pi_path, map_location='cpu', weights_only=False)
if torch.is_tensor(pi_full):
pi_full = pi_full.detach().cpu().numpy()
print(f"Pi tensor shape: {pi_full.shape}") # Should be (N, D, T)
# Average over people (patients) to get (D, T) - average probability for each disease at each age
pi_avg = np.mean(pi_full, axis=0) # Shape: (D, T)
print(f"Average pi shape: {pi_avg.shape}") # Should be (D, T)
# Get age points (assuming age_offset=30, T=52 timepoints)
T = pi_avg.shape[1]
age_points = np.arange(30, 30 + T)
print(f"Age range: {age_points[0]}-{age_points[-1]}")
In [ ]:
# Create Panel B: PSI Heatmap (Disease Index vs Signatures)
# PSI shape: (K, D) where K=21 signatures, D=348 diseases
# We want to show: diseases (rows) vs signatures (columns)
# Signature labels for grouping
SIGNATURE_LABELS = {
0: 'Cardiac Arrhythmias',
1: 'Musculoskeletal',
2: 'Upper GI/Esophageal',
3: 'Mixed/General Medical',
4: 'Upper Respiratory',
5: 'Ischemic cardiovascular',
6: 'Metastatic Cancer',
7: 'Pain/Inflammation',
8: 'Gynecologic',
9: 'Spinal Disorders',
10: 'Ophthalmologic',
11: 'Cerebrovascular',
12: 'Renal/Urologic',
13: 'Male Urogenital',
14: 'Pulmonary/Smoking',
15: 'Metabolic/Diabetes',
16: 'Infectious/Critical Care',
17: 'Lower GI/Colon',
18: 'Hepatobiliary',
19: 'Dermatologic/Oncologic',
20: 'Health'
}
# Transpose PSI to get (D, K) for heatmap
psi_heatmap = psi.T # Shape: (D, K)
# Create figure for PSI heatmap only
# Increased height to accommodate all disease names
fig1, ax1 = plt.subplots(1, 1, figsize=(14, 24))
# Create heatmap
# Note: We'll show a subset of diseases (top diseases per signature) for clarity
# Or show all diseases but cluster them by signature
# Reorder diseases by signature (cluster assignment), then by max psi within each signature
# Group diseases by signature
sig_to_diseases = {}
for d in range(len(clusters)):
sig = clusters[d]
if sig not in sig_to_diseases:
sig_to_diseases[sig] = []
sig_to_diseases[sig].append(d)
# Sort diseases within each signature by their psi value in that signature (most positive first)
# Note: psi is already averaged across batches, so we use psi[sig, d] for disease d in signature sig
disease_order = []
for sig in sorted(sig_to_diseases.keys()):
diseases_in_sig = sig_to_diseases[sig]
# Get psi value for each disease in its assigned signature
sig_psi_values = [psi[sig, d] for d in diseases_in_sig]
# Sort by psi value in this signature (descending - most positive first)
sorted_indices = sorted(range(len(diseases_in_sig)),
key=lambda i: sig_psi_values[i], reverse=True)
disease_order.extend([diseases_in_sig[i] for i in sorted_indices])
# Reorder PSI and disease names
psi_ordered = psi_heatmap[disease_order, :]
disease_names_ordered = [disease_names[i] for i in disease_order]
# Create heatmap with improved aesthetics
im = ax1.imshow(psi_ordered, aspect='auto', cmap='RdBu_r',
vmin=-5, vmax=3, interpolation='nearest')
# Set labels with better formatting
ax1.set_xlabel('Signatures', fontsize=16, fontweight='bold')
ax1.set_ylabel('Disease index (psi kd)', fontsize=16, fontweight='bold')
ax1.set_title('Panel B: Disease-Signature Associations (Log Odds Ratio)',
fontsize=18, fontweight='bold', pad=25)
# Set x-axis ticks (signatures 0-20)
ax1.set_xticks(np.arange(psi.shape[0]))
ax1.set_xticklabels([f'{i}' for i in range(psi.shape[0])], fontsize=11, fontweight='bold')
# Set y-axis ticks (show all disease names - will be hard to read but comprehensive)
n_diseases = len(disease_order)
ax1.set_yticks(np.arange(n_diseases))
ax1.set_yticklabels([disease_names_ordered[i] for i in range(n_diseases)],
fontsize=6) # Small font size to fit all names
# Add colorbar with better formatting
cbar = fig1.colorbar(im, ax=ax1, label='Log odds ratio (psi)', pad=0.02, shrink=0.8)
cbar.ax.tick_params(labelsize=12)
cbar.set_label('Log odds ratio (psi)', fontsize=14, fontweight='bold')
# Add vertical lines to separate signatures
for k in range(psi.shape[0] + 1):
ax1.axvline(k - 0.5, color='white', linewidth=1.5, alpha=0.8)
# Add horizontal lines to separate signature groups (diseases grouped by signature)
# Also add signature labels as brackets on the right side
from matplotlib.patches import FancyBboxPatch
import matplotlib.patches as mpatches
current_y = 0
sig_y_positions = {} # Store y positions for each signature group
for sig in sorted(sig_to_diseases.keys()):
n_diseases_in_sig = len(sig_to_diseases[sig])
if current_y > 0:
ax1.axhline(current_y - 0.5, color='white', linewidth=1.5, alpha=0.8)
# Store the middle y position for this signature group
sig_y_positions[sig] = current_y + n_diseases_in_sig / 2 - 0.5
current_y += n_diseases_in_sig
# Add signature labels on the right side with brackets
ax2 = ax1.twinx() # Create a second y-axis for labels
ax2.set_ylim(ax1.get_ylim())
ax2.set_yticks([])
ax2.set_yticklabels([])
# Get x position for labels (just to the right of the heatmap)
x_label_pos = psi.shape[0] + 0.5
# Add brackets and labels for each signature group
for sig in sorted(sig_to_diseases.keys()):
n_diseases_in_sig = len(sig_to_diseases[sig])
sig_start_y = sum(len(sig_to_diseases[s]) for s in sorted(sig_to_diseases.keys()) if s < sig) - 0.5
sig_end_y = sig_start_y + n_diseases_in_sig
# Get label text
label_text = SIGNATURE_LABELS.get(sig, f'Signature {sig}')
# Draw bracket (curly brace)
# Use annotation with curly brace style
y_mid = (sig_start_y + sig_end_y) / 2
# Draw bracket using FancyBboxPatch or annotation
# Create a bracket shape using lines
bracket_width = 0.3
bracket_x_start = x_label_pos - bracket_width
bracket_x_end = x_label_pos
# Top of bracket
ax2.plot([bracket_x_start, bracket_x_end], [sig_start_y, sig_start_y],
'k-', linewidth=1.5, clip_on=False)
# Bottom of bracket
ax2.plot([bracket_x_start, bracket_x_end], [sig_end_y, sig_end_y],
'k-', linewidth=1.5, clip_on=False)
# Vertical line
ax2.plot([bracket_x_start, bracket_x_start], [sig_start_y, sig_end_y],
'k-', linewidth=1.5, clip_on=False)
# Add label text
ax2.text(x_label_pos + 0.1, y_mid, label_text,
fontsize=9, va='center', ha='left',
rotation=0, fontweight='bold',
clip_on=False)
# Adjust x-axis limits to accommodate labels
ax1.set_xlim(-0.5, x_label_pos + 3)
ax2.set_xlim(-0.5, x_label_pos + 3)
plt.tight_layout()
# Display the figure
from IPython.display import display
#display(fig1)
plt.show()
print("✓ Panel B: PSI heatmap created")
In [ ]:
from generate_psi_heatmap import plot_psi_heatmap_with_arrows
%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures/generate_psi_heatmap.py
In [ ]:
%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures/generate_pi_heatmap2.py