Skip to main content

​
Overview

In warp-md, Plans are the analysis primitives. Each Plan:
  1. Compiles a specific analysis (RMSD, Rg, RDF, etc.)
  2. Validates inputs at construction time
  3. Executes on CPU or GPU via the run() method
Plans are implemented in Rust (traj-engine crate) and exposed to Python via bindings.

​
The Plan Lifecycle

​
1. Create Plan

Instantiate with a Selection and parameters: 
from warp_md import RgPlan

backbone = system.select("backbone")
plan = RgPlan(backbone, mass_weighted=True)
Plan construction is cheap. Validation happens immediately, but no trajectory I/O occurs until run().

​
2. Run Plan

Execute over the trajectory: 
traj = Trajectory.open_xtc("trajectory.xtc", system)
rg = plan.run(traj, system, device="auto", chunk_frames=128)

print(rg.shape)  # (n_frames,)
print(rg.dtype)  # float32

​
Common Plan Interface

All Plans follow this signature: 
plan = SomePlan(
    selection,          # Required: Selection object
    # ...plan-specific parameters
)

result = plan.run(
    traj,               # Trajectory object
    system,             # System object
    device="auto",      # "auto" | "cpu" | "cuda" | "cuda:0"
    chunk_frames=128,   # Frames per chunk (memory tuning)
    frame_indices=None  # Optional: [0, 10, 20] or None for all
)

​
Parameters

​
traj
Trajectory
required
The trajectory to analyze. Must match system.n_atoms().
​
system
System
required
The system providing topology and atom metadata.
​
device
str
default:"auto"
Execution device:
  • "auto": Use CUDA GPU if available, else CPU
  • "cpu": Force CPU execution
  • "cuda": Use GPU device 0
  • "cuda:1": Use GPU device 1
From crates/traj-engine/src/executor.rs:21-52:
pub fn from_spec(spec: &str) -> TrajResult<Self> {
    let spec = spec.trim();
    if spec.eq_ignore_ascii_case("cpu") {
        return Ok(Device::Cpu);
    }
    if spec.eq_ignore_ascii_case("auto") {
        #[cfg(feature = "cuda")]
        {
            if let Ok(ctx) = GpuContext::new(0) {
                return Ok(Device::Cuda(ctx));
            }
        }
        return Ok(Device::Cpu);
    }
    // ...
}
​
chunk_frames
int
default:"128"
Number of frames to read per I/O chunk. Larger values improve throughput but increase memory usage.Tuning advice:
  • Small trajectories: 128-512
  • Large trajectories: 64-128
  • GPU runs: 256-1024 (GPUs amortize transfer overhead)
​
frame_indices
list[int] | None
default:"None"
Optional list of frame indices to analyze. Supports negative indexing.
# Analyze first, middle, last frame
result = plan.run(traj, system, frame_indices=[0, 500, -1])
From crates/traj-engine/src/executor.rs:200-212:
pub fn normalize_frame_indices(frame_indices: Vec<i64>, n_frames: usize) -> Vec<usize> {
    let mut out = Vec::with_capacity(frame_indices.len());
    for raw in frame_indices.into_iter() {
        let mut idx = raw;
        if idx < 0 {
            idx += n_frames as i64;
        }
        if idx >= 0 && (idx as usize) < n_frames {
            out.push(idx as usize);
        }
    }
    out
}

​
Example Plans

​
Radius of Gyration

From python/warp_md/__init__.py:12: 
from warp_md import RgPlan

backbone = system.select("backbone")
plan = RgPlan(backbone, mass_weighted=False)
rg = plan.run(traj, system, device="auto")

print(f"Rg: {rg.mean():.2f} ± {rg.std():.2f} Å")

​
RMSD with Alignment

from warp_md import RmsdPlan

ca_atoms = system.select("name CA")
plan = RmsdPlan(
    ca_atoms,
    reference="topology",  # or "first" or provide coords
    align=True,            # Align before computing RMSD
    mass_weighted=False
)
rmsd = plan.run(traj, system)

​
Distance Between Selections

from warp_md import DistancePlan

sel_a = system.select("resid 10 and name CA")
sel_b = system.select("resid 50 and name CA")

plan = DistancePlan(sel_a, sel_b)
dist = plan.run(traj, system)

print(dist.shape)  # (n_frames,)

​
Clustering Trajectories

From python/warp_md/tests/test_clustering.py:32-76: 
from warp_md import TrajectoryClusterPlan

