Split trajectory frames into chunks

MDANSE/Src/MDANSE/Framework/Configurators/OutputTrajectoryConfigurator.py

@@ -37,7 +37,7 @@ class OutputTrajectoryConfigurator(IConfigurator):
     """
 
     log_options = ("no logs", "DEBUG", "INFO", "WARN", "ERROR", "CRITICAL")
-    _default = ("OUTPUT_TRAJECTORY", 64, "none", "no logs")
+    _default = ("OUTPUT_TRAJECTORY", 64, 128, "none", "no logs")
 
     def __init__(self, name, format=None, **kwargs):
         """
@@ -54,11 +54,12 @@ def __init__(self, name, format=None, **kwargs):
         self._format = "MDTFormat"
         self._dtype = np.float64
         self._compression = "none"
+        self._chunk_limit = 128
 
     def configure(self, value: tuple):
         self._original_input = value
 
-        root, dtype, compression, logs = value
+        root, dtype, chunk_size, compression, logs = value
 
         if logs not in self.log_options:
             self.error_status = "log level option not recognised"
@@ -83,6 +84,8 @@ def configure(self, value: tuple):
         else:
             self._dtype = np.float64
 
+        self._chunk_limit = chunk_size
+
         if compression in TrajectoryWriter.allowed_compression:
             self._compression = compression
         else:
@@ -97,6 +100,7 @@ def configure(self, value: tuple):
         self["file"] = temp_name
         self["dtype"] = self._dtype
         self["compression"] = self._compression
+        self["chunk_size"] = self._chunk_limit
         self["log_level"] = logs
         if logs == "no logs":
             self["write_logs"] = False

MDANSE/Src/MDANSE/Framework/Converters/ASE.py

@@ -44,7 +44,10 @@ class ASETrajectoryFileError(Error):
 
 class ASE(Converter):
     """
-    Converts any trajectory to a HDF trajectory using the ASE io module.
+    Attempts to convert a trajectory file to MDANSE .mdt format (HDF5).
+    The conversion is done using the ase.io module.
+    Please help the ASE format detection mechanism by using
+    standard input file names.
     """
 
     label = "ASE"
@@ -53,8 +56,8 @@ class ASE(Converter):
     settings["trajectory_file"] = (
         "ASEFileConfigurator",
         {
-            "label": "Any MD trajectory file",
-            "default": "INPUT_FILENAME.traj",
+            "label": "An MD trajectory file supported by ASE",
+            "default": "INPUT_FILENAME",
         },
     )
     settings["atom_aliases"] = (
@@ -103,6 +106,7 @@ def initialize(self):
         self._isPeriodic = None
         self._backup_cell = None
         self._keep_running = True
+        self._initial_masses = None
         self._atomicAliases = self.configuration["atom_aliases"]["value"]
 
         # The number of steps of the analysis.
@@ -125,6 +129,7 @@ def initialize(self):
             self._chemicalSystem,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
             initial_charges=self._initial_charges,
         )
@@ -147,6 +152,8 @@ def run_step(self, index):
             LOG.warning(f"Skipping frame {index}")
             return index, None
 
+        variables = {}
+
         try:
             frame = self._input[index]
         except TypeError:
@@ -170,10 +177,23 @@ def run_step(self, index):
         coords = frame.get_positions()
         coords *= unit_conversion_factor
 
+        try:
+            momenta = frame.arrays["momenta"]
+        except KeyError:
+            pass
+        else:
+            if self._initial_masses is not None:
+                velocities = momenta / self._initial_masses.reshape((len(momenta), 1))
+            else:
+                velocities = momenta / np.array(self._chemicalSystem.masses).reshape(
+                    (len(momenta), 1)
+                )
+            variables["velocities"] = velocities * measure(1.0, "ang/fs").toval("nm/ps")
+
         if self._isPeriodic:
             try:
                 realConf = PeriodicRealConfiguration(
-                    self._trajectory.chemical_system, coords, unitCell
+                    self._trajectory.chemical_system, coords, unitCell, **variables
                 )
             except ValueError:
                 self._keep_running = False
@@ -186,8 +206,7 @@ def run_step(self, index):
         else:
             try:
                 realConf = RealConfiguration(
-                    self._trajectory.chemical_system,
-                    coords,
+                    self._trajectory.chemical_system, coords, **variables
                 )
             except ValueError:
                 self._keep_running = False
@@ -204,10 +223,16 @@ def run_step(self, index):
         )
 
         try:
-            charges = frame.get_charges()
+            charges = frame.arrays["charges"]
+        except KeyError:
+            try:
+                charges = frame.get_initial_charges()
+            except:
+                pass
+            else:
+                self._trajectory.write_charges(charges, index)
+        else:
             self._trajectory.write_charges(charges, index)
-        except RuntimeError:
-            pass
 
         return index, None
 
