Removes local MD burn in parameter (#1112)

* Didn't end up necessary
* Can replicate the behavior by calling multiple_steps_local twice if we determine that we need this again.

tests/test_benchmark.py

@@ -284,7 +284,7 @@ def benchmark_local(
 
     local_seed = rng.integers(np.iinfo(np.int32).max)
     # run once before timer starts
-    ctxt.multiple_steps_local(steps_per_batch, ligand_idxs, seed=local_seed, burn_in=0)
+    ctxt.multiple_steps_local(steps_per_batch, ligand_idxs, seed=local_seed)
 
     start = time.time()
 
@@ -293,7 +293,7 @@ def benchmark_local(
         local_seed = rng.integers(np.iinfo(np.int32).max)
         # time the current batch
         batch_start = time.time()
-        _, _ = ctxt.multiple_steps_local(steps_per_batch, ligand_idxs, seed=local_seed, burn_in=0)
+        _, _ = ctxt.multiple_steps_local(steps_per_batch, ligand_idxs, seed=local_seed)
         batch_end = time.time()
 
         delta = batch_end - batch_start

tests/test_local_move.py

@@ -261,7 +261,7 @@ def num_particles_near_ligand(pair):
 
     def local_move(_, steps=500):
         local_seed = rng.integers(np.iinfo(np.int32).max)
-        xs, boxes = ctxt.multiple_steps_local(steps, local_idxs, burn_in=0, k=k, seed=local_seed)
+        xs, boxes = ctxt.multiple_steps_local(steps, local_idxs, k=k, seed=local_seed)
         return (xs[-1], boxes[-1]), None
 
     assert_no_drift(

tests/test_md.py

@@ -300,9 +300,6 @@ def test_multiple_steps_local_validation(freeze_reference):
     with pytest.raises(RuntimeError, match="atom indices must be unique"):
         ctxt.multiple_steps_local(100, np.array([1, 1], dtype=np.int32), radius=radius)
 
-    with pytest.raises(RuntimeError, match="burn in steps must be greater or equal to zero"):
-        ctxt.multiple_steps_local(100, np.array([1], dtype=np.int32), radius=radius, burn_in=-5)
-
     with pytest.raises(RuntimeError, match="radius must be greater or equal to 0.1"):
         ctxt.multiple_steps_local(100, np.array([1], dtype=np.int32), radius=0.01)
 
@@ -362,11 +359,6 @@ def test_multiple_steps_local_selection_validation(freeze_reference):
     with pytest.raises(RuntimeError, match="atom indices must be unique"):
         ctxt.multiple_steps_local_selection(100, reference_idx, np.array([1, 1], dtype=np.int32), radius=radius)
 
-    with pytest.raises(RuntimeError, match="burn in steps must be greater or equal to zero"):
-        ctxt.multiple_steps_local_selection(
-            100, reference_idx, np.array([1], dtype=np.int32), radius=radius, burn_in=-5
-        )
-
     with pytest.raises(RuntimeError, match="radius must be greater or equal to 0.1"):
         ctxt.multiple_steps_local_selection(100, reference_idx, np.array([1], dtype=np.int32), radius=0.01)
 
@@ -386,62 +378,6 @@ def test_multiple_steps_local_selection_validation(freeze_reference):
         ctxt.multiple_steps_local_selection(100, -1, np.array([3], dtype=np.int32))
 
 
-@pytest.mark.parametrize("freeze_reference", [True, False])
-def test_multiple_steps_local_burn_in(freeze_reference):
-    """Verify that burn in steps are identical to regular steps"""
-    seed = 2022
-    np.random.seed(seed)
-
-    N = 8
-    D = 3
-
-    coords = np.random.rand(N, D).astype(dtype=np.float64) * 2
-    box = np.eye(3) * 3.0
-    masses = np.random.rand(N)
-
-    E = 2
-
-    params, potential = prepare_nb_system(
-        coords,
-        E,
-        p_scale=3.0,
-        cutoff=1.0,
-    )
-    nb_pot = potential.to_gpu(np.float32)
-
-    temperature = constants.DEFAULT_TEMP
-    dt = 1.5e-3
-    friction = 0.0
-    radius = 1.2
-
-    # Select a single particle to use as the reference, will be frozen
-    local_idxs = np.array([len(coords) - 1], dtype=np.int32)
-
-    v0 = np.zeros_like(coords)
-    bps = [nb_pot.bind(params).bound_impl]
-
-    intg = LangevinIntegrator(temperature, dt, friction, masses, seed)
-
-    intg_impl = intg.impl()
-
-    steps = 100
-    burn_in = 100
-
-    ctxt = custom_ops.Context(coords, v0, box, intg_impl, bps)
-    ctxt.setup_local_md(temperature, freeze_reference)
-    ref_xs, ref_boxes = ctxt.multiple_steps_local(steps, local_idxs, radius=radius, burn_in=burn_in)
-    assert np.all(ref_xs[-1] < 1000)
-    intg_impl = intg.impl()
-
-    ctxt = custom_ops.Context(coords, v0, box, intg_impl, bps)
-    ctxt.setup_local_md(temperature, freeze_reference)
-    comp_xs, comp_boxes = ctxt.multiple_steps_local(steps + burn_in, local_idxs, radius=radius, burn_in=0)
-
-    # Final frame should be identical
-    np.testing.assert_array_equal(ref_xs, comp_xs)
-    np.testing.assert_array_equal(ref_boxes, comp_boxes)
-
-
 @pytest.mark.parametrize("freeze_reference", [True, False])
 def test_multiple_steps_local_consistency(freeze_reference):
     """Verify that running multiple_steps_local is consistent.
@@ -767,7 +703,7 @@ def test_local_md_with_selection_mask(freeze_reference):
 
     ctxt = custom_ops.Context(coords, v0, box, intg.impl(), bps)
     ctxt.setup_local_md(constants.DEFAULT_TEMP, freeze_reference)
-    xs, boxes = ctxt.multiple_steps_local_selection(steps, reference_idx, free_particles.astype(np.int32), burn_in=0)
+    xs, boxes = ctxt.multiple_steps_local_selection(steps, reference_idx, free_particles.astype(np.int32))
 
     if not freeze_reference:
         free_particles = np.append(free_particles, reference_idx)
@@ -888,11 +824,10 @@ def test_local_md_nonbonded_all_pairs_subset(freeze_reference):
     intg_impl = intg.impl()
 
     steps = 100
-    burn_in = 0
 
     ctxt = custom_ops.Context(coords, v0, box, intg_impl, bps)
     ctxt.setup_local_md(temperature, freeze_reference)
-    ref_xs, ref_boxes = ctxt.multiple_steps_local(steps, local_idxs, radius=radius, burn_in=burn_in)
+    ref_xs, ref_boxes = ctxt.multiple_steps_local(steps, local_idxs, radius=radius)
 
     for (ref_du_dx, ref_u), bp in zip(ref_vals, bps):
         test_du_dx, test_u = bp.execute(coords, box)
@@ -932,8 +867,8 @@ def test_context_invalid_boxes():
 
     ctxt = custom_ops.Context(coords, v0, box, intg.impl(), bps)
     ctxt.multiple_steps(steps)
-    ctxt.multiple_steps_local(steps, ligand_idxs, burn_in=0)
-    ctxt.multiple_steps_local_selection(steps, reference_idx, selection, burn_in=0)
+    ctxt.multiple_steps_local(steps, ligand_idxs)
+    ctxt.multiple_steps_local_selection(steps, reference_idx, selection)
 
     # Make the box way too small, which should trigger the failure
     ctxt.set_box(box * 0.01)
@@ -942,20 +877,18 @@ def test_context_invalid_boxes():
         _, boxes = ctxt.multiple_steps(steps)
         assert len(boxes) == 1
     with pytest.raises(RuntimeError, match=err_msg):
-        _, boxes = ctxt.multiple_steps_local(steps, ligand_idxs, burn_in=0)
+        _, boxes = ctxt.multiple_steps_local(steps, ligand_idxs)
         assert len(boxes) == 1
     with pytest.raises(RuntimeError, match=err_msg):
-        _, boxes = ctxt.multiple_steps_local_selection(steps, reference_idx, selection, burn_in=0)
+        _, boxes = ctxt.multiple_steps_local_selection(steps, reference_idx, selection)
         assert len(boxes) == 1
 
     # Without returning boxes no check will be performed
     _, boxes = ctxt.multiple_steps(steps, store_x_interval=steps + 1)
     assert len(boxes) == 0
-    _, boxes = ctxt.multiple_steps_local(steps, ligand_idxs, burn_in=0, store_x_interval=steps + 1)
+    _, boxes = ctxt.multiple_steps_local(steps, ligand_idxs, store_x_interval=steps + 1)
     assert len(boxes) == 0
-    _, boxes = ctxt.multiple_steps_local_selection(
-        steps, reference_idx, selection, burn_in=0, store_x_interval=steps + 1
-    )
+    _, boxes = ctxt.multiple_steps_local_selection(steps, reference_idx, selection, store_x_interval=steps + 1)
     assert len(boxes) == 0
 
 

timemachine/cpp/src/context.cu

@@ -86,7 +86,6 @@ void Context::_ensure_local_md_intialized() {
 std::array<std::vector<double>, 2> Context::multiple_steps_local(
     const int n_steps,
     const std::vector<int> &local_idxs,
-    const int burn_in,
     const int store_x_interval,
     const double radius,
     const double k,
@@ -122,9 +121,6 @@ std::array<std::vector<double>, 2> Context::multiple_steps_local(
         std::vector<std::shared_ptr<BoundPotential>> local_pots = local_md_pots_->get_potentials();
 
         intg_->initialize(local_pots, d_x_t_, d_v_t_, d_box_t_, d_free_idxs, stream);
-        for (int i = 0; i < burn_in; i++) {
-            this->_step(local_pots, d_free_idxs, stream);
-        }
         for (int i = 1; i <= n_steps; i++) {
             this->_step(local_pots, d_free_idxs, stream);
             if (i % store_x_interval == 0) {
@@ -165,7 +161,6 @@ std::array<std::vector<double>, 2> Context::multiple_steps_local_selection(
     const int n_steps,
     const int reference_idx,
     const std::vector<int> &selection_idxs,
-    const int burn_in,
     const int store_x_interval,
     const double radius,
     const double k) {
@@ -200,9 +195,6 @@ std::array<std::vector<double>, 2> Context::multiple_steps_local_selection(
         std::vector<std::shared_ptr<BoundPotential>> local_pots = local_md_pots_->get_potentials();
 
         intg_->initialize(local_pots, d_x_t_, d_v_t_, d_box_t_, d_free_idxs, stream);
-        for (int i = 0; i < burn_in; i++) {
-            this->_step(local_pots, d_free_idxs, stream);
-        }
         for (int i = 1; i <= n_steps; i++) {
             this->_step(local_pots, d_free_idxs, stream);
             if (i % store_x_interval == 0) {

timemachine/cpp/src/context.hpp

@@ -30,7 +30,6 @@ class Context {
     std::array<std::vector<double>, 2> multiple_steps_local(
         const int n_steps,
         const std::vector<int> &local_idxs,
-        const int burn_in,
         const int store_x_interval,
         const double radius,
         const double k,
@@ -40,7 +39,6 @@ class Context {
         const int n_steps,
         const int reference_idx,
         const std::vector<int> &selection_idxs,
-        const int burn_in,
         const int store_x_interval,
         const double radius,
         const double k);

timemachine/cpp/src/wrap_kernels.cpp

@@ -231,17 +231,13 @@ void declare_context(py::module &m) {
             [](timemachine::Context &ctxt,
                const int n_steps,
                const py::array_t<int, py::array::c_style> &local_idxs,
-               const int burn_in,
                const int store_x_interval,
                const double radius,
                const double k,
                const int seed) -> py::tuple {
                 if (n_steps <= 0) {
                     throw std::runtime_error("local steps must be at least one");
                 }
-                if (burn_in < 0) {
-                    throw std::runtime_error("burn in steps must be greater or equal to zero");
-                }
                 timemachine::verify_local_md_parameters(radius, k);
 
                 const int N = ctxt.num_atoms();
@@ -252,7 +248,7 @@ void declare_context(py::module &m) {
                 timemachine::verify_atom_idxs(N, vec_local_idxs);
 
                 std::array<std::vector<double>, 2> result =
-                    ctxt.multiple_steps_local(n_steps, vec_local_idxs, burn_in, x_interval, radius, k, seed);
+                    ctxt.multiple_steps_local(n_steps, vec_local_idxs, x_interval, radius, k, seed);
                 const int D = 3;
                 const int F = result[0].size() / (N * D);
                 py::array_t<double, py::array::c_style> out_x_buffer({F, N, D});
@@ -264,7 +260,6 @@ void declare_context(py::module &m) {
             },
             py::arg("n_steps"),
             py::arg("local_idxs"),
-            py::arg("burn_in") = 500, // This is arbitrarily selected as a default, TODO make informed choice
             py::arg("store_x_interval") = 0,
             py::arg("radius") = 1.2,
             py::arg("k") = 10000.0,
@@ -295,10 +290,6 @@ void declare_context(py::module &m) {
             selected to be frozen and used as the center of the shell of particles to be simulated. The selected
             idx is constant across all steps.
 
-        burn_in: int
-            How many steps to run prior to storing frames. This is to handle the fact that the local simulation applies a
-            restraint, and burn in helps equilibrate the local simulation. Running with small numbers of steps (< 100) is not recommended.
-
         store_x_interval: int
             How often we store the frames, store after every store_x_interval iterations. Setting to zero collects frames
             at the last step. Setting store_x_interval > n_steps will return no frames and skip runtime validation of box
@@ -333,16 +324,12 @@ void declare_context(py::module &m) {
                const int n_steps,
                const int reference_idx,
                const py::array_t<int, py::array::c_style> &selection_idxs,
-               const int burn_in,
                const int store_x_interval,
                const double radius,
                const double k) -> py::tuple {
                 if (n_steps <= 0) {
                     throw std::runtime_error("local steps must be at least one");
                 }
-                if (burn_in < 0) {
-                    throw std::runtime_error("burn in steps must be greater or equal to zero");
-                }
                 timemachine::verify_local_md_parameters(radius, k);
 
                 const int N = ctxt.num_atoms();
@@ -360,7 +347,7 @@ void declare_context(py::module &m) {
                 }
 
                 std::array<std::vector<double>, 2> result = ctxt.multiple_steps_local_selection(
-                    n_steps, reference_idx, vec_selection_idxs, burn_in, x_interval, radius, k);
+                    n_steps, reference_idx, vec_selection_idxs, x_interval, radius, k);
                 const int D = 3;
                 const int F = result[0].size() / (N * D);
                 py::array_t<double, py::array::c_style> out_x_buffer({F, N, D});
@@ -373,7 +360,6 @@ void declare_context(py::module &m) {
             py::arg("n_steps"),
             py::arg("reference_idx"),
             py::arg("selection_idxs"),
-            py::arg("burn_in") = 500, // This is arbitrarily selected as a default, TODO make informed choice
             py::arg("store_x_interval") = 0,
             py::arg("radius") = 1.2,
             py::arg("k") = 10000.0,
@@ -405,10 +391,6 @@ void declare_context(py::module &m) {
             The idxs of particles that should be free during local MD. Will be restrained to the particle specified by reference_idx particle using a
             flat bottom restraint which is defined by the radius and k values. Can be up to N - 1 particles, IE all particles except the reference_idx.
 
-        burn_in: int
-            How many steps to run prior to storing frames. This is to handle the fact that the local simulation applies a
-            restraint, and burn in helps equilibrate the local simulation. Running with small numbers of steps (< 100) is not recommended.
-
         store_x_interval: int
             How often we store the frames, store after every store_x_interval iterations. Setting to zero collects frames
             at the last step. Setting store_x_interval > n_steps will return no frames and skip runtime validation of box

timemachine/fe/free_energy.py

@@ -405,7 +405,6 @@ def run_production_local_steps(n_steps: int) -> Tuple[NDArray, NDArray]:
                 k=md_params.k,
                 radius=rng.uniform(md_params.min_radius, md_params.max_radius),
                 seed=rng.integers(np.iinfo(np.int32).max),
-                burn_in=0,
             )
 
             if coords is None:

timemachine/lib/custom_ops.pyi

@@ -36,8 +36,8 @@ class Context:
     def get_x_t(self) -> numpy.typing.NDArray[numpy.float64]: ...
     def initialize(self) -> None: ...
     def multiple_steps(self, n_steps: int, store_x_interval: int = ...) -> tuple: ...
-    def multiple_steps_local(self, n_steps: int, local_idxs: numpy.typing.NDArray[numpy.int32], burn_in: int = ..., store_x_interval: int = ..., radius: float = ..., k: float = ..., seed: int = ...) -> tuple: ...
-    def multiple_steps_local_selection(self, n_steps: int, reference_idx: int, selection_idxs: numpy.typing.NDArray[numpy.int32], burn_in: int = ..., store_x_interval: int = ..., radius: float = ..., k: float = ...) -> tuple: ...
+    def multiple_steps_local(self, n_steps: int, local_idxs: numpy.typing.NDArray[numpy.int32], store_x_interval: int = ..., radius: float = ..., k: float = ..., seed: int = ...) -> tuple: ...
+    def multiple_steps_local_selection(self, n_steps: int, reference_idx: int, selection_idxs: numpy.typing.NDArray[numpy.int32], store_x_interval: int = ..., radius: float = ..., k: float = ...) -> tuple: ...
     def set_box(self, box: numpy.typing.NDArray[numpy.float64]) -> None: ...
     def set_v_t(self, velocities: numpy.typing.NDArray[numpy.float64]) -> None: ...
     def set_x_t(self, coords: numpy.typing.NDArray[numpy.float64]) -> None: ...