Merge pull request #655 from ISISNeutronMuon/maciej/simplify-chemical…

…-system

Improve H5MD trajectory support

MDANSE/Src/MDANSE/Framework/InputData/HDFTrajectoryInputData.py

@@ -61,7 +61,7 @@ def info(self):
         val.append("Number of steps:")
         val.append(f"{self._data}\n")
         val.append("Configuration:")
-        val.append(f"\tIs periodic: {'unit_cell' in self._data.file}\n")
+        val.append(f"\tIs periodic: {self._data.unit_cell(0) is not None}\n")
         try:
             val.append(f"First unit cell (nm):\n{self._data.unit_cell(0)._unit_cell}\n")
         except Exception:

MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py

@@ -52,16 +52,34 @@ def __init__(self, h5_filename: Union[Path, str]):
         self._h5_filename = Path(h5_filename)
 
         self._h5_file = h5py.File(self._h5_filename, "r")
-
-        # Load the chemical system
-        try:
+        particle_types = self._h5_file["/particles/all/species"]
+        particle_lookup = h5py.check_enum_dtype(
+            self._h5_file["/particles/all/species"].dtype
+        )
+        if particle_lookup is None:
+            # Load the chemical system
+            try:
+                symbols = self._h5_file["/parameters/atom_symbols"]
+            except KeyError:
+                LOG.error(
+                    f"No information about chemical elements in {self._h5_filename}"
+                )
+                return
+            else:
+                chemical_elements = [byte.decode() for byte in symbols]
+        else:
+            reverse_lookup = {item: key for key, item in particle_lookup.items()}
             chemical_elements = [
-                byte.decode() for byte in self._h5_file["/parameters/atom_symbols"]
+                reverse_lookup[type_number] for type_number in particle_types
             ]
-        except KeyError:
-            chemical_elements = self._h5_file["/particles/all/species"]
         self._chemical_system = ChemicalSystem(self._h5_filename.stem)
-        self._chemical_system.initialise_atoms(chemical_elements)
+        try:
+            self._chemical_system.initialise_atoms(chemical_elements)
+        except Exception:
+            LOG.error(
+                "It was not possible to read chemical element information from an H5MD file."
+            )
+            return
 
         # Load all the unit cells
         self._load_unit_cells()
@@ -73,7 +91,7 @@ def __init__(self, h5_filename: Union[Path, str]):
         except Exception:
             conv_factor = 1.0
         else:
-            if pos_unit == "Ang":
+            if pos_unit in ("Ang", "Angstrom"):
                 pos_unit = "ang"
             conv_factor = measure(1.0, pos_unit).toval("nm")
         coords *= conv_factor
@@ -92,6 +110,7 @@ def file_is_right(self, filename):
                 temp["h5md"]
             except Exception:
                 result = False
+            temp.close()
         return result
 
     def close(self):
@@ -115,7 +134,7 @@ def __getitem__(self, frame):
         except Exception:
             conv_factor = 1.0
         else:
-            if pos_unit == "Ang":
+            if pos_unit in ("Ang", "Angstrom"):
                 pos_unit = "ang"
             conv_factor = measure(1.0, pos_unit).toval("nm")
         configuration = {}
@@ -191,7 +210,7 @@ def coordinates(self, frame):
         except Exception:
             conv_factor = 1.0
         else:
-            if pos_unit == "Ang":
+            if pos_unit in ("Ang", "Angstrom"):
                 pos_unit = "ang"
             conv_factor = measure(1.0, pos_unit).toval("nm")
 
@@ -241,10 +260,10 @@ def _load_unit_cells(self):
         self._unit_cells = []
         try:
             box_unit = self._h5_file["/particles/all/box/edges/value"].attrs["unit"]
-        except Exception:
-            conv_factor = 1.0
+        except (AttributeError, KeyError):
+            conv_factor = 0.1
         else:
-            if box_unit == "Ang":
+            if box_unit in ("Ang", "Angstrom"):
                 box_unit = "ang"
             conv_factor = measure(1.0, box_unit).toval("nm")
         try:
@@ -254,9 +273,16 @@ def _load_unit_cells(self):
         else:
             if len(cells.shape) > 1:
                 for cell in cells:
-                    temp_array = np.array(
-                        [[cell[0], 0.0, 0.0], [0.0, cell[1], 0.0], [0.0, 0.0, cell[2]]]
-                    )
+                    if cell.shape == (3, 3):
+                        temp_array = np.array(cell)
+                    else:
+                        temp_array = np.array(
+                            [
+                                [cell[0], 0.0, 0.0],
+                                [0.0, cell[1], 0.0],
+                                [0.0, 0.0, cell[2]],
+                            ]
+                        )
                     uc = UnitCell(temp_array)
                     self._unit_cells.append(uc)
             else:
@@ -268,14 +294,17 @@ def _load_unit_cells(self):
     def time(self):
         try:
             time_unit = self._h5_file["/particles/all/position/time"].attrs["unit"]
-        except Exception:
+        except KeyError:
             conv_factor = 1.0
         else:
             conv_factor = measure(1.0, time_unit).toval("ps")
         try:
             time = self._h5_file["/particles/all/position/time"] * conv_factor
-        except Exception:
-            time = []
+        except TypeError:
+            try:
+                time = self._h5_file["/particles/all/position/time"][:] * conv_factor
+            except Exception:
+                time = []
         return time
 
     def unit_cell(self, frame):
@@ -367,7 +396,7 @@ def read_com_trajectory(
         except Exception:
             conv_factor = 1.0
         else:
-            if pos_unit == "Ang":
+            if pos_unit in ("Ang", "Angstrom"):
                 pos_unit = "ang"
             conv_factor = measure(1.0, pos_unit).toval("nm")
 
@@ -465,7 +494,7 @@ def read_atomic_trajectory(
         except Exception:
             conv_factor = 1.0
         else:
-            if pos_unit == "Ang":
+            if pos_unit in ("Ang", "Angstrom"):
                 pos_unit = "ang"
             conv_factor = measure(1.0, pos_unit).toval("nm")
         coords = grp[first:last:step, index, :].astype(np.float64) * conv_factor