Support Pint 0.24.4 / openff-units 0.3.0

.github/workflows/CI.yml

@@ -17,15 +17,15 @@ defaults:
 jobs:
   test:
     if: (github.event_name == 'schedule' && github.repository == 'openforcefield/openff-toolkit') || (github.event_name != 'schedule')
-    name: Test on ${{ matrix.os }}, Python ${{ matrix.python-version }}, RDKit=${{ matrix.rdkit }}, OpenEye=${{ matrix.openeye }}
+    name: Test on ${{ matrix.os }}, Python ${{ matrix.python-version }}, RDKit=${{ matrix.rdkit }}, OpenEye=${{ matrix.openeye }}  "new Pint" {{ matrix.new-pint }}
     runs-on: ${{ matrix.os }}
     strategy:
-      fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest]
         python-version: ["3.10", "3.11", "3.12"]
         rdkit: [true, false]
         openeye: [true, false]
+        new-pint: [false, true]
         exclude:
           - rdkit: false
             openeye: false
@@ -79,6 +79,13 @@ jobs:
           OE_LICENSE_TEXT: ${{ secrets.OE_LICENSE }}
         run: echo "${OE_LICENSE_TEXT}" > ${OE_LICENSE}
 
+      - name: Install Pint 0.24.4 and `openff-units` branch
+        if: matrix.new-pint
+        run: |
+          python -m pip install \
+            "pint==0.24.4" \
+            git+https://github.com/openforcefield/[email protected]
+
       - name: Install package
         run: |
           # While Interchange is being installed with pip, there is no need to
@@ -145,7 +152,7 @@ jobs:
           if [[ "$GITHUB_EVENT_NAME" == "schedule" ]]; then
             PYTEST_ARGS+=" -m 'slow or not slow'"
           fi
-
+          python -c "from pint import __version__; print(__version__)"
           python -m pytest --durations=20 $PYTEST_ARGS $COV
 
       - name: Run code snippets in docs

.github/workflows/examples.yml

@@ -63,6 +63,10 @@ jobs:
           OE_LICENSE_TEXT: ${{ secrets.OE_LICENSE }}
         run: echo "${OE_LICENSE_TEXT}" > ${OE_LICENSE}
 
+      - name: Install Pint 0.24.4 and `openff-units` branch
+        run: |
+          pip install "pint==0.24.4" git+https://github.com/openforcefield/[email protected]
+
       - name: Install package
         run: python -m pip install .
 

devtools/conda-envs/openeye.yaml

@@ -41,7 +41,7 @@ dependencies:
   - mdtraj
   - nglview
   - parmed =3
-  - mypy
+  - mypy <1.15
   - typing_extensions
   - pip:
     - types-setuptools

devtools/conda-envs/test_env.yaml

@@ -16,7 +16,7 @@ dependencies:
   - openmm >=7.6
   - openff-forcefields >=2023.05.1
   - smirnoff99Frosst
-  - openff-units =0.2.0
+  - openff-units
   - openff-amber-ff-ports
   - openff-utilities >=0.1.5
   - openff-interchange-base =0.4

openff/toolkit/_tests/test_topology.py

@@ -10,11 +10,10 @@
 import numpy as np
 import pytest
 from openff.units.openmm import from_openmm
-from openff.units.units import Quantity
 from openff.utilities import skip_if_missing
 from openmm import app
 
-from openff.toolkit import unit
+from openff.toolkit import Quantity, unit
 from openff.toolkit._tests.create_molecules import (
     create_ammonia,
     create_cyclohexane,

openff/toolkit/topology/molecule.py

@@ -46,6 +46,7 @@
 )
 
 import numpy as np
+from openff.units import Unit
 from openff.units.elements import MASSES, SYMBOLS
 from openff.utilities.exceptions import MissingOptionalDependencyError
 from typing_extensions import TypeAlias
@@ -110,6 +111,8 @@
 B = TypeVar("B", bound="Bond")
 
 
+_CHARGE_UNITS = set([Unit("elementary_charge")])
+
 class MoleculeDeprecationWarning(UserWarning):
     """Warning for deprecated portions of the Molecule API."""
 
