Collect md steps (#380)

* collect md steps node

* convert to python types

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ipsuite/__init__.pyi

@@ -8,6 +8,7 @@ from .analysis import (
     BoxHeatUp,
     BoxScale,
     CalibrationMetrics,
+    CollectMDSteps,
     ConnectivityCheck,
     DebugCheck,
     DipoleHistogram,
@@ -168,6 +169,7 @@ __all__ = [
     "ThresholdCheck",
     "TemperatureCheck",
     "AnalyseDensity",
+    "CollectMDSteps",
     # Calculators
     "CP2KSinglePoint",
     "ASEGeoOpt",

ipsuite/analysis/__init__.py

@@ -8,7 +8,7 @@
 )
 from ipsuite.analysis.bond_stretch import BondStretchAnalyses
 from ipsuite.analysis.ensemble import ModelEnsembleAnalysis
-from ipsuite.analysis.md import AnalyseDensity
+from ipsuite.analysis.md import AnalyseDensity, CollectMDSteps
 from ipsuite.analysis.model import (
     BoxHeatUp,
     BoxScale,
@@ -62,4 +62,5 @@
     "ForcesUncertaintyHistogram",
     "EnergyUncertaintyHistogram",
     "AnalyseDensity",
+    "CollectMDSteps",
 ]

ipsuite/analysis/md.py

@@ -6,6 +6,7 @@
 import zntrack
 
 from ipsuite import base
+from ipsuite.calculators.ase_md import ASEMD
 from ipsuite.utils.ase_sim import get_density_from_atoms
 
 
@@ -36,3 +37,17 @@ def run(self):
         }
 
         self.results = pd.DataFrame(densities, columns=["density"])
+
+
+class CollectMDSteps(base.IPSNode):
+    mds: list[ASEMD] = zntrack.deps()
+    metrics: dict = zntrack.metrics()
+
+    def run(self):
+        steps: list[int] = [x.steps_before_stopping for x in self.mds]
+
+        self.metrics = {
+            "total": int(np.sum(steps)),
+            "mean": float(np.mean(steps)),
+            "std": float(np.std(steps)),
+        }