README.md

---
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**:
   ```bash
   cd stack-viz
   ```

2. **Install dependencies**:
   ```bash
   pip install -r requirements.txt
   ```
   
   This will install:
   - `gradio>=4.0.0` - Web interface framework

3. **Run the application**:
   ```bash
   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:
```bash
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.