README.md

metadata

title: Stack Viz
emoji: 🧬
colorFrom: indigo
colorTo: purple
sdk: gradio
sdk_version: 6.3.0
app_file: app.py
pinned: false

STACK Model Visualization 🧬

An interactive web-based visualization tool for understanding STACK, a transformer-based model designed for single-cell RNA sequencing analysis with in-context learning capabilities.

📖 What is STACK?

STACK is a deep learning model that learns from gene expression patterns in single-cell data to make zero-shot predictions across different cell types and experimental conditions. The key innovation is its ability to:

• Learn contextual patterns: Understand how gene expression changes under different conditions (e.g., drug treatments, viral infections)
• Transfer knowledge: Apply patterns learned from one cell type to predict behavior in a completely different cell type
• Perform in-context learning: Use "prompt" cells (examples) to guide predictions on "query" cells (targets), similar to how large language models work

How It Works

STACK operates on gene expression matrices where:

• Rows represent gene modules (groups of related genes)
• Columns represent individual cells

The model uses a dual-attention mechanism:

1. Intra-cellular attention: Learns dependencies between genes within each cell (row-wise)
2. Inter-cellular attention: Learns patterns across the cell population (column-wise)

This allows STACK to capture both individual cell biology and population-level context, enabling powerful zero-shot predictions without requiring retraining for new cell types or conditions.

---

🚀 Quick Setup Guide

Prerequisites

• Python 3.8 or higher
• pip package manager

Installation

1. Clone or download this repository:

   cd stack-viz
   

2. Install dependencies:

   pip install -r requirements.txt
   

   This will install:

   • gradio>=4.0.0 - Web interface framework

3. Run the application:

   python app.py
   

4. Access the interface:

   • Open your browser and navigate to http://localhost:7860
   • The Gradio interface will automatically launch

Alternative: Run with auto-reload

For development purposes, you can run with auto-reload:

gradio app.py

---

🎯 Using the Visualization

Architecture Tab 🏗️

Explore how STACK learns from gene expression data:

• Intra-cellular Attention (→): Click to highlight a row, showing how the model learns gene dependencies within individual cells
• Inter-cellular Attention (↓): Click to highlight a column, showing how the model learns context from the cell population
• Masked Pre-training (🔄): Click to see masked gene reconstruction - the model learns by predicting masked gene expression values (two consecutive rows are masked)

Inference Tab 🔮

Test zero-shot predictions:

1. Configure Prompt Cells 🔴 (Known Response):

   • Select a cell type (T-Cell, B-Cell, or Macrophage)
   • Select a condition (Drug A or Viral Infection)
   • These cells provide the "context" for prediction

2. Configure Query Cells 🔵 (To Predict):

   • Select a different cell type
   • Always starts in healthy/baseline state
   • The model will predict how these cells would respond to the condition

3. Run Prediction:

   • Click "Run Zero-Shot Prediction"
   • View predicted gene expression patterns for your query cells

Example Use Case: "If I know how T-Cells respond to Drug A, how would B-Cells respond to the same drug?"

---

📁 Code Organization

Project Structure

stack-viz/
├── app.py              # Main application file
├── requirements.txt    # Python dependencies
├── README.md          # This file
└── .git/              # Git repository

Code Architecture (app.py)

The application is organized into logical sections:

1. Constants & Configuration (Lines 6-17)

• CELL_TYPES: Definitions for T-Cell, B-Cell, and Macrophage with colors
• CONDITIONS: Definitions for Healthy, Drug A, and Viral conditions
• Color schemes for visual consistency

2. Cell Visualization Helpers (Lines 19-68)

• make_cell_svg(): Generates SVG representations of cells
  • Prompt cells: Irregular shape (treated/perturbed state)
  • Query cells: Circular shape with dashed borders (baseline state)
• generate_cell_array(): Creates horizontal arrays of cells for display
• Uses color coding to distinguish cell types and conditions

3. Inference Display Functions (Lines 70-160)

• generate_inference_display(): Creates the main visualization showing prompt and query cells
• generate_mini_matrix(): Generates preview of gene expression matrix
• Combines cell arrays with arrow indicators and matrix representations

4. Architecture Grid Functions (Lines 162-244)

• generate_grid_html(): Creates the 5×5 gene expression matrix visualization
  • Rows = Genes
  • Columns = Cells
  • Interactive highlighting for different attention patterns
• get_step_label(): Returns descriptive labels for each learning step
• update_architecture_view(): Manages state changes between visualization modes

5. Inference Logic (Lines 261-357)

• update_inference_display(): Updates visualization when user changes selections
• run_inference(): Simulates the prediction process with animation
  • Shows "processing" state with spinner
  • Generates predicted cell states
  • Displays gene expression output columns
• reset_inference(): Resets to initial state
• generate_output_column(): Creates visual representation of predicted gene counts

6. Gradio Interface (Lines 359-509)

• create_app(): Main application builder
  • Header: Branding and title
  • Architecture Tab: Grid visualization with control buttons
  • Inference Tab: Interactive prediction interface with controls
  • State management using gr.State()
  • Event handlers for button clicks and selections

7. Styling (Lines 511-545)

• custom_css(): Custom CSS for enhanced UI
  • Animations (spinner for loading)
  • Button styling and transitions
  • Tab styling with gradients
  • Responsive design

8. Application Entry Point (Lines 547-551)

• Launches the Gradio app with custom CSS
• Default port: 7860

Customization Points

To modify the visualization:

• Add new cell types: Edit CELL_TYPES constant
• Add new conditions: Edit CONDITIONS constant
• Change grid size: Modify the grid generation in generate_grid_html()
• Adjust masking logic: Update the masking selection in update_architecture_view()
• Customize styling: Edit custom_css() function

---

🔬 Cell Types & Conditions

Available Cell Types

• 🔵 T-Cell: T lymphocytes (immune cells, blue)
• 🔷 B-Cell: B lymphocytes (antibody producers, indigo)
• � Macrophage: Innate immune cells (orange)

Experimental Conditions

• ✅ Healthy: Normal baseline state
• 💊 Drug A: Pharmaceutical treatment (red indicator)
• 🦠 Viral Infection: Infectious disease state (orange indicator)

---

📄 License

This project is open source and available for educational and research purposes.