Merge pull request i-pi#409 from lab-cosmo/fix_read_traj

Fix traj parsing

ipi/utils/parsing.py

@@ -9,6 +9,7 @@
 import re
 import numpy as np
 from ipi.utils.units import unit_to_user
+from ipi.utils.messages import warning, verbosity
 
 try:
     import ase
@@ -138,8 +139,7 @@ def read_trajectory(
             raise ValueError(f"Unrecognized file format: {format}")
 
     file_handle = open(filename, "r")
-    bohr2angstrom = unit_to_user("length", "angstrom", 1.0)
-    comment_regex = re.compile(r"([^)]+)\{([^}]+)\}")
+    comment_regex = re.compile(r"(\w+)\{([^}]+)\}")
     step_regex = re.compile(r"Step:\s+(\d+)")
 
     frames = []
@@ -200,17 +200,17 @@ def read_trajectory(
 
                 frame = ase.Atoms(
                     ret["atoms"].names,
-                    positions=ret["atoms"].q.reshape((-1, 3)) * bohr2angstrom,
-                    cell=ret["cell"].h.T * bohr2angstrom,
+                    # will apply conversion later!
+                    positions=ret["atoms"].q.reshape((-1, 3)),
+                    cell=ret["cell"].h.T * unit_to_user("length", "ase"),
                     pbc=True,
                 )
 
                 # parse comment to get the property
                 matches = comment_regex.findall(ret["comment"])
-
                 # get what we have found
-                if len(matches) >= 2:
-                    what = matches[-2][0]
+                if len(matches) == 2:
+                    what = matches[0][0]
                 else:  # defaults to reading positions
                     what = "positions"
 
@@ -219,13 +219,25 @@ def read_trajectory(
                 if len(matches) >= 1:
                     frame.info["step"] = int(matches[-1][0])
 
-                # if we have forces, set positions to zero (that data is missing!) and set forces instead
-                if what == "forces":
-                    # set forces and convert to eV/angstrom
-                    frame.positions *= 0
-                    frame.arrays["forces"] = ret["atoms"].q.reshape(
-                        (-1, 3)
-                    ) * unit_to_user("force", "ev/ang", 1.0)
+                # fetch the list of known traj types, cf. `engine/properties.py``
+                from ipi.engine.properties import Trajectories  # avoids circular import
+
+                traj_types = Trajectories().traj_dict
+                if not what in traj_types:
+                    warning(
+                        f"{what} is not a known trajectory type. Will apply no units conversion",
+                        verbosity.low,
+                    )
+                elif traj_types[what]["dimension"] == "length":
+                    # positions is the right place to store, and we just need to convert
+                    frame.positions *= unit_to_user("length", "ase")
+                else:
+                    # if we have another type of value, set positions to zero
+                    # (that data is missing!) and set an array instead
+                    frame.positions *= 0.0
+                    frame.arrays[what] = ret["atoms"].q.reshape((-1, 3)) * unit_to_user(
+                        traj_types[what]["dimension"], "ase"
+                    )
 
                 frames.append(frame)
 

ipi/utils/units.py

@@ -365,7 +365,7 @@ def mass(cls, label):
 #
 
 
-def unit_to_internal(family, unit, number):
+def unit_to_internal(family, unit, number=1.0):
     """Converts a number of given dimensions and units into internal units.
 
     Args:
@@ -411,7 +411,7 @@ def unit_to_internal(family, unit, number):
     return number * UnitMap[family][base.lower()] * UnitPrefix[prefix]
 
 
-def unit_to_user(family, unit, number):
+def unit_to_user(family, unit, number=1.0):
     """Converts a number of given dimensions from internal to user units.
 
     Args:
@@ -423,4 +423,4 @@ def unit_to_user(family, unit, number):
         The number in the user specified units
     """
 
-    return number / unit_to_internal(family, unit, 1.0)
+    return number / unit_to_internal(family, unit)