reading of velocities, forces and charges implemented

PQAnalysis/core/atom.py

@@ -28,13 +28,11 @@
 
 """
 
-import numpy as np
-
 from multimethod import multimethod
 from beartype.typing import Any, Tuple
 from numbers import Real
 
-from PQAnalysis.utils.exceptions import ElementNotFoundError
+from PQAnalysis.exceptions import ElementNotFoundError
 
 
 def guess_element(id: int | str) -> Tuple[str, int, Real]:

PQAnalysis/core/atomicSystem.py

@@ -23,7 +23,7 @@
 
 from .atom import Atom
 from .cell import Cell
-from ..utils.mytypes import Numpy2DFloatArray, Numpy1DFloatArray
+from ..types import Numpy2DFloatArray, Numpy1DFloatArray
 
 
 def check_atoms_pos(func):

PQAnalysis/core/cell.py

@@ -16,7 +16,7 @@
 from beartype.typing import Any
 from numbers import Real
 
-from ..utils.mytypes import Numpy3x3FloatArray, Numpy2DFloatArray, Numpy1DFloatArray
+from ..types import Numpy3x3FloatArray, Numpy2DFloatArray, Numpy1DFloatArray
 
 
 class Cell:

PQAnalysis/exceptions.py

@@ -0,0 +1,52 @@
+"""
+A module containing different exceptions which could be useful.
+
+...
+
+Classes
+-------
+PQException
+    Base class for exceptions in this module.
+ElementNotFoundError
+    Exception raised if the given element id is not valid
+TrajectoryFormatError
+    Exception raised if the given enum is not valid
+"""
+
+from beartype.typing import Any
+
+
+class PQException(Exception):
+    """
+    Base class for exceptions in this module.
+    """
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+        super().__init__(self.message)
+
+
+class ElementNotFoundError(PQException):
+    """
+    Exception raised if the given element id is not valid
+    """
+
+    def __init__(self, id: Any) -> None:
+        self.id = id
+        self.message = f"""Id {self.id} is not a valid element identifier."""
+        super().__init__(self.message)
+
+
+class TrajectoryFormatError(PQException):
+    """
+    Exception raised if the given enum is not valid
+    """
+
+    def __init__(self, value: object, enum: object) -> None:
+        self.enum = enum
+        self.value = value
+        self.message = f"""
+'{self.value}' is not a valid TrajectoryFormat.
+Possible values are: {enum.member_repr()}
+or their case insensitive string representation: {enum.value_repr()}"""
+        super().__init__(self.message)

PQAnalysis/io/trajectoryReader.py

@@ -22,8 +22,8 @@
 from ..core.cell import Cell
 from ..core.atom import Atom
 from ..core.atomicSystem import AtomicSystem
-from ..utils.exceptions import ElementNotFoundError
-from ..utils.mytypes import Numpy2DFloatArray
+from ..exceptions import ElementNotFoundError
+from ..types import Numpy2DFloatArray, Numpy1DFloatArray
 
 
 class TrajectoryReader(BaseReader):
@@ -130,6 +130,10 @@ def read(self, frame_string: str, format: TrajectoryFormat | str = TrajectoryFor
             return self.read_positions(frame_string)
         elif TrajectoryFormat(format) is TrajectoryFormat.VEL:
             return self.read_velocities(frame_string)
+        elif TrajectoryFormat(format) is TrajectoryFormat.FORCE:
+            return self.read_forces(frame_string)
+        elif TrajectoryFormat(format) is TrajectoryFormat.CHARGE:
+            return self.read_charges(frame_string)
 
     def read_positions(self, frame_string: str) -> Frame:
         """
@@ -199,6 +203,74 @@ def read_velocities(self, frame_string: str) -> Frame:
 
         return Frame(AtomicSystem(atoms=atoms, vel=vel, cell=cell))
 
