From f195729c57d7e900779db8823f24ce0ec053a721 Mon Sep 17 00:00:00 2001 From: Forrest York Date: Tue, 8 Aug 2023 11:53:33 -0600 Subject: [PATCH] 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 | 4 +- tests/test_local_move.py | 2 +- tests/test_md.py | 83 +++------------------------- timemachine/cpp/src/context.cu | 8 --- timemachine/cpp/src/context.hpp | 2 - timemachine/cpp/src/wrap_kernels.cpp | 22 +------- timemachine/fe/free_energy.py | 1 - timemachine/lib/custom_ops.pyi | 4 +- 8 files changed, 15 insertions(+), 111 deletions(-) diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 386c55584..1a410fa70 100644 --- a/tests/test_benchmark.py +++ b/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 diff --git a/tests/test_local_move.py b/tests/test_local_move.py index 94267dc00..b168bbb83 100644 --- a/tests/test_local_move.py +++ b/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( diff --git a/tests/test_md.py b/tests/test_md.py index a40f2ac12..826ce4bed 100644 --- a/tests/test_md.py +++ b/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 diff --git a/timemachine/cpp/src/context.cu b/timemachine/cpp/src/context.cu index 35177257c..d2050b49e 100644 --- a/timemachine/cpp/src/context.cu +++ b/timemachine/cpp/src/context.cu @@ -86,7 +86,6 @@ void Context::_ensure_local_md_intialized() { std::array, 2> Context::multiple_steps_local( const int n_steps, const std::vector &local_idxs, - const int burn_in, const int store_x_interval, const double radius, const double k, @@ -122,9 +121,6 @@ std::array, 2> Context::multiple_steps_local( std::vector> 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, 2> Context::multiple_steps_local_selection( const int n_steps, const int reference_idx, const std::vector &selection_idxs, - const int burn_in, const int store_x_interval, const double radius, const double k) { @@ -200,9 +195,6 @@ std::array, 2> Context::multiple_steps_local_selection( std::vector> 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) { diff --git a/timemachine/cpp/src/context.hpp b/timemachine/cpp/src/context.hpp index 2f95044e6..00ece8e69 100644 --- a/timemachine/cpp/src/context.hpp +++ b/timemachine/cpp/src/context.hpp @@ -30,7 +30,6 @@ class Context { std::array, 2> multiple_steps_local( const int n_steps, const std::vector &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 &selection_idxs, - const int burn_in, const int store_x_interval, const double radius, const double k); diff --git a/timemachine/cpp/src/wrap_kernels.cpp b/timemachine/cpp/src/wrap_kernels.cpp index 73a948223..9023a8c70 100644 --- a/timemachine/cpp/src/wrap_kernels.cpp +++ b/timemachine/cpp/src/wrap_kernels.cpp @@ -231,7 +231,6 @@ void declare_context(py::module &m) { [](timemachine::Context &ctxt, const int n_steps, const py::array_t &local_idxs, - const int burn_in, const int store_x_interval, const double radius, const double k, @@ -239,9 +238,6 @@ void declare_context(py::module &m) { 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, 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 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 &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, 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 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 diff --git a/timemachine/fe/free_energy.py b/timemachine/fe/free_energy.py index a8725119b..cce5d1099 100644 --- a/timemachine/fe/free_energy.py +++ b/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: diff --git a/timemachine/lib/custom_ops.pyi b/timemachine/lib/custom_ops.pyi index 3ecfa6839..fb78131ec 100644 --- a/timemachine/lib/custom_ops.pyi +++ b/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: ...