Figure 3: Individual Trajectories

Purpose

Demonstrate individual-specific disease trajectories and how signatures evolve over time.

Panels Required:

• Panel A: Case studies of 2-3 individuals showing signature evolution
• Panel B: Multimorbidity patterns - how theta changes before vs. after diagnosis
• Panel C: Signature response to new diagnoses (real-time updating)
• Panel D: Comparison of trajectories between subtypes of same disease

Key Message:

Show personalized trajectories and dynamic updating capabilities

# 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
sns.set_style("whitegrid")
plt.rcParams['figure.dpi'] = 300

print("Setup complete")

# ============================================================================
# Plot Patient Timeline (Panel A/B/C Style)
# ============================================================================

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

# Load data
initial_clusters = torch.load('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/initial_clusters_400k.pt')
if torch.is_tensor(initial_clusters):
    initial_clusters = initial_clusters.numpy()
K = int(initial_clusters.max() + 1)

checkpoint_path = '/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/enrollment_model_W0.0001_batch_0_10000.pt'
pi_path = '/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/pi_fullmode_batch_0_10000.pt'

ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
lambda_ = ckpt['model_state_dict']['lambda_']
theta = torch.softmax(lambda_, dim=1).numpy()  # (N, K, T)
Y = ckpt['Y']
if torch.is_tensor(Y):
    Y_np = Y.numpy()
else:
    Y_np = Y

pi_full = torch.load(pi_path, map_location='cpu', weights_only=False)
if torch.is_tensor(pi_full):
    pi_np = pi_full.numpy()
else:
    pi_np = pi_full

# Load disease names if available
#disease_names = ckpt.get('disease_names', [f'Disease {i}' for i in range(Y_np.shape[1])])
disease_names=pd.read_csv('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/disease_names.csv')['x'].tolist()
# Find patient with signature spike
N, K_total, T = theta.shape
patient_idx = None
for p_idx in range(N):
    max_theta = theta[p_idx, :, :].max()
    if max_theta > 0.4 and Y_np[p_idx, :, :].sum() > 0:
        patient_idx = p_idx
        break

if patient_idx is None:
    patient_idx = 0

print(f"Selected patient: {patient_idx}")

# Get patient data
patient_theta = theta[patient_idx, :, :]  # (K, T)
patient_pi = pi_np[patient_idx, :, :]  # (D, T)
patient_Y = Y_np[patient_idx, :, :]  # (D, T)

# Find diagnoses
diagnosis_times = {}
for d in range(patient_Y.shape[0]):
    event_times = np.where(patient_Y[d, :] == 1)[0]
    if len(event_times) > 0:
        diagnosis_times[d] = event_times.tolist()

# Convert timepoints to ages
ages = np.arange(30, 30 + T)

# ============================================================================
# Create Plot (Panel A/B/C Style)
# ============================================================================

fig = plt.figure(figsize=(14, 10))
gs = plt.GridSpec(3, 1, height_ratios=[2, 0.8, 1.2], hspace=0.4)

# Panel 1: Signature loadings (θ) vs Age
ax1 = fig.add_subplot(gs[0])
colors = sns.color_palette("tab20", K_total)

for k in range(K_total):
    ax1.plot(ages, patient_theta[k, :], 
             label=f'Signature {k}', linewidth=2, color=colors[k], alpha=0.8)

# Mark diagnosis times
for d, times in diagnosis_times.items():
    for t in times:
        age_at_diag = 30 + t
        ax1.axvline(x=age_at_diag, color='gray', linestyle='--', alpha=0.5, linewidth=1)

ax1.set_ylabel('Signature loadings (θ)', fontsize=12)
ax1.set_title(f'Patient {patient_idx}: Signature Trajectories', fontsize=14, fontweight='bold')
ax1.legend(bbox_to_anchor=(1.02, 1), loc='upper left', fontsize=9, ncol=2)
ax1.grid(True, alpha=0.3)
ax1.set_xlim([30, 81])
ax1.set_ylim([0, None])

