refactor: modernize Molecule class with cached properties and vectorized operations

- Replace manual property updates with @cached_property for total_mass, com, and inertia_tensor
- Introduce invalidate_computed_properties decorator to auto-invalidate cached properties on atom changes
- Vectorize distances_between_atoms, sizes_of_molecule, and minimum_distance calculations using numpy
- Unify and clarify center of mass and standard orientation methods (move_center_of_mass_to_origin, rotate_to_standard_orientation)
- Remove redundant or outdated code, improve typing and error handling
- Update dependent files and tests to use new method names and behaviors

diceplayer/environment/atom.py

@@ -1,4 +1,4 @@
-from diceplayer.utils.ptable import PTable
+from diceplayer.utils.ptable import PTable, AtomInfo
 
 from dataclasses import dataclass
 
@@ -21,3 +21,7 @@ class Atom:
     @property
     def mass(self) -> float:
         return PTable.get_atomic_mass(self.na)
+
+    @property
+    def atom_info(self) -> AtomInfo:
+        return PTable.get_from_atomic_number(self.na)

diceplayer/environment/molecule.py

@@ -1,25 +1,24 @@
 from __future__ import annotations
 
-from typing_extensions import TYPE_CHECKING
+from dataclasses import dataclass, Field
-
+from functools import cached_property
-
-if TYPE_CHECKING:
-    from nptyping import Float, NDArray, Shape
 
 from diceplayer import logger
 from diceplayer.environment import Atom
-from diceplayer.utils.misc import BOHR2ANG
+from diceplayer.utils.cache import invalidate_computed_properties
+from diceplayer.utils.misc import BOHR2ANG, EA_2_DEBYE
 from diceplayer.utils.ptable import GHOST_NUMBER
 
 import numpy as np
 import numpy.typing as npt
 from numpy.linalg import linalg
-from typing_extensions import Any, List, Tuple, Union
+from typing_extensions import List, Tuple, Self
 
 import math
 from copy import deepcopy
 
 
+@dataclass
 class Molecule:
     """
     Molecule class declaration. This class is used throughout the DicePlayer program to represent molecules.
@@ -27,35 +26,59 @@ class Molecule:
     Atributes:
         molname (str): The name of the represented molecule
         atom (List[Atom]): List of atoms of the represented molecule
-        position (npt.NDArray[np.float64]): The position relative to the internal atoms of the represented molecule
-        energy (npt.NDArray[np.float64]): The energy of the represented molecule
-        gradient (npt.NDArray[np.float64]): The first derivative of the energy relative to the position
-        hessian (npt.NDArray[np.float64]): The second derivative of the energy relative to the position
         total_mass (int): The total mass of the molecule
         com (npt.NDArray[np.float64]): The center of mass of the molecule
+        inertia_tensor (npt.NDArray[np.float64]): The inertia tensor of the molecule
     """
+    molname: str
+    atom: List[Atom] = Field(default_factory=list)
 
-    def __init__(self, molname: str) -> None:
+    @cached_property
+    def total_mass(self) -> float:
+        return sum(atom.mass for atom in self.atom)
+
+    @cached_property
+    def com(self) -> npt.NDArray[np.float64]:
+        com = np.zeros(3)
+
+        for atom in self.atom:
+            com += atom.mass * np.array([atom.rx, atom.ry, atom.rz])
+
+        com = com / self.total_mass
+
+        return com
+
+    @cached_property
+    def inertia_tensor(self) -> npt.NDArray[np.float64]:
         """
-        The constructor function __init__ is used to create new instances of the Molecule class.
+        Calculates the inertia tensor of the molecule.
 
-        Args:
+        Returns:
-            molname (str): Molecule name
+            npt.NDArray[np.float64]: inertia tensor of the molecule.
         """
-        self.molname: str = molname
+        inertia_tensor = np.zeros((3, 3), dtype=np.float64)
 
