Add unit class to FlowerMD and allow using real units in Simulation runs

flowermd/__init__.py

@@ -10,5 +10,6 @@
     Simulation,
     System,
 )
+from .internal.units import Units
 
 __version__ = "1.3.1"

flowermd/base/system.py

@@ -3,7 +3,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Union
 
 import gsd
 import mbuild as mb
@@ -14,7 +14,7 @@
 
 from flowermd.base.forcefield import BaseHOOMDForcefield, BaseXMLForcefield
 from flowermd.base.molecule import Molecule
-from flowermd.internal import check_return_iterable, validate_ref_value
+from flowermd.internal import Units, check_return_iterable, validate_unit
 from flowermd.internal.exceptions import ForceFieldError, MoleculeLoadError
 from flowermd.utils import (
     get_target_box_mass_density,
@@ -211,12 +211,10 @@ def reference_length(self, length):
 
         Parameters
         ----------
-        length : string or unyt.unyt_quantity, required
+        length : reference length * `flowermd.Units`, required
             The reference length of the system.
-            It can be provided in the following forms:
-            1) A string with the format of "value unit", for example "1 nm".
-            2) A unyt.unyt_quantity object with the correct dimension. For
-            example, unyt.unyt_quantity(1, "nm").
+            It can be provided in the following form of:
+            value * `flowermd.Units`, for example 1 * `flowermd.Units.angstrom`.
 
         """
         if self.auto_scale:
@@ -225,7 +223,7 @@ def reference_length(self, length):
                 "Setting reference length manually disables auto "
                 "scaling."
             )
-        validated_length = validate_ref_value(length, u.dimensions.length)
+        validated_length = validate_unit(length, u.dimensions.length)
         self._reference_values["length"] = validated_length
 
     @reference_energy.setter
@@ -234,12 +232,10 @@ def reference_energy(self, energy):
 
         Parameters
         ----------
-        energy : string or unyt.unyt_quantity, required
+        energy : reference energy * `flowermd.Units`, required
             The reference energy of the system.
-            It can be provided in the following forms:
-            1) A string with the format of "value unit", for example "1 kJ/mol".
-            2) A unyt.unyt_quantity object with the correct dimension. For
-            example, unyt.unyt_quantity(1, "kJ/mol").
+            It can be provided in the following form of:
+            value * `flowermd.Units`, for example 1 * `flowermd.Units.kcal/mol`.
 
         """
         if self.auto_scale:
@@ -248,7 +244,7 @@ def reference_energy(self, energy):
                 "Setting reference energy manually disables auto "
                 "scaling."
             )
-        validated_energy = validate_ref_value(energy, u.dimensions.energy)
+        validated_energy = validate_unit(energy, u.dimensions.energy)
         self._reference_values["energy"] = validated_energy
 
     @reference_mass.setter
@@ -257,21 +253,18 @@ def reference_mass(self, mass):
 
         Parameters
         ----------
-        mass : string or unyt.unyt_quantity, required
+        mass : reference mass * `flowermd.Units`, required
             The reference mass of the system.