# Panel 2: Static model summary (Average loading bars)
ax2 = fig.add_subplot(gs[1])
avg_theta = patient_theta.mean(axis=1)  # Average across time
bars = ax2.barh(range(K_total), avg_theta, color=colors)
ax2.set_xlabel('Average loading (θ)', fontsize=11)
ax2.set_ylabel('Signature', fontsize=11)
ax2.set_yticks(range(K_total))
ax2.set_yticklabels([f'Sig {k}' for k in range(K_total)], fontsize=9)
ax2.set_xlim([0, 1])
ax2.grid(True, alpha=0.3, axis='x')
ax2.set_title('Static model summary', fontsize=11)

# Panel 3: Disease timeline + Top pi values
ax3 = fig.add_subplot(gs[2])

# Top diseases by max pi
max_pi_per_disease = patient_pi.max(axis=1)
top_diseases = np.argsort(max_pi_per_disease)[::-1][:15]  # Top 15

# Plot pi for top diseases
for i, d in enumerate(top_diseases):
    sig_for_disease = initial_clusters[d] if d < len(initial_clusters) else -1
    color = colors[sig_for_disease] if sig_for_disease < K_total else 'gray'
    ax3.plot(ages, patient_pi[d, :], 
             linewidth=1.5, color=color, alpha=0.6, label=f'D{d}')

# Mark diagnoses
for d, times in diagnosis_times.items():
    if d in top_diseases:
        sig_for_disease = initial_clusters[d] if d < len(initial_clusters) else -1
        color = colors[sig_for_disease] if sig_for_disease < K_total else 'black'
        for t in times:
            age_at_diag = 30 + t
            ax3.scatter(age_at_diag, patient_pi[d, t], 
                       s=100, color=color, edgecolors='black', 
                       linewidth=1.5, zorder=10, marker='o')

ax3.set_xlabel('Age (yr)', fontsize=12)
ax3.set_ylabel('Disease Probability (π)', fontsize=11)
ax3.set_title('Top Disease Probabilities', fontsize=11)
ax3.grid(True, alpha=0.3)
ax3.set_xlim([30, 81])

# Add disease list on the right
disease_list_text = "Diseases:\n"
for d in top_diseases[:10]:
    sig = initial_clusters[d] if d < len(initial_clusters) else -1
    d_name = disease_names[d] if d < len(disease_names) else f'Disease {d}'
    disease_list_text += f"D{d} (sig{sig}): {d_name[:40]}\n"

ax3.text(1.02, 0.5, disease_list_text, transform=ax3.transAxes,
         fontsize=8, verticalalignment='center', 
         bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.3))

plt.suptitle(f'Patient {patient_idx} Timeline', fontsize=16, fontweight='bold', y=0.995)
plt.tight_layout()
plt.savefig(f'/Users/sarahurbut/aladynoulli2/patient_{patient_idx}_timeline_panel_style.pdf', 
            dpi=300, bbox_inches='tight')
print(f"\n✓ Saved plot to: patient_{patient_idx}_timeline_panel_style.pdf")
plt.show()

# ============================================================================
# Plot Patient Timeline - Updated Layout
# ============================================================================

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

# Load data
initial_clusters = torch.load('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/initial_clusters_400k.pt')
if torch.is_tensor(initial_clusters):
    initial_clusters = initial_clusters.numpy()
K = int(initial_clusters.max() + 1)

checkpoint_path = '/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/enrollment_model_W0.0001_batch_0_10000.pt'
pi_path = '/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/pi_fullmode_batch_0_10000.pt'

ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
lambda_ = ckpt['model_state_dict']['lambda_']
theta = torch.softmax(lambda_, dim=1).numpy()  # (N, K, T)
Y = ckpt['Y']
if torch.is_tensor(Y):
    Y_np = Y.numpy()
else:
    Y_np = Y

pi_full = torch.load(pi_path, map_location='cpu', weights_only=False)
if torch.is_tensor(pi_full):
    pi_np = pi_full.numpy()
