Process images with inconsistent keys

tests/test_inconsistent_data.py

@@ -0,0 +1,20 @@
+import numpy.testing as npt
+
+import znh5md
+
+
+def test_keys_missing(tmp_path, s22, s22_energy_forces):
+    io = znh5md.IO(tmp_path / "test.h5")
+
+    images = s22_energy_forces + s22
+    io.extend(images)
+    assert len(io) == len(images)
+    assert len(list(io)) == len(images)
+
+    for a, b in zip(images, io):
+        assert a == b
+        if b.calc is not None:
+            for key in b.calc.results:
+                npt.assert_array_equal(a.calc.results[key], b.calc.results[key])
+        else:
+            assert a.calc is None

znh5md/format.py

@@ -288,10 +288,20 @@ def _combine_dicts(dicts: List[Dict[str, np.ndarray]]) -> Dict[str, np.ndarray]:
     """Helper function to combine dictionaries containing numpy arrays."""
     combined = {}
     for key in dicts[0]:
-        combined[key] = concatenate_varying_shape_arrays(
-            [
-                d[key] if isinstance(d[key], np.ndarray) else np.array(d[key])
-                for d in dicts
-            ]
-        )
+        data = []
+        for d in dicts:
+            if key in d:
+                data.append(d[key])
+            else:
+                dims = dicts[0][key].ndim
+                # Create an array with the appropriate number of dimensions.
+                if dims == 0:
+                    # Handle the case where the number of dimensions is 0
+                    data.append(np.NaN)
+                else:
+                    data.append(np.full_like(dicts[0][key], np.NaN))
+        if data:
+            combined[key] = concatenate_varying_shape_arrays(data)
+        else:
+            raise ValueError(f"Key {key} is missing in one of the data objects.")
     return combined

znh5md/utils.py

@@ -58,6 +58,8 @@ def remove_nan_rows(array: np.ndarray) -> np.ndarray | None:
     1
 
     """
+    if np.isnan(array).all():
+        return None
     if len(np.shape(array)) == 0:
         return array if not np.isnan(array) else None
     return array[~np.isnan(array).all(axis=tuple(range(1, array.ndim)))]
@@ -153,8 +155,9 @@ def build_atoms(args) -> ase.Atoms:
 
     if calc_data is not None:
         if len(calc_data) > 0:
-            atoms.calc = SinglePointCalculator(atoms=atoms)
-            atoms.calc.results = calc_data
+            if not all(val is None for val in calc_data.values()):
+                atoms.calc = SinglePointCalculator(atoms=atoms)
+                atoms.calc.results = calc_data
 
     return atoms