From 14cb6269785367040bb967c07c7f8df46deb27e7 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Fri, 5 Jan 2024 02:40:19 +0000 Subject: [PATCH 1/5] Expand VS tests --- .../test_openmm_equil_rfe_protocols.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py index d961c3b6a..31ef4654e 100644 --- a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py +++ b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py @@ -458,7 +458,31 @@ def test_dry_run_ligand_tip4p(benzene_system, toluene_system, with tmpdir.as_cwd(): sampler = dag_unit.run(dry=True)['debug']['sampler'] assert isinstance(sampler, MultiStateSampler) - assert sampler._factory.hybrid_system + + # Get the htf and check + # 1. The number of particles in the system are as expected + # - tot particles = old particles + new unique + # 2. The number of vs == number of waters + # 3. The vs parameters are as expected + htf = sampler._factory + + # Test 1 + old_particle_count = htf._old_system.getNumParticles() + unique_new_count = len(htf._unique_new_atoms) + hybrid_particle_count = htf.hybrid_system.getNumParticles() + assert old_particle_count + unique_new_count == hybrid_particle_count + + # Test 2 + num_waters = len( + [r for r in htf._old_topology.residues() if r.name =='HOH'] + ) + virtual_sites = [ + ix for ix in range(htf.hybrid_system.getNumParticles()) if + htf.hybrid_system.isVirtualSite(ix) + ] + assert num_waters == len(virtual_sites) + + # Test 3 @pytest.mark.flaky(reruns=3) # bad minimisation can happen From bb92a7f3580350d0b9d201e89533f2fbc7d7feac Mon Sep 17 00:00:00 2001 From: IAlibay Date: Fri, 5 Jan 2024 03:22:26 +0000 Subject: [PATCH 2/5] Add nonbond parameter checks --- .../test_openmm_equil_rfe_protocols.py | 30 ++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py index 31ef4654e..f42ed013d 100644 --- a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py +++ b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py @@ -14,7 +14,10 @@ from importlib import resources import xml.etree.ElementTree as ET -from openmm import app, XmlSerializer, MonteCarloBarostat, NonbondedForce +from openmm import ( + app, XmlSerializer, MonteCarloBarostat, + NonbondedForce, CustomNonbondedForce +) from openmm import unit as omm_unit from openmmtools.multistate.multistatesampler import MultiStateSampler import pathlib @@ -483,6 +486,31 @@ def test_dry_run_ligand_tip4p(benzene_system, toluene_system, assert num_waters == len(virtual_sites) # Test 3 + # get the standard nonbonded force - only every one + nonbond = [f for f in htf.hybrid_system.getForces() + if isinstance(f, NonbondedForce)][0] + + cust_nonbond = [f for f in htf.hybrid_system.getForces() + if isinstance(f, CustomNonbondedForce)][0] + + for entry in virtual_sites: + vs = htf.hybrid_system.getVirtualSite(entry) + vs_mass = htf.hybrid_system.getParticleMass(entry) + assert ensure_quantity(vs_mass, 'openff').m == pytest.approx(0) + vs_weights = [vs.getWeight(ix) for ix in range(vs.getNumParticles())] + np.testing.assert_allclose( + vs_weights, [0.786646558, 0.106676721, 0.106676721] + ) + c, s, e = nonbond.getParticleParameters(entry) + assert ensure_quantity(c, 'openff').m == pytest.approx(-1.04844) + assert ensure_quantity(s, 'openff').m == 1 + assert ensure_quantity(e, 'openff').m == 0 + + s1, e1, s2, e2, i, j = cust_nonbond.getParticleParameters(entry) + + assert i == j == 0 + assert s1 == s2 == 1 + assert e1 == e2 == 0 @pytest.mark.flaky(reruns=3) # bad minimisation can happen From 077456836cbbd19e29c95c32d099a3afaa067f56 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Fri, 5 Jan 2024 03:24:19 +0000 Subject: [PATCH 3/5] some in line comments --- openfe/tests/protocols/test_openmm_equil_rfe_protocols.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py index f42ed013d..6c1f851fb 100644 --- a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py +++ b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py @@ -486,13 +486,15 @@ def test_dry_run_ligand_tip4p(benzene_system, toluene_system, assert num_waters == len(virtual_sites) # Test 3 - # get the standard nonbonded force - only every one + # get the standard and custom nonbonded forces - one of each nonbond = [f for f in htf.hybrid_system.getForces() if isinstance(f, NonbondedForce)][0] cust_nonbond = [f for f in htf.hybrid_system.getForces() if isinstance(f, CustomNonbondedForce)][0] + # loop through every virtual site and check that they have the + # expected tip4p parameters for entry in virtual_sites: vs = htf.hybrid_system.getVirtualSite(entry) vs_mass = htf.hybrid_system.getParticleMass(entry) From e3430783ad2067508db41a96c74e696129be02ff Mon Sep 17 00:00:00 2001 From: IAlibay Date: Mon, 15 Jan 2024 23:00:22 +0000 Subject: [PATCH 4/5] Make fixture --- openfe/tests/protocols/conftest.py | 10 +- .../test_openmm_equil_rfe_protocols.py | 138 +++++++++++------- 2 files changed, 90 insertions(+), 58 deletions(-) diff --git a/openfe/tests/protocols/conftest.py b/openfe/tests/protocols/conftest.py index 45b25723e..050c934b5 100644 --- a/openfe/tests/protocols/conftest.py +++ b/openfe/tests/protocols/conftest.py @@ -16,7 +16,7 @@ def benzene_vacuum_system(benzene_modifications): ) -@pytest.fixture +@pytest.fixture(scope='session') def benzene_system(benzene_modifications): return openfe.ChemicalSystem( {'ligand': benzene_modifications['benzene'], @@ -27,7 +27,7 @@ def benzene_system(benzene_modifications): ) -@pytest.fixture +@pytest.fixture(scope='session') def benzene_complex_system(benzene_modifications, T4_protein_component): return openfe.ChemicalSystem( {'ligand': benzene_modifications['benzene'], @@ -38,14 +38,14 @@ def benzene_complex_system(benzene_modifications, T4_protein_component): ) -@pytest.fixture +@pytest.fixture(scope='session') def toluene_vacuum_system(benzene_modifications): return openfe.ChemicalSystem( {'ligand': benzene_modifications['toluene']}, ) -@pytest.fixture +@pytest.fixture(scope='session') def toluene_system(benzene_modifications): return openfe.ChemicalSystem( {'ligand': benzene_modifications['toluene'], @@ -67,7 +67,7 @@ def toluene_complex_system(benzene_modifications, T4_protein_component): ) -@pytest.fixture +@pytest.fixture(scope='session') def benzene_to_toluene_mapping(benzene_modifications): mapper = openfe.setup.LomapAtomMapper(element_change=False) diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py index 6c1f851fb..8708096b9 100644 --- a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py +++ b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py @@ -431,11 +431,13 @@ def test_confgen_mocked_fail(benzene_system, toluene_system, assert sampler -def test_dry_run_ligand_tip4p(benzene_system, toluene_system, - benzene_to_toluene_mapping, tmpdir): +@pytest.fixture(scope='session') +def tip4p_hybrid_factory( + benzene_system, toluene_system, + benzene_to_toluene_mapping, tmp_path_factory +): """ - Test that we can create a system with virtual sites in the - environment (waters) + Hybrid system with virtual sites in the environment (waters) """ settings = openmm_rfe.RelativeHybridTopologyProtocol.default_settings() settings.forcefield_settings.forcefields = [ @@ -458,61 +460,91 @@ def test_dry_run_ligand_tip4p(benzene_system, toluene_system, ) dag_unit = list(dag.protocol_units)[0] - with tmpdir.as_cwd(): - sampler = dag_unit.run(dry=True)['debug']['sampler'] - assert isinstance(sampler, MultiStateSampler) + shared_temp = tmp_path_factory.mktemp("tip4p_shared") + scratch_temp = tmp_path_factory.mktemp("tip4p_scratch") - # Get the htf and check - # 1. The number of particles in the system are as expected - # - tot particles = old particles + new unique - # 2. The number of vs == number of waters - # 3. The vs parameters are as expected - htf = sampler._factory + dag_unit_result = dag_unit.run( + dry=True, + scratch_basepath=scratch_temp, + shared_basepath=shared_temp, + ) + + return dag_unit_result['debug']['sampler']._factory + + +def test_tip4p_particle_count(tip4p_hybrid_factory): + """ + Check that the total number of particles in the system + are as expected. + """ + + htf = tip4p_hybrid_factory + + old_particle_count = htf._old_system.getNumParticles() + unique_new_count = len(htf._unique_new_atoms) + hybrid_particle_count = htf.hybrid_system.getNumParticles() + + assert old_particle_count + unique_new_count == hybrid_particle_count - # Test 1 - old_particle_count = htf._old_system.getNumParticles() - unique_new_count = len(htf._unique_new_atoms) - hybrid_particle_count = htf.hybrid_system.getNumParticles() - assert old_particle_count + unique_new_count == hybrid_particle_count - # Test 2 - num_waters = len( - [r for r in htf._old_topology.residues() if r.name =='HOH'] +def test_tip4p_num_waters(tip4p_hybrid_factory): + """ + Check that the number of virtual sites is equal the number of waters. + """ + + htf = tip4p_hybrid_factory + + # Test 2 + num_waters = len( + [r for r in htf._old_topology.residues() if r.name =='HOH'] + ) + virtual_sites = [ + ix for ix in range(htf.hybrid_system.getNumParticles()) if + htf.hybrid_system.isVirtualSite(ix) + ] + assert num_waters == len(virtual_sites) + + +def test_tip4p_check_vsite_parameters(tip4p_hybrid_factory): + """ + Check that the virtual site parameters are those expected + as defined by the tip4p-ew parameters in openmmforcefields + """ + + htf = tip4p_hybrid_factory + + virtual_sites = [ + ix for ix in range(htf.hybrid_system.getNumParticles()) if + htf.hybrid_system.isVirtualSite(ix) + ] + + # get the standard and custom nonbonded forces - one of each + nonbond = [f for f in htf.hybrid_system.getForces() + if isinstance(f, NonbondedForce)][0] + + cust_nonbond = [f for f in htf.hybrid_system.getForces() + if isinstance(f, CustomNonbondedForce)][0] + + # loop through every virtual site and check that they have the + # expected tip4p parameters + for entry in virtual_sites: + vs = htf.hybrid_system.getVirtualSite(entry) + vs_mass = htf.hybrid_system.getParticleMass(entry) + assert ensure_quantity(vs_mass, 'openff').m == pytest.approx(0) + vs_weights = [vs.getWeight(ix) for ix in range(vs.getNumParticles())] + np.testing.assert_allclose( + vs_weights, [0.786646558, 0.106676721, 0.106676721] ) - virtual_sites = [ - ix for ix in range(htf.hybrid_system.getNumParticles()) if - htf.hybrid_system.isVirtualSite(ix) - ] - assert num_waters == len(virtual_sites) - - # Test 3 - # get the standard and custom nonbonded forces - one of each - nonbond = [f for f in htf.hybrid_system.getForces() - if isinstance(f, NonbondedForce)][0] - - cust_nonbond = [f for f in htf.hybrid_system.getForces() - if isinstance(f, CustomNonbondedForce)][0] - - # loop through every virtual site and check that they have the - # expected tip4p parameters - for entry in virtual_sites: - vs = htf.hybrid_system.getVirtualSite(entry) - vs_mass = htf.hybrid_system.getParticleMass(entry) - assert ensure_quantity(vs_mass, 'openff').m == pytest.approx(0) - vs_weights = [vs.getWeight(ix) for ix in range(vs.getNumParticles())] - np.testing.assert_allclose( - vs_weights, [0.786646558, 0.106676721, 0.106676721] - ) - c, s, e = nonbond.getParticleParameters(entry) - assert ensure_quantity(c, 'openff').m == pytest.approx(-1.04844) - assert ensure_quantity(s, 'openff').m == 1 - assert ensure_quantity(e, 'openff').m == 0 + c, s, e = nonbond.getParticleParameters(entry) + assert ensure_quantity(c, 'openff').m == pytest.approx(-1.04844) + assert ensure_quantity(s, 'openff').m == 1 + assert ensure_quantity(e, 'openff').m == 0 - s1, e1, s2, e2, i, j = cust_nonbond.getParticleParameters(entry) + s1, e1, s2, e2, i, j = cust_nonbond.getParticleParameters(entry) - assert i == j == 0 - assert s1 == s2 == 1 - assert e1 == e2 == 0 + assert i == j == 0 + assert s1 == s2 == 1 + assert e1 == e2 == 0 @pytest.mark.flaky(reruns=3) # bad minimisation can happen From 21a899f36a0387f7d0e6fcd5f8c79f12bba4cd6b Mon Sep 17 00:00:00 2001 From: IAlibay Date: Mon, 15 Jan 2024 23:09:11 +0000 Subject: [PATCH 5/5] clean up scopes --- openfe/tests/protocols/conftest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openfe/tests/protocols/conftest.py b/openfe/tests/protocols/conftest.py index 050c934b5..0d5fae0fa 100644 --- a/openfe/tests/protocols/conftest.py +++ b/openfe/tests/protocols/conftest.py @@ -27,7 +27,7 @@ def benzene_system(benzene_modifications): ) -@pytest.fixture(scope='session') +@pytest.fixture def benzene_complex_system(benzene_modifications, T4_protein_component): return openfe.ChemicalSystem( {'ligand': benzene_modifications['benzene'], @@ -38,7 +38,7 @@ def benzene_complex_system(benzene_modifications, T4_protein_component): ) -@pytest.fixture(scope='session') +@pytest.fixture def toluene_vacuum_system(benzene_modifications): return openfe.ChemicalSystem( {'ligand': benzene_modifications['toluene']},