else:
    pi_np = pi_full

# Load disease names
disease_names = pd.read_csv('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/disease_names.csv')['x'].tolist()

# Find patient with signature spike
# Find patient with multiple diseases at different times
N, K_total, T = theta.shape
patient_idx = None
best_score = 0

for p_idx in range(N):
    # Count number of diseases with events
    patient_Y_single = Y_np[p_idx, :, :]
    diseases_with_events = np.where(patient_Y_single.sum(axis=1) > 0)[0]
    n_diseases = len(diseases_with_events)
    
    if n_diseases < 5:  # Need at least 5 diseases
        continue
    
    # Get diagnosis times for all diseases
    diagnosis_times_list = []
    for d in diseases_with_events:
        event_times = np.where(patient_Y_single[d, :] == 1)[0]
        if len(event_times) > 0:
            diagnosis_times_list.extend(event_times.tolist())
    
    if len(diagnosis_times_list) == 0:
        continue
    
    # Calculate spread of diagnosis times (want diagnoses spread across time)
    diagnosis_times_array = np.array(diagnosis_times_list)
    time_spread = diagnosis_times_array.max() - diagnosis_times_array.min()
    
    # Calculate diversity: want diagnoses at different timepoints
    unique_times = len(np.unique(diagnosis_times_array))
    
    # Score: prioritize more diseases, more spread, more unique timepoints
    score = n_diseases * 10 + time_spread * 2 + unique_times * 5
    
    if score > best_score:
        best_score = score
        patient_idx = p_idx

if patient_idx is None:
    # Fallback: find any patient with multiple diseases
    for p_idx in range(N):
        patient_Y_single = Y_np[p_idx, :, :]
        n_diseases = np.where(patient_Y_single.sum(axis=1) > 0)[0].shape[0]
        if n_diseases >= 3:
            patient_idx = p_idx
            break

if patient_idx is None:
    patient_idx = 0

print(f"Selected patient: {patient_idx}")

# Get patient data
patient_theta = theta[patient_idx, :, :]  # (K, T)
patient_pi = pi_np[patient_idx, :, :]  # (D, T)
patient_Y = Y_np[patient_idx, :, :]  # (D, T)

# Find diagnoses
diagnosis_times = {}
for d in range(patient_Y.shape[0]):
    event_times = np.where(patient_Y[d, :] == 1)[0]
    if len(event_times) > 0:
        diagnosis_times[d] = event_times.tolist()

# Print summary
n_diseases = len(diagnosis_times)
all_times = []
for times in diagnosis_times.values():
    all_times.extend(times)
time_range = (min(all_times), max(all_times)) if all_times else (0, 0)
print(f"  Number of diseases: {n_diseases}")
print(f"  Diagnosis timepoints: {sorted(set(all_times))}")
print(f"  Time range: {time_range[0]} to {time_range[1]} (ages {30+time_range[0]} to {30+time_range[1]})")

# Get patient data
patient_theta = theta[patient_idx, :, :]  # (K, T)
patient_pi = pi_np[patient_idx, :, :]  # (D, T)
patient_Y = Y_np[patient_idx, :, :]  # (D, T)

# Calculate average theta (K vector) - average across time
avg_theta = patient_theta.mean(axis=1)  # Shape: (K,)

# Find diseases with events
diseases_with_events = []
diagnosis_times = {}
for d in range(patient_Y.shape[0]):
    event_times = np.where(patient_Y[d, :] == 1)[0]
    if len(event_times) > 0:
        diseases_with_events.append(d)
        diagnosis_times[d] = event_times.tolist()

# Convert timepoints to ages
ages = np.arange(30, 30 + T)

# ============================================================================
# Create Plot
# ============================================================================

fig = plt.figure(figsize=(14, 10))
gs = plt.GridSpec(3, 1, height_ratios=[2, 1, 1.2], hspace=0.4)

colors = sns.color_palette("tab20", K_total)
sig_colors = sns.color_palette("tab20", K_total)

