MD-BAX: A general-purpose Bayesian design framework for molecular dynamics simulations with input-dependent noise

Molecular dynamics (MD) simulations are a powerful tool for understanding complex molecular behavior, but exhaustively exploring the large space of input parameters can be computationally prohibitive, especially when the outcomes are noisy and/or expensive to evaluate. In this work, we introduce MD-Bayesian algorithm execution (BAX), a general-purpose, automated design framework that builds on BAX acquisition strategy to efficiently guide simulation campaigns toward learning meaningful features of the system. Unlike optimization-centric Bayesian optimization approaches, MD-BAX seeks to identify broader system properties (e.g., phase transition boundaries, level sets, and threshold crossings) by strategically selecting input/parameter settings based on uncertainty. To accurately represent the variability in simulation outcomes, MD-BAX incorporates a Gaussian process surrogate model with input-dependent noise, estimated directly from MD trajectory statistics at each simulation setting. This enables construction of reliable uncertainty estimates for guiding the next simulation. We demonstrate the approach on a case study involving coil-to-globule transitions in amphiphilic block copolymers, highlighting that explicitly including trajectory-derived noise improves uncertainty calibration and enables our framework to more efficiently map the relationship between polymer structure, solvent quality, and conformational behavior. MD-BAX represents a domain-informed specialization of the BAX framework for MD and is broadly applicable to molecular modeling problems where the goal is to infer key system behaviors from stochastic, trajectory-based simulation outputs rather than to locate a single optimal condition.