-        self.atom: List[Atom] = []
+        for atom in self.atom:
-        self.position: npt.NDArray[np.float64]
+            dx = atom.rx - self.com[0]
-        self.energy: npt.NDArray[np.float64]
+            dy = atom.ry - self.com[1]
-        self.gradient: npt.NDArray[np.float64]
+            dz = atom.rz - self.com[2]
-        self.hessian: npt.NDArray[np.float64]
 
-        self.ghost_atoms: List[Atom] = []
+            inertia_tensor[0, 0] += atom.mass * (dy**2 + dz**2)
-        self.lp_atoms: List[Atom] = []
+            inertia_tensor[1, 1] += atom.mass * (dz**2 + dx**2)
+            inertia_tensor[2, 2] += atom.mass * (dx**2 + dy**2)
 
-        self.total_mass: int = 0
+            inertia_tensor[0, 1] -= atom.mass * dx * dy
-        self.com: Union[None, npt.NDArray[np.float64]] = None
+            inertia_tensor[0, 2] -= atom.mass * dx * dz
+            inertia_tensor[1, 2] -= atom.mass * dy * dz
 
+        # enforce symmetry
+        inertia_tensor[1, 0] = inertia_tensor[0, 1]
+        inertia_tensor[2, 0] = inertia_tensor[0, 2]
+        inertia_tensor[2, 1] = inertia_tensor[1, 2]
+
+        return inertia_tensor
+
+    @invalidate_computed_properties()
     def add_atom(self, a: Atom) -> None:
         """
         Adds Atom instance to the molecule.
@@ -65,25 +88,20 @@ class Molecule:
         """
 
         self.atom.append(a)
-        self.total_mass += a.mass
 
-        self.center_of_mass()
+    @invalidate_computed_properties()
-
+    def remove_atom(self, a: Atom) -> None:
-    def center_of_mass(self) -> npt.NDArray[np.float64]:
         """
-        Calculates the center of mass of the molecule
+        Removes Atom instance from the molecule.
+
+        Args:
+            a (Atom): Atom instance to be removed from atom list.
         """
 
-        self.com = np.zeros(3)
+        self.atom.remove(a)
 
-        for atom in self.atom:
+    @invalidate_computed_properties()
-            self.com += atom.mass * np.array([atom.rx, atom.ry, atom.rz])
+    def move_center_of_mass_to_origin(self) -> None:
-
-        self.com = self.com / self.total_mass
-
-        return self.com
-
-    def center_of_mass_to_origin(self) -> None:
         """
         Updated positions based on the center of mass of the molecule
         """
@@ -92,7 +110,28 @@ class Molecule:
             atom.ry -= self.com[1]
             atom.rz -= self.com[2]
 
-        self.center_of_mass()
+    @invalidate_computed_properties()
+    def rotate_to_standard_orientation(self) -> None:
+        """
+        Rotates the molecule to the standard orientation
+        """
+
+        self.move_center_of_mass_to_origin()
+        evals, evecs = self.principal_axes()
+
+        if np.isclose(linalg.det(evecs), -1):
+            evecs[:, 2] *= -1
+
+        if not np.isclose(linalg.det(evecs), 1):
+            raise RuntimeError(
+                "Error: could not make a rotation matrix while adopting the standard orientation"
+            )
+
+        coords = np.array([(a.rx, a.ry, a.rz) for a in self.atom])
+        rotated = coords @ evecs.T
+
+        for atom, pos in zip(self.atom, rotated):
+            atom.rx, atom.ry, atom.rz = pos
 
     def charges_and_dipole(self) -> List[float]:
         """
@@ -103,7 +142,6 @@ class Molecule:
             second dipole, third dipole and total dipole.
         """
 
-        eA_to_Debye = 1 / 0.20819434
         charge = 0
         dipole = np.zeros(3)
         for atom in self.atom:
@@ -111,7 +149,7 @@ class Molecule:
             dipole += atom.chg * position
             charge += atom.chg
 
-        dipole *= eA_to_Debye
+        dipole *= EA_2_DEBYE
         total_dipole = math.sqrt(dipole[0] ** 2 + dipole[1] ** 2 + dipole[2] ** 2)
 
         return [charge, dipole[0], dipole[1], dipole[2], total_dipole]