# Panel 1: Signature loadings (θ) vs Age
ax1 = fig.add_subplot(gs[0])
colors = sns.color_palette("tab20", K_total)

for k in range(K_total):
    ax1.plot(ages, patient_theta[k, :], 
             label=f'Signature {k}', linewidth=2, color=colors[k], alpha=0.8)

# Add horizontal lines at diagnosis times
for d, times in diagnosis_times.items():
    for t in times:
        age_at_diag = 30 + t
        # Get the signature for this disease
        sig_for_disease = initial_clusters[d] if d < len(initial_clusters) else -1
        if sig_for_disease < K_total:
            # Draw horizontal line at the theta value for this signature at diagnosis time
            theta_at_diag = patient_theta[sig_for_disease, t]
            ax1.axhline(y=theta_at_diag, xmin=(age_at_diag - 30) / (81 - 30), 
                       xmax=(age_at_diag - 30 + 1) / (81 - 30),
                       color=colors[sig_for_disease], linestyle='--', 
                       alpha=0.6, linewidth=1.5)

# Add thin stacked bar showing average theta (single bar, stacked)
# Position it at the top right
ax1_bar = ax1.inset_axes([0.7, 0.7, 0.25, 0.15])  # [x, y, width, height] in axes coordinates

# Sort signatures by average theta (largest first) for better visualization
sorted_indices = np.argsort(avg_theta)[::-1]
sorted_avg_theta = avg_theta[sorted_indices]
sorted_colors = [colors[i] for i in sorted_indices]

# Create stacked bar (single bar)
bottom = 0
for i, (val, color) in enumerate(zip(sorted_avg_theta, sorted_colors)):
    if val > 0.01:  # Only show if > 1% to avoid clutter
        ax1_bar.barh(0, val, left=bottom, color=color, height=0.5, alpha=0.8)
        bottom += val

ax1_bar.set_xlim([0, 1])
ax1_bar.set_ylim([-0.5, 0.5])
ax1_bar.set_xticks([0, 0.5, 1.0])
ax1_bar.set_xticklabels(['0', '0.5', '1'], fontsize=7)
ax1_bar.set_yticks([])
ax1_bar.set_title('Avg θ (stacked)', fontsize=8)
ax1_bar.spines['top'].set_visible(False)
ax1_bar.spines['right'].set_visible(False)
ax1_bar.spines['left'].set_visible(False)

ax1.set_ylabel('Signature loadings (θ)', fontsize=12)
ax1.set_title(f'Patient {patient_idx}: Signature Trajectories', fontsize=14, fontweight='bold')
ax1.legend(bbox_to_anchor=(1.02, 1), loc='upper left', fontsize=9, ncol=2)
ax1.grid(True, alpha=0.3)
ax1.set_xlim([30, 81])
ax1.set_ylim([0, None])

# Panel 2: Disease timeline (scatter plot)
ax2 = fig.add_subplot(gs[1], sharex=ax1)

if len(diagnosis_times) > 0:
    # Sort diagnoses by time
    diag_order = sorted([(d, times[0]) for d, times in diagnosis_times.items()], 
                       key=lambda x: x[1])
    disease_rows = {d: i for i, (d, _) in enumerate(diag_order)}
    
    for d, t_diag in diag_order:
        sig_for_disease = initial_clusters[d] if d < len(initial_clusters) else -1
        color = sig_colors[sig_for_disease] if sig_for_disease < K_total else 'gray'
        age_at_diag = 30 + t_diag
        y = disease_rows[d]
        disease_name = disease_names[d] if d < len(disease_names) else f'Disease {d}'
        ax2.scatter(age_at_diag, y, s=100, color=color, alpha=0.7, zorder=10, 
                   edgecolors='black', linewidths=1)
        ax2.text(age_at_diag + 1, y, f'{disease_name[:30]} (sig{sig_for_disease})', 
                fontsize=8, verticalalignment='center')
    
    ax2.set_yticks(range(len(diag_order)))
    # Label diseases by chronological order (1, 2, 3, ...)
    ax2.set_yticklabels([f'{i+1}' for i in range(len(diag_order))], fontsize=8)
