Test AHFE user charges

openfe/protocols/openmm_afe/base.py

@@ -324,16 +324,15 @@ def _get_modeller(
         # Note by default this is cached to ctx.shared/db.json which should
         # reduce some of the costs.
         for mol in smc_components.values():
-            if mol.partial_charges is not None and np.any(mol.partial_charges):
-                # skip if we have existing partial charges
-                continue
-            try:
-                # try and follow official spec method
-                mol.assign_partial_charges('am1bcc')
-            except ValueError:  # this is what a confgen failure yields
-                # but fallback to using existing conformer
-                mol.assign_partial_charges('am1bcc',
-                                           use_conformers=mol.conformers)
+            # don't do this if we have user charges
+            if not (mol.partial_charges is not None and np.any(mol.partial_charges)):
+                try:
+                    # try and follow official spec method
+                    mol.assign_partial_charges('am1bcc')
+                except ValueError:  # this is what a confgen failure yields
+                    # but fallback to using existing conformer
+                    mol.assign_partial_charges('am1bcc',
+                                               use_conformers=mol.conformers)
 
             system_generator.create_system(
                 mol.to_topology().to_openmm(), molecules=[mol]

openfe/tests/protocols/test_openmm_afe_solvation_protocol.py

@@ -4,8 +4,10 @@
 import json
 import pytest
 from unittest import mock
+from openmm import NonbondedForce, CustomNonbondedForce
 from openmmtools.multistate.multistatesampler import MultiStateSampler
 from openff.units import unit as offunit
+from openff.units.openmm import ensure_quantity
 import mdtraj as mdt
 import numpy as np
 import gufe
@@ -357,6 +359,92 @@ def test_dry_run_solv_benzene_tip4p(benzene_modifications, tmpdir):
         assert sol_sampler.is_periodic
 
 
+def test_dry_run_solv_user_charges_benzene(benzene_modifications, tmpdir):
+    """
+    Create a test system with fictitious user supplied charges and
+    ensure that they are properly passed through to the constructed
+    alchemical system.
+    """
+    s = openmm_afe.AbsoluteSolvationProtocol.default_settings()
+    s.alchemsampler_settings.n_repeats = 1
+
+    protocol = openmm_afe.AbsoluteSolvationProtocol(
+            settings=s,
+    )
+
+    def assign_fictitious_charges(offmol):
+        """
+        Get a random array of fake partial charges for your offmol.
+        """
+        rand_arr = np.random.randint(1, 10, size=offmol.n_atoms) / 100
+        rand_arr[-1] = -sum(rand_arr[:-1])
+        return rand_arr * offunit.elementary_charge
+
+    benzene_offmol = benzene_modifications['benzene'].to_openff()
+    offmol_pchgs = assign_fictitious_charges(benzene_offmol)
+    benzene_offmol.partial_charges = offmol_pchgs
+    benzene_smc = openfe.SmallMoleculeComponent.from_openff(benzene_offmol)
+
+    # check propchgs
+    prop_chgs = benzene_smc.to_dict()['molprops']['atom.dprop.PartialCharge']
+    prop_chgs = np.array(prop_chgs.split(), dtype=float)
+    np.testing.assert_allclose(prop_chgs, offmol_pchgs)
+
+    # Create ChemicalSystems
+    stateA = ChemicalSystem({
+        'benzene': benzene_smc,
+        'solvent': SolventComponent()
+    })
+
+    stateB = ChemicalSystem({
+        'solvent': SolventComponent(),
+    })
+
+    # Create DAG from protocol, get the vacuum and solvent units
+    # and eventually dry run the first solvent unit
+    dag = protocol.create(stateA=stateA, stateB=stateB, mapping=None,)
+    prot_units = list(dag.protocol_units)
+
+    vac_unit = [u for u in prot_units
+                if isinstance(u, AbsoluteSolvationVacuumUnit)][0]
+    sol_unit = [u for u in prot_units
+                if isinstance(u, AbsoluteSolvationSolventUnit)][0]
+
+    # check sol_unit charges
+    with tmpdir.as_cwd():
+        sampler = sol_unit.run(dry=True)['debug']['sampler']
+        system = sampler._thermodynamic_states[0].system
+        nonbond = [f for f in system.getForces()
+                   if isinstance(f, NonbondedForce)]
+
+        assert len(nonbond) == 1
+
+        # loop through the 12 benzene atoms
+        # partial charge is stored in the offset
+        for i in range(12):
+            offsets = nonbond[0].getParticleParameterOffset(i)
+            c = ensure_quantity(offsets[2], 'openff')
+            assert pytest.approx(c) == prop_chgs[i]
+
+    # check vac_unit charges
+    with tmpdir.as_cwd():
+        sampler = vac_unit.run(dry=True)['debug']['sampler']
+        system = sampler._thermodynamic_states[0].system
+        nonbond = [f for f in system.getForces()
+                   if isinstance(f, CustomNonbondedForce)]
+        assert len(nonbond) == 4
+
+        custom_elec = [
+            n for n in nonbond if
+            n.getGlobalParameterName(0) == 'lambda_electrostatics'][0]
+
+        # loop through the 12 benzene atoms
+        for i in range(12):
+            c, s = custom_elec.getParticleParameters(i)
+            c = ensure_quantity(c, 'openff')
+            assert pytest.approx(c) == prop_chgs[i]
+
+
 def test_nreplicas_lambda_mismatch(benzene_modifications, tmpdir):
     s = AbsoluteSolvationProtocol.default_settings()
     s.alchemsampler_settings.n_repeats = 1

openfe/tests/protocols/test_openmm_equil_rfe_protocols.py

@@ -481,7 +481,7 @@ def assign_fictitious_charges(offmol):
         create a molecule that has a total charge that is different from
         the expected formal charge, hence we enforce a zero charge here.
         """
-        rand_arr = np.random.randint(1, 10, size=offmol.n_atoms)
+        rand_arr = np.random.randint(1, 10, size=offmol.n_atoms) / 100
         rand_arr[-1] = -sum(rand_arr[:-1])
         return rand_arr * unit.elementary_charge
 
@@ -497,12 +497,10 @@ def check_propchgs(smc, charge_array):
     # Create new smc with overriden charges
     benzene_offmol = benzene_modifications['benzene'].to_openff()
     toluene_offmol = benzene_modifications['toluene'].to_openff()
-    benzene_offmol.assign_partial_charges(partial_charge_method='am1bcc')
-    toluene_offmol.assign_partial_charges(partial_charge_method='am1bcc')
     benzene_rand_chg = assign_fictitious_charges(benzene_offmol)
     toluene_rand_chg = assign_fictitious_charges(toluene_offmol)
-    benzene_offmol.partial_charges[:] = benzene_rand_chg
-    toluene_offmol.partial_charges[:] = toluene_rand_chg
+    benzene_offmol.partial_charges = benzene_rand_chg
+    toluene_offmol.partial_charges = toluene_rand_chg
     benzene_smc = openfe.SmallMoleculeComponent.from_openff(benzene_offmol)
     toluene_smc = openfe.SmallMoleculeComponent.from_openff(toluene_offmol)