Merge pull request #76 from MolarVerse/dev

Release 1.0.6

.github/workflows/pylint.yml

@@ -93,7 +93,7 @@ jobs:
 
     #add changes of .github/.pylint_cache to the commit if it is based on push event
     - name: Add .github/.pylint_cache to the commit
-      if: github.event_name == 'push' && github.ref_name != 'main'
+      if: (github.event_name == 'push' || github.event.pull_request.merged == true) && github.ref_name != 'main'
       run: |
         git config --global user.email "github-actions[bot]@users.noreply.github.com"
         git config --global user.name "github-actions[bot]"

PQAnalysis/atomic_system/atomic_system.py

@@ -15,11 +15,12 @@
 
 from PQAnalysis.core import Atom, Atoms, Cell, distance
 from PQAnalysis.topology import Topology
-from PQAnalysis.types import PositiveReal, PositiveInt
+from PQAnalysis.types import PositiveReal, PositiveInt, Bool
 from PQAnalysis.type_checking import runtime_type_checking
 from PQAnalysis.exceptions import PQNotImplementedError
 from PQAnalysis.utils.random import get_random_seed
 from PQAnalysis.utils.custom_logging import setup_logger
+from PQAnalysis.utils.math import allclose_vectorized
 from PQAnalysis import __package_name__
 
 from PQAnalysis.types import (
@@ -533,7 +534,7 @@ def copy(self) -> "AtomicSystem":
             topology=self.topology
         )
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: Any) -> Bool:
         """
         Checks whether the AtomicSystem is equal to another AtomicSystem.
 
@@ -547,6 +548,34 @@ def __eq__(self, other: Any) -> bool:
         bool
             Whether the AtomicSystem is equal to the other AtomicSystem.
         """
+
+        return self.isclose(other)
+
+    @runtime_type_checking
+    def isclose(
+        self,
+        other: Any,
+        rtol: PositiveReal = 1e-9,
+        atol: PositiveReal = 0.0,
+    ) -> Bool:
+        """
+        Checks whether the AtomicSystem is close to another AtomicSystem.
+
+        Parameters
+        ----------
+        other : AtomicSystem
+            The other AtomicSystem to compare to.
+        rtol : PositiveReal, optional
+            The relative tolerance, by default 1e-9
+        atol : PositiveReal, optional
+            The absolute tolerance, by default 0.0
+
+        Returns
+        -------
+        bool
+            Whether the AtomicSystem is close to the other AtomicSystem.
+        """
+
         if not isinstance(other, AtomicSystem):
             return False
 
@@ -556,19 +585,23 @@ def __eq__(self, other: Any) -> bool:
         if self.topology != other.topology:
             return False
 
-        if self.cell != other.cell:
+        if not self.cell.isclose(other.cell, rtol=rtol, atol=atol):
             return False
 
-        if not np.allclose(self.pos, other.pos):
+        if not allclose_vectorized(self.pos, other.pos, rtol=rtol, atol=atol):
             return False
 
-        if not np.allclose(self.vel, other.vel):
+        if not allclose_vectorized(self.vel, other.vel, rtol=rtol, atol=atol):
             return False
 
-        if not np.allclose(self.forces, other.forces):
+        if not allclose_vectorized(
+            self.forces, other.forces, rtol=rtol, atol=atol
+        ):
             return False
 
-        if not np.allclose(self.charges, other.charges):
+        if not allclose_vectorized(
+            self.charges, other.charges, rtol=rtol, atol=atol
+        ):
             return False
 
         return True

PQAnalysis/core/cell/cell.py

@@ -13,7 +13,8 @@
 from beartype.vale import Is
 
 from PQAnalysis.type_checking import runtime_type_checking
-from PQAnalysis.types import Np3x3NumberArray, Np2DNumberArray, NpnDNumberArray
+from PQAnalysis.types import Np3x3NumberArray, Np2DNumberArray, NpnDNumberArray, PositiveReal, Bool
+from PQAnalysis.utils.math import allclose_vectorized
 
 from ._standard_properties import _StandardPropertiesMixin
 