else:
    ax2.text(0.5, 0.5, 'No diagnoses', transform=ax2.transAxes, 
            ha='center', va='center', fontsize=12)

ax2.set_ylabel('Disease', fontsize=11)
ax2.set_title('Disease timeline', fontsize=11)
ax2.grid(True, alpha=0.3, axis='x')
ax2.set_xlim([30, 81])

# Panel 3: Disease probabilities (π) - stop after diagnosis
ax3 = fig.add_subplot(gs[2], sharex=ax1)

# Plot diseases with events, colored by signature, stopping after diagnosis
for d in diseases_with_events:
    disease_name = disease_names[d] if d < len(disease_names) else f'Disease {d}'
    sig_for_disease = initial_clusters[d] if d < len(initial_clusters) else -1
    color = sig_colors[sig_for_disease] if sig_for_disease < K_total else 'gray'
    
    # Get diagnosis timepoint
    if d in diagnosis_times:
        first_diag_t = min(diagnosis_times[d])
        # Only plot up to and including diagnosis timepoint
        plot_ages = ages[:first_diag_t + 1]
        plot_pi = patient_pi[d, :first_diag_t + 1]
    else:
        plot_ages = ages
        plot_pi = patient_pi[d, :]
    
    # Plot probability curve (stops after diagnosis)
    ax3.plot(plot_ages, plot_pi, 
             label=f"{disease_name} (Sig {sig_for_disease})",
             color=color, linewidth=2, alpha=0.7)
    
    # Mark diagnosis timepoint
    if d in diagnosis_times:
        for t in diagnosis_times[d]:
            age_at_diag = 30 + t
            ax3.scatter(age_at_diag, patient_pi[d, t], 
                       color=color, s=100, zorder=10, marker='o', 
                       edgecolors='black', linewidths=1.5)

ax3.set_xlabel('Age (yr)', fontsize=12)
ax3.set_ylabel('Disease Probability (π)', fontsize=11)
ax3.set_title('Disease Probabilities Over Time (colored by primary signature)', fontsize=11)
ax3.legend(bbox_to_anchor=(1.02, 1), loc='upper left', fontsize=8)
ax3.grid(True, alpha=0.3)
ax3.set_xlim([30, 81])

plt.suptitle(f'Patient {patient_idx} Timeline', fontsize=16, fontweight='bold', y=0.995)
plt.tight_layout()
plt.savefig(f'/Users/sarahurbut/aladynoulli2/patient_{patient_idx}_timeline_panel_style.pdf', 
            dpi=300, bbox_inches='tight')
print(f"\n✓ Saved plot to: patient_{patient_idx}_timeline_panel_style.pdf")
plt.show()

# ============================================================================
# Plot Patient Timeline - Updated Layout
# ============================================================================

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from pathlib import Path

# Load data
initial_clusters = torch.load('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/initial_clusters_400k.pt')
if torch.is_tensor(initial_clusters):
    initial_clusters = initial_clusters.numpy()
K = int(initial_clusters.max() + 1)

# Load theta from the specified file
theta_path = '/Users/sarahurbut/aladynoulli2/pyScripts/new_thetas_with_pcs_retrospective_correctE.pt'
theta_full = torch.load(theta_path, map_location='cpu', weights_only=False)

# Check structure of theta file
if isinstance(theta_full, dict):
    # If it's a dict, try common keys
    if 'theta' in theta_full:
        theta = theta_full['theta']
    elif 'thetas' in theta_full:
        theta = theta_full['thetas']
    elif 'lambda_' in theta_full:
        theta = torch.softmax(theta_full['lambda_'], dim=1)
    else:
        # Try first tensor value
        theta = list(theta_full.values())[0]
        if torch.is_tensor(theta) and theta.dim() == 3:
            theta = torch.softmax(theta, dim=1)
else:
    theta = theta_full

