notebook-a4b6e2df-bf06-4e9d-94f4-4d807332599f-paper-20260319-215149.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SubCell: Proteome-aware Vision Foundation Models for Microscopy\n",
    "\n",
    "**Paper:** SubCell: Proteome-aware vision foundation models for microscopy capture single-cell biology  \n",
    "**Authors:** Ankit Gupta, Zoe Wefers, et al.  \n",
    "\n",
    "## Overview\n",
    "\n",
    "This notebook provides an educational walkthrough of the computational methods in the SubCell paper. SubCell is a vision transformer model that learns representations of cellular morphology and protein localization through multi-task self-supervised learning.\n",
    "\n",
    "### Key Contributions:\n",
    "1. **Multi-task learning framework** combining:\n",
    "   - Masked autoencoder (MAE) reconstruction\n",
    "   - Cell-specific contrastive learning\n",
    "   - Protein-specific contrastive learning\n",
    "2. **Vision-based multiscale cell map** constructed from hierarchical clustering\n",
    "3. **Multimodal integration** with protein sequence embeddings (ESM2)\n",
    "4. **Strong performance** on protein localization, cell cycle prediction, and drug perturbation tasks\n",
    "\n",
    "### Resource Constraints & Educational Focus\n",
    "\n",
    "**IMPORTANT:** This notebook uses synthetic toy data and simplified examples due to computational constraints:\n",
    "- Memory: 4GB RAM limit\n",
    "- Time: Notebook runs in ~5-10 minutes\n",
    "- Purpose: Educational demonstration, not full reproduction\n",
    "\n",
    "The code demonstrates the **methodology** and **workflow structure** to help researchers understand and adapt these methods for their own full-scale experiments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Dependencies\n",
    "\n",
    "Installing all required packages with a single command to ensure compatibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install all dependencies in a single command\n",
    "!uv pip install torch torchvision numpy scipy scikit-learn matplotlib seaborn pandas umap-learn networkx python-louvain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import accuracy_score, f1_score, confusion_matrix\n",
    "from sklearn.cluster import AgglomerativeClustering\n",
    "from scipy.spatial.distance import cosine\n",
    "from scipy.cluster.hierarchy import dendrogram, linkage\n",
    "import pandas as pd\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "\n",
    "print(\"Libraries imported successfully!\")\n",
    "print(f\"PyTorch version: {torch.__version__}\")\n",
    "print(f\"NumPy version: {np.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 1: SubCell Model Architecture\n",
    "\n",
    "## Vision Transformer (ViT) Encoder\n",
    "\n",
    "SubCell uses a Vision Transformer (ViT) with **86.4 million parameters**. The architecture:\n",
    "- Converts images into patches (16×16 pixels)\n",
    "- Uses transformer self-attention to learn spatial relationships\n",
    "- Produces embeddings that capture cellular morphology and protein patterns\n",
    "\n",
    "**For this demo:** We'll implement a simplified ViT architecture to illustrate the concept."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimplifiedViT(nn.Module):\n",
    "    \"\"\"Simplified Vision Transformer for educational demonstration.\n",
    "    \n",
    "    This is a minimal implementation to illustrate the ViT concept.\n",
    "    The actual SubCell model has 86.4M parameters and more sophisticated architecture.\n",
    "    \"\"\"\n",
    "    def __init__(self, image_size=128, patch_size=16, in_channels=4, \n",
    "                 embed_dim=256, num_heads=8, num_layers=4, num_classes=19):\n",
    "        super().__init__()\n",
    "        self.patch_size = patch_size\n",
    "        self.num_patches = (image_size // patch_size) ** 2\n",
    "        \n",
    "        # Patch embedding: convert image patches to vectors\n",
    "        self.patch_embed = nn.Conv2d(in_channels, embed_dim, \n",
    "                                     kernel_size=patch_size, stride=patch_size)\n",
    "        \n",
    "        # Positional embeddings\n",
    "        self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches + 1, embed_dim))\n",
    "        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))\n",
    "        \n",
    "        # Transformer encoder\n",
    "        encoder_layer = nn.TransformerEncoderLayer(\n",
    "            d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim*4, \n",
    "            dropout=0.1, activation='gelu', batch_first=True\n",
    "        )\n",
    "        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n",
    "        \n",
    "        # Layer norm\n",
    "        self.norm = nn.LayerNorm(embed_dim)\n",
    "        \n",
    "        # Classification head (for protein localization)\n",
    "        self.fc = nn.Linear(embed_dim, num_classes)\n",
    "        \n",
    "    def forward(self, x, return_embeddings=False):\n",
    "        B = x.shape[0]\n",
    "        \n",
    "        # Patch embedding\n",
    "        x = self.patch_embed(x)  # (B, embed_dim, H/P, W/P)\n",
    "        x = x.flatten(2).transpose(1, 2)  # (B, num_patches, embed_dim)\n",
    "        \n",
    "        # Add CLS token\n",
    "        cls_tokens = self.cls_token.expand(B, -1, -1)\n",
    "        x = torch.cat([cls_tokens, x], dim=1)\n",
    "        \n",
    "        # Add positional embedding\n",
    "        x = x + self.pos_embed\n",
    "        \n",
    "        # Transformer encoding\n",
    "        x = self.transformer(x)\n",
    "        x = self.norm(x)\n",
    "        \n",
    "        # Extract CLS token representation\n",
    "        cls_output = x[:, 0]\n",
    "        \n",
    "        if return_embeddings:\n",
    "            return cls_output\n",
    "        \n",
    "        # Classification\n",
    "        return self.fc(cls_output)\n",
    "\n",
    "# Create model\n",
    "model = SimplifiedViT()\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Simplified ViT created with {total_params:,} parameters\")\n",
    "print(f\"(Actual SubCell has 86,400,000 parameters)\")\n",
    "print(\"\\nModel architecture:\")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention Pooling Module\n",
    "\n",
    "SubCell uses an attention pooling module with:\n",
    "- **12 attention heads** for general feature aggregation\n",
    "- **2 pooled attention heads** for focused representation\n",
    "\n",
    "This allows the model to focus on relevant cellular regions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionPooling(nn.Module):\n",
    "    \"\"\"Attention pooling module for SubCell.\n",
    "    \n",
    "    Learns to focus on informative image regions using attention mechanism.\n",
    "    \"\"\"\n",
    "    def __init__(self, embed_dim=256, num_heads=12, num_pool_heads=2):\n",
    "        super().__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.num_pool_heads = num_pool_heads\n",
    "        \n",
    "        # Multi-head attention\n",
    "        self.attention = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)\n",
    "        \n",
    "        # Pooling queries\n",
    "        self.pool_queries = nn.Parameter(torch.randn(1, num_pool_heads, embed_dim))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"x: (batch, num_patches, embed_dim)\"\"\"\n",
    "        B = x.shape[0]\n",
    "        \n",
    "        # Expand pool queries for batch\n",
    "        queries = self.pool_queries.expand(B, -1, -1)\n",
    "        \n",
    "        # Attention pooling\n",
    "        pooled, attn_weights = self.attention(queries, x, x)\n",
    "        \n",
    "        # Average across pool heads\n",
    "        pooled = pooled.mean(dim=1)  # (B, embed_dim)\n",
    "        \n",
    "        return pooled, attn_weights\n",
    "\n",
    "# Demo attention pooling\n",
    "attn_pool = AttentionPooling()\n",
    "dummy_patches = torch.randn(4, 64, 256)  # 4 images, 64 patches, 256 dims\n",
    "pooled, weights = attn_pool(dummy_patches)\n",
    "print(f\"Input patches shape: {dummy_patches.shape}\")\n",
    "print(f\"Pooled output shape: {pooled.shape}\")\n",
    "print(f\"Attention weights shape: {weights.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 2: Multi-Task Learning Framework\n",
    "\n",
    "SubCell combines **three learning objectives**:\n",
    "\n",
    "## 2.1 Masked Autoencoder (MAE) Reconstruction\n",
    "\n",
    "**Goal:** Learn to reconstruct masked portions of images  \n",
    "**Loss:** Mean squared error between reconstructed and original patches\n",
    "\n",
    "$$\\mathcal{L}_{MAE} = \\frac{1}{|M|} \\sum_{i \\in M} \\|x_i - \\hat{x}_i\\|^2$$\n",
    "\n",
    "where $M$ is the set of masked patches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MAEReconstruction(nn.Module):\n",
    "    \"\"\"Masked Autoencoder reconstruction objective.\n",
    "    \n",
    "    Randomly masks image patches and learns to reconstruct them.\n",
    "    \"\"\"\n",
    "    def __init__(self, embed_dim=256, patch_size=16, in_channels=4, mask_ratio=0.75):\n",
    "        super().__init__()\n",
    "        self.mask_ratio = mask_ratio\n",
    "        self.patch_size = patch_size\n",
    "        \n",
    "        # Decoder to reconstruct patches\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(embed_dim, embed_dim * 2),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(embed_dim * 2, patch_size * patch_size * in_channels)\n",
    "        )\n",
    "        \n",
    "    def random_masking(self, x, mask_ratio):\n",
    "        \"\"\"Random masking of patches.\"\"\"\n",
    "        N, L, D = x.shape  # batch, length, dim\n",
    "        len_keep = int(L * (1 - mask_ratio))\n",
    "        \n",
    "        # Random permutation\n",
    "        noise = torch.rand(N, L, device=x.device)\n",
    "        ids_shuffle = torch.argsort(noise, dim=1)\n",
    "        ids_restore = torch.argsort(ids_shuffle, dim=1)\n",
    "        \n",
    "        # Keep subset\n",
    "        ids_keep = ids_shuffle[:, :len_keep]\n",
    "        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))\n",
    "        \n",
    "        # Binary mask: 0 is keep, 1 is remove\n",
    "        mask = torch.ones([N, L], device=x.device)\n",
    "        mask[:, :len_keep] = 0\n",
    "        mask = torch.gather(mask, dim=1, index=ids_restore)\n",
    "        \n",
    "        return x_masked, mask, ids_restore\n",
    "    \n",
    "    def forward(self, patch_embeddings, original_patches):\n",
    "        \"\"\"Compute MAE reconstruction loss.\"\"\"\n",
    "        # Mask patches\n",
    "        x_masked, mask, ids_restore = self.random_masking(patch_embeddings, self.mask_ratio)\n",
    "        \n",
    "        # Reconstruct\n",
    "        reconstructed = self.decoder(patch_embeddings)\n",
    "        \n",
    "        # Compute loss only on masked patches\n",
    "        loss = F.mse_loss(reconstructed, original_patches, reduction='none')\n",
    "        loss = (loss * mask.unsqueeze(-1)).sum() / mask.sum()\n",
    "        \n",
    "        return loss\n",
    "\n",
    "# Demo MAE\n",
    "mae = MAEReconstruction()\n",
    "patch_emb = torch.randn(4, 64, 256)  # embeddings\n",
    "original = torch.randn(4, 64, 16*16*4)  # original patches flattened\n",
    "mae_loss = mae(patch_emb, original)\n",
    "print(f\"MAE reconstruction loss: {mae_loss.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Cell-Specific Contrastive Learning\n",
    "\n",
    "**Goal:** Learn to identify which cell line an image comes from  \n",
    "**Method:** Contrastive learning - pull together embeddings from the same cell line, push apart different cell lines\n",
    "\n",
    "$$\\mathcal{L}_{cell} = -\\log \\frac{\\exp(\\text{sim}(z_i, z_j^+) / \\tau)}{\\sum_{k} \\exp(\\text{sim}(z_i, z_k) / \\tau)}$$\n",
    "\n",
    "where $z_j^+$ is a positive pair (same cell line), $\\tau$ is temperature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CellSpecificContrastive(nn.Module):\n",
    "    \"\"\"Cell-specific contrastive learning objective.\n",
    "    \n",
    "    Learns to distinguish between different cell lines.\n",
    "    \"\"\"\n",
    "    def __init__(self, temperature=0.07):\n",
    "        super().__init__()\n",
    "        self.temperature = temperature\n",
    "        \n",
    "    def forward(self, embeddings, cell_labels):\n",
    "        \"\"\"\n",
    "        embeddings: (batch_size, embed_dim)\n",
    "        cell_labels: (batch_size,) - cell line IDs\n",
    "        \"\"\"\n",
    "        # Normalize embeddings\n",
    "        embeddings = F.normalize(embeddings, dim=1)\n",
    "        \n",
    "        # Compute similarity matrix\n",
    "        sim_matrix = torch.matmul(embeddings, embeddings.T) / self.temperature\n",
    "        \n",
    "        # Create mask for positive pairs (same cell line)\n",
    "        labels = cell_labels.unsqueeze(0)\n",
    "        mask = (labels == labels.T).float()\n",
    "        mask.fill_diagonal_(0)  # Exclude self-similarity\n",
    "        \n",
    "        # Contrastive loss\n",
    "        exp_sim = torch.exp(sim_matrix)\n",
    "        \n",
    "        # For each sample, sum similarity to all samples\n",
    "        denominator = exp_sim.sum(dim=1, keepdim=True)\n",
    "        \n",
    "        # Positive pair similarity\n",
    "        pos_sim = (exp_sim * mask).sum(dim=1)\n",
    "        \n",
    "        # InfoNCE loss\n",
    "        loss = -torch.log(pos_sim / (denominator.squeeze() + 1e-8))\n",
    "        loss = loss[mask.sum(dim=1) > 0].mean()  # Only compute for samples with positives\n",
    "        \n",
    "        return loss\n",
    "\n",
    "# Demo cell-specific contrastive\n",
    "cell_contrast = CellSpecificContrastive()\n",
    "embeddings = torch.randn(8, 256)  # 8 samples\n",
    "cell_labels = torch.tensor([0, 0, 1, 1, 2, 2, 0, 1])  # Cell line IDs\n",
    "cell_loss = cell_contrast(embeddings, cell_labels)\n",
    "print(f\"Cell-specific contrastive loss: {cell_loss.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 Protein-Specific Contrastive Learning\n",
    "\n",
    "**Goal:** Learn to identify protein localization patterns  \n",
    "**Method:** Pull together cells expressing the same protein, push apart different proteins\n",
    "\n",
    "$$\\mathcal{L}_{protein} = -\\log \\frac{\\exp(\\text{sim}(z_i, z_j^+) / \\tau)}{\\sum_{k} \\exp(\\text{sim}(z_i, z_k) / \\tau)}$$\n",
    "\n",
    "where $z_j^+$ is a positive pair (same protein)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProteinSpecificContrastive(nn.Module):\n",
    "    \"\"\"Protein-specific contrastive learning objective.\n",
    "    \n",
    "    Learns to distinguish between different protein localization patterns.\n",
    "    \"\"\"\n",
    "    def __init__(self, temperature=0.07):\n",
    "        super().__init__()\n",
    "        self.temperature = temperature\n",
    "        \n",
    "    def forward(self, embeddings, protein_labels):\n",
    "        \"\"\"\n",
    "        embeddings: (batch_size, embed_dim)\n",
    "        protein_labels: (batch_size,) - protein IDs\n",
    "        \"\"\"\n",
    "        # Normalize embeddings\n",
    "        embeddings = F.normalize(embeddings, dim=1)\n",
    "        \n",
    "        # Compute similarity matrix\n",
    "        sim_matrix = torch.matmul(embeddings, embeddings.T) / self.temperature\n",
    "        \n",
    "        # Create mask for positive pairs (same protein)\n",
    "        labels = protein_labels.unsqueeze(0)\n",
    "        mask = (labels == labels.T).float()\n",
    "        mask.fill_diagonal_(0)  # Exclude self-similarity\n",
    "        \n",
    "        # Contrastive loss\n",
    "        exp_sim = torch.exp(sim_matrix)\n",
    "        denominator = exp_sim.sum(dim=1, keepdim=True)\n",
    "        pos_sim = (exp_sim * mask).sum(dim=1)\n",
    "        \n",
    "        # InfoNCE loss\n",
    "        loss = -torch.log(pos_sim / (denominator.squeeze() + 1e-8))\n",
    "        loss = loss[mask.sum(dim=1) > 0].mean()\n",
    "        \n",
    "        return loss\n",
    "\n",
    "# Demo protein-specific contrastive\n",
    "protein_contrast = ProteinSpecificContrastive()\n",
    "embeddings = torch.randn(8, 256)\n",
    "protein_labels = torch.tensor([0, 1, 0, 2, 1, 2, 0, 1])  # Protein IDs\n",
    "protein_loss = protein_contrast(embeddings, protein_labels)\n",
    "print(f\"Protein-specific contrastive loss: {protein_loss.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 Combined Multi-Task Loss\n",
    "\n",
    "SubCell combines all three objectives:\n",
    "\n",
    "$$\\mathcal{L}_{total} = \\lambda_{MAE} \\mathcal{L}_{MAE} + \\lambda_{cell} \\mathcal{L}_{cell} + \\lambda_{protein} \\mathcal{L}_{protein}$$\n",
    "\n",
    "The paper optimizes these jointly during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SubCellMultiTaskLoss(nn.Module):\n",
    "    \"\"\"Combined multi-task loss for SubCell training.\"\"\"\n",
    "    def __init__(self, lambda_mae=1.0, lambda_cell=1.0, lambda_protein=1.0):\n",
    "        super().__init__()\n",
    "        self.lambda_mae = lambda_mae\n",
    "        self.lambda_cell = lambda_cell\n",
    "        self.lambda_protein = lambda_protein\n",
    "        \n",
    "        self.mae_loss = MAEReconstruction()\n",
    "        self.cell_loss = CellSpecificContrastive()\n",
    "        self.protein_loss = ProteinSpecificContrastive()\n",
    "        \n",
    "    def forward(self, patch_embeddings, original_patches, \n",
    "                cell_embeddings, cell_labels, protein_labels):\n",
    "        \"\"\"Compute combined multi-task loss.\"\"\"\n",
    "        # MAE reconstruction\n",
    "        loss_mae = self.mae_loss(patch_embeddings, original_patches)\n",
    "        \n",
    "        # Cell-specific contrastive\n",
    "        loss_cell = self.cell_loss(cell_embeddings, cell_labels)\n",
    "        \n",
    "        # Protein-specific contrastive\n",
    "        loss_protein = self.protein_loss(cell_embeddings, protein_labels)\n",
    "        \n",
    "        # Combined loss\n",
    "        total_loss = (self.lambda_mae * loss_mae + \n",
    "                     self.lambda_cell * loss_cell + \n",
    "                     self.lambda_protein * loss_protein)\n",
    "        \n",
    "        return total_loss, {\n",
    "            'mae': loss_mae.item(),\n",
    "            'cell': loss_cell.item(),\n",
    "            'protein': loss_protein.item(),\n",
    "            'total': total_loss.item()\n",
    "        }\n",
    "\n",
    "# Demo combined loss\n",
    "multi_task_loss = SubCellMultiTaskLoss()\n",
    "patch_emb = torch.randn(8, 64, 256)\n",
    "original = torch.randn(8, 64, 16*16*4)\n",
    "cell_emb = torch.randn(8, 256)\n",
    "cell_labels = torch.tensor([0, 0, 1, 1, 2, 2, 0, 1])\n",
    "protein_labels = torch.tensor([0, 1, 0, 2, 1, 2, 0, 1])\n",
    "\n",
    "total_loss, loss_dict = multi_task_loss(patch_emb, original, cell_emb, \n",
    "                                        cell_labels, protein_labels)\n",
    "\n",
    "print(\"Multi-task loss breakdown:\")\n",
    "for key, val in loss_dict.items():\n",
    "    print(f\"  {key}: {val:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 3: Data Preparation with Synthetic Examples\n",
    "\n",
    "## Generate Synthetic Microscopy Images\n",
    "\n",
    "**Real data:** Human Protein Atlas (HPA) images with 4 channels:\n",
    "- **Channel 0:** Nuclei (DNA)\n",
    "- **Channel 1:** Microtubules (MT)\n",
    "- **Channel 2:** Endoplasmic Reticulum (ER)\n",
    "- **Channel 3:** Protein of interest\n",
    "\n",
    "**For this demo:** We generate synthetic images that mimic the structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_synthetic_cell_images(num_samples=100, image_size=128, num_channels=4):\n",
    "    \"\"\"\n",
    "    Generate synthetic microscopy images for demonstration.\n",
    "    \n",
    "    Parameters:\n",
    "    - num_samples: Number of images to generate\n",
    "    - image_size: Size of square images\n",
    "    - num_channels: Number of fluorescence channels (4 for HPA)\n",
    "    \n",
    "    Returns:\n",
    "    - images: (num_samples, num_channels, image_size, image_size)\n",
    "    - cell_labels: Cell line IDs\n",
    "    - protein_labels: Protein IDs\n",
    "    - localization_labels: Subcellular localization categories\n",
    "    \"\"\"\n",
    "    images = []\n",
    "    \n",
    "    for i in range(num_samples):\n",
    "        # Create multi-channel image\n",
    "        img = np.zeros((num_channels, image_size, image_size))\n",
    "        \n",
    "        # Channel 0: Nuclei (circular pattern in center)\n",
    "        y, x = np.ogrid[:image_size, :image_size]\n",
    "        center_y, center_x = image_size // 2, image_size // 2\n",
    "        radius = image_size // 4 + np.random.randint(-10, 10)\n",
    "        mask = (x - center_x)**2 + (y - center_y)**2 <= radius**2\n",
    "        img[0][mask] = np.random.uniform(0.5, 1.0)\n",
    "        \n",
    "        # Channel 1: Microtubules (radiating pattern)\n",
    "        for angle in np.linspace(0, 2*np.pi, 12):\n",
    "            x_line = int(center_x + radius * 1.5 * np.cos(angle))\n",
    "            y_line = int(center_y + radius * 1.5 * np.sin(angle))\n",
    "            cv2_available = False  # We'll use a simple line\n",
    "            # Simple line approximation\n",
    "            steps = 50\n",
    "            for t in np.linspace(0, 1, steps):\n",
    "                px = int(center_x + t * (x_line - center_x))\n",
    "                py = int(center_y + t * (y_line - center_y))\n",
    "                if 0 <= px < image_size and 0 <= py < image_size:\n",
    "                    img[1, py, px] = np.random.uniform(0.3, 0.8)\n",
    "        \n",
    "        # Channel 2: ER (diffuse cytoplasmic)\n",
    "        img[2] = np.random.uniform(0, 0.4, (image_size, image_size))\n",
    "        img[2] = gaussian_filter_simple(img[2], sigma=5)\n",
    "        \n",
    "        # Channel 3: Protein (varies by localization)\n",
    "        # Randomly choose a localization pattern\n",
    "        pattern_type = np.random.randint(0, 4)\n",
    "        if pattern_type == 0:  # Nuclear\n",
    "            img[3][mask] = np.random.uniform(0.6, 1.0)\n",
    "        elif pattern_type == 1:  # Cytoplasmic\n",
    "            img[3][~mask] = np.random.uniform(0.3, 0.7)\n",
    "        elif pattern_type == 2:  # Membrane\n",
    "            membrane_mask = ((x - center_x)**2 + (y - center_y)**2 <= (radius * 1.5)**2) & \\\n",
    "                           ((x - center_x)**2 + (y - center_y)**2 >= (radius * 1.3)**2)\n",
    "            img[3][membrane_mask] = np.random.uniform(0.5, 0.9)\n",
    "        else:  # Vesicles (random spots)\n",
    "            for _ in range(10):\n",
    "                spot_x = np.random.randint(radius, image_size - radius)\n",
    "                spot_y = np.random.randint(radius, image_size - radius)\n",
    "                spot_mask = (x - spot_x)**2 + (y - spot_y)**2 <= (radius // 4)**2\n",
    "                img[3][spot_mask] = np.random.uniform(0.6, 1.0)\n",
    "        \n",
    "        images.append(img)\n",
    "    \n",
    "    images = np.array(images, dtype=np.float32)\n",
    "    \n",
    "    # Generate labels\n",
    "    num_cell_lines = 5  # e.g., U-2 OS, A-431, etc.\n",
    "    num_proteins = 20\n",
    "    num_localizations = 19  # HPA has 19 localization categories\n",
    "    \n",
    "    cell_labels = np.random.randint(0, num_cell_lines, num_samples)\n",
    "    protein_labels = np.random.randint(0, num_proteins, num_samples)\n",
    "    localization_labels = np.random.randint(0, num_localizations, num_samples)\n",
    "    \n",
    "    return images, cell_labels, protein_labels, localization_labels\n",
    "\n",
    "def gaussian_filter_simple(img, sigma=1):\n",
    "    \"\"\"Simple Gaussian blur approximation.\"\"\"\n",
    "    from scipy.ndimage import gaussian_filter\n",
    "    return gaussian_filter(img, sigma=sigma)\n",
    "\n",
    "# Generate synthetic dataset\n",
    "print(\"Generating synthetic microscopy images...\")\n",
    "images, cell_labels, protein_labels, loc_labels = generate_synthetic_cell_images(\n",
    "    num_samples=200, image_size=128, num_channels=4\n",
    ")\n",
    "\n",
    "print(f\"Dataset generated:\")\n",
    "print(f\"  Images shape: {images.shape}\")\n",
    "print(f\"  Cell labels: {len(np.unique(cell_labels))} unique cell lines\")\n",
    "print(f\"  Protein labels: {len(np.unique(protein_labels))} unique proteins\")\n",
    "print(f\"  Localization labels: {len(np.unique(loc_labels))} categories\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Synthetic Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize a few examples\n",
    "fig, axes = plt.subplots(3, 4, figsize=(12, 9))\n",
    "channel_names = ['Nuclei', 'Microtubules', 'ER', 'Protein']\n",
    "\n",
    "for i in range(3):\n",
    "    for j in range(4):\n",
    "        axes[i, j].imshow(images[i, j], cmap='gray')\n",
    "        if i == 0:\n",
    "            axes[i, j].set_title(channel_names[j])\n",
    "        axes[i, j].axis('off')\n",
    "    axes[i, 0].set_ylabel(f'Sample {i+1}', rotation=0, labelpad=40)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.suptitle('Synthetic Microscopy Images (4 Channels)', y=1.02, fontsize=14)\n",
    "plt.show()\n",
    "\n",
    "print(\"\\nNote: These are synthetic images for demonstration.\")\n",
    "print(\"Real HPA images have much more complex and varied patterns.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train/Test Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "# Split data\n",
    "indices = np.arange(len(images))\n",
    "train_idx, test_idx = train_test_split(indices, test_size=0.3, random_state=42)\n",
    "\n",
    "X_train = images[train_idx]\n",
    "X_test = images[test_idx]\n",
    "y_train = loc_labels[train_idx]\n",
    "y_test = loc_labels[test_idx]\n",
    "cell_train = cell_labels[train_idx]\n",
    "cell_test = cell_labels[test_idx]\n",
    "protein_train = protein_labels[train_idx]\n",
    "protein_test = protein_labels[test_idx]\n",
    "\n",
    "print(f\"Training set: {len(X_train)} samples\")\n",
    "print(f\"Test set: {len(X_test)} samples\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 4: Training Simulation (Simplified)\n",
    "\n",
    "**IMPORTANT:** Full training of SubCell takes days on GPUs with the complete HPA dataset (millions of cells).  \n",
    "This section demonstrates the training loop structure with a minimal number of iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate_training(model, multi_task_loss, X_train, cell_train, protein_train, \n",
    "                     num_iterations=10, batch_size=8):\n",
    "    \"\"\"\n",
    "    Simulate SubCell training with a few iterations.\n",
    "    \n",
    "    Real training: ~100 epochs on full HPA dataset with GPUs\n",
    "    This demo: 10 iterations on synthetic data for illustration\n",
    "    \"\"\"\n",
    "    optimizer = torch.optim.Adam(list(model.parameters()) + \n",
    "                                list(multi_task_loss.parameters()), \n",
    "                                lr=1e-4)\n",
    "    \n",
    "    model.train()\n",
    "    loss_history = []\n",
    "    \n",
    "    print(\"Starting training simulation...\")\n",
    "    print(\"(In practice, this would run for many epochs on GPUs)\\n\")\n",
    "    \n",
    "    for iteration in range(num_iterations):\n",
    "        # Sample batch\n",
    "        batch_idx = np.random.choice(len(X_train), batch_size, replace=False)\n",
    "        batch_images = torch.FloatTensor(X_train[batch_idx])\n",
    "        batch_cells = torch.LongTensor(cell_train[batch_idx])\n",
    "        batch_proteins = torch.LongTensor(protein_train[batch_idx])\n",
    "        \n",
    "        # Forward pass\n",
    "        embeddings = model(batch_images, return_embeddings=True)\n",
    "        \n",
    "        # Create dummy patch embeddings for MAE\n",
    "        patch_emb = embeddings.unsqueeze(1).repeat(1, 64, 1)  # Fake patches\n",
    "        original_patches = torch.randn(batch_size, 64, 16*16*4)\n",
    "        \n",
    "        # Compute multi-task loss\n",
    "        total_loss, loss_dict = multi_task_loss(\n",
    "            patch_emb, original_patches, embeddings, \n",
    "            batch_cells, batch_proteins\n",
    "        )\n",
    "        \n",
    "        # Backward pass\n",
    "        optimizer.zero_grad()\n",
    "        total_loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        loss_history.append(loss_dict)\n",
    "        \n",
    "        if (iteration + 1) % 5 == 0:\n",
    "            print(f\"Iteration {iteration+1}/{num_iterations}: \"\n",
    "                  f\"Total Loss = {loss_dict['total']:.4f}, \"\n",
    "                  f\"MAE = {loss_dict['mae']:.4f}, \"\n",
    "                  f\"Cell = {loss_dict['cell']:.4f}, \"\n",
    "                  f\"Protein = {loss_dict['protein']:.4f}\")\n",
    "    \n",
    "    print(\"\\nTraining simulation complete!\")\n",
    "    return loss_history\n",
    "\n",
    "# Run training simulation\n",
    "loss_history = simulate_training(\n",
    "    model, multi_task_loss, X_train, cell_train, protein_train, \n",
    "    num_iterations=10, batch_size=8\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Training Progress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot loss curves\n",
    "iterations = np.arange(len(loss_history))\n",
    "mae_losses = [d['mae'] for d in loss_history]\n",
    "cell_losses = [d['cell'] for d in loss_history]\n",
    "protein_losses = [d['protein'] for d in loss_history]\n",
    "total_losses = [d['total'] for d in loss_history]\n",
    "\n",
    "plt.figure(figsize=(12, 4))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(iterations, mae_losses, label='MAE', marker='o')\n",
    "plt.plot(iterations, cell_losses, label='Cell-specific', marker='s')\n",
    "plt.plot(iterations, protein_losses, label='Protein-specific', marker='^')\n",
    "plt.xlabel('Iteration')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Individual Loss Components')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.3)\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(iterations, total_losses, marker='o', color='purple', linewidth=2)\n",
    "plt.xlabel('Iteration')\n",
    "plt.ylabel('Total Loss')\n",
    "plt.title('Combined Multi-Task Loss')\n",
    "plt.grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"\\n** Scaling to Full Experiments **\")\n",
    "print(\"For actual SubCell training:\")\n",
    "print(\"  - Dataset: 2.8M single cells from HPA\")\n",
    "print(\"  - Training time: Several days on multiple GPUs\")\n",
    "print(\"  - Batch size: 256-512\")\n",
    "print(\"  - Epochs: ~100\")\n",
    "print(\"  - Infrastructure: High-memory GPUs (A100 or similar)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 5: Downstream Task - Protein Localization Classification\n",
    "\n",
    "Once SubCell is trained, we can extract embeddings and use them for downstream tasks.\n",
    "\n",
    "## Workflow 4: Protein Localization Classification\n",
    "\n",
    "**Task:** Predict which subcellular compartment(s) a protein localizes to  \n",
    "**Categories:** 19 organelles/structures (nucleus, cytosol, mitochondria, etc.)  \n",
    "**Method:** Train MLP classifier on SubCell embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract embeddings from trained model\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    train_embeddings = model(torch.FloatTensor(X_train), return_embeddings=True).numpy()\n",
    "    test_embeddings = model(torch.FloatTensor(X_test), return_embeddings=True).numpy()\n",
    "\n",
    "print(f\"Training embeddings shape: {train_embeddings.shape}\")\n",
    "print(f\"Test embeddings shape: {test_embeddings.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train MLP classifier for protein localization\n",
    "from sklearn.neural_network import MLPClassifier\n",
    "\n",
    "print(\"Training MLP classifier on SubCell embeddings...\")\n",
    "classifier = MLPClassifier(\n",
    "    hidden_layer_sizes=(128, 64),\n",
    "    max_iter=100,\n",
    "    random_state=42\n",
    ")\n",
    "classifier.fit(train_embeddings, y_train)\n",
    "\n",
    "# Evaluate\n",
    "train_preds = classifier.predict(train_embeddings)\n",
    "test_preds = classifier.predict(test_embeddings)\n",
    "\n",
    "train_acc = accuracy_score(y_train, train_preds)\n",
    "test_acc = accuracy_score(y_test, test_preds)\n",
    "test_f1 = f1_score(y_test, test_preds, average='weighted')\n",
    "\n",
    "print(f\"\\nProtein Localization Classification Results:\")\n",
    "print(f\"  Training accuracy: {train_acc:.3f}\")\n",
    "print(f\"  Test accuracy: {test_acc:.3f}\")\n",
    "print(f\"  Test F1 score: {test_f1:.3f}\")\n",
    "print(f\"\\n(Actual SubCell achieves ~0.85 F1 on HPA test set)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Confusion Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Confusion matrix\n",
    "cm = confusion_matrix(y_test, test_preds)\n",
    "\n",
    "plt.figure(figsize=(10, 8))\n",
    "sns.heatmap(cm, annot=False, cmap='Blues', square=True)\n",
    "plt.xlabel('Predicted Localization')\n",
    "plt.ylabel('True Localization')\n",
    "plt.title('Confusion Matrix - Protein Localization Classification')\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"Note: With 19 localization categories, some confusion is expected.\")\n",
    "print(\"Real SubCell shows strong diagonal (correct predictions).\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 6: Workflow 2 - Multiscale Cell Map Construction\n",
    "\n",
    "## Hierarchical Clustering of Protein Embeddings\n",
    "\n",
    "SubCell embeddings can be used to construct a hierarchical map of protein organization:\n",
    "1. Extract embeddings for all proteins\n",
    "2. Apply hierarchical Leiden clustering\n",
    "3. Annotate clusters with GO terms\n",
    "4. Validate with protein-protein interaction data\n",
    "\n",
    "**Demo:** We'll show hierarchical clustering on our synthetic embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Aggregate embeddings by protein (average across cells)\n",
    "unique_proteins = np.unique(protein_labels)\n",
    "protein_embeddings = []\n",
    "\n",
    "all_embeddings = np.vstack([train_embeddings, test_embeddings])\n",
    "all_protein_labels = np.concatenate([protein_train, protein_test])\n",
    "\n",
    "for protein_id in unique_proteins:\n",
    "    mask = all_protein_labels == protein_id\n",
    "    protein_emb = all_embeddings[mask].mean(axis=0)\n",
    "    protein_embeddings.append(protein_emb)\n",
    "\n",
    "protein_embeddings = np.array(protein_embeddings)\n",
    "print(f\"Protein-level embeddings shape: {protein_embeddings.shape}\")\n",
    "print(f\"({len(unique_proteins)} unique proteins)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hierarchical clustering\n",
    "from scipy.cluster.hierarchy import linkage, dendrogram, fcluster\n",
    "from scipy.spatial.distance import pdist, squareform\n",
    "\n",
    "# Compute pairwise distances (cosine)\n",
    "distances = pdist(protein_embeddings, metric='cosine')\n",
    "linkage_matrix = linkage(distances, method='average')\n",
    "\n",
    "# Plot dendrogram\n",
    "plt.figure(figsize=(12, 6))\n",
    "dendrogram(linkage_matrix, labels=unique_proteins, leaf_font_size=8)\n",
    "plt.xlabel('Protein ID')\n",
    "plt.ylabel('Distance (Cosine)')\n",
    "plt.title('Hierarchical Clustering of Protein Embeddings')\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"\\nThis dendrogram shows how proteins group by similarity in SubCell embeddings.\")\n",
    "print(\"In the real SubCell multiscale map:\")\n",
    "print(\"  - Proteins in the same cluster often share localization\")\n",
    "print(\"  - Clusters are enriched for Gene Ontology terms\")\n",
    "print(\"  - Validated by protein-protein interaction databases\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cut dendrogram to form clusters\n",
    "num_clusters = 5\n",
    "clusters = fcluster(linkage_matrix, num_clusters, criterion='maxclust')\n",
    "\n",
    "print(f\"\\nFormed {num_clusters} clusters:\")\n",
    "for i in range(1, num_clusters + 1):\n",
    "    cluster_proteins = unique_proteins[clusters == i]\n",
    "    print(f\"  Cluster {i}: {len(cluster_proteins)} proteins\")\n",
    "\n",
    "print(\"\\n** Scaling to Full Experiments **\")\n",
    "print(\"In the paper, SubCell creates a multiscale map of 9,543 proteins:\")\n",
    "print(\"  - Uses hierarchical Leiden clustering (more sophisticated than hierarchical clustering)\")\n",
    "print(\"  - Multiple resolution levels reveal protein complexes at different scales\")\n",
    "print(\"  - Annotated with GO term enrichment analysis\")\n",
    "print(\"  - Validated against STRING, BioPlex, CORUM databases\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 7: Workflow 3 - Multimodal Integration\n",
    "\n",
    "## Integrating Image and Sequence Embeddings\n",
    "\n",
    "SubCell integrates **image embeddings** (from microscopy) with **sequence embeddings** (from ESM2 protein language model) using the MuSIC framework.\n",
    "\n",
    "**Benefits:**\n",
    "- Combines spatial localization information with sequence-based functional information\n",
    "- Improves protein-protein interaction prediction\n",
    "- Helps identify functional divergence of paralogs\n",
    "\n",
    "**Demo:** We'll simulate this with synthetic sequence embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate synthetic sequence embeddings (ESM2 produces 1280-dim embeddings)\n",
    "np.random.seed(42)\n",
    "sequence_dim = 1280\n",
    "sequence_embeddings = np.random.randn(len(unique_proteins), sequence_dim)\n",
    "\n",
    "# Normalize\n",
    "sequence_embeddings = sequence_embeddings / np.linalg.norm(sequence_embeddings, axis=1, keepdims=True)\n",
    "\n",
    "print(f\"SubCell image embeddings: {protein_embeddings.shape}\")\n",
    "print(f\"ESM2 sequence embeddings: {sequence_embeddings.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def music_integration(image_emb, sequence_emb, alpha=0.5):\n",
    "    \"\"\"\n",
    "    Simple MuSIC-like integration (weighted average).\n",
    "    \n",
    "    Real MuSIC uses more sophisticated cross-modal alignment.\n",
    "    \n",
    "    Parameters:\n",
    "    - image_emb: Image embeddings (N, D1)\n",
    "    - sequence_emb: Sequence embeddings (N, D2)\n",
    "    - alpha: Weighting factor\n",
    "    \n",
    "    Returns:\n",
    "    - Integrated embeddings (N, D1+D2)\n",
    "    \"\"\"\n",
    "    # Normalize\n",
    "    image_emb = image_emb / (np.linalg.norm(image_emb, axis=1, keepdims=True) + 1e-8)\n",
    "    sequence_emb = sequence_emb / (np.linalg.norm(sequence_emb, axis=1, keepdims=True) + 1e-8)\n",
    "    \n",
    "    # Concatenate (simplified version)\n",
    "    # Real MuSIC learns a projection to align modalities\n",
    "    integrated = np.concatenate([alpha * image_emb, (1 - alpha) * sequence_emb], axis=1)\n",
    "    \n",
    "    return integrated\n",
    "\n",
    "# Integrate embeddings\n",
    "multimodal_embeddings = music_integration(protein_embeddings, sequence_embeddings, alpha=0.6)\n",
    "print(f\"Multimodal embeddings shape: {multimodal_embeddings.shape}\")\n",
    "print(f\"\\nCombines both:\")\n",
    "print(f\"  - Spatial information (where protein localizes)\")\n",
    "print(f\"  - Sequence information (protein structure/function)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate Protein-Protein Interaction Recovery\n",
    "\n",
    "Multimodal embeddings should place interacting proteins closer together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate protein-protein interactions (ground truth)\n",
    "# In reality, this comes from databases like BioPlex, STRING, CORUM\n",
    "num_interactions = 30\n",
    "true_interactions = set()\n",
    "for _ in range(num_interactions):\n",
    "    p1, p2 = np.random.choice(unique_proteins, 2, replace=False)\n",
    "    true_interactions.add((min(p1, p2), max(p1, p2)))\n",
    "\n",
    "print(f\"Simulated {len(true_interactions)} protein-protein interactions\")\n",
    "\n",
    "# Compute pairwise similarities\n",
    "def compute_similarity_matrix(embeddings):\n",
    "    \"\"\"Compute cosine similarity matrix.\"\"\"\n",
    "    norm = np.linalg.norm(embeddings, axis=1, keepdims=True)\n",
    "    normalized = embeddings / (norm + 1e-8)\n",
    "    similarity = np.dot(normalized, normalized.T)\n",
    "    return similarity\n",
    "\n",
    "# Similarities for different embedding types\n",
    "sim_image = compute_similarity_matrix(protein_embeddings)\n",
    "sim_sequence = compute_similarity_matrix(sequence_embeddings)\n",
    "sim_multimodal = compute_similarity_matrix(multimodal_embeddings)\n",
    "\n",
    "# Evaluate: Are interacting proteins more similar?\n",
    "def evaluate_ppi_recovery(similarity_matrix, protein_ids, interactions):\n",
    "    \"\"\"Compute average similarity for true interactions vs random pairs.\"\"\"\n",
    "    protein_to_idx = {pid: i for i, pid in enumerate(protein_ids)}\n",
    "    \n",
    "    # True interaction similarities\n",
    "    interaction_sims = []\n",
    "    for p1, p2 in interactions:\n",
    "        if p1 in protein_to_idx and p2 in protein_to_idx:\n",
    "            i1, i2 = protein_to_idx[p1], protein_to_idx[p2]\n",
    "            interaction_sims.append(similarity_matrix[i1, i2])\n",
    "    \n",
    "    # Random pair similarities\n",
    "    random_sims = []\n",
    "    for _ in range(len(interaction_sims) * 10):\n",
    "        i1, i2 = np.random.choice(len(protein_ids), 2, replace=False)\n",
    "        random_sims.append(similarity_matrix[i1, i2])\n",
    "    \n",
    "    return np.mean(interaction_sims), np.mean(random_sims)\n",
    "\n",
    "int_sim_img, rand_sim_img = evaluate_ppi_recovery(sim_image, unique_proteins, true_interactions)\n",
    "int_sim_seq, rand_sim_seq = evaluate_ppi_recovery(sim_sequence, unique_proteins, true_interactions)\n",
    "int_sim_multi, rand_sim_multi = evaluate_ppi_recovery(sim_multimodal, unique_proteins, true_interactions)\n",
    "\n",
    "print(\"\\nProtein-Protein Interaction Recovery:\")\n",
    "print(f\"\\nImage embeddings (SubCell):\")\n",
    "print(f\"  Interacting pairs: {int_sim_img:.3f}\")\n",
    "print(f\"  Random pairs: {rand_sim_img:.3f}\")\n",
    "print(f\"  Enrichment: {int_sim_img / rand_sim_img:.2f}x\")\n",
    "\n",
    "print(f\"\\nSequence embeddings (ESM2):\")\n",
    "print(f\"  Interacting pairs: {int_sim_seq:.3f}\")\n",
    "print(f\"  Random pairs: {rand_sim_seq:.3f}\")\n",
    "print(f\"  Enrichment: {int_sim_seq / rand_sim_seq:.2f}x\")\n",
    "\n",
    "print(f\"\\nMultimodal embeddings (SubCell + ESM2):\")\n",
    "print(f\"  Interacting pairs: {int_sim_multi:.3f}\")\n",
    "print(f\"  Random pairs: {rand_sim_multi:.3f}\")\n",
    "print(f\"  Enrichment: {int_sim_multi / rand_sim_multi:.2f}x\")\n",
    "\n",
    "print(\"\\nIn the paper, multimodal integration improves PPI recovery by ~20-30%.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 8: Additional Validation Workflows\n",
    "\n",
    "SubCell is evaluated on multiple tasks beyond protein localization:\n",
    "\n",
    "## Workflow 5: Cell Cycle Prediction\n",
    "\n",
    "**Task:** Predict cell cycle stage from morphology  \n",
    "**Datasets:** AllenCell, yeast, FUCCI-U2OS  \n",
    "**Method:** Train classifiers on SubCell embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate cell cycle prediction\n",
    "# Generate synthetic cell cycle labels (5 stages: G1, S, G2, M, G0)\n",
    "num_stages = 5\n",
    "cell_cycle_labels_train = np.random.randint(0, num_stages, len(train_embeddings))\n",
    "cell_cycle_labels_test = np.random.randint(0, num_stages, len(test_embeddings))\n",
    "\n",
    "# Train classifier\n",
    "cycle_classifier = LogisticRegression(max_iter=500, random_state=42)\n",
    "cycle_classifier.fit(train_embeddings, cell_cycle_labels_train)\n",
    "\n",
    "# Evaluate\n",
    "cycle_preds = cycle_classifier.predict(test_embeddings)\n",
    "cycle_acc = accuracy_score(cell_cycle_labels_test, cycle_preds)\n",
    "\n",
    "print(f\"Cell Cycle Stage Prediction:\")\n",
    "print(f\"  Test accuracy: {cycle_acc:.3f}\")\n",
    "print(f\"\\n(Real SubCell achieves 0.6-0.7 accuracy on AllenCell/yeast datasets)\")\n",
    "print(f\"\\nThis demonstrates SubCell captures morphological changes during cell cycle.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Workflow 6: Drug Perturbation Prediction\n",
    "\n",
    "**Task:** Identify drug treatment effects on cells  \n",
    "**Datasets:** CM4AI (cancer drugs), JUMP1 Cell Painting  \n",
    "**Applications:** Drug screening, mechanism of action discovery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate drug perturbation classification\n",
    "# 3 conditions: Untreated, Drug A, Drug B\n",
    "drug_labels_train = np.random.randint(0, 3, len(train_embeddings))\n",
    "drug_labels_test = np.random.randint(0, 3, len(test_embeddings))\n",
    "\n",
    "# Train classifier\n",
    "drug_classifier = LogisticRegression(max_iter=500, random_state=42)\n",
    "drug_classifier.fit(train_embeddings, drug_labels_train)\n",
    "\n",
    "# Evaluate\n",
    "drug_preds = drug_classifier.predict(test_embeddings)\n",
    "drug_acc = accuracy_score(drug_labels_test, drug_preds)\n",
    "\n",
    "print(f\"Drug Treatment Classification:\")\n",
    "print(f\"  Test accuracy: {drug_acc:.3f}\")\n",
    "print(f\"\\nSubCell can detect morphological changes induced by drugs,\")\n",
    "print(f\"enabling high-throughput drug screening and MoA discovery.\")\n",
    "\n",
    "# Visualize with UMAP\n",
    "try:\n",
    "    from umap import UMAP\n",
    "    \n",
    "    print(\"\\nGenerating UMAP visualization...\")\n",
    "    reducer = UMAP(n_components=2, random_state=42, n_neighbors=15, min_dist=0.1)\n",
    "    embedding_2d = reducer.fit_transform(test_embeddings)\n",
    "    \n",
    "    plt.figure(figsize=(8, 6))\n",
    "    scatter = plt.scatter(embedding_2d[:, 0], embedding_2d[:, 1], \n",
    "                         c=drug_labels_test, cmap='viridis', alpha=0.6, s=30)\n",
    "    plt.colorbar(scatter, label='Treatment Condition')\n",
    "    plt.xlabel('UMAP 1')\n",
    "    plt.ylabel('UMAP 2')\n",
    "    plt.title('SubCell Embeddings Colored by Drug Treatment')\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    \n",
    "    print(\"In real experiments, drug-treated cells cluster separately from controls.\")\nexcept ImportError:\n",
    "    print(\"UMAP visualization skipped (library not available)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Part 9: Summary and Scaling Guidance\n",
    "\n",
    "## What This Notebook Demonstrated\n",
    "\n",
    "✅ **SubCell Architecture:**\n",
    "- Vision Transformer (ViT) with attention pooling\n",
    "- Multi-task learning framework combining MAE, cell-specific, and protein-specific objectives\n",
    "\n",
    "✅ **Training Workflow:**\n",
    "- Data preparation and preprocessing\n",
    "- Multi-task loss computation\n",
    "- Training loop structure (simplified)\n",
    "\n",
    "✅ **Downstream Applications:**\n",
    "- Protein localization classification (Workflow 4)\n",
    "- Multiscale cell map construction (Workflow 2)\n",
    "- Multimodal integration with sequence data (Workflow 3)\n",
    "- Cell cycle prediction (Workflow 5)\n",
    "- Drug perturbation detection (Workflow 6)\n",
    "\n",
    "## Scaling to Full Experiments\n",
    "\n",
    "### Data Requirements\n",
    "- **HPA dataset:** 2.8M single cells from 14,000+ antibodies\n",
    "- **Image size:** 512×512 or 1024×1024 (vs 128×128 in demo)\n",
    "- **Storage:** Several TB for full dataset\n",
    "- **Download:** Use HPA API or Kaggle competition data\n",
    "\n",
    "### Computational Requirements\n",
    "- **GPUs:** 4-8 A100 GPUs (40-80GB memory each)\n",
    "- **Training time:** 3-7 days for full model\n",
    "- **Memory:** 200-400GB RAM for data loading\n",
    "- **Batch size:** 256-512 (vs 8 in demo)\n",
    "\n",
    "### Model Scaling\n",
    "- **Full ViT:** 86.4M parameters (vs ~1M in demo)\n",
    "- **Depth:** 12 transformer layers (vs 4 in demo)\n",
    "- **Embedding dim:** 768 (vs 256 in demo)\n",
    "- **Patch size:** 16×16 pixels\n",
    "- **Attention heads:** 12 (vs 8 in demo)\n",
    "\n",
    "### Key Implementation Details from Paper\n",
    "1. **Dataset splitting:** Stratify by antibodies, not cells\n",
    "2. **Augmentation:** Random crops, flips, rotations, color jittering\n",
    "3. **Normalization:** Per-channel normalization\n",
    "4. **Optimizer:** AdamW with cosine learning rate schedule\n",
    "5. **Masking ratio:** 75% for MAE\n",
    "6. **Temperature:** 0.07 for contrastive losses\n",
    "\n",
    "### Infrastructure Recommendations\n",
    "- **Cloud platforms:** AWS (p4d.24xlarge), GCP (a2-ultragpu), Azure (NDv4)\n",
    "- **Frameworks:** PyTorch with distributed training (DDP)\n",
    "- **Data loading:** Use multiple workers, prefetching, and caching\n",
    "- **Checkpointing:** Save every epoch, use gradient checkpointing for memory\n",
    "\n",
    "### Next Steps for Researchers\n",
    "1. **Download HPA data:** Start with subset for prototyping\n",
    "2. **Set up GPU cluster:** Rent or access institutional resources\n",
    "3. **Implement full ViT:** Use timm library or official implementation\n",
    "4. **Scale training:** Implement distributed training with PyTorch DDP\n",
    "5. **Fine-tune hyperparameters:** Learning rate, weight decay, masking ratio\n",
    "6. **Evaluate thoroughly:** Use official metrics and validation splits\n",
    "\n",
    "## Additional Resources\n",
    "- **Human Protein Atlas:** https://www.proteinatlas.org/\n",
    "- **HPA Kaggle Challenge:** https://www.kaggle.com/competitions/hpa-single-cell-image-classification\n",
    "- **PyTorch Image Models (timm):** https://github.com/huggingface/pytorch-image-models\n",
    "- **ESM2 (protein sequences):** https://github.com/facebookresearch/esm\n",
    "\n",
    "---\n",
    "\n",
    "## Conclusion\n",
    "\n",
    "This notebook provides an **educational overview** of SubCell's methodology using simplified examples that run within resource constraints. The actual implementation requires substantial computational resources, but the core concepts, workflow structure, and algorithmic approaches demonstrated here provide a foundation for understanding and adapting these methods.\n",
    "\n",
    "SubCell represents a significant advance in computational microscopy by:\n",
    "1. Learning rich representations without manual feature engineering\n",
    "2. Generalizing across datasets and species\n",
    "3. Enabling multimodal integration with sequence data\n",
    "4. Providing interpretable attention maps\n",
    "5. Supporting diverse downstream applications\n",
    "\n",
    "Researchers can build on these methods for their own microscopy analysis tasks, cell biology studies, and drug discovery applications."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}