diff --git a/MDANSE/Src/MDANSE/Framework/Configurators/MockTrajectoryConfigurator.py b/MDANSE/Src/MDANSE/Framework/Configurators/MockTrajectoryConfigurator.py index 6cce4e767..858ca25cc 100644 --- a/MDANSE/Src/MDANSE/Framework/Configurators/MockTrajectoryConfigurator.py +++ b/MDANSE/Src/MDANSE/Framework/Configurators/MockTrajectoryConfigurator.py @@ -23,7 +23,9 @@ class MockTrajectoryConfigurator(IConfigurator): """ This is a replacement for a trajectory stored in and HDF5 file. It is intended to be a drop-in replacement for HDFTrajectoryConfigurator, - even though it is NOT file-based. + even though it is NOT based on an HDF5 file. + It can use a JSON file with MockTrajectory parameters to create + a trajectory entirely in the RAM. """ _default = None diff --git a/MDANSE/Src/MDANSE/Framework/InputData/MockTrajectoryInputData.py b/MDANSE/Src/MDANSE/Framework/InputData/MockTrajectoryInputData.py index df8f849c1..edf77daae 100644 --- a/MDANSE/Src/MDANSE/Framework/InputData/MockTrajectoryInputData.py +++ b/MDANSE/Src/MDANSE/Framework/InputData/MockTrajectoryInputData.py @@ -19,6 +19,9 @@ class MockTrajectoryInputData(InputFileData): + """Imitates the HDFTrajectoryInputData, + but builds a MockTrajectory out of a JSON file instead. + """ extension = "json" def load(self): @@ -41,16 +44,6 @@ def info(self): val.append("%s\n" % self._name) val.append("Number of steps:") val.append("%s\n" % len(self._data)) - val.append("Configuration:") - val.append( - "\tIs periodic: {}\n".format( - "unit_cell" in self._data.file["/configuration"] - ) - ) - val.append("Variables:") - for k, v in self._data.file["/configuration"].items(): - val.append("\t- {}: {}".format(k, v.shape)) - mol_types = {} val.append("\nMolecular types found:") for ce in self._data.chemical_system.chemical_entities: @@ -76,4 +69,11 @@ def chemical_system(self): @property def hdf(self): + """There is no HDF5 file for a mock trajectory + + Returns + ------- + str + name of a nonexistent file + """ return self._data.file diff --git a/MDANSE/Src/MDANSE/MolecularDynamics/MockTrajectory.py b/MDANSE/Src/MDANSE/MolecularDynamics/MockTrajectory.py index e9aa910b2..d5631d05c 100644 --- a/MDANSE/Src/MDANSE/MolecularDynamics/MockTrajectory.py +++ b/MDANSE/Src/MDANSE/MolecularDynamics/MockTrajectory.py @@ -14,6 +14,7 @@ import math import json +from typing import TypeVar import numpy as np from icecream import ic @@ -26,6 +27,7 @@ from MDANSE.MolecularDynamics.Configuration import ( PeriodicRealConfiguration, RealConfiguration, + _Configuration, ) from MDANSE.MolecularDynamics.TrajectoryUtils import ( build_connectivity, @@ -35,7 +37,16 @@ from MDANSE.MolecularDynamics.UnitCell import UnitCell +Self = TypeVar("Self", bound="MockTrajectory") + + class MockTrajectory: + """For testing purposes, MockTrajectory can replace a trajectory. +It acts as a trajectory of predefined composition and size, +while taking only a necessary minimum of resources. +The main goal is performance testing of different analysis types +without the need to run Molecular Dynamics simulations beforehand. + """ def __init__( self, number_of_frames: int = 1000, @@ -73,6 +84,22 @@ def __init__( self._chemicalSystem.add_chemical_entity(Atom(symbol=atom)) def set_coordinates(self, coords: np.ndarray): + """Sets the initial (equlibrium) positions of atoms from + the input array. + + The array must have as many rows as there are atoms in a single box, + and the positions will be replicated between boxes. + + Parameters + ---------- + coords : np.ndarray + positions of atoms in a single box + + Returns + ------- + bool + False if the number of elements was wrong + """ if len(coords) != self._num_atoms_in_box: return False coords_nm = coords * measure(1.0, "ang").toval("nm") @@ -93,6 +120,31 @@ def modulate_structure( period: int = 10, amplitude: float = 0.1, ): + """Creates a number of frames in the trajectory which contain + coordinates displaced by a mechanical wave. The atom positions + oscillate around the equlibrium positions. + + Parameters + ---------- + polarisation : np.ndarray, optional + direction of atom displacements + propagation_vector : np.ndarray, optional + propagation vector of the wave (phonon). All zeros for standing wave + period : int, optional + Number of frames corresponding to a total 2pi period of the wave. + amplitude : float, optional + Maximum displacement along the polarisation vector (in Angstrom) + + Returns + ------- + bool + False on wrong size of the input array + + Raises + ------ + ValueError + if the period of the new modulation is incommensurate with the number of frames + """ if len(polarisation) * self._multiplier == len(self._start_coordinates): polarisation = np.row_stack(self._multiplier * [polarisation]) if len(polarisation) != len(self._start_coordinates): @@ -140,17 +192,28 @@ def modulate_structure( self._real_length = n_steps def close(self): - """Close the trajectory.""" - - def __getitem__(self, frame): - """Return the configuration at a given frame - - :param frame: the frame - :type frame: int - - :return: the configuration - :rtype: dict of ndarray + """Present for compatibility with Trajectory""" + + def __getitem__(self, frame: int): + """Returns the configuration of the system at the Nth frame. + + Parameters + ---------- + frame : int + number of the frame at which to get the configuration + + Returns + ------- + dict + coordinates, time and unit cell at the specified frame + + Raises + ------ + IndexError + if frame is outside of range """ + if frame < 0 or frame >= len(self): + raise IndexError(f"Invalid frame number: {frame}") configuration = {} configuration["coordinates"] = self.coordinates(frame).astype(np.float64) @@ -160,19 +223,32 @@ def __getitem__(self, frame): return configuration def __getstate__(self): + """Only added for compatibility with Trajectory + """ pass def __setstate__(self, state): + """Only added for compatibility with Trajectory + """ pass - def coordinates(self, frame): - """Return the coordinates at a given frame. + def coordinates(self, frame: int) -> np.ndarray: + """Returns the atom coordinates at the specified frame - :param frame: the frame - :type frame: int + Parameters + ---------- + frame : int + number of the simulation step (frame) - :return: the coordinates - :rtype: ndarray + Returns + ------- + np.ndarray + an array (N,3) of atom coordinates. N is the numer of atoms. + + Raises + ------ + IndexError + if frame is out of range """ if frame < 0 or frame >= len(self): @@ -185,14 +261,23 @@ def coordinates(self, frame): return self._coordinates[scaled_index].astype(np.float64) - def configuration(self, frame): - """Build and return a configuration at a given frame. + def configuration(self, frame: int) -> '_Configuration': + """An MDANSE Configuration at the specified frame number. - :param frame: the frame - :type frame: int + Parameters + ---------- + frame : int + number of the simulation step (frame) - :return: the configuration - :rtype: MDANSE.MolecularDynamics.Configuration.Configuration + Returns + ------- + _Configuration + An object holding the atom positions, unit cell, etc. + + Raises + ------ + IndexError + if frame is out of range """ if frame < 0 or frame >= len(self): @@ -214,25 +299,31 @@ def configuration(self, frame): return conf def _load_unit_cells(self): - """Load all the unit cells.""" + """Only added for compatibility with Trajectory.""" - def unit_cell(self, frame): - """Return the unit cell at a given frame. If no unit cell is defined, returns None. + def unit_cell(self, frame: int) -> UnitCell: + """Returns the UnitCell the size of the system. - :param frame: the frame number - :type frame: int + Parameters + ---------- + frame : int + ignored - :return: the unit cell - :rtype: ndarray + Returns + ------- + UnitCell + Object defining the system size """ return UnitCell(self._full_box_size) - def __len__(self): - """Returns the length of the trajectory. + def __len__(self) -> int: + """Length of the mock trajectory - :return: the number of frames of the trajectory - :rtype: int + Returns + ------- + int + number of frames that can be returned by MockTrajectory """ return self._number_of_frames @@ -432,8 +523,9 @@ def chemical_system(self): return self._chemicalSystem @property - def file(self): - """Return the trajectory file object. + def file(self) -> str: + """There is no trajectory file. + A string is returned instead. :return: the trajectory file object :rtype: HDF5 file object @@ -442,8 +534,8 @@ def file(self): return "nonexistent_file.h5" @property - def filename(self): - """Return the trajectory filename. + def filename(self) -> str: + """Returns a file name, but the file does not exist. :return: the trajectory filename :rtype: str @@ -452,11 +544,20 @@ def filename(self): return "nonexistent_file.h5" @property - def has_velocities(self): + def has_velocities(self) -> bool: + """True if the trajectory contains atom velocities, + False otherwise. + + Returns + ------- + bool + True if velocities are stored in MockTrajectory + """ return "velocities" in self._variables.keys() def variables(self): """Return the configuration variables stored in this trajectory. + Most likely empty for MockTrajectory, but does not have to be. :return; the configuration variable :rtype: list @@ -467,7 +568,21 @@ def variables(self): return list(grp.keys()) @classmethod - def from_json(cls, filename: str): + def from_json(cls, filename: str) -> Self: + """Builds and returns an instance of MockTrajectory + using the parameters in a JSON file. + + Parameters + ---------- + filename : str + must be a valid JSON file with "parameters", + "coordinates" and "modulations" sections. + + Returns + ------- + Self + a MockTrajectory instance + """ with open(filename, "r") as source: struct = json.load(source) temp = struct["parameters"]