# Convert to numpy if needed
if torch.is_tensor(theta):
    theta = theta.numpy()
elif isinstance(theta, list):
    theta = np.array(theta)

# Load Y and pi from checkpoint (for patient 148746)
# Load Y and pi from checkpoint (for patient 148746)
checkpoint_path = '/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/enrollment_model_W0.0001_batch_0_10000.pt'
pi_path = '/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/pi_fullmode_400k.pt'

ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
Y = torch.load('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/Y_tensor.pt')
if torch.is_tensor(Y):
    Y_np = Y.numpy()
else:
    Y_np = Y

pi_full = torch.load(pi_path, map_location='cpu', weights_only=False)
if torch.is_tensor(pi_full):
    pi_np = pi_full.numpy()
else:
    pi_np = pi_full

# Load disease names
disease_names = pd.read_csv('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/disease_names.csv')['x'].tolist()

# Set patient index to 148746
patient_idx = 148745

# Check dimensions
N_theta, K_total, T_theta = theta.shape
N_y, D, T_y = Y_np.shape
N_pi, D_pi, T_pi = pi_np.shape

print(f"Theta shape: ({N_theta}, {K_total}, {T_theta})")
print(f"Y shape: ({N_y}, {D}, {T_y})")
print(f"Pi shape: ({N_pi}, {D_pi}, {T_pi})")
print(f"\nRequested patient index: {patient_idx}")

# Verify patient index is valid
if patient_idx >= N_theta:
    print(f"WARNING: Patient {patient_idx} not in theta array (max index: {N_theta-1})")
    patient_idx = min(patient_idx, N_theta - 1)
    print(f"Using patient index: {patient_idx}")

if patient_idx >= N_y:
    print(f"WARNING: Patient {patient_idx} not in Y array (max index: {N_y-1})")
    # We'll handle this below

if patient_idx >= N_pi:
    print(f"WARNING: Patient {patient_idx} not in pi array (max index: {N_pi-1})")
    # We'll handle this below

# Get patient data
patient_theta = theta[patient_idx, :, :]  # (K, T)
patient_Y = Y_np[patient_idx, :, :] if patient_idx < N_y else np.zeros((D, T_y))  # (D, T)
patient_pi = pi_np[patient_idx, :, :] if patient_idx < N_pi else np.zeros((D_pi, T_pi))  # (D, T)

# Ensure T dimensions match (use minimum)
T = min(T_theta, T_y, T_pi)
if T_theta != T:
    patient_theta = patient_theta[:, :T]
if T_y != T:
    patient_Y = patient_Y[:, :T]
if T_pi != T:
    patient_pi = patient_pi[:, :T]

print(f"\nUsing T = {T} (aligned across all arrays)")

# Find diagnoses
diagnosis_times = {}
for d in range(patient_Y.shape[0]):
    event_times = np.where(patient_Y[d, :] == 1)[0]
    if len(event_times) > 0:
        diagnosis_times[d] = event_times.tolist()

# Print summary
n_diseases = len(diagnosis_times)
all_times = []
for times in diagnosis_times.values():
    all_times.extend(times)
time_range = (min(all_times), max(all_times)) if all_times else (0, 0)
print(f"\nPatient {patient_idx} Summary:")
print(f"  Number of diseases: {n_diseases}")
print(f"  Diagnosis timepoints: {sorted(set(all_times))}")
print(f"  Time range: {time_range[0]} to {time_range[1]} (ages {30+time_range[0]} to {30+time_range[1]})")

# Calculate average theta (K vector) - average across time
avg_theta = patient_theta.mean(axis=1)  # Shape: (K,)

# Find diseases with events
diseases_with_events = []
for d in range(patient_Y.shape[0]):
    event_times = np.where(patient_Y[d, :] == 1)[0]
    if len(event_times) > 0:
        diseases_with_events.append(d)
        if d not in diagnosis_times:
            diagnosis_times[d] = event_times.tolist()

# Convert timepoints to ages
ages = np.arange(30, 30 + T)