@@ -98,7 +99,8 @@ def bounding_edges(self) -> Np2DNumberArray:
         for i, x in enumerate([-0.5, 0.5]):
             for j, y in enumerate([-0.5, 0.5]):
                 for k, z in enumerate([-0.5, 0.5]):
-                    edges[i*4+j*2+k, :] = self.box_matrix @ np.array([x, y, z])
+                    edges[i * 4 + j * 2 +
+                          k, :] = self.box_matrix @ np.array([x, y, z])
 
         return edges
 
@@ -107,8 +109,7 @@ def volume(self) -> Real:
         """volume: The volume of the unit cell."""
         with warnings.catch_warnings():
             warnings.filterwarnings(
-                "ignore",
-                message="overflow encountered in det"
+                "ignore", message="overflow encountered in det"
             )
             volume = np.linalg.det(self.box_matrix)
 
@@ -154,7 +155,7 @@ def image(self, pos: NpnDNumberArray) -> NpnDNumberArray:
 
         return np.reshape(pos, original_shape)
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: Any) -> Bool:
         """
         Checks if the Cell is equal to another Cell.
 
@@ -165,16 +166,54 @@ def __eq__(self, other: Any) -> bool:
 
         Returns
         -------
-        bool
+        Bool
             True if the Cells are equal, False otherwise.
         """
 
+        return self.isclose(other)
+
+    @runtime_type_checking
+    def isclose(
+        self,
+        other: Any,
+        rtol: PositiveReal = 1e-9,
+        atol: PositiveReal = 0.0,
+    ) -> Bool:
+        """
+        Checks if the Cell is close to another Cell.
+
+        Parameters
+        ----------
+        other : Cell
+            The Cell to compare with.
+        rtol : PositiveReal, optional
+            The relative tolerance parameter. Default is 1e-9.
+        atol : PositiveReal, optional
+            The absolute tolerance parameter. Default is 0.0.
+
+        Returns
+        -------
+        Bool
+            True if the Cells are close, False otherwise.
+        """
+
         if not isinstance(other, Cell):
             return False
 
-        is_equal = True
-        is_equal &= np.allclose(self.box_lengths, other.box_lengths)
-        is_equal &= np.allclose(self.box_angles, other.box_angles)
+        is_equal = allclose_vectorized(
+            self.box_lengths,
+            other.box_lengths,
+            rtol=rtol,
+            atol=atol,
+        )
+
+        is_equal &= allclose_vectorized(
+            self.box_angles,
+            other.box_angles,
+            rtol=rtol,
+            atol=atol,
+        )
+
         return is_equal
 
     def __str__(self) -> str:
@@ -233,21 +272,16 @@ def init_from_box_matrix(cls, box_matrix: Np3x3NumberArray) -> "Cell":
         beta = np.arccos(box_matrix[0][2] / z)
         alpha = np.arccos(
             (
-            box_matrix[0][1] * box_matrix[0][2] +
-            box_matrix[1][1] * box_matrix[1][2]
+                box_matrix[0][1] * box_matrix[0][2] +
+                box_matrix[1][1] * box_matrix[1][2]
             ) / (y * z)
         )
 
         print(alpha, beta, gamma)
         print(np.rad2deg(alpha), np.rad2deg(beta), np.rad2deg(gamma))
 
         return cls(
-            x,
-            y,
-            z,
-            np.rad2deg(alpha),
-            np.rad2deg(beta),
-            np.rad2deg(gamma)
+            x, y, z, np.rad2deg(alpha), np.rad2deg(beta), np.rad2deg(gamma)
         )
 
 
@@ -256,6 +290,5 @@ def init_from_box_matrix(cls, box_matrix: Np3x3NumberArray) -> "Cell":
 Cells = NewType(
     "Cells",
     Annotated[list,
-    Is[lambda list: all(isinstance(atom,
-    Cell) for atom in list)]]
+              Is[lambda list: all(isinstance(atom, Cell) for atom in list)]]
 )

PQAnalysis/topology/shake_topology.py

