Remove NaN results

janus_core/single_point.py

@@ -236,18 +236,23 @@ def _get_stress(self) -> MaybeList[ndarray]:
     @staticmethod
     def _remove_invalid_props(
         struct: Atoms,
+        results: CalcResults = None,
         properties: Collection[str] = (),
     ) -> None:
         """
-        Remove any invalid properties from calculator results.
+        Remove any invalid properties from calculated results.
 
         Parameters
         ----------
         struct : Atoms
             ASE Atoms structure with attached calculator results.
+        results : CalcResults
+            Dictionary of calculated results. Default is {}.
         properties : Collection[str]
             Physical properties requested to be calculated. Default is ().
         """
+        results = results if results else {}
+
         # Find any properties with non-finite values
         rm_keys = [
             prop
@@ -262,21 +267,29 @@ def _remove_invalid_props(
                     f"'{prop}' contains non-finite values for this structure."
                 )
             del struct.calc.results[prop]
+            if prop in results:
+                del results[prop]
 
-    def _clean_results(self, properties: Collection[str] = ()) -> None:
+    def _clean_results(
+        self, results: CalcResults = None, properties: Collection[str] = ()
+    ) -> None:
         """
-        Remove results with NaN or inf values from calc.results dictionary.
+        Remove NaN and inf values from results and calc.results dictionaries.
 
         Parameters
         ----------
+        results : CalcResults
+            Dictionary of calculated results. Default is {}.
         properties : Collection[str]
             Physical properties requested to be calculated. Default is ().
         """
+        results = results if results else {}
+
         if isinstance(self.struct, list):
             for image in self.struct:
-                self._remove_invalid_props(image, properties)
+                self._remove_invalid_props(image, results, properties)
         else:
-            self._remove_invalid_props(self.struct, properties)
+            self._remove_invalid_props(self.struct, results, properties)
 
     def run_single_point(
         self,
@@ -327,7 +340,7 @@ def run_single_point(
         if "stress" in properties or len(properties) == 0:
             results["stress"] = self._get_stress()
 
-        self._clean_results(properties=properties)
+        self._clean_results(results, properties=properties)
 
         if self.logger:
             self.logger.info("Single point calculation complete")

tests/test_single_point.py

@@ -64,6 +64,20 @@ def test_single_point_none():
         assert prop in results
 
 
+def test_single_point_clean():
+    """Test single point stress using MACE calculator."""
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "H2O.cif",
+        architecture="mace",
+        calc_kwargs={"model_paths": MODEL_PATH},
+    )
+
+    results = single_point.run_single_point()
+    for prop in ["energy", "forces"]:
+        assert prop in results
+    assert "stress" not in results
+
+
 def test_single_point_traj():
     """Test single point stress using MACE calculator."""
     single_point = SinglePoint(