@@ -264,9 +289,18 @@ def parse_first_step(self, mapping):
         if self._isPeriodic:
             self._backup_cell = first_frame.cell.array
 
+        LOG.info(
+            f"The following arrays were found in the trajectory: {list(first_frame.arrays.keys())}"
+        )
+
+        if "masses" in first_frame.arrays.keys():
+            self._initial_masses = first_frame.arrays["masses"]
+        else:
+            self._initial_masses = None
+
         try:
-            self._initial_charges = first_frame.get_charges()
-        except RuntimeError:
+            self._initial_charges = first_frame.arrays["charges"]
+        except KeyError:
             LOG.warning("ASE converter could not read partial charges from file.")
             self._initial_charges = None
 

MDANSE/Src/MDANSE/Framework/Converters/CASTEP.py

@@ -101,6 +101,7 @@ def initialize(self):
             self._chemical_system,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Converters/CP2K.py

@@ -206,6 +206,7 @@ def initialize(self):
             self._chemical_system,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Converters/DCD.py

@@ -334,6 +334,7 @@ def initialize(self):
             self._chemical_system,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Converters/DL_POLY.py

@@ -189,6 +189,7 @@ def initialize(self):
             self._chemicalSystem,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
             initial_charges=self._fieldFile.get_atom_charges(),
         )

MDANSE/Src/MDANSE/Framework/Converters/Discover.py

@@ -338,6 +338,7 @@ def initialize(self):
             self._chemicalSystem,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
             initial_charges=self.configuration["xtd_file"].get_atom_charges(),
         )

MDANSE/Src/MDANSE/Framework/Converters/Forcite.py

@@ -352,6 +352,7 @@ def initialize(self):
             self._chemicalSystem,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
             initial_charges=self.configuration["xtd_file"].get_atom_charges(),
         )

MDANSE/Src/MDANSE/Framework/Converters/Gromacs.py

@@ -130,6 +130,7 @@ def initialize(self):
             chemical_system,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Converters/ImprovedASE.py

@@ -135,6 +135,7 @@ def initialize(self):
             self._chemicalSystem,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Converters/LAMMPS.py

@@ -828,6 +828,7 @@ def initialize(self):
             self._chemicalSystem,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
             initial_charges=charges,
         )

MDANSE/Src/MDANSE/Framework/Converters/MDAnalysis.py

@@ -144,6 +144,7 @@ def initialize(self):
 
         kwargs = {
             "positions_dtype": self.configuration["output_files"]["dtype"],
+            "chunking_limit": self.configuration["output_files"]["chunk_size"],
             "compression": self.configuration["output_files"]["compression"],
         }
         if hasattr(self.u.atoms, "charges"):

MDANSE/Src/MDANSE/Framework/Converters/VASP.py

@@ -98,6 +98,7 @@ def initialize(self):
             self._chemicalSystem,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Jobs/CenterOfMassesTrajectory.py

@@ -91,6 +91,7 @@ def initialize(self):
             chemical_system,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Jobs/CroppedTrajectory.py

@@ -79,6 +79,7 @@ def initialize(self):
             self.numberOfSteps,
             self._selectedAtoms,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
             initial_charges=[
                 self.configuration["trajectory"]["instance"].charges(0)[ind]

MDANSE/Src/MDANSE/Framework/Jobs/GlobalMotionFilteredTrajectory.py

@@ -108,6 +108,7 @@ def initialize(self):
             self.numberOfSteps,
             self._selected_atoms.atom_list,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Jobs/MoleculeFinder.py

@@ -102,6 +102,7 @@ def initialize(self):
             chemical_system,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
             initial_charges=self.configuration["trajectory"]["instance"].charges(0),
         )

MDANSE/Src/MDANSE/Framework/Jobs/RigidBodyTrajectory.py

@@ -153,6 +153,7 @@ def initialize(self):
             self.configuration["frames"]["number"],
             selectedAtoms,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 

MDANSE/Src/MDANSE/Framework/Jobs/TrajectoryEditor.py

@@ -122,6 +122,7 @@ def initialize(self):
             new_chemical_system,
             self.numberOfSteps,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
         )
 
@@ -161,7 +162,9 @@ def run_step(self, index):
             )
         else:
             com_conf = RealConfiguration(
-                self._output_trajectory.chemical_system, coords, **variables
+                self._output_trajectory.chemical_system,
+                coords[self._indices],
+                **variables,
             )
 
         self._output_trajectory.chemical_system.configuration = com_conf

MDANSE/Src/MDANSE/Framework/Jobs/UnfoldedTrajectory.py

@@ -82,6 +82,7 @@ def initialize(self):
             self.numberOfSteps,
             self._selectedAtoms,
             positions_dtype=self.configuration["output_files"]["dtype"],
+            chunking_limit=self.configuration["output_files"]["chunk_size"],
             compression=self.configuration["output_files"]["compression"],
             initial_charges=[
                 self.configuration["trajectory"]["instance"].charges(0)[ind]