-            It can be provided in the following forms:
-            1) A string with the format of "value unit", for example "1 amu".
-            2) A unyt.unyt_quantity object with the correct dimension. For
-            example, unyt.unyt_quantity(1, "amu").
-
+            It can be provided in the following form of:
+            value * `flowermd.Units`, for example 1 * `flowermd.Units.amu`.
         """
         if self.auto_scale:
             warnings.warn(
                 "`auto_scale` was set to True for this system. "
                 "Setting reference mass manually disables auto "
                 "scaling."
             )
-        validated_mass = validate_ref_value(mass, u.dimensions.mass)
+        validated_mass = validate_unit(mass, u.dimensions.mass)
         self._reference_values["mass"] = validated_mass
 
     @reference_values.setter
@@ -407,6 +400,23 @@ def to_gsd(self, file_name):
         with gsd.hoomd.open(file_name, "w") as traj:
             traj.append(self.hoomd_snapshot)
 
+    def save_reference_values(self, file_path="reference_values.pickle"):
+        """Save the reference values of the system to a pickle file.
+
+        Parameters
+        ----------
+        file_path : str, default "reference_values.pickle"
+            The path to save the pickle file to.
+
+        """
+        if not self.reference_values:
+            raise ValueError(
+                "Reference values have not been set. "
+                "See System.reference_values"
+            )
+        f = open(file_path, "wb")
+        pickle.dump(self.reference_values, f)
+
     def _convert_to_gmso(self):
         """Convert the mbuild system to a gmso system."""
         topology = from_mbuild(self.system)
@@ -618,11 +628,12 @@ class Pack(System):
 
     Parameters
     ----------
-    density : float, required
-        The desired density of the system (g/cm^3). Used to set the
+    density : float or unyt_quantity or flowermd.internal.Units, required
+        The desired density of the system. Used to set the
         target_box attribute. Can be useful when initializing
         systems at low density and running a shrink simulation
-        to achieve a target density.
+        to achieve a target density. If no unit is provided, assuming the
+        density is in g/cm**3.
     packing_expand_factor : int, default 5
         The factor by which to expand the box for packing.
     edge : float, default 0.2
@@ -657,7 +668,7 @@ class Pack(System):
     def __init__(
         self,
         molecules,
-        density: float,
+        density: Union[int, float, u.unyt_quantity, u.unyt_array, Units],
         base_units=dict(),
         packing_expand_factor=5,
         edge=0.2,
@@ -666,11 +677,11 @@ def __init__(
         fix_orientation=False,
         **kwargs,
     ):
-        if not isinstance(density, u.array.unyt_quantity):
-            self.density = density * u.Unit("g") / u.Unit("cm**3")
+        if isinstance(density, (int, float)):
             warnings.warn(
                 "Units for density were not given, assuming units of g/cm**3."
             )
+            self.density = density * Units.g_cm3
         else:
             self.density = density
         self.packing_expand_factor = packing_expand_factor
@@ -683,8 +694,8 @@ def __init__(
         )
 
     def _build_system(self, **kwargs):
-        mass_density = u.Unit("kg") / u.Unit("m**3")
-        number_density = u.Unit("m**-3")
+        mass_density = Units.kg_m3
+        number_density = Units.n_m3
         if self.density.units.dimensions == mass_density.dimensions:
             target_box = get_target_box_mass_density(
                 density=self.density, mass=self.mass

flowermd/internal/__init__.py

@@ -1,3 +1,4 @@
 # ruff: noqa: F401
 from .ff_utils import xml_to_gmso_ff
-from .utils import check_return_iterable, validate_ref_value
+from .units import Units
+from .utils import check_return_iterable, validate_unit

flowermd/internal/exceptions.py

@@ -53,7 +53,7 @@ def __init__(self, msg):
         super().__init__(msg)
 
 
-class ReferenceUnitError(Exception):
+class UnitError(Exception):
     def __init__(self, msg):
         super().__init__(msg)
 

flowermd/internal/units.py

@@ -0,0 +1,86 @@
+"""FlowerMD Unit class."""
+
+import unyt as u
+
+
+class Units:
+    """FlowerMD Unit class.
+
+    Example usage:
+    --------------
+
+    ::
+        length = 1.0 * flowermd.internal.Units.angstrom
+        energy = 1.0 * flowermd.internal.Units.kcal_mol
+        mass = 1.0 * flowermd.internal.Units.amu
+        time = 1.0 * flowermd.internal.Units.ps
+        temperature = 1.0 * flowermd.internal.Units.K
+
+
+    """
+
+    # length units
+    m = u.m
+    meter = u.m
+    cm = u.cm
+    centimeter = u.cm
+    nm = u.nm
+    nanometer = u.nm
+    pm = u.pm
+    picometer = u.pm
+    ang = u.angstrom
+    angstrom = u.angstrom
+
+    # energy units
+    J = u.J
+    Joule = u.J
+    kJ = u.kJ
+    cal = u.cal
+    kcal = u.kcal
+    kcal_mol = u.kcal / u.mol
+    kJ_mol = u.kJ / u.mol
+
+    # mass units
+    g = u.g
+    gram = u.g
+    amu = u.amu
+    kg = u.kg
+
+    mol = u.mol
+
+    # time units
+    s = u.s
+    second = u.s
+    ns = u.ns
+    nanosecond = u.ns
+    ps = u.ps
+    picosecond = u.ps
+    fs = u.fs
+    femtosecond = u.fs
+
+    # temperature units
+    K = u.K
+    Kelvin = u.K
+    C = u.degC
+    Celsius = u.degC
+    F = u.degF
+    Fahrenheit = u.degF
+
+    # mass density units
+    g_cm3 = u.g / u.cm**3
+    kg_m3 = u.kg / u.m**3
+    amu_A3 = u.amu / u.angstrom**3
+    # number density units
+    n_m3 = u.Unit("m**-3")
+    n_cm3 = u.Unit("cm**-3")
+    n_A3 = u.Unit("angstrom**-3")
+    n_nm3 = u.Unit("nm**-3")
+
+    # pressure units
+    atm = u.atm
+    bar = u.bar
+    Pa = u.Pa
+    kPa = u.kPa
+    MPa = u.MPa
+    GPa = u.GPa
+    psi = u.psi

flowermd/internal/utils.py

@@ -1,6 +1,8 @@
+import unyt
 import unyt as u
 
-from flowermd.internal.exceptions import ReferenceUnitError
+from flowermd.internal import Units
+from flowermd.internal.exceptions import UnitError
 
 """utils.py
    Internal utility methods for flowerMD.