backbone = system.select("backbone")
plan = TrajectoryClusterPlan(
    backbone,
    method="dbscan",  # or "kmeans"
    eps=2.0,
    min_samples=5,
    memory_budget_bytes=1024 * 1024 * 512  # 512 MB
)

result = plan.run(traj, system, device="cuda")
print(result["labels"])     # Cluster assignments per frame
print(result["centroids"])  # Centroid frame indices
print(result["sizes"])      # Cluster sizes

​
Plan Catalog

From python/warp_md/__init__.py:8-107, 96 Plans are available:
  • RmsdPlan — RMSD with optional alignment
  • RmsdPerResPlan — Per-residue RMSD
  • RmsfPlan — Root mean square fluctuation
  • SymmRmsdPlan — RMSD with symmetry equivalents
  • DistanceRmsdPlan — Distance-based RMSD (dRMSD)
  • PairwiseRmsdPlan — All-vs-all RMSD matrix
  • RgPlan — Radius of gyration
  • RadgyrTensorPlan — Gyration tensor eigenvalues
  • DistancePlan — Distance between two selections
  • PairDistPlan — Distance between atom pairs
  • AnglePlan — Bond angles
  • DihedralPlan — Dihedral angles
  • VectorPlan — Geometric vectors
  • AlignPlan — Align trajectory to reference
  • SuperposePlan — Superpose structures
  • TranslatePlan — Translate coordinates
  • RotatePlan — Rotate coordinates
  • ScalePlan — Scale coordinates
  • TransformPlan — Apply transformation matrix
  • TrajectoryClusterPlan — Cluster frames (DBSCAN/k-means)
  • PcaPlan — Principal component analysis
  • ProjectionPlan — Project onto modes
  • MatrixPlan — Covariance/correlation matrix
  • MsdPlan — Mean square displacement
  • AtomicCorrPlan — Atomic correlation functions
  • VelocityAutoCorrPlan — Velocity autocorrelation
  • XcorrPlan — Cross-correlation
  • RotAcfPlan — Rotational autocorrelation
  • RdfPlan — Radial distribution function
  • DensityPlan — 3D density grid
  • VolmapPlan — Volumetric map
  • GistGridPlan — GIST analysis (grid-based)
  • GistDirectPlan — GIST analysis (direct)
  • HbondPlan — Hydrogen bond analysis
  • NmrIredPlan — NMR relaxation (iRED)
  • WaveletPlan — Wavelet transform
  • EndToEndPlan — End-to-end distance
  • ContourLengthPlan — Contour length
  • ChainRgPlan — Radius of gyration per chain
  • PersistenceLengthPlan — Persistence length
  • BondLengthDistributionPlan — Bond length distribution
  • BondAngleDistributionPlan — Bond angle distribution
See the API Reference for complete documentation of all 96 Plans.

​
Execution Model

From crates/traj-engine/src/executor.rs:74-78: 
pub struct Executor {
    system: System,
    chunk_frames: usize,
    device: Device,
}
Plans execute via the Executor:
  1. Chunk reading: Trajectory read in chunks (size = chunk_frames)
  2. Device dispatch: Coords copied to CPU/GPU
  3. Kernel execution: Plan-specific compute (Rust or CUDA)
  4. Accumulation: Results aggregated across chunks
  5. Finalization: Return NumPy array
GPU execution is automatic when device="auto" and CUDA is available. No code changes needed.

​
Streaming vs Batch

Most Plans stream data: 
# Streams trajectory in chunks, low memory
rg = RgPlan(sel).run(traj, system, chunk_frames=128)
Some Plans require full trajectory in memory:
  • PairwiseRmsdPlan (all-vs-all matrix)
  • TrajectoryClusterPlan (DBSCAN/k-means)
  • PcaPlan (covariance matrix)
Batch Plans may OOM on large trajectories (>10k frames, >10k atoms). Use memory_budget_bytes to limit memory.

​
Error Handling

Plans raise Python exceptions on errors: 
try:
    plan = RmsdPlan(selection, reference="topology")
    rmsd = plan.run(traj, system)
except RuntimeError as e:
    print(f"Plan failed: {e}")
Common errors:
  • Atom count mismatch: Trajectory has 1000 atoms, System has 999
  • Invalid selection: Selection contains out-of-bounds indices
  • Device error: CUDA unavailable; rebuild with --features cuda

​
See Also

Build docs developers (and LLMs) love

Get started for freeTalk to us