Fix the handling of H5MD trajectories

ISISNeutronMuon · Jan 29, 2025 · 2cf3247 · 2cf3247
1 parent fde578b
commit 2cf3247
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py b/MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py
@@ -87,12 +87,12 @@ def file_is_right(self, filename):
         result = True
         try:
             temp = h5py.File(filename)
-        except:
+        except FileNotFoundError:
             result = False
         else:
             try:
                 temp["h5md"]
-            except:
+            except KeyError:
                 result = False
         return result
 
@@ -338,20 +338,33 @@ def read_com_trajectory(
         if last is None:
             last = len(self)
 
+        if len(atom_indices) == 1:
+            return self.read_atomic_trajectory(
+                atom_indices[0],
+                first=first,
+                last=last,
+                step=step,
+                box_coordinates=box_coordinates,
+            )
+
         atoms = self.chemical_system.atom_list
 
         try:
-            masses = self._h5_file["/particles/all/mass/value"][:].astype(np.float64)
+            masses = self._h5_file["/particles/all/mass/value"][atom_indices].astype(
+                np.float64
+            )
         except KeyError:
             try:
-                masses = self._h5_file["/particles/all/mass"][:].astype(np.float64)
+                masses = self._h5_file["/particles/all/mass"][atom_indices].astype(
+                    np.float64
+                )
             except KeyError:
                 masses = np.array(
                     [
                         ATOMS_DATABASE.get_atom_property(at, "atomic_weight")
                         for at in atoms
                     ]
-                )
+                )[atom_indices]
         grp = self._h5_file["/particles/all/position/value"]
         try:
             pos_unit = self._h5_file["/particles/all/position/value"].attrs["unit"]
@@ -362,7 +375,7 @@ def read_com_trajectory(
                 pos_unit = "ang"
             conv_factor = measure(1.0, pos_unit).toval("nm")
 
-        coords = grp[first:last:step, :, :].astype(np.float64) * conv_factor
+        coords = grp[first:last:step, atom_indices, :].astype(np.float64) * conv_factor
 
         if coords.ndim == 2:
             coords = coords[np.newaxis, :, :]

diff --git a/MDANSE/Src/MDANSE/Trajectory/MdanseTrajectory.py b/MDANSE/Src/MDANSE/Trajectory/MdanseTrajectory.py
@@ -74,9 +74,9 @@ def file_is_right(self, filename):
                     os.path.splitext(os.path.basename(filename))[0]
                 )
                 temp_cs.load(file_object)
-            except Exception as ex:
-                LOG.error(
-                    f"Loading the ChemicalSystem from {filename} failed: {ex}, {traceback.format_exc()}"
+            except Exception:
+                LOG.warning(
+                    f"Could not load ChemicalSystem from {filename}. MDANSE will try to read it as H5MD next."
                 )
                 result = False
             file_object.close()