@@ -338,12 +341,10 @@ def formal_charge(self, other):
         Set the atom's formal charge. Accepts either ints or unit-wrapped ints with units of charge.
         """
         if isinstance(other, int):
-            self._formal_charge = Quantity(other, unit.elementary_charge)
+            self._formal_charge = Quantity(other, "elementary_charge")
         elif isinstance(other, Quantity):
             # Faster to check equality than convert, so short-circuit
-            if other.units is unit.elementary_charge:
-                self.formal_charge = other
-            elif other.units in unit.elementary_charge.compatible_units():
+            if other.units in _CHARGE_UNITS:
                 self._formal_charge = other
             else:
                 raise IncompatibleUnitError(
@@ -361,7 +362,7 @@ def formal_charge(self, other):
             from openff.units.openmm import from_openmm
 
             converted = from_openmm(other)
-            if converted.units in unit.elementary_charge.compatible_units():
+            if converted.units in _CHARGE_UNITS:
                 self._formal_charge = converted
             else:
                 raise IncompatibleUnitError(
@@ -3073,14 +3074,14 @@ def _add_conformer(self, coordinates: Quantity):
         index
             The index of this conformer
         """
-        if coordinates.shape != (self.n_atoms, 3):
+        if coordinates.shape != (self.n_atoms, 3):  # type: ignore[attr-defined]
             raise InvalidConformerError(
                 "molecule.add_conformer given input of the wrong shape: "
-                f"Given {coordinates.shape}, expected {(self.n_atoms, 3)}"
+                f"Given {coordinates.shape}, expected {(self.n_atoms, 3)}"  # type: ignore[attr-defined]
             )
 
         if isinstance(coordinates, Quantity):
-            if not coordinates.units.is_compatible_with(unit.angstrom):
+            if not coordinates.units.is_compatible_with(unit.angstrom):  # type: ignore[attr-defined]
                 raise IncompatibleUnitError(
                     "Coordinates passed to Molecule._add_conformer with incompatible units. "
                     "Ensure that units are dimension of length."
@@ -3115,7 +3116,7 @@ def _add_conformer(self, coordinates: Quantity):
             np.zeros(shape=(self.n_atoms, 3), dtype=float), unit.angstrom
         )
         try:
-            tmp_conf[:] = coordinates
+            tmp_conf[:] = coordinates  # type: ignore[index]
         except AttributeError as e:
             # TODO: Make this a warning, log it, or do something other than print
             print(e)
@@ -3168,12 +3169,12 @@ def partial_charges(self, charges):
             )
 
         if isinstance(charges, Quantity):
-            if charges.units in unit.elementary_charge.compatible_units():
+            if charges.units in _CHARGE_UNITS:
                 self._partial_charges = charges.astype(float)
             else:
                 raise IncompatibleUnitError(
                     "Unsupported unit passed to partial_charges setter. "
-                    f"Found unit {charges.units}, expected {unit.elementary_charge}"
+                    f"Found unit {charges.units}, expected elementary_charge"
                 )
 
         elif hasattr(charges, "unit"):
@@ -3189,12 +3190,12 @@ def partial_charges(self, charges):
                 from openff.units.openmm import from_openmm
 
                 converted = from_openmm(charges)
-                if converted.units in unit.elementary_charge.compatible_units():
+                if converted.units in _CHARGE_UNITS:
                     self._partial_charges = converted.astype(float)
                 else:
                     raise IncompatibleUnitError(
                         "Unsupported unit passed to partial_charges setter. "
-                        f"Found unit {converted.units}, expected {unit.elementary_charge}"
+                        f"Found unit {converted.units}, expected elementary_charge"
                     )
 
         else:
@@ -4108,7 +4109,7 @@ def title(frame):
         # add the data to the xyz_data list
         for i, geometry in enumerate(conformers, 1):
             xyz_data.write(f"{self.n_atoms}\n" + title(end))
