3D Protein Simulation

Molecular dynamics simulation and visualization of protein structures using Warp and Polyscope. The codebase provides a modular framework for simulating and visualizing protein dynamics with GPU acceleration. The example demonstrates visualization of the tumor suppressor protein 1TUP.

Project Structure

3D_protein/
├── run_simulation.py       # Main entry point (run from here)
├── requirements.txt
├── data/
│   └── 1TUP.cif            # Example protein structure (Tumor Suppressor P53)
├── src/
│   ├── simulation/         # Simulation engine
│   │   ├── __init__.py
│   │   └── engine.py       # Langevin dynamics, Warp kernels, analyze_structure, render_frame
│   └── visualization/     # Polyscope visualization
│       ├── __init__.py
│       └── polyscope_viz.py
└── scripts/
    └── run_simulation_only.py   # Run simulation without visualization

• src/simulation/ – Core simulation engine

  • run_simulation(), analyze_structure(), render_frame()
  • GPU-accelerated Warp kernels for bond, non-bonded, and Langevin integration

• src/visualization/ – Polyscope-based visualization

  • visualize_simulation(), prepare_visualization_data(), setup_polyscope_visualization()
  • Before/after comparison and animated trajectory playback

• run_simulation.py – Main entry point; run from project root to execute the full pipeline

Installation

Create and activate a virtual environment, then install dependencies:

python -m venv .venv
source .venv/bin/activate   # On Windows: .venv\Scripts\activate
pip install -r requirements.txt

Important: Run all commands below with this venv activated (or use .venv/bin/python explicitly). Otherwise you may get ModuleNotFoundError: No module named 'warp' because the system Python doesn’t have warp-lang installed.

Usage

Quick Start

Run the complete simulation and visualization pipeline:

python run_simulation.py

This will:

1. Load the protein structure from data/1TUP.cif
2. Run the molecular dynamics simulation
3. Generate protein_simulation.usd in the project root
4. Launch an interactive Polyscope visualization window

Running Components Separately

Run only the simulation (no visualization):

python scripts/run_simulation_only.py

This generates the USD file but does not launch the visualization.

Run only the visualization (requires pre-computed trajectory):

from src.simulation import run_simulation
from src.visualization import visualize_simulation

trajectory, forces, chains_data = run_simulation(
    filename='data/1TUP.cif',
    n_steps=200,
    dt=0.002
)
visualize_simulation(trajectory, forces, chains_data)

Programmatic Usage

Import from the src package:

import warp as wp
from src.simulation import run_simulation, analyze_structure
from src.visualization import visualize_simulation, prepare_visualization_data

wp.init()

trajectory, forces, chains_data = run_simulation(
    filename='data/your_protein.cif',
    n_steps=500,
    dt=0.001,
    output_usd='output.usd',
    device='cuda'
)

visualize_simulation(
    trajectory=trajectory,
    forces=forces,
    chains_data=chains_data,
    screenshot_path='comparison.png'
)

API Reference

src.simulation.run_simulation()

def run_simulation(
    filename: str = '1TUP.cif',
    n_steps: int = 200,
    dt: float = 0.002,
    output_usd: str = "protein_simulation.usd",
    device: str = "cpu"
) -> tuple:
    """Returns (trajectory, forces_history, chains_data). Call wp.init() first."""

src.visualization.visualize_simulation()

def visualize_simulation(
    trajectory: np.ndarray,
    forces: np.ndarray,
    chains_data: List[Dict[str, Any]],
    screenshot_path: str = None
) -> None:
    """Main function to visualize a protein simulation."""

Theoretical Background

This project implements a basic coarse-grained molecular dynamics simulation to demonstrate high-performance physics simulation using NVIDIA Warp. The primary goal is to showcase how Warp can be used to write differentiable, GPU-accelerated simulation kernels in Python.

1. Langevin Dynamics (Overdamped)

The motion of particles is governed by overdamped Langevin dynamics, which approximates the behavior of particles in a solvent where viscous drag dominates inertia:

$$ r_{i}(t+\Delta t) = r_{i}(t) + \frac{F_{i}(t)}{\gamma} \Delta t + \sqrt{2 k_B T \frac{\Delta t}{\gamma}} \xi(t) $$

Where:

• $r_i$ is the position of particle $i$
• $\gamma$ is the friction coefficient
• $k_B T$ is the thermal energy
• $\xi(t)$ is a Gaussian random noise vector (Brownian motion)

2. Interaction Forces

The simulation includes several force terms to model molecular interactions:

Harmonic Bonds: Maintains the connectivity of the protein chain. $$ F_{bond} = -k_{bond} (r - r_0) \hat{r} $$

Lennard-Jones Potential: Models the van der Waals forces (short-range repulsion and long-range attraction). $$ F_{LJ} = \frac{24\epsilon}{r} \left[ 2\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6 \right] \hat{r} $$

Electrostatics (Coulomb): Models the attraction and repulsion between charged residues (e.g., DNA- and Protein+). $$ F_{elec} = k_{coulomb} \frac{q_i q_j}{r^2} \hat{r} $$

Note: This is a simplified coarse-grained model (1 bead per residue) intended for visualization and performance demonstration, not for rigorous biophysical predictions.

Code Organization

The codebase follows best practices for modular Python development:

• Clear APIs: All functions have docstrings with type hints
• Reusable Components: Functions can be imported and used independently
• Separation of Concerns: Simulation, visualization, and execution are separate modules
• No Side Effects: Modules don't execute code on import
• Proper Initialization: wp.init() is called by the user or main script, not at module level

Output Files

• protein_simulation.usd: USD format file containing the full simulation trajectory
• output images/before_after_comparison.png: Screenshot comparing initial and final conformations