BAX4MD

This repository contains code supporting the paper "MD-BAX: A General-Purpose Bayesian Design Framework for Molecular Dynamics Simulations with Input-Dependent Noise".

Overview

• /branin contains UQ calibration results based on a synthetic problem with varying levels of input-dependent noise.
• /data contains the estimated Rg value and noise data for all the points in the parameter grid. We used the precomputed diblock dataset for both UQ calibration analysis and case studies on BAX applications.
• /examples contains notebooks illustrating two MD-BAX application: level set estimation (diblock) and manifold crawling (both diblock and triblock).
• /src contains the source code for reproducing the results in our paper. UQtools.py was adapted from https://github.com/jensengroup/UQ_validation_methods.
• /template provides a template file for interested readers to apply MD-BAX (requires users to integrate their own MD code and design goals) on their own system.
• /uq_calibration contains UQ calibration results based on MD simulation data.

Packages

• botorch 0.10.0
• gpytorch 1.11
• matplotlib 3.9.2
• numpy 1.26.4
• pandas 2.2.3
• pytorch 2.4.0
• scikit-learn 1.1.3
• scipy 1.13.1

Usage

To efficiently obtain meaningful statistical performance results, we used a pseudo simulator that directly reads pre-computed real simulation outcomes (both Rg value and estimated noise) from data files when given specified inputs on a predefined grid. The Simulator class in copolymer_simulator.py can be adapted to a real simulator and seamlessly integrated with our framework.

from copolymer_simulator import Simulator

simulator = Simulator(polymer='diblock') # diblock case
# simulator = Simulator(polymer='triblock') # triblock case

For closed-loop implementation with real simulators, readers can refer to run_bax.py. Users will need to replace the following dummy code with their own MD code:

    def _run_simulation(self, X):
        """
        Run a single MD simulation at input condition X.

        Parameters
        ----------
        X : np.ndarray or torch.Tensor, shape (1, d)
            - The simulation input parameters (e.g., polymer composition, temperature, etc.).
            - Currently unused in this dummy implementation, but in a real MD code you
            would pass these values into your simulator.

        Returns
        -------
        traj : np.ndarray, shape (T, 2)
            - Two-column array representing the trajectory.
            - Column 0: timesteps (monotonically increasing, dtype float).
            - Column 1: property of interest (e.g., radius of gyration Rg, float).
            - This is the format expected by `_mean_sem_from_traj_original` downstream.

        Notes
        -----
        - The downstream code expects the trajectory to have *at least two columns*:
        time in the first column, and the property value in the second column.
        - Time should be increasing, ideally in the same units as `acf_window` if you
        use one (e.g., ps, steps). If your MD engine outputs irregular time points,
        you should still provide them sorted in ascending order.
        - The property can be anything scalar (Rg, energy, etc.) but should be sampled
        at each time point.
        """

        # ------------------ Dummy Example ------------------
        n = 100000                          # number of timesteps
        t = np.arange(0, n*10, 10, dtype=float)  # (T,) time vector: 0, 10, 20, ... 999,990

        rng = self.rng                      # random generator for reproducibility
        # Toy signal: mean ~3.0, slow sinusoidal oscillation + Gaussian noise
        y = 3.0 + 0.2*np.sin(2*np.pi*t/5000.0) + rng.normal(0, 0.05, size=n)

        # Combine into trajectory array (T, 2)
        traj = np.column_stack([t, y])
        return traj

Our general-purpose framework can extend beyond tasks of level set estimation and manifold crawling shown in the paper. To do this, users need to update TaskHandler class in run_bax.py to apply their own (algorithms).

class TaskHandler:
    def __init__(self, task_kwargs):
        self.task_kwargs = task_kwargs

    def get_valid_indices(self, X, sample_Y, pred_Ymean, pred_Yvar, max_var_indices):
        raise NotImplementedError

Citation

@article{
author = {Tianhong Tan, Ting-Yeh Chen, Jacob R. Breese, Lisa M. Hall, Joel A. Paulson},
title = {MD-BAX: A General-Purpose Bayesian Design Framework for Molecular Dynamics Simulations with Input-Dependent Noise},
journal = {},
volume = {},
number = {},
pages = {},
doi = {},
abstract = {}}