@@ -123,52 +161,11 @@ class Molecule:
         Returns:
            NDArray[Shape["Any,Any"],Float]: distances between the atoms.
         """
+        coords = np.array([(a.rx, a.ry, a.rz) for a in self.atom], dtype=np.float64)
+        diff = coords[:, None, :] - coords[None, :, :]
+        return np.linalg.norm(diff, axis=-1)
 
-        distances = []
+    def principal_axes(self) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
-        dim = len(self.atom)
-        for index1, atom1 in enumerate(self.atom):
-            for index2, atom2 in enumerate(self.atom):
-                if index1 != index2:
-                    dx = atom1.rx - atom2.rx
-                    dy = atom1.ry - atom2.ry
-                    dz = atom1.rz - atom2.rz
-                    distances.append(math.sqrt(dx**2 + dy**2 + dz**2))
-
-        return np.array(distances).reshape(dim, dim - 1)
-
-    def inertia_tensor(self) -> npt.NDArray[np.float64]:
-        """
-        Calculates the inertia tensor of the molecule.
-
-        Returns:
-            npt.NDArray[np.float64]: inertia tensor of the molecule.
-        """
-
-        self.center_of_mass()
-
-        Ixx = 0.0
-        Ixy = 0.0
-        Ixz = 0.0
-        Iyy = 0.0
-        Iyz = 0.0
-        Izz = 0.0
-
-        for atom in self.atom:
-            dx = atom.rx - self.com[0]
-            dy = atom.ry - self.com[1]
-            dz = atom.rz - self.com[2]
-
-            Ixx += atom.mass * (dy**2 + dz**2)
-            Iyy += atom.mass * (dz**2 + dx**2)
-            Izz += atom.mass * (dx**2 + dy**2)
-
-            Ixy += atom.mass * dx * dy * -1
-            Ixz += atom.mass * dx * dz * -1
-            Iyz += atom.mass * dy * dz * -1
-
-        return np.array([[Ixx, Ixy, Ixz], [Ixy, Iyy, Iyz], [Ixz, Iyz, Izz]])
-
-    def principal_axes(self) -> Tuple[np.ndarray, np.ndarray]:
         """
         Calculates the principal axes of the molecule
 
@@ -178,7 +175,11 @@ class Molecule:
         """
 
         try:
-            evals, evecs = linalg.eigh(self.inertia_tensor())
+            evals, evecs = linalg.eigh(self.inertia_tensor)
+
+            idx = np.argsort(evals)
+            evals = evals[idx]
+            evecs = evecs[:, idx]
         except ValueError:
             raise RuntimeError(
                 "Error: diagonalization of inertia tensor did not converge"
@@ -186,22 +187,16 @@ class Molecule:
 
         return evals, evecs
 
-    def read_position(self) -> np.ndarray:
+    def read_position(self) -> npt.NDArray[np.float64]:
         """Reads the position of the molecule from the position values of the atoms
 
         Returns:
             np.ndarray: internal position relative to atoms of the molecule
         """
+        coords = np.array([(a.rx, a.ry, a.rz) for a in self.atom], dtype=np.float64)
+        return coords.ravel() * BOHR2ANG
 
-        position_list = []
+    def update_charges(self, charges: npt.NDArray[np.float64]) -> int:
-        for atom in self.atom:
-            position_list.extend([atom.rx, atom.ry, atom.rz])
-        position = np.array(position_list)
-        position *= BOHR2ANG
-
-        return position
-
-    def update_charges(self, charges: NDArray) -> int:
         """
         Updates the charges of the atoms of the molecule and
         returns the max difference between the new and old charges
@@ -213,34 +208,6 @@ class Molecule:
 
         return diff
 
