Ref (#25)

* removing openff dependency
* Add functions to generate PDB file from molecule and update test system generation
* add heavy atom threshold
* relax convergence criteria
* remove phosphorus

guardowl/analysis.py

@@ -5,49 +5,40 @@
 
 
 class PropertyCalculator:
-    """
-    A class for calculating various properties for a molecular dynamics trajectory.
-
-    Attributes
-    ----------
-    md_traj : md.Trajectory
-        The molecular dynamics trajectory.
-
-    Methods
-    -------
-    calculate_water_rdf()
-        Calculate the radial distribution function of water molecules.
-    monitor_water_bond_length()
-        Monitor the bond length between water molecules.
-    monitor_water_angle()
-        Monitor the angle between water molecules.
-    monitor_bond_length(bond_pairs)
-        Monitor the bond length for specific atom pairs.
-    monitor_angle_length(angle_list)
-        Monitor the angles for specific sets of atoms.
-
-    """
 
     def __init__(self, md_traj: md.Trajectory) -> None:
         """
-        Initialize the PropertyCalculator with a molecular dynamics trajectory.
+        Calculates various properties for a molecular dynamics trajectory.
 
         Parameters
         ----------
         md_traj : md.Trajectory
-            The molecular dynamics trajectory to be analyzed.
-
+            The molecular dynamics trajectory for analysis.
         """
         self.md_traj = md_traj
 
     def calculate_heat_capacity(
         self, total_energy: np.array, volumn: np.array
     ) -> float:
         """
-        Calculate the heat capacity of the trajectory.
-        C_p = <\Delta E^2> / k_B T^2 V
+        Calculates the heat capacity of the system using the trajectory data.
+
+        Parameters
+        ----------
+        total_energy : np.array
+            The total energy for each frame of the trajectory.
+        volume : np.array
+            The volume for each frame of the trajectory.
+
+        Returns
+        -------
+        float
+            The calculated heat capacity of the system.
+
+        Notes
+        -----
+        The formula used for calculation is C_p = <\Delta E^2> / (kB * T^2 * V).
         """
-        from openmm.unit import kelvin, nanometer, kilojoule_per_mole
         from .constants import kB, temperature
 
         mean_energy = np.mean(total_energy)
@@ -68,31 +59,32 @@ def calculate_isothermal_compressability_kappa_T(self):
         kappa_T = md.isothermal_compressability_kappa_T(
             self.md_traj, temperature.value_in_unit(kelvin)
         )
-        print(kappa_T)
         log.debug(f"isothermal_compressability_kappa_T: {kappa_T}")
         return kappa_T
 
-    def calculate_water_rdf(self):  # type: ignore
+    def calculate_water_rdf(self) -> np.ndarray:
         """
-        Calculate the radial distribution function (RDF) for water molecules in the trajectory.
+        Calculates the radial distribution function (RDF) for water molecules within the trajectory.
 
         Returns
         -------
         np.ndarray
-            The RDF values for the water molecules.
+            The RDF values for water molecules.
         """
-        oxygen_pairs = self.md_traj.top.select_pairs(
+        oxygen_pairs = self.md_traj.topology.select_pairs(
             "name O and water", "name O and water"
         )
-        bins = 300
-        r_max = 1
+        r_max = 1.0
         r_min = 0.01
+        bins = 300
 
-        mdtraj_rdf = md.compute_rdf(
-            self.md_traj, oxygen_pairs, (r_min, r_max), n_bins=bins
+        rdf_result = md.compute_rdf(
+            self.md_traj,
+            pairs=oxygen_pairs,
+            r_range=(r_min, r_max),
+            bin_width=(r_max - r_min) / bins,
         )
-
-        return mdtraj_rdf
+        return rdf_result
 
     def _extract_water_bonds(self) -> List[Tuple[int, int]]:
         bond_list = []
@@ -108,81 +100,77 @@ def _extract_bonds_except_water(self) -> List[Tuple[int, int]]:
                 bond_list.append((bond.atom1.index, bond.atom2.index))
         return bond_list
 
-    def monitor_water_bond_length(self):  # type: ignore
+    def monitor_water_bond_length(self) -> np.ndarray:
         """
-        Monitor the bond length between water molecules in the trajectory.
+        Monitors the bond length between water molecules throughout the trajectory.
 
         Returns
         -------
         np.ndarray
-            The bond lengths between water molecules.
-
+            The bond lengths between water molecules across all frames of the trajectory.
         """
-
         bond_list = self._extract_water_bonds()
         return self.monitor_bond_length(bond_list)
 
-    def monitor_bond_length_except_water(self):  # type: ignore
+    def monitor_bond_length_except_water(self):
+        """
+        This method monitors the bond deviation in molecules that are *not* water molecules throughout the trajectory.
+        """
         bond_list = self._extract_bonds_except_water()
         bond_length = self.monitor_bond_length(bond_list)
         compare_to = bond_length[0]
         bond_diff = np.abs(bond_length - compare_to)
         return bond_diff
 
-    def monitor_water_angle(self):  # type: ignore
+    def monitor_water_angle(self) -> np.ndarray:
         """
-        Monitor the angle between water molecules in the trajectory.
+        Monitors the angle between water molecules throughout the trajectory.
 
         Returns
         -------
         np.ndarray
-            The angles between water molecules.
-
+            The angles between water molecules across all frames of the trajectory.
         """
+        angle_list = self._extract_water_angles()
+        return self.monitor_angle_length(angle_list)
 
-        def _extract_angles() -> list:
-            """
-            Helper function to extract angles between water molecules.
-
-            Returns
-            -------
-            List[List[int]]
-                A list of atom index triplets representing the angles to monitor.
+    def _extract_water_angles(self) -> List[List[int]]:
+        """
+        Extracts sets of atom indices to compute angles within water molecules.
 
-            """
+        Returns
+        -------
+        List[List[int]]
+            A list of lists, where each inner list contains indices of three atoms forming an angle.
+        """
 
-            angle_list = []
-            for bond_1 in self.md_traj.top.bonds:
+        angle_list = []
+        for bond_1 in self.md_traj.top.bonds:
+            # skip if bond is not a water molecule
+            if bond_1.atom1.residue.name != "HOH":
+                continue
+            for bond_2 in self.md_traj.top.bonds:
                 # skip if bond is not a water molecule
-                if bond_1.atom1.residue.name != "HOH":
+                if bond_2.atom1.residue.name != "HOH":
                     continue
-                for bond_2 in self.md_traj.top.bonds:
-                    # skip if bond is not a water molecule
-                    if bond_2.atom1.residue.name != "HOH":
-                        continue
-                    water = {}
-                    for bond in [bond_1, bond_2]:
-                        water[bond.atom1.index] = bond.atom1
-                        water[bond.atom2.index] = bond.atom2
-                    # skip if atoms are not part of the same molecule
-                    if len(water.keys()) != 3:
-                        continue
-
-                    sorted_water = [
-                        water[key].index for key in sorted(water.keys(), reverse=True)
-                    ]  # oxygen is first
-                    angle_list.append(
-                        [sorted_water[1], sorted_water[2], sorted_water[0]]
-                    )
-
-            return [list(x) for x in set(tuple(x) for x in angle_list)]
-
-        angle_list = _extract_angles()
-        return self.monitor_angle_length(angle_list)
+                water = {}
+                for bond in [bond_1, bond_2]:
+                    water[bond.atom1.index] = bond.atom1
+                    water[bond.atom2.index] = bond.atom2
+                # skip if atoms are not part of the same molecule
+                if len(water.keys()) != 3:
+                    continue
+
+                sorted_water = [
+                    water[key].index for key in sorted(water.keys(), reverse=True)
+                ]  # oxygen is first
+                angle_list.append([sorted_water[1], sorted_water[2], sorted_water[0]])
 
-    def monitor_bond_length(self, bond_pairs: list):  # type: ignore
+        return [list(x) for x in set(tuple(x) for x in angle_list)]
+
+    def monitor_bond_length(self, bond_pairs: List[Tuple[int, int]]) -> np.ndarray:
         """
-        Monitor the bond length between specific atom pairs in the trajectory.
+        Monitors the bond length for specified pairs of atoms across the trajectory
 
         Parameters
         ----------
@@ -192,31 +180,42 @@ def monitor_bond_length(self, bond_pairs: list):  # type: ignore
         Returns
         -------
         np.ndarray
-            The bond lengths for the specified atom pairs.
+            An array of bond lengths for the specified atom pairs across all frames of the trajectory.
 
         """
-        bond_length = md.compute_distances(self.md_traj, bond_pairs)
-        return bond_length
+        bond_lengths = md.compute_distances(self.md_traj, atom_pairs=bond_pairs)
+        return bond_lengths
 
-    def monitor_angle_length(self, angle_list: list):  # type: ignore
+    def monitor_angle_length(self, angle_list: List[List[int]]) -> np.ndarray:
         """
-        Monitor the angles for specific sets of atoms in the trajectory.
+        Monitors the angles for specified sets of atoms throughout the trajectory.
 
         Parameters
         ----------
         angle_list : List[List[int]]
-            A list of atom index triplets whose angles are to be monitored.
+            A list of lists, where each inner list contains indices of three atoms forming an angle to be monitored.
 
         Returns
         -------
         np.ndarray
-            The angles for the specified sets of atoms.
+            An array of angles in degrees for the specified sets of atoms across all frames of the trajectory.
 
         """
-        angles = md.compute_angles(self.md_traj, angle_list) * (180 / np.pi)
+        angles = md.compute_angles(self.md_traj, angle_indices=angle_list) * (
+            180 / np.pi
+        )
         return angles
 
     def monitor_phi_psi(self) -> Tuple[np.ndarray, np.ndarray]:
-        phi = md.compute_phi(self.md_traj)[1]
-        psi = md.compute_psi(self.md_traj)[1]
-        return (phi, psi)
+        """
+        Monitors the phi and psi dihedral angles throughout the trajectory.
+
+        Returns
+        -------
+        Tuple[np.ndarray, np.ndarray]
+            Two arrays representing the phi and psi angles in radians for each frame of the trajectory.
+
+        """
+        _, phi_angles = md.compute_phi(self.md_traj)
+        _, psi_angle = md.compute_psi(self.md_traj)
+        return (phi_angles, psi_angles)

guardowl/benchmark.py

@@ -100,7 +100,7 @@ def run(self) -> None:
         print(f"{self.implementation=} {self.platform=}")
         potential = MLPotential(self.nnp)
 
-        system = self.system_factory.initialize_ml_system(
+        system = self.system_factory.initialize_system(
             potential,
             self.testsystem.topology,
             self.remove_constraints,
@@ -139,7 +139,6 @@ def run(self) -> None:
 
 
 class Benchmark:
-
     """
     A class to benchmark the performance of a neural network potential.
     It creates two processes, one that tracks the GPU memory usage and one that performs the benchmark.

guardowl/constants.py

@@ -1,7 +1,6 @@
 from openmm import unit
 from openmm.unit import Quantity
 from openmmtools.constants import kB
-
 # define units
 distance_unit = unit.angstrom
 time_unit = unit.femto * unit.seconds

guardowl/parameters.py

@@ -1,8 +1,9 @@
 from dataclasses import dataclass, field
+from typing import List, Literal, Optional, Union
+
 from openmm import Platform, System, unit
 from openmm.app import StateDataReporter
 from openmmtools.testsystems import TestSystem
-from typing import List, Union, Optional
 
 
 @dataclass
@@ -41,7 +42,7 @@ class MinimizationTestParameters(BaseParameters):
     ----------
     convergence_criteria : unit.Quantity
         The energy convergence criteria for the minimization process.
-    env : str
+    env : Literal["vacuum", "solution"]
         The environment of the simulation (e.g., 'vacuum', 'solution').
     device_index : int
         Index of the GPU device to use for the simulation.
@@ -54,10 +55,10 @@ class MinimizationTestParameters(BaseParameters):
 
     convergence_criteria: unit.Quantity = field(
         default_factory=lambda: unit.Quantity(
-            0.5, unit.kilojoule_per_mole / unit.angstrom
+            1 * unit.kilojoule_per_mole / unit.angstrom
         )
     )
-    env: str = "vacuum"
+    env: Literal["vacuum"] = "vacuum"
     device_index: int = 0
     temperature: unit.Quantity = field(
         default_factory=lambda: unit.Quantity(300, unit.kelvin)
@@ -66,6 +67,9 @@ class MinimizationTestParameters(BaseParameters):
     ensemble: str = "NVT"
 
 
+from typing import Literal
+
+
 @dataclass
 class StabilityTestParameters(BaseParameters):
     """Parameters for a stability test.
@@ -78,7 +82,7 @@ class StabilityTestParameters(BaseParameters):
         The duration of the test protocol.
     temperature : unit.Quantity
         The temperature at which the test is conducted.
-    env : str
+    env : Literal["vacuum", "solution"]
         The environment for the simulation (e.g., 'vacuum', 'solution').
     simulated_annealing : bool
         Flag to indicate if simulated annealing is used.
@@ -93,11 +97,11 @@ class StabilityTestParameters(BaseParameters):
 
     protocol_length: int
     temperature: unit.Quantity
-    env: str
+    env: Literal["vacuum", "solution"]
     simulated_annealing: bool
     state_data_reporter: StateDataReporter
     device_index: int = 0
-    ensemble: Optional[str] = None
+    ensemble: Literal["NVT", "NPT", "NVE", None] = None
 
 
 @dataclass