change behavior of reference evaluate function to never return a NullState

Author: svandenhaute

psiflow/reference/_cp2k.py

@@ -117,7 +117,6 @@ def parse_cp2k_output(
             energy = float(line.split()[-1]) * Ha
     if energy is None:
         return NullState
-    # atoms.reference_status = True
     geometry.energy = energy
     geometry.per_atom.forces[:] = np.nan
 
@@ -136,7 +135,6 @@ def parse_cp2k_output(
             forces[j, :] = np.array([float(f) for f in line.split()[3:6]])
         forces *= Ha / Bohr
         geometry.per_atom.forces[:] = forces
-    # atoms.info.pop("stress", None)  # remove if present for some reason
     geometry.stress = None
     return geometry
 
@@ -185,9 +183,6 @@ def cp2k_singlepoint_post(
     from psiflow.geometry import NullState, new_nullstate
     from psiflow.reference._cp2k import parse_cp2k_output
 
-    if geometry == NullState:
-        return NullState.copy()  # copy?
-
     with open(inputs[0], "r") as f:
         cp2k_output_str = f.read()
     geometry = parse_cp2k_output(cp2k_output_str, properties, geometry)

psiflow/reference/gpaw_.py

@@ -56,7 +56,7 @@ def gpaw_singlepoint_post(
     with open(inputs[0], "r") as f:
         lines = f.read().split("\n")
 
-    geometry = new_nullstate()
+    geometry = new_nullstate()  # GPAW parsing doesn't require initial geometry
     for i, line in enumerate(lines):
         if "CALCULATION SUCCESSFUL" in line:
             natoms = int(lines[i + 1])

psiflow/reference/reference.py

@@ -20,7 +20,7 @@
 
 @typeguard.typechecked
 def _extract_energy(state: Geometry):
-    if state == NullState:
+    if state.energy is None:
         return 1e10
     else:
         return state.energy
@@ -40,6 +40,24 @@ def get_minimum_energy(element, configs, *energies):
     return copy_app_future(energy)
 
 
+@typeguard.typechecked
+def _nan_if_unsuccessful(
+    geometry: Geometry,
+    result: Geometry,
+) -> Geometry:
+    if result == NullState:
+        geometry.energy = None
+        geometry.per_atom.forces[:] = np.nan
+        geometry.per_atom.stress = None
+        geometry.stdout = result.stdout
+        return geometry
+    else:
+        return result
+
+
+nan_if_unsuccessful = python_app(_nan_if_unsuccessful, executors=['default_threads'])
+
+
 @join_app
 @typeguard.typechecked
 def evaluate(
@@ -54,10 +72,11 @@ def evaluate(
             stdout=parsl.AUTO_LOGNAME,
             stderr=parsl.AUTO_LOGNAME,
         )
-        return reference.app_post(
+        result = reference.app_post(
             geometry=geometry.copy(),
             inputs=[future.stdout, future.stderr, future],
         )
+        return nan_if_unsuccessful(geometry, result)
 
 
 @join_app

tests/test_reference.py

@@ -320,7 +320,8 @@ def test_cp2k_failure(context, tmp_path):
     evaluated = evaluate(geometry, reference)
     assert isinstance(evaluated, AppFuture)
     state = evaluated.result()
-    assert state == NullState
+    assert state.energy is None
+    assert np.all(np.isnan(state.per_atom.forces))
     with open(state.stdout, "r") as f:
         log = f.read()
     assert "ABORT" in log  # verify error is captured
@@ -445,4 +446,4 @@ def test_gpaw_single(dataset, dataset_h2):
     assert gpaw.compute_atomic_energy("Zr", box_size=9).result() == 0.0
 
     gpaw = GPAW(askdfj="asdfk")  # invalid input
-    assert evaluate(dataset_h2[1], gpaw).result() == NullState
+    assert evaluate(dataset_h2[1], gpaw).result().energy is None