-    # @staticmethod
-    # def update_hessian(
-    #         step: np.ndarray,
-    #         cur_gradient: np.ndarray,
-    #         old_gradient: np.ndarray,
-    #         hessian: np.ndarray,
-    # ) -> np.ndarray:
-    #     """
-    #     Updates the Hessian of the molecule based on the current hessian, the current gradient and the previous gradient
-    #
-    #     Args:
-    #         step (np.ndarray): step value of the iteration
-    #         cur_gradient (np.ndarray): current gradient
-    #         old_gradient (np.ndarray): previous gradient
-    #         hessian (np.ndarray): current hessian
-    #
-    #     Returns:
-    #         np.ndarray: updated hessian of the molecule
-    #     """
-    #
-    #     dif_gradient = cur_gradient - old_gradient
-    #
-    #     mat1 = 1 / np.dot(dif_gradient, step) * np.matmul(dif_gradient.T, dif_gradient)
-    #     mat2 = 1 / np.dot(step, np.matmul(hessian, step.T).T)
-    #     mat2 *= np.matmul(np.matmul(hessian, step.T), np.matmul(step, hessian))
-    #
-    #     return hessian + mat1 - mat2
-
     def sizes_of_molecule(self) -> List[float]:
         """
         Calculates sides of the smallest box that the molecule could fit
@@ -248,56 +215,10 @@ class Molecule:
         Returns:
             List[float]: list of the sizes of the molecule
         """
+        coords = np.array([(a.rx, a.ry, a.rz) for a in self.atom], dtype=np.float64)
+        return (coords.max(axis=0) - coords.min(axis=0)).tolist()
 
-        x_list = []
+    def translate(self, vector: np.ndarray) -> Self:
-        y_list = []
-        z_list = []
-
-        for atom in self.atom:
-            x_list.append(atom.rx)
-            y_list.append(atom.ry)
-            z_list.append(atom.rz)
-
-        x_max = max(x_list)
-        x_min = min(x_list)
-        y_max = max(y_list)
-        y_min = min(y_list)
-        z_max = max(z_list)
-        z_min = min(z_list)
-
-        sizes = [x_max - x_min, y_max - y_min, z_max - z_min]
-
-        return sizes
-
-    def standard_orientation(self) -> None:
-        """
-        Rotates the molecule to the standard orientation
-        """
-
-        self.center_of_mass_to_origin()
-        evals, evecs = self.principal_axes()
-
-        if round(linalg.det(evecs)) == -1:
-            evecs[0, 2] *= -1
-            evecs[1, 2] *= -1
-            evecs[2, 2] *= -1
-
-        if round(linalg.det(evecs)) != 1:
-            raise RuntimeError(
-                "Error: could not make a rotation matrix while adopting the standard orientation"
-            )
-
-        rot_matrix = evecs.T
-
-        for atom in self.atom:
-            position = np.array([atom.rx, atom.ry, atom.rz])
-            new_position = np.matmul(rot_matrix, position.T).T
-
-            atom.rx = new_position[0]
-            atom.ry = new_position[1]
-            atom.rz = new_position[2]
-
-    def translate(self, vector: np.ndarray) -> "Molecule":
         """
         Creates a new Molecule object where its' atoms has been translated by a vector
 
@@ -307,6 +228,9 @@ class Molecule:
         Returns:
             Molecule: new Molecule object translated by a vector
         """
+        vec = np.asarray(vector, dtype=np.float64)
+        if vec.shape != (3,):
+            raise ValueError("translation vector must be shape (3,)")
 
         new_molecule = deepcopy(self)
 
@@ -327,7 +251,6 @@ class Molecule:
                 self.com[0], self.com[1], self.com[2]
             )
         )
-        self.inertia_tensor()
         evals, evecs = self.principal_axes()
 
         logger.info(
@@ -368,7 +291,7 @@ class Molecule:
             )
         )
 
-    def minimum_distance(self, molec: "Molecule") -> float:
+    def minimum_distance(self, molec: Self) -> float:
         """
         Return the minimum distance between two molecules
 
@@ -378,15 +301,12 @@ class Molecule:
         Returns:
             float: minimum distance between the two molecules
         """
+        coords_a = np.array([(a.rx, a.ry, a.rz) for a in self.atom if a.na != GHOST_NUMBER])
+        coords_b = np.array([(a.rx, a.ry, a.rz) for a in molec.atom if a.na != GHOST_NUMBER])
 
