Merge pull request #284 from ISISNeutronMuon/fix-ase-converter

Fix ase converter

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

@@ -20,6 +20,7 @@
 
 from ase.io import iread, read
 from ase.atoms import Atoms as ASEAtoms
+from ase.io.trajectory import Trajectory as ASETrajectory
 import numpy as np
 import h5py
 
@@ -54,14 +55,7 @@ class ASE(Converter):
         "InputFileConfigurator",
         {
             "label": "Any MD trajectory file",
-            "default": "INPUT_FILENAME.lammps",
-        },
-    )
-    settings["configuration_file"] = (
-        "InputFileConfigurator",
-        {
-            "label": "An optional structure/configuration file",
-            "default": "INPUT_FILENAME.lammps",
+            "default": "INPUT_FILENAME.traj",
         },
     )
     settings["time_step"] = (
@@ -97,7 +91,15 @@ def initialize(self):
         # The number of steps of the analysis.
         self.numberOfSteps = self.configuration["n_steps"]["value"]
 
+        self._timestep = float(self.configuration["time_step"]["value"]) * measure(
+            1.0, self.configuration["time_unit"]["value"]
+        ).toval("ps")
+
         self.parse_first_step()
+        self._start = 0
+
+        if self.numberOfSteps < 1:
+            self.numberOfSteps = self._total_number_of_steps
 
         # A trajectory is opened for writing.
         self._trajectory = TrajectoryWriter(
@@ -110,11 +112,7 @@ def initialize(self):
             [(at.name, at.index) for at in self._trajectory.chemical_system.atom_list]
         )
 
-        self._start = 0
-
-        float(self.configuration["time_step"]["value"]) * measure(
-            1.0, self.configuration["time_unit"]["value"]
-        )
+        print(f"total steps: {self.numberOfSteps}")
 
     def run_step(self, index):
         """Runs a single step of the job.
@@ -125,114 +123,37 @@ def run_step(self, index):
         @note: the argument index is the index of the loop note the index of the frame.
         """
 
-        self._input = iread(self.configuration["trajectory_file"]["value"])
-
-        for stepnum, frame in enumerate(self._input):
-            time = stepnum * self._timestep
-
-        for _ in range(self._itemsPosition["TIMESTEP"][0]):
-            line = self._lammps.readline()
-            if not line:
-                return index, None
-
-        time = (
-            float(self._lammps.readline())
-            * self.configuration["time_step"]["value"]
-            * measure(1.0, "fs").toval("ps")
-        )
-
-        for _ in range(
-            self._itemsPosition["TIMESTEP"][1], self._itemsPosition["BOX BOUNDS"][0]
-        ):
-            self._lammps.readline()
-
-        unitCell = np.zeros((9), dtype=np.float64)
-        temp = [float(v) for v in self._lammps.readline().split()]
-        if len(temp) == 2:
-            xlo, xhi = temp
-            xy = 0.0
-        elif len(temp) == 3:
-            xlo, xhi, xy = temp
-        else:
-            raise ASETrajectoryFileError("Bad format for A vector components")
-
-        temp = [float(v) for v in self._lammps.readline().split()]
-        if len(temp) == 2:
-            ylo, yhi = temp
-            xz = 0.0
-        elif len(temp) == 3:
-            ylo, yhi, xz = temp
-        else:
-            raise ASETrajectoryFileError("Bad format for B vector components")
-
-        temp = [float(v) for v in self._lammps.readline().split()]
-        if len(temp) == 2:
-            zlo, zhi = temp
-            yz = 0.0
-        elif len(temp) == 3:
-            zlo, zhi, yz = temp
-        else:
-            raise ASETrajectoryFileError("Bad format for C vector components")
-
-        # The ax component.
-        unitCell[0] = xhi - xlo
-
-        # The bx and by components.
-        unitCell[3] = xy
-        unitCell[4] = yhi - ylo
-
-        # The cx, cy and cz components.
-        unitCell[6] = xz
-        unitCell[7] = yz
-        unitCell[8] = zhi - zlo
+        try:
+            frame = self._input[index]
+        except TypeError:
+            frame = read(self.configuration["trajectory_file"]["value"], index=index)
+        time = self._timeaxis[index]
 
-        unitCell = np.reshape(unitCell, (3, 3))
+        unitCell = frame.cell.array
 
         unitCell *= measure(1.0, "ang").toval("nm")
         unitCell = UnitCell(unitCell)
 
-        for _ in range(
-            self._itemsPosition["BOX BOUNDS"][1], self._itemsPosition["ATOMS"][0]
-        ):
-            self._lammps.readline()
-
-        coords = np.empty(
-            (self._trajectory.chemical_system.number_of_atoms, 3), dtype=np.float64
-        )
-
-        for i, _ in enumerate(
-            range(self._itemsPosition["ATOMS"][0], self._itemsPosition["ATOMS"][1])
-        ):
-            temp = self._lammps.readline().split()
-            idx = self._nameToIndex[self._rankToName[int(temp[0]) - 1]]
-            coords[idx, :] = np.array(
-                [temp[self._x], temp[self._y], temp[self._z]], dtype=np.float64
-            )
+        coords = frame.get_positions()
+        coords *= measure(1.0, "ang").toval("nm")
 
         if self._fractionalCoordinates:
             conf = PeriodicBoxConfiguration(
                 self._trajectory.chemical_system, coords, unitCell
             )
             realConf = conf.to_real_configuration()
         else:
-            coords *= measure(1.0, "ang").toval("nm")
             realConf = PeriodicRealConfiguration(
                 self._trajectory.chemical_system, coords, unitCell
             )
 
-        if self.configuration["fold"]["value"]:
-            # The whole configuration is folded in to the simulation box.
-            realConf.fold_coordinates()
-
         self._trajectory.chemical_system.configuration = realConf
 
         # A snapshot is created out of the current configuration.
         self._trajectory.dump_configuration(
             time, units={"time": "ps", "unit_cell": "nm", "coordinates": "nm"}
         )
 
-        self._start += self._last
-
         return index, None
 
     def combine(self, index, x):
@@ -251,24 +172,39 @@ def finalize(self):
         Finalize the job.
         """
 
-        self._lammps.close()
+        self._input.close()
 
         # Close the output trajectory.
         self._trajectory.close()
 
         super(ASE, self).finalize()
 
     def parse_first_step(self):
-        self._input = iread(self.configuration["trajectory_file"]["value"], index=0)
+        try:
+            self._input = ASETrajectory(self.configuration["trajectory_file"]["value"])
+        except:
+            self._input = iread(
+                self.configuration["trajectory_file"]["value"], index="[:]"
+            )
+            first_frame = read(self.configuration["trajectory_file"]["value"], index=0)
+            last_iterator = 0
+            generator = iread(self.configuration["trajectory_file"]["value"])
+            for _ in generator:
+                last_iterator += 1
+            generator.close()
+            self._total_number_of_steps = last_iterator
+        else:
+            first_frame = self._input[0]
+            self._total_number_of_steps = len(self._input)
 
-        for i in self._input[:1]:
-            frame = i
+        self._timeaxis = self._timestep * np.arange(self._total_number_of_steps)
+
+        self._fractionalCoordinates = np.all(first_frame.get_pbc())
 
         g = Graph()
 
         element_count = {}
-        element_list = frame.get_chemical_symbols()
-        id_list = np.arange(len(element_list)) + 1
+        element_list = first_frame.get_chemical_symbols()
 
         self._nAtoms = len(element_list)
 
@@ -287,7 +223,7 @@ def parse_first_step(self):
                 try:
                     obj = Atom(node.element, name=node.atomName)
                 except TypeError:
-                    print("EXCEPTION in LAMMPS loader")
+                    print("EXCEPTION in ASE loader")
                     print(f"node.element = {node.element}")
                     print(f"node.atomName = {node.atomName}")
                     print(f"rankToName = {self._rankToName}")

MDANSE/Src/MDANSE/MolecularDynamics/Trajectory.py

@@ -127,7 +127,7 @@ def coordinates(self, frame):
         """
 
         if frame < 0 or frame >= len(self):
-            raise TrajectoryError("Invalid frame number")
+            raise TrajectoryError(f"Invalid frame number: {frame}")
 
         grp = self._h5_file["/configuration"]
 
@@ -144,7 +144,7 @@ def configuration(self, frame):
         """
 
         if frame < 0 or frame >= len(self):
-            raise TrajectoryError("Invalid frame number")
+            raise TrajectoryError(f"Invalid frame number: {frame}")
 
         if self._unit_cells is not None:
             unit_cell = self._unit_cells[frame]
@@ -189,7 +189,7 @@ def unit_cell(self, frame):
         """
 
         if frame < 0 or frame >= len(self):
-            raise TrajectoryError("Invalid frame number")
+            raise TrajectoryError(f"Invalid frame number: {frame}")
 
         if self._unit_cells is not None:
             return self._unit_cells[frame]
@@ -504,7 +504,7 @@ def dump_configuration(self, time, units=None):
 
         if self._current_index >= self._n_steps:
             raise TrajectoryError(
-                "The current steps is greater than the actual number of steps of the trajectory"
+                f"The current index {self._current_index} is greater than the actual number of steps of the trajectory {self._n_steps}"
             )
 
         configuration = self._chemical_system.configuration