Abstract

This project delivers a high-performance, GPU-accelerated framework for the classic 2D Müller-Brown potential, designed as a ground-truth testbed for novel Machine Learning in Molecular Dynamics (ML-MD) algorithms. The architecture is built around a JIT-compiled PyTorch module that provides efficient force computation via analytical Jacobian derivation. This optimization achieves a validated 3-10x performance gain over standard PyTorch Autograd, ensuring microsecond-level computation for large batch simulations.

Core Computational Capabilities

  • Dual Force Kernels: Implemented analytically-derived force kernel for peak simulation performance, complemented by an Autograd-enabled mode for seamless ML training and differentiation
  • Accelerated Execution (JIT/CUDA): Native PyTorch module with JIT compilation (@torch.jit.script) for sub-microsecond latency, ensuring GPU compatibility and low-latency scaling
  • Robust Numerical Integrator: Features a production-grade Langevin Dynamics (LD) simulator using the Velocity-Verlet integration scheme for numerically stable and accurate trajectory generation
  • Modular & Extensible Architecture: Clean separation of physics (Potential class), numerics (LangevinSimulator class), and data I/O, establishing a template for simulating larger, complex systems

Implementation & Technical Deliverables

Architecture

Core implementation built with PyTorch best practices:

  • Physics Module: Implemented the energy surface as a JIT-compiled torch.nn.Module, correctly initializing potential parameters as Buffers
  • Force Kernel Optimization: Calculated and implemented the analytical Jacobian for the potential, offering a guaranteed path for peak performance, bypassing autograd overhead during long simulations
  • Vectorized Execution: All kernel operations are fully vectorized, enabling efficient batch processing and concurrent multi-particle simulations
  • Hardware Agnostic: Full support for moving all operations (potential, force, simulation) to CUDA devices via native PyTorch tensor handling

Performance Results

Benchmarking across different batch sizes showed:

  • Analytical derivatives consistently outperform autograd by 3-10x
  • Per-particle computation stays under 1 microsecond for large systems
  • Linear scaling enables processing 100,000+ coordinates efficiently
  • JIT compilation provides significant speedups over pure Python

Molecular Dynamics

Included Langevin dynamics simulations demonstrate the potential’s behavior:

  • Particle motion in different energy basins
  • Thermal fluctuations around local minima
  • Transitions between energy wells
  • Trajectory visualization on the potential surface

Why It Matters

Core Engineering & Scientific Skills Demonstrated

This project demonstrates:

  • High-Performance Kernel Design: Expertise in designing and implementing fast, differentiable mathematical kernels for scientific computing problems using PyTorch’s performance tools (JIT, vectorization)
  • Performance Engineering Trade-offs: The ability to identify bottlenecks (Autograd) and implement analytical solutions for targeted 3-10x speedup, proving competency in computational efficiency
  • Bridging Scientific Computing and ML: Successfully creating a native PyTorch component that is immediately ready for integration into machine learning pipelines (e.g., training a surrogate force field or using enhanced sampling methods)

Value Proposition for Scientific Development

  • ML Algorithm Testbed: Provides a known, high-quality ground truth surface for rapidly benchmarking and validating new ML/MD methods without the complexity of a full-scale molecular simulator
  • Educational Resource: Simple 2D system for learning molecular dynamics concepts with production-quality code
  • Foundation for Advanced Methods: An ideal starting point for implementing and testing rare event sampling, path finding, or advanced kinetic analysis algorithms

Technical Notes

Implementation focuses on clarity and performance:

  • Clean separation between mathematical formulation and computational details
  • Comprehensive documentation and examples
  • Reproducible benchmarking with all code included
  • Interactive visualizations to understand the energy landscape

This work demonstrates how classic scientific computing problems can benefit from modern frameworks while maintaining computational efficiency.

Simulation Videos

These videos demonstrate Langevin dynamics simulations on the Müller-Brown potential surface:

A Basin Dynamics: Particle motion and thermal fluctuations around the A minimum.

B Basin Dynamics: Exploration of the deeper B minimum energy well.

Transition Path: Particle transitioning between energy basins, demonstrating barrier crossing.

This implementation is documented in detail in: