Merge pull request #118 from MolarVerse/feature/xyz2rst

CLI: xyz2rst tool added

Author: galjos

PQAnalysis/atomic_system/atomic_system.py

@@ -36,7 +36,6 @@
 from .exceptions import AtomicSystemError
 
 
-
 class AtomicSystem(
     _PropertiesMixin,
     _StandardPropertiesMixin,
@@ -194,7 +193,7 @@ def __init__(
         self._stress = stress
         self._cell = cell
 
-    #TODO: check why dynamic formatting does
+    # TODO: check why dynamic formatting does
     #      not work here for "AtomicSystem"
     #
     # @runtime_type_checking
@@ -403,15 +402,15 @@ def center_of_mass_residues(self) -> "AtomicSystem":
         -------
         AtomicSystem
             The center of mass of the residues in the system.
-            
+
         Raises:
         -------
         AtomicSystemError
             If the number of residues in the system is not a 
             multiple of the number of atoms.
         PQNotImplementedError
             if system has forces, velocities or charges.
-            
+
         TODO:
         -----
         Include also center of mass velocities, forces and so on...
@@ -519,6 +518,34 @@ def compute_com_atomic_system(self, group=None) -> "AtomicSystem":
 
         return AtomicSystem(pos=np.array(pos), atoms=names, cell=self.cell)
 
+    def randomize_positions(
+        self,
+        stdev: float = 0.001,
+        seed: int | None = None,
+    ) -> "AtomicSystem":
+        """
+        Randomizes the positions of the atoms in the system with a normal distribution.
+
+        Parameters
+        ----------
+        stdev : float, optional
+            The standard deviation of the random displacements, by default 0.001
+        seed : int, optional
+            The seed for the random number generator, by default None
+
+        Returns
+        -------
+        AtomicSystem
+            The AtomicSystem with the randomized positions.
+        """
+
+        if seed is not None:
+            np.random.seed(seed)
+
+        self.pos += np.random.normal(scale=stdev, size=self.pos.shape)
+
+        return self
+
     def copy(self) -> "AtomicSystem":
         """
         Returns a copy of the AtomicSystem.

PQAnalysis/cli/main.py

@@ -8,6 +8,7 @@
 from .traj2qmcfc import Traj2QMCFCCLI
 from .traj2box import Traj2BoxCLI
 from .rst2xyz import Rst2XYZCLI
+from .xyz2rst import XYZ2RstCLI
 from .rdf import RDFCLI
 from .gen2xyz import GEN2XYZCLI
 from .continue_input import ContinueInputCLI
@@ -43,6 +44,7 @@ def main():
         GEN2XYZCLI.program_name(): GEN2XYZCLI,
         RDFCLI.program_name(): RDFCLI,
         Rst2XYZCLI.program_name(): Rst2XYZCLI,
+        XYZ2RstCLI.program_name(): XYZ2RstCLI,
         Traj2BoxCLI.program_name(): Traj2BoxCLI,
         Traj2QMCFCCLI.program_name(): Traj2QMCFCCLI,
         XYZ2GENCLI.program_name(): XYZ2GENCLI,

PQAnalysis/cli/xyz2rst.py

@@ -0,0 +1,123 @@
+"""
+.. _cli.rst2xyz:
+
+Command Line Tool for Converting XYZ Files to Restart Files
+-----------------------------------------------------------
+
+
+"""
+
+from PQAnalysis.io import xyz2rst
+from PQAnalysis.config import code_base_url
+from ._argument_parser import _ArgumentParser
+from ._cli_base import CLIBase
+
+__outputdoc__ = """
+
+This command line tool can be used to convert xyz files to restart files. 
+
+"""
+
+__epilog__ = "\n"
+__epilog__ += "For more information on required and optional input file keys please visit "
+__epilog__ += f"{code_base_url}PQAnalysis.cli.xyz2rst.html."
+__epilog__ += "\n"
+__epilog__ += "\n"
+
+__doc__ += __outputdoc__
+
+
+class XYZ2RstCLI(CLIBase):
+
+    """
+    Command Line Tool for Converting XYZ Files to Restart Files
+    """
+
+    @classmethod
+    def program_name(cls) -> str:
+        """
+        Returns the name of the program.
+
+        Returns
+        -------
+        str
+            The name of the program.
+        """
+        return 'xyz2rst'
+
+    @classmethod
+    def add_arguments(cls, parser: _ArgumentParser) -> None:
+        """
+        Adds the arguments to the parser.
+
+        Parameters
+        ----------
+        parser : _ArgumentParser
+            The parser to which the arguments should be added.
+        """
+        parser.parse_output_file()
+
+        parser.add_argument(
+            'xyz_file',
+            type=str,
+            help='The xyz file to be converted.'
+        )
+
+        parser.add_argument(
+            '--velocity_file',
+            type=str,
+            default=None,
+            help='The velocity file to be converted.'
+        )
+
+        parser.add_argument(
+            '--force_file',
+            type=str,
+            default=None,
+            help='The force file to be converted.'
+        )
+
+        parser.add_argument(
+            '--randomize',
+            type=float,
+            default=0.0,
+            help='Randomize the atom order. Default is 0.0.'
+        )
+
+        parser.parse_engine()
+        parser.parse_mode()
+
+    @classmethod
+    def run(cls, args):
+        """
+        Runs the command line tool.
+
+        Parameters
+        ----------
+        args : argparse.Namespace
+            The arguments parsed by the parser.
+        """
+        xyz2rst(
+            xyz_file=args.xyz_file,
+            velocity_file=args.velocity_file,
+            force_file=args.force_file,
+            randomize=args.randomize,
+            output=args.output,
+            md_format=args.engine,
+            mode=args.mode,
+        )
+
+
+def main():
+    """
+    Main function of the xyz2rst command line tool, which is basically just a wrapper 
+    for the xyz2rst function. For more information on the xyz2rst function 
+    please visit :py:func:`PQAnalysis.io.api.xyz2rst`.
+    """
+    parser = _ArgumentParser(description=__outputdoc__, epilog=__epilog__)
+
+    XYZ2RstCLI.add_arguments(parser)
+
+    args = parser.parse_args()
+
+    XYZ2RstCLI.run(args)

PQAnalysis/io/__init__.py

@@ -48,5 +48,6 @@
 from .input_file_reader import InputFileFormat
 
 from .api import continue_input_file
-from .conversion_api import (gen2xyz, xyz2gen, rst2xyz, traj2box, traj2qmcfc)
+from .conversion_api import (
+    gen2xyz, xyz2gen, rst2xyz, xyz2rst, traj2box, traj2qmcfc)
 from .write_api import write, write_box

PQAnalysis/io/conversion_api.py

@@ -19,11 +19,13 @@
     write_gen_file,
     read_gen_file,
 )
-from .restart_file.api import read_restart_file
+from .restart_file.api import (
+    read_restart_file,
+    write_restart_file
+)
 from .traj_file.api import write_trajectory
 
 
-
 @runtime_type_checking
 def gen2xyz(
     gen_file: str,
@@ -65,7 +67,6 @@ def gen2xyz(
     )
 
 
-
 @runtime_type_checking
 def xyz2gen(
     xyz_file: str,
@@ -110,7 +111,6 @@ def xyz2gen(
     )
 
 
-
 @runtime_type_checking
 def rst2xyz(
     restart_file: str,
@@ -152,6 +152,75 @@ def rst2xyz(
     )
 
 
+@runtime_type_checking
+def xyz2rst(
+    xyz_file: str,
+    velocity_file: str | None = None,
+    force_file: str | None = None,
+    randomize: float = 0.0,
+    random_seed: int | None = 0,
+    output: str | None = None,
+    md_format: MDEngineFormat | str = MDEngineFormat.PQ,
+    mode: FileWritingMode | str = "w"
+) -> None:
+    """
+    Converts a xyz file to a restart file and prints it to stdout or writes it to a file.
+
+    When the print_box flag is set to True, the box is printed as well.
+    This means that after the number of atoms the box is printed in the
+    same line in the format a b c alpha beta gamma.
+
+    Parameters
+    ----------
+    xyz_file : str
+        The xyz file to be converted.
+    velocity_file : str | None
+        The velocity file to be converted. Default is None.
+    force_file : str | None
+        The force file to be converted. Default is None.
+    randomize : float, optional
+        Randomize the atom order. Default is 0.0.
+    output : str | None
+        The output file. If not specified, the output is printed to stdout.
+    md_format : MDEngineFormat | str, optional
+        The format of the md engine for the output file. The default is MDEngineFormat.PQ.
+    mode : FileWritingMode | str, optional
+        The writing mode, by default "w". The following modes are available:
+        - "w": write
+        - "a": append
+        - "o": overwrite
+    """
+
+    system = TrajectoryReader(
+        xyz_file,
+        md_format=md_format,
+        traj_format="xyz",
+    ).read()[-1]
+
+    if velocity_file is not None:
+        system.vel = TrajectoryReader(
+            velocity_file,
+            md_format=md_format,
+            traj_format="vel",
+        ).read()[-1].vel
+
+    if force_file is not None:
+        system.forces = TrajectoryReader(
+            force_file,
+            md_format=md_format,
+            traj_format="force",
+        ).read()[-1].forces
+
+    if randomize != 0.0:
+        system.randomize_positions(stdev=randomize, seed=random_seed)
+
+    write_restart_file(
+        atomic_system=system,
+        filename=output,
+        md_engine_format=md_format,
+        mode=mode,
+    )
+
 
 @runtime_type_checking
 def traj2box(
@@ -201,7 +270,6 @@ def traj2box(
         writer.write(trajectory, reset_counter=False)
 
 
-
 @runtime_type_checking
 def traj2qmcfc(
     trajectory_files: List[str],

PQAnalysis/io/restart_file/api.py

@@ -4,11 +4,12 @@
 
 from PQAnalysis.core import Residues
 from PQAnalysis.traj import MDEngineFormat
+from PQAnalysis.io.formats import FileWritingMode
 from PQAnalysis.atomic_system import AtomicSystem
 from PQAnalysis.type_checking import runtime_type_checking
 
 from .restart_reader import RestartFileReader
-
+from .restart_writer import RestartFileWriter
 
 
 @runtime_type_checking
@@ -51,3 +52,36 @@ def read_restart_file(
     )
 
     return reader.read()
+
+
+@runtime_type_checking
+def write_restart_file(
+    atomic_system: AtomicSystem,
+    filename: str | None = None,
+    md_engine_format: MDEngineFormat | str = MDEngineFormat.PQ,
+    mode: FileWritingMode | str = 'w'
+) -> None:
+    """
+    API function for reading a restart file.
+
+    Parameters
+    ----------
+    atomic_system : AtomicSystem
+        The AtomicSystem object to write.
+    filename : str
+        The filename of the restart file.
+    md_engine_format : MDEngineFormat | str, optional
+        The format of the restart file,
+        by default MDEngineFormat.PQ
+    mode : FileWritingMode | str, optional
+        The mode of the file. Either 'w' for write, 
+        'a' for append or 'o' for overwrite. The default is 'w'.
+    """
+
+    writer = RestartFileWriter(
+        filename=filename,
+        md_engine_format=md_engine_format,
+        mode=mode
+    )
+
+    writer.write(atomic_system)