# ============================================================================
# Create Plot
# ============================================================================

fig = plt.figure(figsize=(14, 10))
gs = plt.GridSpec(3, 1, height_ratios=[2, 1, 1.2], hspace=0.4)

colors = sns.color_palette("tab20", K_total)
sig_colors = sns.color_palette("tab20", K_total)

# Panel 1: Signature loadings (θ) vs Age
ax1 = fig.add_subplot(gs[0])

for k in range(K_total):
    ax1.plot(ages, patient_theta[k, :], 
             label=f'Signature {k}', linewidth=2, color=colors[k], alpha=0.8)

# Add horizontal lines at diagnosis times
for d, times in diagnosis_times.items():
    for t in times:
        if t >= T:
            continue
        age_at_diag = 30 + t
        # Get the signature for this disease
        sig_for_disease = initial_clusters[d] if d < len(initial_clusters) else -1
        if sig_for_disease < K_total:
            # Draw horizontal line at the theta value for this signature at diagnosis time
            theta_at_diag = patient_theta[sig_for_disease, t]
            ax1.axhline(y=theta_at_diag, xmin=(age_at_diag - 30) / (81 - 30), 
                       xmax=(age_at_diag - 30 + 1) / (81 - 30),
                       color=colors[sig_for_disease], linestyle='--', 
                       alpha=0.6, linewidth=1.5)

# Add thin stacked bar showing average theta (single bar, stacked)
# Position it at the top right
ax1_bar = ax1.inset_axes([0.7, 0.7, 0.25, 0.15])  # [x, y, width, height] in axes coordinates

# Sort signatures by average theta (largest first) for better visualization
sorted_indices = np.argsort(avg_theta)[::-1]
sorted_avg_theta = avg_theta[sorted_indices]
sorted_colors = [colors[i] for i in sorted_indices]

# Create stacked bar (single bar)
bottom = 0
for i, (val, color) in enumerate(zip(sorted_avg_theta, sorted_colors)):
    if val > 0.01:  # Only show if > 1% to avoid clutter
        ax1_bar.barh(0, val, left=bottom, color=color, height=0.5, alpha=0.8)
        bottom += val

ax1_bar.set_xlim([0, 1])
ax1_bar.set_ylim([-0.5, 0.5])
ax1_bar.set_xticks([0, 0.5, 1.0])
ax1_bar.set_xticklabels(['0', '0.5', '1'], fontsize=7)
ax1_bar.set_yticks([])
ax1_bar.set_title('Avg θ (stacked)', fontsize=8)
ax1_bar.spines['top'].set_visible(False)
ax1_bar.spines['right'].set_visible(False)
ax1_bar.spines['left'].set_visible(False)

ax1.set_ylabel('Signature loadings (θ)', fontsize=12)
ax1.set_title(f'Patient {patient_idx}: Signature Trajectories', fontsize=14, fontweight='bold')
ax1.legend(bbox_to_anchor=(1.02, 1), loc='upper left', fontsize=9, ncol=2)
ax1.grid(True, alpha=0.3)
ax1.set_xlim([30, 81])
ax1.set_ylim([0, None])

# Panel 2: Disease timeline (scatter plot)
ax2 = fig.add_subplot(gs[1], sharex=ax1)

if len(diagnosis_times) > 0:
    # Sort diagnoses by time
    diag_order = sorted([(d, times[0]) for d, times in diagnosis_times.items()], 
                       key=lambda x: x[1])
    disease_rows = {d: i for i, (d, _) in enumerate(diag_order)}
    
    for d, t_diag in diag_order:
        if t_diag >= T:
            continue
        sig_for_disease = initial_clusters[d] if d < len(initial_clusters) else -1
        color = sig_colors[sig_for_disease] if sig_for_disease < K_total else 'gray'
        age_at_diag = 30 + t_diag
        y = disease_rows[d]
        disease_name = disease_names[d] if d < len(disease_names) else f'Disease {d}'
        ax2.scatter(age_at_diag, y, s=100, color=color, alpha=0.7, zorder=10, 
                   edgecolors='black', linewidths=1)
        ax2.text(age_at_diag + 1, y, f'{disease_name[:30]} (sig{sig_for_disease})', 
                fontsize=8, verticalalignment='center')
    
    ax2.set_yticks(range(len(diag_order)))
    # Label diseases by chronological order (1, 2, 3, ...)
    ax2.set_yticklabels([f'{i+1}' for i in range(len(diag_order))], fontsize=8)
