diff --git a/ipsuite/__init__.pyi b/ipsuite/__init__.pyi index 27735998..20c269fd 100644 --- a/ipsuite/__init__.pyi +++ b/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", diff --git a/ipsuite/analysis/__init__.py b/ipsuite/analysis/__init__.py index 77f4f0e6..28a74dd8 100644 --- a/ipsuite/analysis/__init__.py +++ b/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", ] diff --git a/ipsuite/analysis/md.py b/ipsuite/analysis/md.py index d8d19bfe..f09c94f2 100644 --- a/ipsuite/analysis/md.py +++ b/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)), + }