@@ -2,6 +2,7 @@
 A module containing the ShakeTopologyGenerator class.
 """
 
+import logging
 import numpy as np
 
 from beartype.typing import List
@@ -10,6 +11,9 @@
 from PQAnalysis.types import Np1DIntArray, Np2DIntArray
 from PQAnalysis.io import BaseWriter, FileWritingMode
 from PQAnalysis.type_checking import runtime_type_checking
+from PQAnalysis.utils.custom_logging import setup_logger
+from PQAnalysis.exceptions import PQValueError
+from PQAnalysis import __package_name__
 
 from .selection import SelectionCompatible, Selection
 
@@ -21,6 +25,9 @@ class ShakeTopologyGenerator:
     A class for generating the shake topology for a given trajectory
     """
 
+    logger = logging.getLogger(__package_name__).getChild(__qualname__)
+    logger = setup_logger(logger)
+
     @runtime_type_checking
     def __init__(
         self,
@@ -46,6 +53,7 @@ def __init__(
         self.target_indices = None
         self.distances = None
         self._topology = None
+        self.line_comments = None
 
     @runtime_type_checking
     def generate_topology(self, trajectory: Trajectory) -> None:
@@ -69,8 +77,7 @@ def generate_topology(self, trajectory: Trajectory) -> None:
         self._topology = trajectory.topology
 
         indices = self.selection.select(
-            self._topology,
-            self._use_full_atom_info
+            self._topology, self._use_full_atom_info
         )
 
         target_indices, distances = atomic_system.nearest_neighbours(
@@ -96,7 +103,8 @@ def generate_topology(self, trajectory: Trajectory) -> None:
     @runtime_type_checking
     def average_equivalents(
         self,
-        indices: List[Np1DIntArray] | Np2DIntArray
+        indices: List[Np1DIntArray] | Np2DIntArray,
+        comments: List[str] | None = None
     ) -> None:
         """
         Averages the distances for equivalent atoms.
@@ -109,6 +117,8 @@ def average_equivalents(
         ----------
         indices : List[Np1DIntArray] | Np2DIntArray
             The indices of the equivalent atoms.
+        comments : List[str], optional
+            The line comments for the averaged distances, by default None
         """
 
         for equivalent_indices in indices:
@@ -118,6 +128,46 @@ def average_equivalents(
 
             self.distances[_indices] = mean_distance
 
+        if comments is not None:
+            if len(comments) != len(indices):
+                self.logger.error(
+                    "The number of comments does not match the number of indices.",
+                    exception=PQValueError
+                )
+
+            self.line_comments = [""] * len(self.indices)
+
+            for i, equivalent_indices in enumerate(indices):
+                _indices = np.nonzero(
+                    np.in1d(self.indices, equivalent_indices)
+                )[0]
+
+                for index in _indices:
+                    self.line_comments[index] = comments[i]
+
+    @runtime_type_checking
+    def add_comments(self, comments: List[str]) -> None:
+        """
+        Adds comments to the topology.
+
+        Parameters
+        ----------
+        comments : List[str]
+            The comments to add to each line of the shake topology.
+        """
+
+        if self.indices is None or len(comments) != len(self.indices):
+            self.logger.error(
+                "The number of comments does not match the number of indices.",
+                exception=PQValueError
+            )
+
+        if self.line_comments is None:
+            self.line_comments = [""] * len(self.indices)
+
+        for i, comment in enumerate(comments):
+            self.line_comments[i] = comment
+
     @runtime_type_checking
     def write_topology(
         self,
@@ -147,8 +197,8 @@ def write_topology(
 
         print(
             (
-            f"SHAKE {len(self.indices)}  "
-            f"{len(np.unique(self.target_indices))}  0"
+                f"SHAKE {len(self.indices)}  "
+                f"{len(np.unique(self.target_indices))}  0"
             ),
             file=writer.file
         )
@@ -157,7 +207,16 @@ def write_topology(
             target_index = self.target_indices[i]
             distance = self.distances[i]
 
-            print(f"{index+1} {target_index+1} {distance}", file=writer.file)
+            print(
+                f"{index+1} {target_index+1} {distance}",
+                end="",
+                file=writer.file
+            )
+
+            if self.line_comments is not None:
+                print(f"  # {self.line_comments[i]}", file=writer.file)
+            else:
+                print("", file=writer.file)
 
         print("END", file=writer.file)