make sure we don't install the newest toolkit by mistake (#237)

* make sure we don't install the newest toolkit by mistake * make sure we don't pull in a really old version of the toolkit * forgot to add a verb there * Start to make internals compatible with both unit packages * Try different logic * Improve logic in test * Fix some tests * Fix * Fix * Better units logic in tests * Remove temporary code * That's too semantic a variable name Co-authored-by: Matthew W. Thompson <[email protected]>
openmm · Sep 20, 2022 · 33d500e · 33d500e
1 parent 35496e5
commit 33d500e
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 41 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -57,6 +57,11 @@ jobs:
       run: |
         micromamba update -y -c conda-forge "openff-toolkit >=0.11.0"
 
+    - name: "Install the older  openff-toolkit"
+      if: ${{ matrix.latest-openff-toolkit == false }}
+      run: |
+        micromamba install -y -c conda-forge "openff-toolkit==0.10.6" "openff-toolkit-base==0.10.6"
+
     - name: Install Package
       run: |
         pip list

diff --git a/devtools/conda-envs/test_env.yaml b/devtools/conda-envs/test_env.yaml
@@ -19,6 +19,8 @@ dependencies:
     # Requirements for converted force field installer
   - openmm >=7.6.0
 
+  - openff-units >=0.1.8
+
     # Requirements for conversion tools
   - pyyaml
   - ambertools >=18.0 # contains sufficiently recent ParmEd

diff --git a/openmmforcefields/generators/template_generators.py b/openmmforcefields/generators/template_generators.py
@@ -219,13 +219,13 @@ def _molecule_has_user_charges(self, molecule):
 
         """
         import numpy as np
-        from openff.units import unit
+        from openff.units.openmm import ensure_quantity
 
         if molecule.partial_charges is None:
             return False
 
         if np.allclose(
-            molecule.partial_charges.m_as(unit.elementary_charge),
+            ensure_quantity(molecule.partial_charges, "openff").m,
             np.zeros([molecule.n_atoms]),
         ):
             return False
@@ -575,6 +575,14 @@ def generate_residue_template(self, molecule, residue_atoms=None):
         """
         import numpy as np
         from openff.units import unit
+        from openff.units.openmm import ensure_quantity
+
+        uses_old_api = hasattr(molecule.atoms[0], "element")
+
+        if uses_old_api:
+            unit_solution = "openmm"
+        else:
+            unit_solution = "openff"
 
         # Use the canonical isomeric SMILES to uniquely name the template
         smiles = molecule.to_smiles()
@@ -586,10 +594,6 @@ def generate_residue_template(self, molecule, residue_atoms=None):
         # Compute net formal charge
         net_charge = molecule.total_charge
 
-        if type(net_charge) != unit.Quantity:
-            # openforcefield toolkit < 0.7.0 did not return unit-bearing quantity
-            # how long should openmmforcefields support < 0.7.0?
-            net_charge = float(net_charge) * unit.elementary_charge
         _logger.debug(f'Total charge is {net_charge}')
 
         # Compute partial charges if required
@@ -637,11 +641,26 @@ def generate_residue_template(self, molecule, residue_atoms=None):
         #       or pure numbers.
         _logger.debug(f'Fixing partial charges...')
         _logger.debug(f'{molecule.partial_charges}')
-        residue_charge = 0.0 * unit.elementary_charge
+        residue_charge = ensure_quantity(0.0 * unit.elementary_charge, unit_solution)
         total_charge = molecule.partial_charges.sum()
+
         sum_of_absolute_charge = np.sum(np.abs(molecule.partial_charges))
+
+        if uses_old_api:
+            from openmm import unit as openmm_unit
+
+            redistribute = sum_of_absolute_charge > 0.0
+
+            sum_of_absolute_charge = openmm_unit.Quantity(
+                sum_of_absolute_charge,
+                openmm_unit.elementary_charge,
+            )
+        else:
+            redistribute = sum_of_absolute_charge.m > 0.0
+
         charge_deficit = net_charge - total_charge
-        if (sum_of_absolute_charge / unit.elementary_charge).m > 0.0:
+
+        if redistribute:
             # Redistribute excess charge proportionally to absolute charge
             molecule.partial_charges = molecule.partial_charges + charge_deficit * abs(molecule.partial_charges) / sum_of_absolute_charge
         _logger.debug(f'{molecule.partial_charges}')
@@ -675,7 +694,24 @@ def generate_residue_template(self, molecule, residue_atoms=None):
         residues = etree.SubElement(root, "Residues")
         residue = etree.SubElement(residues, "Residue", name=smiles)
         for atom in molecule.atoms:
