diff --git a/cvpack/__init__.py b/cvpack/__init__.py index 7ac0ee68..b0d4ea8a 100644 --- a/cvpack/__init__.py +++ b/cvpack/__init__.py @@ -1,8 +1,5 @@ """Useful Collective Variables for OpenMM""" -import yaml - -# Add imports here from ._version import __version__ # noqa: F401 from .angle import Angle # noqa: F401 from .atomic_function import AtomicFunction # noqa: F401 @@ -24,29 +21,3 @@ from .sheet_rmsd_content import SheetRMSDContent # noqa: F401 from .torsion import Torsion # noqa: F401 from .torsion_similarity import TorsionSimilarity # noqa: F401 - -for _cv in [ - Angle, - AtomicFunction, - AttractionStrength, - CentroidFunction, - CompositeRMSD, - Distance, - HelixAngleContent, - HelixHBondContent, - HelixRMSDContent, - HelixTorsionContent, - NumberOfContacts, - OpenMMForceWrapper, - PathInCVSpace, - RadiusOfGyration, - RadiusOfGyrationSq, - ResidueCoordination, - RMSD, - SheetRMSDContent, - Torsion, - TorsionSimilarity, -]: - yaml.SafeDumper.add_representer(_cv, _cv.to_yaml) - yaml.SafeLoader.add_constructor(_cv.yaml_tag, _cv.from_yaml) -del _cv diff --git a/cvpack/angle.py b/cvpack/angle.py index f532f494..aea52139 100644 --- a/cvpack/angle.py +++ b/cvpack/angle.py @@ -60,11 +60,12 @@ class Angle(openmm.CustomAngleForce, BaseCollectiveVariable): """ - yaml_tag = "!cvpack.Angle" - def __init__(self, atom1: int, atom2: int, atom3: int, pbc: bool = False) -> None: super().__init__("theta") self.addAngle(atom1, atom2, atom3, []) self.setUsesPeriodicBoundaryConditions(pbc) self._registerCV(mmunit.radians, atom1, atom2, atom3, pbc) self._registerPeriod(2 * math.pi) + + +Angle.registerTag("!cvpack.Angle") diff --git a/cvpack/atomic_function.py b/cvpack/atomic_function.py index 3f67e6cd..a29c7467 100644 --- a/cvpack/atomic_function.py +++ b/cvpack/atomic_function.py @@ -107,8 +107,6 @@ class AtomicFunction(openmm.CustomCompoundBondForce, BaseCustomFunction): 429.479... kJ/mol """ - yaml_tag = "!cvpack.AtomicFunction" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -354,3 +352,6 @@ def fromOpenMMForce( if isinstance(force, openmm.PeriodicTorsionForce): return cls._fromPeriodicTorsionForce(force, unit, pbc) raise TypeError(f"Force {force} is not convertible to an AtomicFunction") + + +AtomicFunction.registerTag("!cvpack.AtomicFunction") diff --git a/cvpack/attraction_strength.py b/cvpack/attraction_strength.py index 15e872fd..e57b61c6 100644 --- a/cvpack/attraction_strength.py +++ b/cvpack/attraction_strength.py @@ -166,8 +166,6 @@ class AttractionStrength(openmm.CustomNonbondedForce, BaseCollectiveVariable): 3880.8... """ - yaml_tag = "!cvpack.AttractionStrength" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -236,3 +234,6 @@ def __init__( # pylint: disable=too-many-arguments reference, contrastScaling, ) + + +AttractionStrength.registerTag("!cvpack.AttractionStrength") diff --git a/cvpack/centroid_function.py b/cvpack/centroid_function.py index 1476a1a5..77a18c2a 100644 --- a/cvpack/centroid_function.py +++ b/cvpack/centroid_function.py @@ -151,8 +151,6 @@ class CentroidFunction(openmm.CustomCentroidBondForce, BaseCustomFunction): 33.0 dimensionless """ - yaml_tag = "!cvpack.CentroidFunction" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -195,3 +193,6 @@ def __init__( # pylint: disable=too-many-arguments ) if period is not None: self._registerPeriod(period) + + +CentroidFunction.registerTag("!cvpack.CentroidFunction") diff --git a/cvpack/composite_rmsd.py b/cvpack/composite_rmsd.py index e5c21d3f..14de945d 100644 --- a/cvpack/composite_rmsd.py +++ b/cvpack/composite_rmsd.py @@ -123,8 +123,6 @@ class CompositeRMSD(CompositeRMSDForce, BaseCollectiveVariable): 0.0 nm """ - yaml_tag = "!cvpack.CompositeRMSD" - @mmunit.convert_quantities def __init__( self, @@ -149,3 +147,6 @@ def __init__( for group in groups: self.addGroup(group) self._registerCV(mmunit.nanometers, defined_coords, groups, num_atoms) + + +CompositeRMSD.registerTag("!cvpack.CompositeRMSD") diff --git a/cvpack/cvpack.py b/cvpack/cvpack.py index 08f29194..3a23f9c6 100644 --- a/cvpack/cvpack.py +++ b/cvpack/cvpack.py @@ -17,17 +17,16 @@ from cvpack import unit as mmunit +from .serializer import Serializable from .unit import value_in_md_units from .utils import compute_effective_mass, get_single_force_state -class SerializableAtom(yaml.YAMLObject): +class SerializableAtom(Serializable): r""" A serializable version of OpenMM's Atom class. """ - yaml_tag = "!cvpack.Atom" - def __init__( # pylint: disable=super-init-not-called self, atom: t.Union[mmapp.topology.Atom, "SerializableAtom"] ) -> None: @@ -48,16 +47,11 @@ def __setstate__(self, keywords: t.Dict[str, t.Any]) -> None: self.__dict__.update(keywords) -yaml.SafeDumper.add_representer(SerializableAtom, SerializableAtom.to_yaml) -yaml.SafeLoader.add_constructor(SerializableAtom.yaml_tag, SerializableAtom.from_yaml) - +SerializableAtom.registerTag("!cvpack.Atom") -class SerializableResidue(yaml.YAMLObject): - r""" - A serializable version of OpenMM's Residue class. - """ - yaml_tag = "!cvpack.Residue" +class SerializableResidue(Serializable): + r"""A serializable version of OpenMM's Residue class.""" def __init__( # pylint: disable=super-init-not-called self, residue: t.Union[mmapp.topology.Residue, "SerializableResidue"] @@ -85,13 +79,10 @@ def atoms(self): return iter(self._atoms) -yaml.SafeDumper.add_representer(SerializableResidue, SerializableResidue.to_yaml) -yaml.SafeLoader.add_constructor( - SerializableResidue.yaml_tag, SerializableResidue.from_yaml -) +SerializableResidue.registerTag("!cvpack.Residue") -class BaseCollectiveVariable(openmm.Force, yaml.YAMLObject): +class BaseCollectiveVariable(openmm.Force, Serializable): r""" An abstract class with common attributes and method for all CVs. """ diff --git a/cvpack/distance.py b/cvpack/distance.py index 895df0bc..741441e1 100644 --- a/cvpack/distance.py +++ b/cvpack/distance.py @@ -49,10 +49,11 @@ class Distance(openmm.CustomBondForce, BaseCollectiveVariable): """ - yaml_tag = "!cvpack.Distance" - def __init__(self, atom1: int, atom2: int, pbc: bool = False) -> None: super().__init__("r") self.addBond(atom1, atom2, []) self.setUsesPeriodicBoundaryConditions(pbc) self._registerCV(mmunit.nanometers, atom1, atom2, pbc) + + +Distance.registerTag("!cvpack.Distance") diff --git a/cvpack/helix_angle_content.py b/cvpack/helix_angle_content.py index dcf89754..19478e9f 100644 --- a/cvpack/helix_angle_content.py +++ b/cvpack/helix_angle_content.py @@ -96,8 +96,6 @@ class HelixAngleContent(openmm.CustomAngleForce, BaseCollectiveVariable): 18.7605... dimensionless """ - yaml_tag = "!cvpack.HelixAngleContent" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -139,3 +137,6 @@ def find_alpha_carbon(residue: mmapp.topology.Residue) -> int: halfExponent, normalize, ) + + +HelixAngleContent.registerTag("!cvpack.HelixAngleContent") diff --git a/cvpack/helix_hbond_content.py b/cvpack/helix_hbond_content.py index 5b495f91..7574d856 100644 --- a/cvpack/helix_hbond_content.py +++ b/cvpack/helix_hbond_content.py @@ -84,8 +84,6 @@ class HelixHBondContent(openmm.CustomBondForce, BaseCollectiveVariable): 15.880... dimensionless """ - yaml_tag = "!cvpack.HelixHBondContent" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -126,3 +124,6 @@ def find_atom(residue: mmapp.topology.Residue, pattern: t.Pattern) -> int: halfExponent, normalize, ) + + +HelixHBondContent.registerTag("!cvpack.HelixHBondContent") diff --git a/cvpack/helix_rmsd_content.py b/cvpack/helix_rmsd_content.py index 216206da..c561e98e 100644 --- a/cvpack/helix_rmsd_content.py +++ b/cvpack/helix_rmsd_content.py @@ -123,8 +123,6 @@ class HelixRMSDContent(BaseRMSDContent): 15.98... dimensionless """ - yaml_tag = "!cvpack.HelixRMSDContent" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -155,3 +153,6 @@ def __init__( # pylint: disable=too-many-arguments stepFunction, normalize, ) + + +HelixRMSDContent.registerTag("!cvpack.HelixRMSDContent") diff --git a/cvpack/helix_torsion_content.py b/cvpack/helix_torsion_content.py index 8ae9ced6..133fad55 100644 --- a/cvpack/helix_torsion_content.py +++ b/cvpack/helix_torsion_content.py @@ -106,8 +106,6 @@ class HelixTorsionContent(openmm.CustomTorsionForce, BaseCollectiveVariable): 17.452... dimensionless """ - yaml_tag = "!cvpack.HelixTorsionContent" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -157,3 +155,6 @@ def find_atom(residue: mmapp.topology.Residue, name: str) -> int: tolerance, halfExponent, ) + + +HelixTorsionContent.registerTag("!cvpack.HelixTorsionContent") diff --git a/cvpack/number_of_contacts.py b/cvpack/number_of_contacts.py index 9111bab4..bccd344b 100644 --- a/cvpack/number_of_contacts.py +++ b/cvpack/number_of_contacts.py @@ -126,8 +126,6 @@ class NumberOfContacts(openmm.CustomNonbondedForce, BaseCollectiveVariable): 0.99999... dimensionless """ - yaml_tag = "!cvpack.NumberOfContacts" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -175,3 +173,6 @@ def __init__( # pylint: disable=too-many-arguments cutoffFactor, switchFactor, ) + + +NumberOfContacts.registerTag("!cvpack.NumberOfContacts") diff --git a/cvpack/openmm_force_wrapper.py b/cvpack/openmm_force_wrapper.py index c15967b6..32627af1 100644 --- a/cvpack/openmm_force_wrapper.py +++ b/cvpack/openmm_force_wrapper.py @@ -62,8 +62,6 @@ class OpenMMForceWrapper(BaseCollectiveVariable): 0.00538... nm**2 Da/(rad**2) """ - yaml_tag = "!cvpack.OpenMMForceWrapper" - def __init__( # pylint: disable=too-many-arguments, super-init-not-called self, openmmForce: t.Union[openmm.Force, str], @@ -79,3 +77,6 @@ def __init__( # pylint: disable=too-many-arguments, super-init-not-called self._registerCV(unit, openmmForce, unit, period) if period is not None: self._registerPeriod(period) + + +OpenMMForceWrapper.registerTag("!cvpack.OpenMMForceWrapper") diff --git a/cvpack/path.py b/cvpack/path.py index 7944f742..4d48fb93 100644 --- a/cvpack/path.py +++ b/cvpack/path.py @@ -7,10 +7,10 @@ """ -import yaml +from .serializer import Serializable -class Metric(yaml.YAMLObject): +class Metric(Serializable): """ A measure of progress or deviation with respect to a path in CV space """ @@ -26,9 +26,14 @@ def __repr__(self) -> str: def __eq__(self, other: object) -> bool: return isinstance(other, Metric) and self.name == other.name + def __getstate__(self) -> dict: + return {"name": self.name} -yaml.SafeDumper.add_representer(Metric, Metric.to_yaml) -yaml.SafeLoader.add_constructor(Metric.yaml_tag, Metric.from_yaml) + def __setstate__(self, state: dict) -> None: + self.name = state["name"] + + +Metric.registerTag("!cvpack.path.Metric") progress: Metric = Metric("progress") diff --git a/cvpack/path_in_cv_space.py b/cvpack/path_in_cv_space.py index 1149b3df..36b9dde7 100644 --- a/cvpack/path_in_cv_space.py +++ b/cvpack/path_in_cv_space.py @@ -118,8 +118,6 @@ class PathInCVSpace(openmm.CustomCVForce, BaseCollectiveVariable): z = 0.25... dimensionless """ - yaml_tag = "!cvpack.PathInCVSpace" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -177,3 +175,6 @@ def __init__( # pylint: disable=too-many-arguments sigma, scales, ) + + +PathInCVSpace.registerTag("!cvpack.PathInCVSpace") diff --git a/cvpack/radius_of_gyration.py b/cvpack/radius_of_gyration.py index 10358fd4..23ae56ae 100644 --- a/cvpack/radius_of_gyration.py +++ b/cvpack/radius_of_gyration.py @@ -75,8 +75,6 @@ class RadiusOfGyration(BaseRadiusOfGyration): """ - yaml_tag = "!cvpack.RadiusOfGyration" - def __init__( self, group: t.Iterable[int], pbc: bool = False, weighByMass: bool = False ) -> None: @@ -91,3 +89,6 @@ def __init__( ) self.addBond(list(range(num_groups))) self._registerCV(mmunit.nanometers, group, pbc, weighByMass) + + +RadiusOfGyration.registerTag("!cvpack.RadiusOfGyration") diff --git a/cvpack/radius_of_gyration_sq.py b/cvpack/radius_of_gyration_sq.py index 092e5a8e..99325715 100644 --- a/cvpack/radius_of_gyration_sq.py +++ b/cvpack/radius_of_gyration_sq.py @@ -77,8 +77,6 @@ class RadiusOfGyrationSq(BaseRadiusOfGyration): """ - yaml_tag = "!cvpack.RadiusOfGyrationSq" - def __init__( self, group: t.Iterable[int], pbc: bool = False, weighByMass: bool = False ) -> None: @@ -88,3 +86,6 @@ def __init__( for atom in group: self.addBond([atom, num_atoms]) self._registerCV(mmunit.nanometers**2, group, pbc, weighByMass) + + +RadiusOfGyrationSq.registerTag("!cvpack.RadiusOfGyrationSq") diff --git a/cvpack/residue_coordination.py b/cvpack/residue_coordination.py index 08de58c5..b38fab30 100644 --- a/cvpack/residue_coordination.py +++ b/cvpack/residue_coordination.py @@ -105,8 +105,6 @@ class ResidueCoordination(openmm.CustomCentroidBondForce, BaseCollectiveVariable 0.99999... dimensionless """ - yaml_tag = "!cvpack.ResidueCoordination" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -184,3 +182,6 @@ def setReferenceValue(self, value: mmunit.ScalarQuantity) -> None: expression.replace(f"refval={self._ref_val}", f"refval={value}") ) self._ref_val = value + + +ResidueCoordination.registerTag("!cvpack.ResidueCoordination") diff --git a/cvpack/rmsd.py b/cvpack/rmsd.py index e1a66f39..1c7fb323 100644 --- a/cvpack/rmsd.py +++ b/cvpack/rmsd.py @@ -84,8 +84,6 @@ class RMSD(openmm.RMSDForce, BaseCollectiveVariable): """ - yaml_tag = "!cvpack.RMSD" - @mmunit.convert_quantities def __init__( self, @@ -149,3 +147,6 @@ def getNullBondForce(self) -> openmm.HarmonicBondForce: for i, j in zip(group[:-1], group[1:]): force.addBond(i, j, 0.0, 0.0) return force + + +RMSD.registerTag("!cvpack.RMSD") diff --git a/cvpack/serializer/__init__.py b/cvpack/serializer/__init__.py index 2b9c56f7..d4aaef2e 100644 --- a/cvpack/serializer/__init__.py +++ b/cvpack/serializer/__init__.py @@ -3,6 +3,4 @@ """ -# Add imports here -from .serializer import deserialize # noqa: F401 -from .serializer import serialize # noqa: F401 +from .serializer import Serializable, deserialize, serialize # noqa: F401 diff --git a/cvpack/serializer/serializer.py b/cvpack/serializer/serializer.py index 434ae125..1a602a0f 100644 --- a/cvpack/serializer/serializer.py +++ b/cvpack/serializer/serializer.py @@ -12,6 +12,26 @@ import yaml +class Serializable(yaml.YAMLObject): + """ + A mixin class that allows serialization and deserialization of objects with PyYAML. + """ + + @classmethod + def registerTag(cls, tag: str) -> None: + """ + Register a class for serialization and deserialization with PyYAML. + + Parameters + ---------- + tag + The YAML tag to be used for this class. + """ + cls.yaml_tag = tag + yaml.SafeDumper.add_representer(cls, cls.to_yaml) + yaml.SafeLoader.add_constructor(tag, cls.from_yaml) + + def serialize(obj: t.Any, iostream: t.IO) -> None: """ Serializes a cvpack object. diff --git a/cvpack/sheet_rmsd_content.py b/cvpack/sheet_rmsd_content.py index f710cfa8..017006cc 100644 --- a/cvpack/sheet_rmsd_content.py +++ b/cvpack/sheet_rmsd_content.py @@ -181,8 +181,6 @@ class SheetRMSDContent(BaseRMSDContent): 0.9859... dimensionless """ - yaml_tag = "!cvpack.SheetRMSDContent" - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, @@ -235,3 +233,6 @@ def __init__( # pylint: disable=too-many-arguments stepFunction, normalize, ) + + +SheetRMSDContent.registerTag("!cvpack.SheetRMSDContent") diff --git a/cvpack/torsion.py b/cvpack/torsion.py index ad116376..c431e6a2 100644 --- a/cvpack/torsion.py +++ b/cvpack/torsion.py @@ -67,8 +67,6 @@ class Torsion(openmm.CustomTorsionForce, BaseCollectiveVariable): """ - yaml_tag = "!cvpack.Torsion" - def __init__( # pylint: disable=too-many-arguments self, atom1: int, atom2: int, atom3: int, atom4: int, pbc: bool = False ) -> None: @@ -77,3 +75,6 @@ def __init__( # pylint: disable=too-many-arguments self.setUsesPeriodicBoundaryConditions(pbc) self._registerCV(mmunit.radians, atom1, atom2, atom3, atom4, pbc) self._registerPeriod(2 * math.pi) + + +Torsion.registerTag("!cvpack.Torsion") diff --git a/cvpack/torsion_similarity.py b/cvpack/torsion_similarity.py index d0323d7a..371d9e27 100644 --- a/cvpack/torsion_similarity.py +++ b/cvpack/torsion_similarity.py @@ -79,8 +79,6 @@ class TorsionSimilarity(openmm.CustomCompoundBondForce, BaseCollectiveVariable): 18.659... dimensionless """ - yaml_tag = "!cvpack.TorsionSimilarity" - def __init__( self, firstList: t.Iterable[t.Tuple[int, int, int, int]], @@ -96,3 +94,6 @@ def __init__( self.addBond([*first, *second], []) self.setUsesPeriodicBoundaryConditions(pbc) self._registerCV(mmunit.dimensionless, firstList, secondList) + + +TorsionSimilarity.registerTag("!cvpack.TorsionSimilarity") diff --git a/cvpack/unit/__init__.py b/cvpack/unit/__init__.py index b72d6b02..c78f2a44 100644 --- a/cvpack/unit/__init__.py +++ b/cvpack/unit/__init__.py @@ -18,9 +18,10 @@ import numpy as np import openmm -import yaml from openmm import unit as _mmunit +from ..serializer import Serializable + ScalarQuantity = t.Union[_mmunit.Quantity, Real] VectorQuantity = t.Union[_mmunit.Quantity, np.ndarray, openmm.Vec3] MatrixQuantity = t.Union[ @@ -46,7 +47,7 @@ def visit_Name( # pylint: disable=invalid-name return ast.Attribute(value=mod, attr=node.id, ctx=ast.Load()) -class SerializableUnit(_mmunit.Unit, yaml.YAMLObject): +class SerializableUnit(_mmunit.Unit, Serializable): r""" A child class of openmm.unit.Unit that allows for serialization/deserialization. @@ -72,8 +73,6 @@ class SerializableUnit(_mmunit.Unit, yaml.YAMLObject): Quantity(value=2, unit=nanometer/picosecond) """ - yaml_tag = "!cvpack.Unit" - def __init__(self, base_or_scaled_units): if isinstance(base_or_scaled_units, _mmunit.Unit): self.__dict__ = base_or_scaled_units.__dict__ @@ -96,11 +95,10 @@ def __setstate__(self, kwds) -> None: self.__init__(kwds["description"]) -yaml.SafeDumper.add_representer(SerializableUnit, SerializableUnit.to_yaml) -yaml.SafeLoader.add_constructor(SerializableUnit.yaml_tag, SerializableUnit.from_yaml) +SerializableUnit.registerTag("!cvpack.Unit") -class SerializableQuantity(_mmunit.Quantity, yaml.YAMLObject): +class SerializableQuantity(_mmunit.Quantity, Serializable): r""" A child class of openmm.unit.Quantity that allows for serialization/deserialization. @@ -131,8 +129,6 @@ class SerializableQuantity(_mmunit.Quantity, yaml.YAMLObject): Quantity(value=1.0, unit=nanometer) """ - yaml_tag = "!cvpack.Quantity" - def __init__(self, value, unit=None): # pylint: disable=redefined-outer-name if unit is None: super().__init__(value._value, SerializableUnit(value.unit)) @@ -156,10 +152,7 @@ def value_in_md_units(self): # pylint: disable=invalid-name return value_in_md_units(self) -yaml.SafeDumper.add_representer(SerializableQuantity, SerializableQuantity.to_yaml) -yaml.SafeLoader.add_constructor( - SerializableQuantity.yaml_tag, SerializableQuantity.from_yaml -) +SerializableQuantity.registerTag("!cvpack.Quantity") def value_in_md_units( # pylint: disable=redefined-outer-name diff --git a/cvpack/utils.py b/cvpack/utils.py index c2c6cf29..38d7c4a1 100644 --- a/cvpack/utils.py +++ b/cvpack/utils.py @@ -12,20 +12,19 @@ import numpy as np import openmm -import yaml from numpy import typing as npt from openmm import XmlSerializer from cvpack import unit as mmunit +from .serializer import Serializable + class NonbondedForceSurrogate( - yaml.YAMLObject + Serializable ): # pylint: disable=too-many-instance-attributes """A surrogate class for the NonbondedForce class in OpenMM.""" - yaml_tag = "!cvpack.NonbondedForce" - def __init__(self, other: openmm.NonbondedForce) -> None: self._cutoff = other.getCutoffDistance() self._uses_pbc = other.usesPeriodicBoundaryConditions() @@ -102,12 +101,7 @@ def getSwitchingDistance(self) -> float: return mmunit.value_in_md_units(self._switching_distance) -yaml.SafeDumper.add_representer( - NonbondedForceSurrogate, NonbondedForceSurrogate.to_yaml -) -yaml.SafeLoader.add_constructor( - NonbondedForceSurrogate.yaml_tag, NonbondedForceSurrogate.from_yaml -) +NonbondedForceSurrogate.registerTag("!cvpack.NonbondedForce") def evaluate_in_context(