From 154d34ceac9d1c2312ff3d83f118dce40af8d98c Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Wed, 6 Nov 2024 09:26:53 -0600 Subject: [PATCH 1/2] Defer/fix some minor import issues --- openff/toolkit/topology/topology.py | 5 ++++- openff/toolkit/utils/_viz.py | 11 ++++++----- openff/toolkit/utils/ambertools_wrapper.py | 3 ++- openff/toolkit/utils/rdkit_wrapper.py | 3 ++- openff/toolkit/utils/utils.py | 9 ++------- 5 files changed, 16 insertions(+), 15 deletions(-) diff --git a/openff/toolkit/topology/topology.py b/openff/toolkit/topology/topology.py index 8e618e671..5598ce7c3 100644 --- a/openff/toolkit/topology/topology.py +++ b/openff/toolkit/topology/topology.py @@ -31,7 +31,6 @@ import numpy as np from numpy.typing import NDArray -from openff.units import ensure_quantity from typing_extensions import TypeAlias from openff.toolkit import Quantity, unit @@ -618,6 +617,8 @@ def box_vectors(self, box_vectors): return if not hasattr(box_vectors, "units"): if hasattr(box_vectors, "unit"): + from openff.units import ensure_quantity + # this is probably an openmm.unit.Quantity; we should gracefully import OpenMM but # the chances of this being an object with the two previous conditions met is low box_vectors = ensure_quantity(box_vectors, "openff") @@ -1538,6 +1539,8 @@ def from_openmm( topology.box_vectors = from_openmm(openmm_topology.getPeriodicBoxVectors()) if positions is not None: + from openff.units import ensure_quantity + topology.set_positions(ensure_quantity(positions, "openff")) # TODO: How can we preserve metadata from the openMM topology when creating the OFF topology? diff --git a/openff/toolkit/utils/_viz.py b/openff/toolkit/utils/_viz.py index 254d0fa89..91286f917 100644 --- a/openff/toolkit/utils/_viz.py +++ b/openff/toolkit/utils/_viz.py @@ -1,15 +1,16 @@ import uuid from io import StringIO +from typing import TYPE_CHECKING from nglview.base_adaptor import Structure, Trajectory -from openff.toolkit import Molecule, Topology, unit +if TYPE_CHECKING: + from openff.toolkit import Molecule, Topology MOLECULE_DEFAULT_REPS = [ dict(type="licorice", params=dict(radius=0.25, multipleBond=True)) ] - class MoleculeNGLViewTrajectory(Structure, Trajectory): """ OpenFF Molecule adaptor. @@ -40,7 +41,7 @@ class MoleculeNGLViewTrajectory(Structure, Trajectory): def __init__( self, - molecule: Molecule, + molecule: "Molecule", ext: str = "MOL2", ): if not molecule.conformers: @@ -53,7 +54,7 @@ def __init__( self.id = str(uuid.uuid4()) def get_coordinates(self, index: int = 0): - return self.molecule.conformers[index].m_as(unit.angstrom) + return self.molecule.conformers[index].m_as("angstrom") @property def n_frames(self): @@ -91,7 +92,7 @@ class TopologyNGLViewStructure(Structure): def __init__( self, - topology: Topology, + topology: "Topology", ext: str = "PDB", ): self.topology = topology diff --git a/openff/toolkit/utils/ambertools_wrapper.py b/openff/toolkit/utils/ambertools_wrapper.py index ff5b16ada..64167fc4a 100644 --- a/openff/toolkit/utils/ambertools_wrapper.py +++ b/openff/toolkit/utils/ambertools_wrapper.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Optional import numpy as np -from openff.utilities.provenance import get_ambertools_version from openff.toolkit import Quantity, unit from openff.toolkit.utils import base_wrapper, rdkit_wrapper @@ -62,6 +61,8 @@ class AmberToolsToolkitWrapper(base_wrapper.ToolkitWrapper): SUPPORTED_CHARGE_METHODS = _supported_charge_methods def __init__(self): + from openff.utilities.provenance import get_ambertools_version + super().__init__() if not self.is_available(): diff --git a/openff/toolkit/utils/rdkit_wrapper.py b/openff/toolkit/utils/rdkit_wrapper.py index 2a4366c16..c7d8a70e7 100644 --- a/openff/toolkit/utils/rdkit_wrapper.py +++ b/openff/toolkit/utils/rdkit_wrapper.py @@ -18,6 +18,7 @@ import numpy as np from cachetools import LRUCache, cached +from numpy.typing import NDArray from openff.units.elements import SYMBOLS from openff.toolkit import Quantity, unit @@ -1985,7 +1986,7 @@ def _elf_compute_electrostatic_energy( return 0.5 * interaction_energies.sum() @classmethod - def _elf_compute_rms_matrix(cls, molecule: "Molecule") -> np.ndarray: + def _elf_compute_rms_matrix(cls, molecule: "Molecule") -> NDArray: """Computes the symmetric RMS matrix of all conformers in a molecule taking only heavy atoms into account. diff --git a/openff/toolkit/utils/utils.py b/openff/toolkit/utils/utils.py index f03d27e11..b2d760459 100644 --- a/openff/toolkit/utils/utils.py +++ b/openff/toolkit/utils/utils.py @@ -28,9 +28,9 @@ import logging from typing import TYPE_CHECKING, Any, Iterable, TypeVar, Union, overload -import numpy import numpy as np import pint +from numpy.typing import NDArray from openff.units import Quantity, Unit, unit from openff.utilities import requires_package @@ -397,8 +397,6 @@ def serialize_numpy(np_array) -> tuple[bytes, tuple[int]]: shape The shape of the serialized array """ - import numpy as np - bigendian_float = np.dtype(float).newbyteorder(">") bigendian_array = np_array.astype(bigendian_float) serialized = bigendian_array.tobytes() @@ -409,7 +407,7 @@ def serialize_numpy(np_array) -> tuple[bytes, tuple[int]]: def deserialize_numpy( serialized_np: Union[bytes, list], shape: tuple[int, ...], -) -> numpy.ndarray: +) -> NDArray: """ Deserializes a numpy array from a bytestring or list. The input, if a bytestring, is assumed to be in big-endian byte order. @@ -426,9 +424,6 @@ def deserialize_numpy( np_array The deserialized numpy array """ - - import numpy as np - if isinstance(serialized_np, list): np_array = np.array(serialized_np) if isinstance(serialized_np, bytes): From c0939be9b19fc7f91d51847e60ec15a686391811 Mon Sep 17 00:00:00 2001 From: Jeff Wagner Date: Thu, 7 Nov 2024 17:15:35 -0800 Subject: [PATCH 2/2] update releasehistory --- docs/releasehistory.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/releasehistory.md b/docs/releasehistory.md index febf27c32..bb218edd0 100644 --- a/docs/releasehistory.md +++ b/docs/releasehistory.md @@ -13,6 +13,7 @@ Releases follow the `major.minor.micro` scheme recommended by [PEP440](https://w ### Behavior changes - [PR #1954](https://github.com/openforcefield/openff-toolkit/pull/1954): `Topology.from_openmm` no longer casts residue numbers to int. +- [PR #1959](https://github.com/openforcefield/openff-toolkit/pull/1959): Speeds up import times. ### Bugfixes