else:
    ax2.text(0.5, 0.5, 'No diagnoses', transform=ax2.transAxes, 
            ha='center', va='center', fontsize=12)

ax2.set_ylabel('Disease', fontsize=11)
ax2.set_title('Disease timeline', fontsize=11)
ax2.grid(True, alpha=0.3, axis='x')
ax2.set_xlim([30, 81])

# Panel 3: Disease probabilities (π) - stop after diagnosis
ax3 = fig.add_subplot(gs[2], sharex=ax1)

# Plot diseases with events, colored by signature, stopping after diagnosis
for d in diseases_with_events:
    disease_name = disease_names[d] if d < len(disease_names) else f'Disease {d}'
    sig_for_disease = initial_clusters[d] if d < len(initial_clusters) else -1
    color = sig_colors[sig_for_disease] if sig_for_disease < K_total else 'gray'
    
    # Get diagnosis timepoint
    if d in diagnosis_times:
        first_diag_t = min(diagnosis_times[d])
        if first_diag_t >= T:
            first_diag_t = T - 1
        # Only plot up to and including diagnosis timepoint
        plot_ages = ages[:first_diag_t + 1]
        plot_pi = patient_pi[d, :first_diag_t + 1]
    else:
        plot_ages = ages
        plot_pi = patient_pi[d, :]
    
    # Plot probability curve (stops after diagnosis)
    ax3.plot(plot_ages, plot_pi, 
             label=f"{disease_name} (Sig {sig_for_disease})",
             color=color, linewidth=2, alpha=0.7)
    
    # Mark diagnosis timepoint
    if d in diagnosis_times:
        for t in diagnosis_times[d]:
            if t >= T:
                continue
            age_at_diag = 30 + t
            ax3.scatter(age_at_diag, patient_pi[d, t], 
                       color=color, s=100, zorder=10, marker='o', 
                       edgecolors='black', linewidths=1.5)

ax3.set_xlabel('Age (yr)', fontsize=12)
ax3.set_ylabel('Disease Probability (π)', fontsize=11)
ax3.set_title('Disease Probabilities Over Time (colored by primary signature)', fontsize=11)
ax3.legend(bbox_to_anchor=(1.02, 1), loc='upper left', fontsize=8)
ax3.grid(True, alpha=0.3)
ax3.set_xlim([30, 81])

plt.suptitle(f'Patient {patient_idx} Timeline', fontsize=16, fontweight='bold', y=0.995)
plt.tight_layout()
plt.savefig(f'/Users/sarahurbut/aladynoulli2/patient_{patient_idx}_timeline_panel_style.pdf', 
            dpi=300, bbox_inches='tight')
print(f"\n✓ Saved plot to: patient_{patient_idx}_timeline_panel_style.pdf")
plt.show()

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures/plot_patient_timeline.py

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures/generate_prs_signature_plots.py --batch_dir "/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized_noPCS/" --pattern "enrollment_model_VECTORIZED_W0.0001_batch_*_*.pt" --output_dir "/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/paper_figs/prs_signatures_nopPCS_E" --n_top 30

from analyze_age_onset import analyze_age_onset_patterns

early_indices, late_indices, stats = analyze_age_onset_patterns(
    results_base_dir='/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized',
    disease_index=113,  # MI
    early_threshold=55,
    late_threshold=70,
    output_path='/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/paper_figs/fig3/mi_onset_patterns_all_batches.pdf',  # optional
    return_stats=True
)