-            for j, atom_coords in enumerate(geometry.m_as(unit.angstrom)):
+            for j, atom_coords in enumerate(geometry.m_as(unit.angstrom)):  # type: ignore[arg-type]
                 x, y, z = atom_coords
                 xyz_data.write(
                     f"{SYMBOLS[self.atoms[j].atomic_number]}       {x: .10f}   {y: .10f}   {z: .10f}\n"

openff/toolkit/topology/topology.py

@@ -28,6 +28,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from openff.units import Unit
 from typing_extensions import TypeAlias
 
 from openff.toolkit import Quantity, unit
@@ -76,6 +77,20 @@
 TKR: TypeAlias = Union["ToolkitRegistry", "ToolkitWrapper"]
 MoleculeLike: TypeAlias = Union["Molecule", "FrozenMolecule", "_SimpleMolecule"]
 
+_DISTANCE_UNITS = set(
+    [
+        Unit(val) for val in
+        [
+            'nanometer',
+            'angstrom',
+            'micron',
+            'bohr',
+            'meter',
+            'fermi',
+        ]
+    ]
+)
+
 
 class _TransformedDict(MutableMapping):
     """A dictionary that transform and sort keys.
@@ -623,10 +638,7 @@ def box_vectors(self, box_vectors):
             else:
                 raise InvalidBoxVectorsError("Given unitless box vectors")
 
-        # Unit.compatible_units() returns False with itself, for some reason
-        if (box_vectors.units != unit.nm) and (
-            box_vectors.units not in unit.nm.compatible_units()
-        ):
+        if box_vectors.units not in _DISTANCE_UNITS:
             raise InvalidBoxVectorsError(
                 f"Cannot set box vectors with quantities with unit {box_vectors.units}"
             )
@@ -1236,7 +1248,7 @@ def to_dict(self) -> dict:
         return_dict["aromaticity_model"] = self._aromaticity_model
 
         return_dict["constrained_atom_pairs"] = {
-            constrained_atom_pair: quantity_to_string(distance)
+            constrained_atom_pair: quantity_to_string(distance)  # type: ignore[arg-type]
             for constrained_atom_pair, distance in self._constrained_atom_pairs.items()
         }
 
@@ -1916,7 +1928,7 @@ def to_openmm(
                     (last_residue.name == atom_residue_name)
                     and (int(last_residue.id) == int(atom_residue_number))
                     and (last_residue.insertionCode == atom_insertion_code)
-                    and (chain.id == last_chain.id)
+                    and (chain.id == last_chain.id)  # type: ignore[union-attr]
                 ):
                     residue = last_residue
                 else:
@@ -2045,7 +2057,7 @@ def to_file(
         if isinstance(positions, openmm.unit.Quantity):
             openmm_positions: openmm.unit.Quantity = positions
         elif isinstance(positions, Quantity):
-            openmm_positions = positions.to_openmm()
+            openmm_positions = positions.to_openmm()  # type: ignore[attr-defined]
         elif isinstance(positions, np.ndarray):
             openmm_positions = openmm.unit.Quantity(positions, openmm.unit.angstroms)
         elif positions is None:
@@ -2152,18 +2164,18 @@ def set_positions(self, array: Quantity):
 
         # Copy the array in nanometers and make it an OpenFF Quantity
         array = Quantity(np.asarray(array.to(unit.nanometer).magnitude), unit.nanometer)
-        if array.shape != (self.n_atoms, 3):
+        if array.shape != (self.n_atoms, 3):  # type: ignore[attr-defined]
             raise WrongShapeError(
-                f"Array has shape {array.shape} but should have shape {self.n_atoms, 3}"
+                f"Array has shape {array.shape} but should have shape {self.n_atoms, 3}"  # type: ignore[attr-defined]
             )
 
         start = 0
         for molecule in self.molecules:
             stop = start + molecule.n_atoms
             if molecule.conformers is None:
-                molecule._conformers = [array[start:stop]]
+                molecule._conformers = [array[start:stop]]  # type: ignore[index]
             else:
-                molecule.conformers[0:1] = [array[start:stop]]
+                molecule.conformers[0:1] = [array[start:stop]]  # type: ignore[index]
             start = stop
 
     @classmethod

openff/toolkit/typing/engines/smirnoff/forcefield.py

@@ -803,7 +803,7 @@ def parse_sources(self, sources, allow_cosmetic_attributes: bool = True):
                 smirnoff_data, allow_cosmetic_attributes=allow_cosmetic_attributes
             )
 
-    def _to_smirnoff_data(self, discard_cosmetic_attributes: bool = False) -> dict:
+    def _to_smirnoff_data(self, discard_cosmetic_attributes: bool = False) -> list | dict:
         """
         Convert this ForceField and all related ParameterHandlers to a dict representing a SMIRNOFF
         data object.

openff/toolkit/typing/engines/smirnoff/parameters.py

@@ -213,7 +213,7 @@ def _validate_units(attr, value: Union[str, Quantity], units: Unit):
     value = object_to_quantity(value)
 
     try:
-        if not units.is_compatible_with(value.units):
+        if not units.is_compatible_with(value.units):  # type: ignore
             raise IncompatibleUnitError(
                 f"{attr.name}={value} should have units of {units}"
             )

openff/toolkit/utils/ambertools_wrapper.py

@@ -265,12 +265,11 @@ def assign_partial_charges(
             with open(f"{tmpdir}/charges.txt") as infile:
                 contents = infile.read()
             text_charges = contents.split()
-            charges = np.zeros([molecule.n_atoms], np.float64)
+            charges_array = np.zeros([molecule.n_atoms], np.float64)
             for index, token in enumerate(text_charges):
-                charges[index] = float(token)
+                charges_array[index] = float(token)
             # TODO: Ensure that the atoms in charged.mol2 are in the same order as in molecule.sdf
-        charges = Quantity(charges, unit.elementary_charge)
-        molecule.partial_charges = charges
+        molecule.partial_charges = Quantity(charges_array, unit.elementary_charge)
 
         if normalize_partial_charges:
             molecule._normalize_partial_charges()

openff/toolkit/utils/builtin_wrapper.py

@@ -117,7 +117,7 @@ def assign_partial_charges(
             )
         else:
             mol_copy._conformers = None
-            for conformer in use_conformers:
+            for conformer in use_conformers:  # type: ignore[attr-defined]
                 mol_copy._add_conformer(conformer)
             self._check_n_conformers(
                 mol_copy,

openff/toolkit/utils/openeye_wrapper.py

@@ -23,15 +23,10 @@
 
 import numpy as np
 from cachetools import LRUCache, cached
+from openff.units.elements import SYMBOLS
 from typing_extensions import TypeAlias
 
 from openff.toolkit import Quantity, unit
-
-if TYPE_CHECKING:
-    from openff.toolkit.topology.molecule import Atom, Bond, FrozenMolecule, Molecule
-
-from openff.units.elements import SYMBOLS
-
 from openff.toolkit.utils.base_wrapper import (
     ToolkitWrapper,
     _ChargeSettings,
@@ -64,6 +59,12 @@
 )
 from openff.toolkit.utils.utils import inherit_docstrings
 
+if TYPE_CHECKING:
+    import openmm.app
+
+    from openff.toolkit.topology.molecule import Atom, Bond, FrozenMolecule, Molecule
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -332,8 +333,8 @@ def _polymer_openmm_topology_to_offmol(
 
     def _polymer_openmm_topology_to_oemol(
         self,
-        omm_top,
-        substructure_library,
+        omm_top: "openmm.app.Topology",
+        substructure_library: dict[str, dict],
     ):
         """
         Parameters
@@ -1674,7 +1675,7 @@ def to_openeye(
                 # OE needs a 1 x (3*n_Atoms) double array as input
                 flat_coords = np.zeros(shape=oemol.NumAtoms() * 3, dtype=np.float64)
                 for index, oe_idx in off_to_oe_idx.items():
-                    (x, y, z) = conf[index, :].m_as(unit.angstrom)
+                    (x, y, z) = conf[index, :].m_as(unit.angstrom)  # type: ignore[index]
                     flat_coords[(3 * oe_idx)] = x
                     flat_coords[(3 * oe_idx) + 1] = y
                     flat_coords[(3 * oe_idx) + 2] = z
@@ -2560,7 +2561,7 @@ def assign_partial_charges(
             index = oeatom.GetIdx()
             charge = oeatom.GetPartialCharge()
             charge = charge * unit.elementary_charge
-            charges[index] = charge
+            charges[index] = charge  # type: ignore[index]
 
         molecule.partial_charges = charges