Expand water virtual site tests

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'],
@@ -45,7 +45,7 @@ def toluene_vacuum_system(benzene_modifications):
     )
 
 
-@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)
 

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
@@ -428,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 = [
@@ -455,10 +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)
-        assert sampler._factory.hybrid_system
+    shared_temp = tmp_path_factory.mktemp("tip4p_shared")
+    scratch_temp = tmp_path_factory.mktemp("tip4p_scratch")
+
+    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
+
+
+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]
+        )
+        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