correct attention patterns

README.md

@@ -1,70 +1,206 @@
 ---
 title: Stack Viz
-emoji: 🌍
-colorFrom: green
-colorTo: pink
+emoji: 🧬
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
 sdk_version: 6.3.0
 app_file: app.py
 pinned: false
 ---
 
-# STACK Model Visualization
+# STACK Model Visualization 🧬
 
-An interactive Gradio-based visualization of the STACK (Structured and Contextualized Knowledge) model architecture and inference capabilities.
+An interactive web-based visualization tool for understanding **STACK** (Single-cell Transcriptomic Analysis with Contextual Knowledge), a transformer-based model designed for single-cell RNA sequencing analysis with in-context learning capabilities.
 
-## Features
+## 📖 What is STACK?
 
-### Architecture View
-- **Visual Grid Representation**: 5x5 grid showing cells and gene modules
-- **Interactive Steps**:
-  - **Intra-cellular**: Visualize gene dependencies within cells (row-wise attention)
-  - **Inter-cellular**: Show population context across cells (column-wise attention)
-  - **Pre-training**: Demonstrate masked reconstruction for model training
+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:
 
-### Inference View
-- **Context Configuration**: Select cell type and condition for the prompt
-- **Target Configuration**: Choose query cell type (always healthy initial state)
-- **Zero-shot Prediction**: Run predictions using context from one cell type to predict another
-- **Visual Results**: See predicted cell states with animations
+- **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
 
-## Installation
+### 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
-pip install -r requirements.txt
+gradio app.py
 ```
 
-## Usage
+---
 
-Run the application:
+## 🎯 Using the Visualization
 
-```bash
-python app.py
+### 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"
+   - Watch the model process the context
+   - 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
 ```
 
-The app will launch in your browser at `http://localhost:7860`
+### 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
 
-## How It Works
+---
+
+## 🔬 Cell Types & Conditions
 
-1. **Architecture Tab**: Explore how STACK learns dependencies:
-   - Click "Intra-cellular" to see within-cell attention patterns
-   - Click "Inter-cellular" to see across-cell attention patterns
-   - Click "Pre-training" to see masked reconstruction in action
+### Available Cell Types
+- 🔵 **T-Cell**: T lymphocytes (immune cells, blue)
+- 🔷 **B-Cell**: B lymphocytes (antibody producers, indigo)  
+- � **Macrophage**: Innate immune cells (orange)
 
-2. **Inference Tab**: Test zero-shot predictions:
-   - Configure the context (prompt) with a cell type and condition
-   - Configure the target (query) with a different cell type
-   - Click "Run Prediction" to see the predicted cell state
-   - The model uses patterns from the prompt cells to predict query cell behavior
+### Experimental Conditions
+- ✅ **Healthy**: Normal baseline state
+- 💊 **Drug A**: Pharmaceutical treatment (red indicator)
+- 🦠 **Viral Infection**: Infectious disease state (orange indicator)
 
-## Cell Types
 
-- 🔵 **T-Cell**: T lymphocytes
-- 🔷 **B-Cell**: B lymphocytes  
-- 🟢 **Macro**: Macrophages
+---
 
-## Conditions
+## 📄 License
 
-- **Healthy**: Normal cell state
-- **Drug A**: Drug-treated state (indicated by red marker)
-- **Viral**: Virus-infected state (indicated by dashed border)
+This project is open source and available for educational and research purposes.