-        distances = []
+        if len(coords_a) == 0 or len(coords_b) == 0:
-        for atom1 in self.atom:
+            raise ValueError("No real atoms to compare")
-            if atom1.na != GHOST_NUMBER:
-                for atom2 in molec.atom:
-                    if atom2.na != GHOST_NUMBER:
-                        dx = atom1.rx - atom2.rx
-                        dy = atom1.ry - atom2.ry
-                        dz = atom1.rz - atom2.rz
-                        distances.append(math.sqrt(dx**2 + dy**2 + dz**2))
 
-        return min(distances)
+        diff = coords_a[:, None, :] - coords_b[None, :, :]
+        d2 = np.sum(diff ** 2, axis=-1)
+        return np.sqrt(d2.min())

diceplayer/environment/system.py

@@ -33,7 +33,7 @@ class System:
         Args:
             m (Molecule): The instance of the new type of molecule
         """
-        if isinstance(m, Molecule) is False:
+        if not isinstance(m, Molecule):
             raise TypeError("Error: molecule is not a Molecule instance")
         self.molecule.append(m)
 
@@ -71,8 +71,8 @@ class System:
         new_projecting_mol = deepcopy(projecting_mol)
         new_reference_mol = deepcopy(reference_mol)
 
-        new_projecting_mol.center_of_mass_to_origin()
+        new_projecting_mol.move_center_of_mass_to_origin()
-        new_reference_mol.center_of_mass_to_origin()
+        new_reference_mol.move_center_of_mass_to_origin()
 
         x = []
         y = []
@@ -129,90 +129,10 @@ class System:
             new_projecting_mol.atom[i].ry = x[i, 1]
             new_projecting_mol.atom[i].rz = x[i, 2]
 
-        reference_mol.center_of_mass()
-
         projected_mol = new_projecting_mol.translate(reference_mol.com)
 
         return rmsd, projected_mol
 
-    # def center_of_mass_distance(self, a: int, b: int) -> float:
-    #     """
-    #     Calculates the distance between the center of mass of two molecules
-    #
-    #     Args:
-    #         a (Molecule): First Molecule Instance
-    #         b (Molecule): Second Molecule Instance
-    #
-    #     Returns:
-    #         float: module of the distance between the two center of masses
-    #     """
-    #
-    #     com1 = self.molecule[a].center_of_mass()
-    #     com2 = self.molecule[b].center_of_mass()
-    #     dx = com1[0] - com2[0]
-    #     dy = com1[1] - com2[1]
-    #     dz = com1[2] - com2[2]
-    #     distance = math.sqrt(dx**2 + dy**2 + dz**2)
-    #
-    #     return distance
-
-    # def nearest_image(
-    #     self,
-    #     index_r: int,
-    #     index_m: int,
-    #     lx: float,
-    #     ly: float,
-    #     lz: float,
-    #     criterium=None,
-    # ) -> Tuple[float, Molecule]:
-    #
-    #     if criterium in None:
-    #         criterium = "com"
-    #
-    #     if criterium != "com" and criterium != "min":
-    #         raise RuntimeError("Error in value passed to function nearest_image")
-    #
-    #     min_dist = 1e20
-    #
-    #     for i in range(-1, 2):
-    #         for j in range(-1, 2):
-    #             for k in range(-1, 2):
-    #
-    #                 tr_vector = [i * lx, j * ly, k * lz]
-    #                 self.add_molecule(self.molecule[index_m].translate(tr_vector))
-    #
-    #                 if criterium == "com":
-    #                     dist = self.center_of_mass_distance(index_r, -1)
-    #                 else:
-    #                     dist = self.minimum_distance(index_r, -1)
-    #
-    #                 if dist < min_dist:
-    #                     min_dist = dist
-    #                     nearestmol = deepcopy(self.molecule[-1])
-    #
-    #     self.molecule.pop(-1)
-    #
-    #     return min_dist, nearestmol
-
-    # def print_geom(self, cycle: int, fh: TextIO) -> None:
-    #     """
-    #             Print the geometry of the molecule in the Output file
-    #
-    #             Args:
-    #                 cycle (int): Number of the cycle
-    #                 fh (TextIO): Output file
-    #             """
-    #
-    #     fh.write("Cycle # {}\n".format(cycle))
-    #     fh.write("Number of site: {}\n".format(len(self.molecule[0].atom)))
-    #     for atom in self.molecule[0].atom:
-    #         symbol = atomsymb[atom.na]
-    #         fh.write(
-    #             "{:<2s}    {:>10.6f}  {:>10.6f}  {:>10.6f}\n".format(
-    #                 symbol, atom.rx, atom.ry, atom.rz
-    #             )
-    #         )
-    #
     def print_charges_and_dipole(self, cycle: int) -> None:
         """
         Print the charges and dipole of the molecule in the Output file

diceplayer/player.py

@@ -76,7 +76,7 @@ class Player:
                 "\n    Translating and rotating molecule to standard orientation..."
             )
 
-            mol.standard_orientation()
+            mol.rotate_to_standard_orientation()
             logger.info("\n Done")
             logger.info("\nNew values:\n")
             mol.print_mol_info()

diceplayer/utils/cache.py

@@ -0,0 +1,29 @@
+from functools import cached_property
+
+
+def invalidate_computed_properties():
+    """
+    Decorator function to invalidate the cached properties of the molecule when a new atom is added
+
+    Args:
+        properties (list[str]): list of the names of the properties to be invalidated
+    """
+
+    def get_cached_properies(cls: type) -> set[str]:
+        return {
+            name
+            for name, value in cls.__dict__.items()
+            if isinstance(value, cached_property)
+        }
+
+    def decorator(func):
+        def wrapper(self, *args, **kwargs):
+            result = func(self, *args, **kwargs)
+            for prop in get_cached_properies(self.__class__):
+                if hasattr(self, prop):
+                    delattr(self, prop)
+            return result
+
+        return wrapper
+
+    return decorator

diceplayer/utils/misc.py

@@ -13,6 +13,8 @@ import time
 BOHR2ANG: Final[float] = 0.52917721092
 ANG2BOHR: Final[float] = 1 / BOHR2ANG
 
+EA_2_DEBYE = 1 / 0.20819434
+
 
 #######################################  functions  ######################################
 

tests/shared/environment/test_molecule.py

@@ -41,7 +41,7 @@ class TestMolecule(unittest.TestCase):
             Atom(lbl=1, na=1, rx=1.0, ry=1.0, rz=1.0, chg=1.0, eps=1.0, sig=1.0)
         )
 
-        mol.center_of_mass_to_origin()
+        mol.move_center_of_mass_to_origin()
 
         npt.assert_equal(mol.com, [0, 0, 0])
 
@@ -68,11 +68,14 @@ class TestMolecule(unittest.TestCase):
             Atom(lbl=1, na=1, rx=1.0, ry=1.0, rz=1.0, chg=1.0, eps=1.0, sig=1.0)
         )
 
-        expected_distance_between_atoms = [[1.73205081], [1.73205081]]
+        expected = [
-        actual_distance_between_atoms = mol.distances_between_atoms()
+            [0.0, 1.73205081],
+            [1.73205081, 0.0]
+        ]
+        actual = mol.distances_between_atoms()
 
         npt.assert_almost_equal(
-            expected_distance_between_atoms, actual_distance_between_atoms
+            expected, actual
         )
 
     def test_inertia_tensor(self):
@@ -91,7 +94,7 @@ class TestMolecule(unittest.TestCase):
             [-0.50395, -0.50395, 1.0079],
         ]
 
-        actual_inertia_tensor = mol.inertia_tensor()
+        actual_inertia_tensor = mol.inertia_tensor
 
         npt.assert_equal(expected_inertia_tensor, actual_inertia_tensor)
 
@@ -163,7 +166,7 @@ class TestMolecule(unittest.TestCase):
             Atom(lbl=1, na=1, rx=1.0, ry=1.0, rz=1.0, chg=1.0, eps=1.0, sig=1.0)
         )
 
-        mol.standard_orientation()
+        mol.rotate_to_standard_orientation()
 
         expected_position = [0.0, 0.0, 0.0]