From bccfb233c6403de00cc681bc2713d6c84b5a18d1 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 20 Nov 2023 15:25:40 +0100 Subject: [PATCH 1/6] Change get_simsteps to handle NVT equilibration in plain MD protocol --- openfe/protocols/openmm_afe/base.py | 10 +++-- .../protocols/openmm_rfe/equil_rfe_methods.py | 11 +++-- .../openmm_utils/settings_validation.py | 43 ++++++++----------- openfe/tests/protocols/test_openmmutils.py | 19 ++++---- 4 files changed, 41 insertions(+), 42 deletions(-) diff --git a/openfe/protocols/openmm_afe/base.py b/openfe/protocols/openmm_afe/base.py index 662f049cd..dd1616c7f 100644 --- a/openfe/protocols/openmm_afe/base.py +++ b/openfe/protocols/openmm_afe/base.py @@ -746,9 +746,13 @@ def _run_simulation( # Get the relevant simulation steps mc_steps = settings['integrator_settings'].n_steps.m - equil_steps, prod_steps = settings_validation.get_simsteps( - equil_length=settings['simulation_settings'].equilibration_length, - prod_length=settings['simulation_settings'].production_length, + equil_steps = settings_validation.get_simsteps( + sim_length=settings['simulation_settings'].equilibration_length, + timestep=settings['integrator_settings'].timestep, + mc_steps=mc_steps, + ) + prod_steps = settings_validation.get_simsteps( + sim_length=settings['simulation_settings'].production_length, timestep=settings['integrator_settings'].timestep, mc_steps=mc_steps, ) diff --git a/openfe/protocols/openmm_rfe/equil_rfe_methods.py b/openfe/protocols/openmm_rfe/equil_rfe_methods.py index 8086aaa72..97035386e 100644 --- a/openfe/protocols/openmm_rfe/equil_rfe_methods.py +++ b/openfe/protocols/openmm_rfe/equil_rfe_methods.py @@ -637,10 +637,13 @@ def run(self, *, dry=False, verbose=True, settings_validation.validate_timestep( forcefield_settings.hydrogen_mass, timestep ) - equil_steps, prod_steps = settings_validation.get_simsteps( - equil_length=sim_settings.equilibration_length, - prod_length=sim_settings.production_length, - timestep=timestep, mc_steps=mc_steps + equil_steps = settings_validation.get_simsteps( + sim_length=sim_settings.equilibration_length, + timestep=timestep, mc_steps=mc_steps, + ) + prod_steps = settings_validation.get_simsteps( + sim_length=sim_settings.production_length, + timestep=timestep, mc_steps=mc_steps, ) solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) diff --git a/openfe/protocols/openmm_utils/settings_validation.py b/openfe/protocols/openmm_utils/settings_validation.py index 5745a178d..2deb2f90d 100644 --- a/openfe/protocols/openmm_utils/settings_validation.py +++ b/openfe/protocols/openmm_utils/settings_validation.py @@ -37,17 +37,15 @@ def validate_timestep(hmass: float, timestep: unit.Quantity): raise ValueError(errmsg) -def get_simsteps(equil_length: unit.Quantity, prod_length: unit.Quantity, - timestep: unit.Quantity, mc_steps: int) -> Tuple[int, int]: +def get_simsteps(sim_length: unit.Quantity, + timestep: unit.Quantity, mc_steps: int) -> int: """ - Gets and validates the number of equilibration and production steps. + Gets and validates the number of simulation steps. Parameters ---------- - equil_length : unit.Quantity - Simulation equilibration length. - prod_length : unit.Quantity - Simulation production length. + sim_length : unit.Quantity + Simulation length. timestep : unit.Quantity Integration timestep. mc_steps : int @@ -55,29 +53,22 @@ def get_simsteps(equil_length: unit.Quantity, prod_length: unit.Quantity, Returns ------- - equil_steps : int - The number of equilibration timesteps. - prod_steps : int - The number of production timesteps. + sim_steps : int + The number of simulation timesteps. """ - equil_time = round(equil_length.to('attosecond').m) - prod_time = round(prod_length.to('attosecond').m) + sim_time = round(sim_length.to('attosecond').m) ts = round(timestep.to('attosecond').m) - equil_steps, mod = divmod(equil_time, ts) + sim_steps, mod = divmod(sim_time, ts) if mod != 0: - raise ValueError("Equilibration time not divisible by timestep") - prod_steps, mod = divmod(prod_time, ts) - if mod != 0: - raise ValueError("Production time not divisible by timestep") + raise ValueError("Simulation time not divisible by timestep") - for var in [("Equilibration", equil_steps, equil_time), - ("Production", prod_steps, prod_time)]: - if (var[1] % mc_steps) != 0: - errmsg = (f"{var[0]} time {var[2]/1000000} ps should contain a " - "number of steps divisible by the number of integrator " - f"timesteps between MC moves {mc_steps}") - raise ValueError(errmsg) + var = ["Simulation", sim_steps, sim_time] + if (var[1] % mc_steps) != 0: + errmsg = (f"{var[0]} time {var[2]/1000000} ps should contain a " + "number of steps divisible by the number of integrator " + f"timesteps between MC moves {mc_steps}") + raise ValueError(errmsg) - return equil_steps, prod_steps + return sim_steps diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py index 7147f82d4..72be4136f 100644 --- a/openfe/tests/protocols/test_openmmutils.py +++ b/openfe/tests/protocols/test_openmmutils.py @@ -27,17 +27,18 @@ def test_validate_timestep(): settings_validation.validate_timestep(2.0, 4.0 * unit.femtoseconds) -@pytest.mark.parametrize('e,p,ts,mc,es,ps', [ - [1 * unit.nanoseconds, 5 * unit.nanoseconds, 4 * unit.femtoseconds, - 250, 250000, 1250000], - [1 * unit.picoseconds, 1 * unit.picoseconds, 2 * unit.femtoseconds, - 250, 500, 500], +@pytest.mark.parametrize('s,ts,mc,es', [ + [5 * unit.nanoseconds, 4 * unit.femtoseconds, + 250, 1250000], + [1 * unit.nanoseconds, 4 * unit.femtoseconds, + 250, 250000], + [1 * unit.picoseconds, 2 * unit.femtoseconds, + 250, 500], ]) -def test_get_simsteps(e, p, ts, mc, es, ps): - equil_steps, prod_steps = settings_validation.get_simsteps(e, p, ts, mc) +def test_get_simsteps(s, ts, mc, es): + sim_steps = settings_validation.get_simsteps(s, ts, mc) - assert equil_steps == es - assert prod_steps == ps + assert sim_steps == es @pytest.mark.parametrize('nametype, timelengths', [ From de12b3c1b18220506070a1972d3aacf6025d3637 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 20 Nov 2023 15:28:22 +0100 Subject: [PATCH 2/6] Small fix --- openfe/tests/protocols/test_openmmutils.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py index 72be4136f..f554ba55d 100644 --- a/openfe/tests/protocols/test_openmmutils.py +++ b/openfe/tests/protocols/test_openmmutils.py @@ -28,12 +28,9 @@ def test_validate_timestep(): @pytest.mark.parametrize('s,ts,mc,es', [ - [5 * unit.nanoseconds, 4 * unit.femtoseconds, - 250, 1250000], - [1 * unit.nanoseconds, 4 * unit.femtoseconds, - 250, 250000], - [1 * unit.picoseconds, 2 * unit.femtoseconds, - 250, 500], + [5 * unit.nanoseconds, 4 * unit.femtoseconds, 250, 1250000], + [1 * unit.nanoseconds, 4 * unit.femtoseconds, 250, 250000], + [1 * unit.picoseconds, 2 * unit.femtoseconds, 250, 500], ]) def test_get_simsteps(s, ts, mc, es): sim_steps = settings_validation.get_simsteps(s, ts, mc) From 2ac7bf12e8c58aa87f4d1836c6008fcfcf92fc2b Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 21 Nov 2023 09:56:28 +0100 Subject: [PATCH 3/6] Fix get_simsteps tests --- openfe/tests/protocols/test_openmmutils.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py index f554ba55d..4f641e37f 100644 --- a/openfe/tests/protocols/test_openmmutils.py +++ b/openfe/tests/protocols/test_openmmutils.py @@ -39,29 +39,25 @@ def test_get_simsteps(s, ts, mc, es): @pytest.mark.parametrize('nametype, timelengths', [ - ['Equilibration', [1.003 * unit.picoseconds, 1 * unit.picoseconds]], - ['Production', [1 * unit.picoseconds, 1.003 * unit.picoseconds]], + ['Simulation', [1.003 * unit.picoseconds, 1 * unit.picoseconds]], ]) def test_get_simsteps_indivisible_simtime(nametype, timelengths): errmsg = f"{nametype} time not divisible by timestep" with pytest.raises(ValueError, match=errmsg): settings_validation.get_simsteps( timelengths[0], - timelengths[1], 2 * unit.femtoseconds, 100) @pytest.mark.parametrize('nametype, timelengths', [ - ['Equilibration', [1 * unit.picoseconds, 10 * unit.picoseconds]], - ['Production', [10 * unit.picoseconds, 1 * unit.picoseconds]], + ['Simulation', [1 * unit.picoseconds, 10 * unit.picoseconds]], ]) def test_mc_indivisible(nametype, timelengths): errmsg = f"{nametype} time 1.0 ps should contain" with pytest.raises(ValueError, match=errmsg): settings_validation.get_simsteps( - timelengths[0], timelengths[1], - 2 * unit.femtoseconds, 1000) + timelengths[0], 2 * unit.femtoseconds, 1000) def test_get_alchemical_components(benzene_modifications, From 85115ff0a4d98eb73c3a56ba3f5cd71bd247563c Mon Sep 17 00:00:00 2001 From: Hannah Baumann <43765638+hannahbaumann@users.noreply.github.com> Date: Wed, 22 Nov 2023 15:18:10 +0100 Subject: [PATCH 4/6] Update openfe/protocols/openmm_utils/settings_validation.py Co-authored-by: Irfan Alibay --- openfe/protocols/openmm_utils/settings_validation.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/openfe/protocols/openmm_utils/settings_validation.py b/openfe/protocols/openmm_utils/settings_validation.py index 2deb2f90d..976bc4963 100644 --- a/openfe/protocols/openmm_utils/settings_validation.py +++ b/openfe/protocols/openmm_utils/settings_validation.py @@ -64,9 +64,8 @@ def get_simsteps(sim_length: unit.Quantity, if mod != 0: raise ValueError("Simulation time not divisible by timestep") - var = ["Simulation", sim_steps, sim_time] - if (var[1] % mc_steps) != 0: - errmsg = (f"{var[0]} time {var[2]/1000000} ps should contain a " + if (sim_steps % mc_steps) != 0: + errmsg = (f"Simulation time {sim_time/1000000} ps should contain a " "number of steps divisible by the number of integrator " f"timesteps between MC moves {mc_steps}") raise ValueError(errmsg) From e88dfa5b591a1ef3e2329a4f4ec2d7dae504db22 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 15:30:13 +0100 Subject: [PATCH 5/6] remove unnecessary fixture --- openfe/tests/protocols/test_openmmutils.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py index 4f641e37f..b96245b40 100644 --- a/openfe/tests/protocols/test_openmmutils.py +++ b/openfe/tests/protocols/test_openmmutils.py @@ -38,16 +38,11 @@ def test_get_simsteps(s, ts, mc, es): assert sim_steps == es -@pytest.mark.parametrize('nametype, timelengths', [ - ['Simulation', [1.003 * unit.picoseconds, 1 * unit.picoseconds]], -]) -def test_get_simsteps_indivisible_simtime(nametype, timelengths): - errmsg = f"{nametype} time not divisible by timestep" +def test_get_simsteps_indivisible_simtime(): + errmsg = "Simulation time not divisible by timestep" + timelengths = 1.003 * unit.picosecond with pytest.raises(ValueError, match=errmsg): - settings_validation.get_simsteps( - timelengths[0], - 2 * unit.femtoseconds, - 100) + settings_validation.get_simsteps(timelengths, 2 * unit.femtoseconds, 100) @pytest.mark.parametrize('nametype, timelengths', [ @@ -84,7 +79,7 @@ def test_get_alchemical_components(benzene_modifications, def test_duplicate_chemical_components(benzene_modifications): stateA = openfe.ChemicalSystem({'A': benzene_modifications['toluene'], - 'B': benzene_modifications['toluene'],}) + 'B': benzene_modifications['toluene'], }) stateB = openfe.ChemicalSystem({'A': benzene_modifications['toluene']}) errmsg = "state A components B:" @@ -133,7 +128,7 @@ def test_multiple_proteins(T4_protein_component): def test_get_components_gas(benzene_modifications): state = openfe.ChemicalSystem({'A': benzene_modifications['benzene'], - 'B': benzene_modifications['toluene'],}) + 'B': benzene_modifications['toluene'], }) s, p, mols = system_validation.get_components(state) @@ -146,7 +141,7 @@ def test_components_solvent(benzene_modifications): state = openfe.ChemicalSystem({'S': openfe.SolventComponent(), 'A': benzene_modifications['benzene'], - 'B': benzene_modifications['toluene'],}) + 'B': benzene_modifications['toluene'], }) s, p, mols = system_validation.get_components(state) From c45f85aa46a3134537782d509777ce78bfaba6a7 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 22 Nov 2023 15:33:21 +0100 Subject: [PATCH 6/6] Remove unnessesary fixtures 2 --- openfe/tests/protocols/test_openmmutils.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py index b96245b40..fb80e90fc 100644 --- a/openfe/tests/protocols/test_openmmutils.py +++ b/openfe/tests/protocols/test_openmmutils.py @@ -40,19 +40,17 @@ def test_get_simsteps(s, ts, mc, es): def test_get_simsteps_indivisible_simtime(): errmsg = "Simulation time not divisible by timestep" - timelengths = 1.003 * unit.picosecond + timelength = 1.003 * unit.picosecond with pytest.raises(ValueError, match=errmsg): - settings_validation.get_simsteps(timelengths, 2 * unit.femtoseconds, 100) + settings_validation.get_simsteps(timelength, 2 * unit.femtoseconds, 100) -@pytest.mark.parametrize('nametype, timelengths', [ - ['Simulation', [1 * unit.picoseconds, 10 * unit.picoseconds]], -]) -def test_mc_indivisible(nametype, timelengths): - errmsg = f"{nametype} time 1.0 ps should contain" +def test_mc_indivisible(): + errmsg = "Simulation time 1.0 ps should contain" + timelength = 1 * unit.picoseconds with pytest.raises(ValueError, match=errmsg): settings_validation.get_simsteps( - timelengths[0], 2 * unit.femtoseconds, 1000) + timelength, 2 * unit.femtoseconds, 1000) def test_get_alchemical_components(benzene_modifications,