@@ -19,77 +21,41 @@
         return [obj]
 
 
-def validate_ref_value(ref_value, dimension):
-    """Validate the reference value and checks the unit dimension.
-    This function validates the reference value. The reference value can be
-    provided in three ways:
-        1. An unyt_quantity instance.
-        2. A string with the value and unit , for example "1.0 g".
-        3. A string with the value and unit separated by a "/", for example
-        "1.0 kcal/mol".
+def validate_unit(value, dimension):
+    """Validate the unit and checks the unit dimension.
+
     Parameters
     ----------
-    ref_value : unyt_quantity or str; required
-        The reference value.
+    value : `unit value * flowermd.Units`; required
+        The unit value to be validated.
     dimension : unyt_dimension; required
-        The dimension of the reference value.
-
-    Returns
-    -------
-    The validated reference value as an unyt.unyt_quantity instance.
+        The dimension of the unit.
     """
 
-    def _is_valid_dimension(ref_unit):
-        if ref_unit.dimensions != dimension:
-            raise ReferenceUnitError(
-                f"Invalid unit dimension. The reference "
-                f"value must be in {dimension} "
-                f"dimension."
+    def _sample_unit_str(dimension):
+        if dimension == u.dimensions.temperature:
+            return "temperature = 300 * flowermd.Units.K"
+        elif dimension == u.dimensions.mass:
+            return "mass = 1.0 * flowermd.Units.g"
+        elif dimension == u.dimensions.length:
+            return "length = 1.0 * flowermd.Units.angstrom"
+        elif dimension == u.dimensions.time:
+            return "time = 1.0 * flowermd.Units.ps"
+        elif dimension == u.dimensions.energy:
+            return "energy = 1.0 * flowermd.Units.kcal_mol"
+
+    if isinstance(value, (u.unyt_quantity, unyt.unyt_array, Units)):
+        if value.units.dimensions != dimension:
+            raise UnitError(
+                f"Invalid unit dimension. The unit must be in "
+                f"{dimension} dimension. Check `flowermd.Units` for "
+                f"valid units."
             )
-        return True
-
-    def _parse_and_validate_unit(value, unit_str):
-        if hasattr(u, unit_str):
-            if unit_str == "amu":
-                u_unit = u.Unit("amu")
-            else:
-                u_unit = getattr(u, unit_str)
-            if _is_valid_dimension(u_unit):
-                return float(value) * u_unit
-        # if the unit contains "/" character, for example "g/mol", check if
-        # the unit is a valid unit and has the correct dimension.
-        if len(unit_str.split("/")) == 2:
-            unit1, unit2 = unit_str.split("/")
-            if hasattr(u, unit1) and hasattr(u, unit2):
-                comb_unit = getattr(u, unit1) / getattr(u, unit2)
-                if _is_valid_dimension(comb_unit):
-                    return float(value) * comb_unit
-        raise ReferenceUnitError(
-            f"Invalid reference value. Please provide "
-            f"a reference value with unit of "
-            f"{dimension} dimension."
-        )
-
-    def _is_float(num):
-        try:
-            return float(num)
-        except ValueError:
-            raise ValueError("The reference value is not a number.")
-
-    # if ref_value is an instance of unyt_quantity, check the dimension.
-    if isinstance(ref_value, u.unyt_quantity) and _is_valid_dimension(
-        ref_value.units
-    ):
-        return ref_value
-    # if ref_value is a string, check if it is a number and if it is, check if
-    # the unit exists in unyt and has the correct dimension.
-    elif isinstance(ref_value, str) and len(ref_value.split()) == 2:
-        value, unit_str = ref_value.split()
-        value = _is_float(value)
-        return _parse_and_validate_unit(value, unit_str)
+        return value
     else:
-        raise ReferenceUnitError(
-            f"Invalid reference value. Please provide "
-            f"a reference value with unit of "
-            f"{dimension} dimension."
+        raise UnitError(
+            "The unit value must be provided from the "
+            "`flowermd.Units` class. For example, "
+            f"{_sample_unit_str(dimension)}. Check "
+            "`flowermd.Units` for valid units."
         )