+    def read_forces(self, frame_string: str) -> Frame:
+        """
+        Reads the forces of the atoms in a frame from a string.
+
+        Parameters
+        ----------
+        frame_string : str
+            The string to read the frame from.
+
+        Returns
+        -------
+        Frame
+            The frame read from the string.
+
+        Raises
+        ------
+        TypeError
+            If the given frame_string is not a string.
+        """
+
+        splitted_frame_string = frame_string.split('\n')
+        header_line = splitted_frame_string[0]
+
+        n_atoms, cell = self._read_header_line(header_line)
+
+        forces, atoms = self._read_xyz(splitted_frame_string, n_atoms)
+
+        try:
+            atoms = [Atom(atom) for atom in atoms]
+        except ElementNotFoundError:
+            atoms = [Atom(atom, use_guess_element=False) for atom in atoms]
+
+        return Frame(AtomicSystem(atoms=atoms, forces=forces, cell=cell))
+
+    def read_charges(self, frame_string: str) -> Frame:
+        """
+        Reads the charge values of the atoms in a frame from a string.
+
+        Parameters
+        ----------
+        frame_string : str
+            The string to read the frame from.
+
+        Returns
+        -------
+        Frame
+            The frame read from the string.
+
+        Raises
+        ------
+        TypeError
+            If the given frame_string is not a string.
+        """
+
+        splitted_frame_string = frame_string.split('\n')
+        header_line = splitted_frame_string[0]
+
+        n_atoms, cell = self._read_header_line(header_line)
+
+        charges, atoms = self._read_scalar(splitted_frame_string, n_atoms)
+
+        try:
+            atoms = [Atom(atom) for atom in atoms]
+        except ElementNotFoundError:
+            atoms = [Atom(atom, use_guess_element=False) for atom in atoms]
+
+        return Frame(AtomicSystem(atoms=atoms, charges=charges, cell=cell))
+
     def _read_header_line(self, header_line: str) -> Tuple[int, Cell | None]:
         """
         Reads the header line of a frame.
@@ -281,3 +353,41 @@ def _read_xyz(self, splitted_frame_string: List[str], n_atoms: int) -> Tuple[Num
             atoms.append(line.split()[0])
 
         return xyz, atoms
+
+    def _read_scalar(self, splitted_frame_string: List[str], n_atoms: int) -> Tuple[Numpy1DFloatArray, List[str]]:
+        """
+        Reads the scalar values and the atom names from the given string.
+
+        Parameters
+        ----------
+        splitted_frame_string : str
+            The string to read the scalar values and the atom names from.
+        n_atoms : int
+            The number of atoms in the frame.
+
+        Returns
+        -------
+        scalar : np.array
+            The scalar values of the atoms.
+        atoms : list of str
+            The names of the atoms.
+
+        Raises
+        ------
+        ValueError
+            If the given string does not contain the correct number of lines.
+        """
+
+        scalar = np.zeros((n_atoms))
+        atoms = []
+        for i in range(n_atoms):
+            line = splitted_frame_string[2+i]
+
+            if len(line.split()) != 2:
+                raise ValueError(
+                    'Invalid file format in scalar values of Frame.')
+
+            scalar[i] = float(line.split()[1])
+            atoms.append(line.split()[0])
+
+        return scalar, atoms

PQAnalysis/io/trajectoryWriter.py

@@ -17,7 +17,7 @@
 from ..traj.trajectory import Trajectory
 from ..core.cell import Cell
 from ..core.atom import Atom
-from ..utils.mytypes import Numpy2DFloatArray
+from ..types import Numpy2DFloatArray
 
 
 def write_trajectory(traj, filename: str | None = None, format: str | None = None) -> None:

PQAnalysis/physicalData/energy.py

@@ -14,7 +14,7 @@
 from beartype.typing import Dict
 from collections import defaultdict
 
-from ..utils.mytypes import Numpy2DFloatArray, Numpy1DFloatArray
+from ..types import Numpy2DFloatArray, Numpy1DFloatArray
 
 
 class Energy():

PQAnalysis/traj/frame.py

@@ -19,7 +19,7 @@
 from ..core.atomicSystem import AtomicSystem
 from ..core.atom import Atom
 from ..core.cell import Cell
-from ..utils.mytypes import Numpy2DFloatArray, Numpy1DFloatArray
+from ..types import Numpy2DFloatArray, Numpy1DFloatArray
 
 
 class Frame:
@@ -202,6 +202,30 @@ def vel(self) -> Numpy2DFloatArray:
         """
         return self.system.vel
 
+    @property
+    def forces(self) -> Numpy2DFloatArray:
+        """
+        The forces on the atoms in the system.
+
+        Returns
+        -------
+        Numpy2DFloatArray
+            The forces on the atoms in the system.
+        """
+        return self.system.forces
+
+    @property
+    def charges(self) -> Numpy1DFloatArray:
+        """
+        The charges of the atoms in the system.
+
+        Returns
+        -------
+        Numpy1DFloatArray
+            The charges of the atoms in the system.
+        """
+        return self.system.charges
+
     @property
     def atoms(self) -> List[Atom]:
         """

PQAnalysis/traj/trajectory.py

@@ -5,6 +5,8 @@
 
 Classes
 -------
+TrajectoryFormat
+    An enumeration of the supported trajectory formats.
 Trajectory
     A trajectory is a sequence of frames.
 """
@@ -17,6 +19,7 @@
 from enum import Enum
 
 from .frame import Frame
+from ..exceptions import TrajectoryFormatError
 
 
 class TrajectoryFormat(Enum):
@@ -31,10 +34,16 @@ class TrajectoryFormat(Enum):
         The XYZ format.
     VEL : str
         The VEL format.
+    FORCE : str
+        The FORCE format.
+    CHARGE : str
+        The CHARGE format.
     """
 
     XYZ = "XYZ"
     VEL = "VEL"
+    FORCE = "FORCE"
+    CHARGE = "CHARGE"
 
     @classmethod
     def _missing_(cls, value: object) -> Any:
@@ -55,7 +64,34 @@ def _missing_(cls, value: object) -> Any:
         for member in cls:
             if member.value.lower() == value:
                 return member
-        return None
+
+        raise TrajectoryFormatError(value, cls)
+
+    @classmethod
+    def member_repr(cls) -> str:
+        """
+        This method returns a string representation of the members of the enumeration.
+
+        Returns
+        -------
+        str
+            A string representation of the members of the enumeration.
+        """
+
+        return ', '.join([str(member) for member in cls])
+
+    @classmethod
+    def value_repr(cls) -> str:
+        """
+        This method returns a string representation of the values of the members of the enumeration.
+
+        Returns
+        -------
+        str
+            A string representation of the values of the members of the enumeration.
+        """
+
+        return ', '.join([str(member.value) for member in cls])
 
 
 class Trajectory:

PQAnalysis/utils/mytypes.py → PQAnalysis/types.py

@@ -3,9 +3,6 @@
 from beartype.vale import Is
 from typing import Annotated
 
-# Import the requisite machinery.
-from beartype import beartype, BeartypeConf
-
 
 Numpy2DFloatArray = Annotated[np.ndarray, Is[lambda array:
                                              array.ndim == 2 and

tests/core/test_atom.py

@@ -4,7 +4,7 @@
 from multimethod import DispatchError
 
 from PQAnalysis.core.atom import Atom, guess_element
-from PQAnalysis.utils.exceptions import ElementNotFoundError
+from PQAnalysis.exceptions import ElementNotFoundError
 
 
 def test_guess_element():