From a5504b043558cb10217e5905fba683310bf27145 Mon Sep 17 00:00:00 2001 From: Vitor Hideyoshi Date: Sat, 28 Feb 2026 09:58:42 -0300 Subject: [PATCH] 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 | 6 +- diceplayer/environment/molecule.py | 294 ++++++++-------------- diceplayer/environment/system.py | 86 +------ diceplayer/player.py | 2 +- diceplayer/utils/cache.py | 29 +++ diceplayer/utils/misc.py | 2 + tests/shared/environment/test_molecule.py | 15 +- 7 files changed, 156 insertions(+), 278 deletions(-) create mode 100644 diceplayer/utils/cache.py diff --git a/diceplayer/environment/atom.py b/diceplayer/environment/atom.py index 5aeeff8..34adea5 100644 --- a/diceplayer/environment/atom.py +++ b/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) diff --git a/diceplayer/environment/molecule.py b/diceplayer/environment/molecule.py index 4df0409..95f975c 100644 --- a/diceplayer/environment/molecule.py +++ b/diceplayer/environment/molecule.py @@ -1,25 +1,24 @@ from __future__ import annotations -from typing_extensions import TYPE_CHECKING - - -if TYPE_CHECKING: - from nptyping import Float, NDArray, Shape +from dataclasses import dataclass, Field +from functools import cached_property 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: - molname (str): Molecule name + Returns: + 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] = [] - self.position: npt.NDArray[np.float64] - self.energy: npt.NDArray[np.float64] - self.gradient: npt.NDArray[np.float64] - self.hessian: npt.NDArray[np.float64] + for atom in self.atom: + dx = atom.rx - self.com[0] + dy = atom.ry - self.com[1] + dz = atom.rz - self.com[2] - self.ghost_atoms: List[Atom] = [] - self.lp_atoms: List[Atom] = [] + inertia_tensor[0, 0] += atom.mass * (dy**2 + dz**2) + 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 - self.com: Union[None, npt.NDArray[np.float64]] = None + inertia_tensor[0, 1] -= atom.mass * dx * dy + 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() - - def center_of_mass(self) -> npt.NDArray[np.float64]: + @invalidate_computed_properties() + def remove_atom(self, a: Atom) -> None: """ - 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: - self.com += atom.mass * np.array([atom.rx, atom.ry, atom.rz]) - - self.com = self.com / self.total_mass - - return self.com - - def center_of_mass_to_origin(self) -> None: + @invalidate_computed_properties() + def move_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 = [] - 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]: + def principal_axes(self) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: """ 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 = [] - 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: + def update_charges(self, charges: npt.NDArray[np.float64]) -> 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 = [] - 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": + def translate(self, vector: np.ndarray) -> Self: """ 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 = [] - for atom1 in self.atom: - 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)) + if len(coords_a) == 0 or len(coords_b) == 0: + raise ValueError("No real atoms to compare") - return min(distances) + diff = coords_a[:, None, :] - coords_b[None, :, :] + d2 = np.sum(diff ** 2, axis=-1) + return np.sqrt(d2.min()) \ No newline at end of file diff --git a/diceplayer/environment/system.py b/diceplayer/environment/system.py index ecb4aad..f84e656 100644 --- a/diceplayer/environment/system.py +++ b/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_reference_mol.center_of_mass_to_origin() + new_projecting_mol.move_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 diff --git a/diceplayer/player.py b/diceplayer/player.py index 75ee643..b646fb4 100644 --- a/diceplayer/player.py +++ b/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() diff --git a/diceplayer/utils/cache.py b/diceplayer/utils/cache.py new file mode 100644 index 0000000..e8d9228 --- /dev/null +++ b/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 diff --git a/diceplayer/utils/misc.py b/diceplayer/utils/misc.py index 848c48d..ea6c33b 100644 --- a/diceplayer/utils/misc.py +++ b/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 ###################################### diff --git a/tests/shared/environment/test_molecule.py b/tests/shared/environment/test_molecule.py index c910390..ad1670f 100644 --- a/tests/shared/environment/test_molecule.py +++ b/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]] - actual_distance_between_atoms = mol.distances_between_atoms() + expected = [ + [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]