NEB bug fixes

emmet-core/emmet/core/mobility/migrationgraph.py

@@ -690,7 +690,7 @@ def _get_energy_struct_info(npr: NebPathwayResult) -> dict[str, dict[str, Any]]:
         for hop_key, data in npr.hops.items():
             energy_struct_info[hop_key] = {
                 "hop_key": hop_key,
-                "max_barrier": npr.max_barriers[hop_key],
+                "max_barrier": npr.max_barriers.get(hop_key),
                 "energy_range": data.barrier_energy_range,
                 "energies": data.energies,
                 "state": data.state,

emmet-core/emmet/core/neb.py

@@ -651,20 +651,16 @@ def set_top_levels(self) -> Self:
     @property
     def max_barriers(self) -> dict[str, float | None] | None:
         """Retrieve the maximum barrier along each hop."""
-
-        barriers: dict[str, list[float]] = {}
-        for b in [self.forward_barriers, self.reverse_barriers]:
-            if b:
-                for idx, barrier in b.items():
-                    if barrier:
-                        if idx not in barriers:
-                            barriers[idx] = []
-                        barriers[idx].append(barrier)
-
-        if len(barriers) == 0:
-            return None
-
-        return {idx: max(v) for idx, v in barriers.items()}
+        fwd_barr = self.forward_barriers or {}
+        rev_barr = self.reverse_barriers or {}
+        return {
+            hop_key: max(
+                fwd_barr.get(hop_key) or float("-inf"),
+                rev_barr.get(hop_key) or float("-inf"),
+            )
+            for hop_key in self.hops
+            if fwd_barr.get(hop_key) or rev_barr.get(hop_key)
+        } or None
 
     @property
     def barrier_ranges(self) -> dict[str, float | None]:

emmet-core/tests/test_neb.py

@@ -9,8 +9,10 @@
 from pymatgen.core import Structure
 
 from emmet.core.neb import (
+    NebResult,
     NebTaskDoc,
     NebMethod,
+    NebPathwayResult,
     BarrierAnalysis,
     NebIntermediateImagesDoc,
 )
@@ -120,3 +122,49 @@ def test_from_directories(neb_test_dir):
         == getattr(neb_doc.barrier_analysis, f"{direction}_barrier", None)
         for direction in ("forward", "reverse")
     )
+
+
+def test_top_level_assignment(test_dir):
+
+    structure = Structure.from_file(test_dir / "Si_mp_149.cif")
+    fictive_doc = NebPathwayResult(
+        hops={
+            k: NebResult(forward_barrier=barriers[0], reverse_barrier=barriers[1])
+            for k, barriers in {
+                "a": (1.0, None),
+                "b": (5.0, 10.0),
+                "c": (None, 0.1),
+            }.items()
+        },
+        host_structure=structure,
+    )
+    assert fictive_doc.forward_barriers == {
+        "a": 1.0,
+        "b": 5.0,
+        "c": None,
+    }
+    assert fictive_doc.reverse_barriers == {
+        "a": None,
+        "b": 10.0,
+        "c": 0.1,
+    }
+    assert fictive_doc.max_barriers == {
+        "a": 1.0,
+        "b": 10.0,
+        "c": 0.1,
+    }
+
+    assert fictive_doc.host_formula == structure.formula
+    assert fictive_doc.host_formula_reduced == structure.reduced_formula
+    assert fictive_doc.host_chemsys == structure.chemical_system
+
+    energies = {
+        "0+1": [0.0, 0.5, 0.1],
+        "0+2": [-0.5, 1.25, 0.75],
+    }
+    fictive_doc = NebPathwayResult(
+        hops={k: NebResult(energies=ens) for k, ens in energies.items()}
+    )
+    assert fictive_doc.barrier_ranges == {
+        k: max(ens) - min(ens) for k, ens in energies.items()
+    }