-            atom = etree.SubElement(residue, "Atom", name=atom.name, type=atom.gaff_type, charge=str(atom.partial_charge.m_as(unit.elementary_charge)))
+
+            if uses_old_api:
+                charge_string =str(
+                    atom.partial_charge.value_in_unit(openmm_unit.elementary_charge)
+                )
+            else:
+                charge_string = str(
+                    atom.partial_charge.m_as(unit.elementary_charge)
+                )
+
+            atom = etree.SubElement(
+                residue,
+                "Atom",
+                name=atom.name,
+                type=atom.gaff_type,
+                charge=charge_string,
+            )
+
         for bond in molecule.bonds:
             if (bond.atom1 in residue_atoms) and (bond.atom2 in residue_atoms):
                 bond = etree.SubElement(residue, "Bond", atomName1=bond.atom1.name, atomName2=bond.atom2.name)

diff --git a/openmmforcefields/tests/test_template_generators.py b/openmmforcefields/tests/test_template_generators.py
@@ -6,10 +6,8 @@
 
 import numpy as np
 import openmm
-from openff.toolkit import __version__ as toolkit_version
 from openff.toolkit.topology import Molecule
 from openmm.app import PME, ForceField, Modeller, NoCutoff, PDBFile
-from packaging.version import Version
 
 from openmmforcefields.generators import (
     EspalomaTemplateGenerator,
@@ -58,6 +56,7 @@ def filter_molecules(self, molecules):
 
     def setUp(self):
         from openff.units import unit
+        from openff.units.openmm import ensure_quantity
 
         # TODO: Harmonize with test_system_generator.py infrastructure
         # Read test molecules
@@ -68,7 +67,13 @@ def setUp(self):
         molecule = Molecule.from_smiles('C=O')
         molecule.generate_conformers(n_conformers=1)
 
-        molecule.conformers[0][0,0] += 0.1*unit.angstroms
+        uses_old_api = hasattr(molecule.atoms[0], "element")
+
+        molecule.conformers[0][0,0] += ensure_quantity(
+            unit.Quantity(0.1, unit.angstroms),
+            "openmm" if uses_old_api else "openff",
+        )
+
         molecules.insert(0, molecule)
         # DEBUG END
 
@@ -137,8 +142,11 @@ def test_add_molecules(self):
             system = forcefield.createSystem(openmm_topology, nonbondedMethod=NoCutoff)
         except Exception as e:
             print(forcefield._atomTypes.keys())
-            from openff.units.openmm import to_openmm
-            PDBFile.writeFile(openmm_topology, to_openmm(molecule.conformers[0]))
+            from openff.units.openmm import ensure_quantity
+            PDBFile.writeFile(
+                openmm_topology,
+                ensure_quantity(molecule.conformers[0], "openmm"),
+            )
             raise e
         assert system.getNumParticles() == molecule.n_atoms
 
@@ -191,11 +199,17 @@ def charges_are_equal(self, system, molecule):
             True if the partial charges are equal, False if not
         """
         from openff.units import unit
+        from openff.units.openmm import ensure_quantity
 
         assert system.getNumParticles() == molecule.n_atoms
 
+        # charges_from_system returns a NumPy array that we trust to be implicitly e
         system_charges: np.ndarray = self.charges_from_system(system)
-        molecule_charges: np.ndarray = molecule.partial_charges.m_as(unit.elementary_charge)
+
+        # type(molecule.partial_charges) depends on the toolkit version
+        molecule_charges: np.ndarray = ensure_quantity(
+            molecule.partial_charges, "openff",
+        ).m_as(unit.elementary_charge)
 
         result = np.allclose(system_charges, molecule_charges)
 
@@ -230,10 +244,7 @@ def test_charge(self):
 
     def test_charge_from_molecules(self):
         """Test that user-specified partial charges are used if requested"""
-        from openff.units import unit
-
-        if Version(toolkit_version) < Version("0.11.0"):
-            self.skipTest("Test written with new toolkit API")
+        from openff.units.openmm import ensure_quantity
 
         # Create a generator that does not know about any molecules
         generator = self.TEMPLATE_GENERATOR()
@@ -245,14 +256,22 @@ def test_charge_from_molecules(self):
         # Check that parameterizing a molecule using user-provided charges produces expected charges
 
         molecule = self.molecules[0]
+        uses_old_api = hasattr(molecule.atoms[0], "element")
+
+        if uses_old_api:
+            from openmm import unit
+        else:
+            from openff.units import unit
 
         # Populate the molecule with arbitrary partial charges that still sum to 0.0
         molecule.partial_charges = unit.Quantity(
             np.linspace(-0.5, 0.5, molecule.n_atoms),
             unit.elementary_charge,
         )
 
-        assert (molecule.partial_charges is not None) and not np.all(molecule.partial_charges.m_as(unit.elementary_charge) == 0)
+        assert (molecule.partial_charges is not None)
+
+        assert not np.all(ensure_quantity(molecule.partial_charges, "openff").m == 0)
 
         generator.add_molecules(molecule)
 
@@ -350,11 +369,11 @@ def check_cache(generator, n_expected):
     def test_add_solvent(self):
         """Test using openmm.app.Modeller to add solvent to a small molecule parameterized by template generator"""
         # Select a molecule to add solvent around
-        from openff.units.openmm import to_openmm
+        from openff.units.openmm import ensure_quantity
         from openmm import unit
         molecule = self.molecules[0]
         openmm_topology = molecule.to_topology().to_openmm()
-        openmm_positions = to_openmm(molecule.conformers[0])
+        openmm_positions = ensure_quantity(molecule.conformers[0], "openmm")
         # Try adding solvent without residue template generator; this will fail
         forcefield = ForceField('tip3p.xml')
         # Add solvent to a system containing a small molecule
@@ -613,16 +632,16 @@ def compare_energies(cls, molecule, template_generated_system, reference_system)
             System generated by reference parmaeterization engine
 
         """
-        from openff.units.openmm import to_openmm
+        from openff.units.openmm import ensure_quantity
 
         # Compute energies
         reference_energy, reference_forces = cls.compute_energy(
             template_generated_system,
-            to_openmm(molecule.conformers[0]),
+            ensure_quantity(molecule.conformers[0], "openmm"),
         )
         template_energy, template_forces = cls.compute_energy(
             reference_system,
-            to_openmm(molecule.conformers[0]),
+            ensure_quantity(molecule.conformers[0], "openmm"),
         )
 
         from openmm import unit
@@ -708,21 +727,29 @@ def propagate_dynamics(self, molecule, system):
 
         """
         # Run some dynamics
-        from openff.units.openmm import from_openmm as from_openmm_quantity
-        from openff.units.openmm import to_openmm as to_openmm_quantity
+        from openff.units.openmm import ensure_quantity
         from openmm import unit
+
+        uses_old_api = hasattr(molecule.atoms[0], "element")
+
         temperature = 300 * unit.kelvin
         collision_rate = 1.0 / unit.picoseconds
         timestep = 1.0 * unit.femtoseconds
         nsteps = 100
         integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
         platform = openmm.Platform.getPlatformByName('Reference')
         context = openmm.Context(system, integrator, platform)
-        context.setPositions(to_openmm_quantity(molecule.conformers[0]))
+        context.setPositions(ensure_quantity(molecule.conformers[0], "openmm"))
         integrator.step(nsteps)
         # Copy the molecule, storing new conformer
         new_molecule = copy.deepcopy(molecule)
-        new_molecule.conformers[0] = from_openmm_quantity(context.getState(getPositions=True).getPositions())
+        new_positions = context.getState(getPositions=True).getPositions()
+
+        new_molecule.conformers[0] = ensure_quantity(
+            new_positions,
+            "openmm" if uses_old_api else "openff",
+        )
+
 
         del context, integrator
 
@@ -750,9 +777,6 @@ def test_INSTALLED_FORCEFIELDS(self):
     def test_energies(self):
         """Test potential energies match between openff-toolkit and OpenMM ForceField"""
 
-        if Version(toolkit_version) < Version("0.11.0"):
-            self.skipTest("Test written with new toolkit API")
-
         # Test all supported SMIRNOFF force fields
         for small_molecule_forcefield in SMIRNOFFTemplateGenerator.INSTALLED_FORCEFIELDS:
             if "ff14sb" in small_molecule_forcefield:
@@ -835,12 +859,10 @@ def propagate_dynamics(self, molecule, system):
 
         """
         # Run some dynamics
-        from openff.units.openmm import from_openmm as from_openmm_quantity
-        from openff.units.openmm import to_openmm as to_openmm_quantity
+        from openff.units.openmm import ensure_quantity
         from openmm import unit
 
-        if Version(toolkit_version) < Version("0.11.0"):
-            self.skipTest("Test written with new toolkit API")
+        uses_old_api = hasattr(molecule.atoms[0], "element")
 
         temperature = 300 * unit.kelvin
         collision_rate = 1.0 / unit.picoseconds
@@ -849,11 +871,17 @@ def propagate_dynamics(self, molecule, system):
         integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
         platform = openmm.Platform.getPlatformByName('Reference')
         context = openmm.Context(system, integrator, platform)
-        context.setPositions(to_openmm_quantity(molecule.conformers[0]))
+        context.setPositions(ensure_quantity(molecule.conformers[0], "openmm"))
         integrator.step(nsteps)
         # Copy the molecule, storing new conformer
         new_molecule = copy.deepcopy(molecule)
-        new_molecule.conformers[0] = from_openmm_quantity(context.getState(getPositions=True).getPositions())
+        new_positions = context.getState(getPositions=True).getPositions()
+
+        new_molecule.conformers[0] = ensure_quantity(
+            new_positions,
+            "openmm" if uses_old_api else "openff",
+        )
+
         # Clean up
         del context, integrator
 
@@ -880,9 +908,6 @@ def test_retrieve_forcefields(self):
     def test_energies(self):
         """Test potential energies match between openff-toolkit and OpenMM ForceField"""
 
-        if Version(toolkit_version) < Version("0.11.0"):
-            self.skipTest("Test written with new toolkit API")
-
         # Test all supported SMIRNOFF force fields
         for small_molecule_forcefield in EspalomaTemplateGenerator.INSTALLED_FORCEFIELDS:
             print(f'Testing energies for {small_molecule_forcefield}...')