From 4d81148f358db7ff122a7e66f3942e9b5f3c1e47 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matti=20Hellstr=C3=B6m?= Date: Thu, 13 Aug 2026 11:31:55 +0200 Subject: [PATCH 1/2] add plot_molecule_counts and AMSResults.get_molecule_count_history() SO-- --- CHANGELOG.md | 2 +- src/scm/plams/interfaces/adfsuite/ams.py | 70 ++++++++++++++++++++++++ src/scm/plams/tools/plot.py | 61 +++++++++++++++++++++ unit_tests/test_tools_plot.py | 26 +++++++++ 4 files changed, 158 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5925914dc..44df63a2a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ This changelog is effective from the 2025 releases. ### Added * MultiJob now supports generic Job types for the self.children attribute * Function `view_orbital` to visualize orbitals of completed AMS jobs +* `AMSResults.get_molecule_count_history` to extract molecular-formula populations from reactive MD trajectories, and `plot_molecule_counts` in `scm.plams.tools.plot` to visualize them ### Changed * `Molecule.readmol2` can now read non-integer bond orders @@ -138,4 +139,3 @@ This changelog is effective from the 2025 releases. * Legacy `BANDJob`, `DFTBJob`, `UFFJob`, `MOPACJob`, `ReaxFFJob`, `CSHessianADFJob` and `ADFJob` have been removed * Exception classes `AMSPipeDecodeError`, `AMSPipeError`, `AMSPipeInvalidArgumentError`, `AMSPipeLogicError`, `AMSPipeRuntimeError`, `AMSPipeUnknownArgumentError`, `AMSPipeUnknownMethodError`, `AMSPipeUnknownVersionError`, were moved from scm.plams to scm.amspipe. - diff --git a/src/scm/plams/interfaces/adfsuite/ams.py b/src/scm/plams/interfaces/adfsuite/ams.py index 08586982e..75f200daa 100644 --- a/src/scm/plams/interfaces/adfsuite/ams.py +++ b/src/scm/plams/interfaces/adfsuite/ams.py @@ -1,5 +1,6 @@ import os import re +from collections import Counter from typing import ( Dict, @@ -532,6 +533,75 @@ def get_history_property(self, varname: str, history_section: str = "History") - values = [main.read(history_section, f"{varname}({step})") for step in range(1, nentries + 1)] # type: ignore[misc] return values + def get_molecule_count_history( + self, species: Optional[Sequence[str]] = None + ) -> Tuple[np.ndarray, np.ndarray, Dict[str, np.ndarray]]: + """Return molecular formula populations for every saved MD frame. + + The MD trajectory must have been run with + ``MolecularDynamics%Trajectory%WriteMolecules=True``. Returns one-based frame indices, + MD times in fs, and a dictionary with one population array per molecular formula. When + *species* is ``None``, all detected formulas are returned; otherwise, only the requested + formulas are included. + """ + n_molecule_types = self.readrkf("Molecules", "Num molecules") + if n_molecule_types is None: + raise KeyError( + "Molecule information is not present in ams.rkf. " + "Run MD with MolecularDynamics%Trajectory%WriteMolecules=True." + ) + + # The Molecules section defines the lookup table: molecule type N has the formula + # stored in ``Molecule name N``. History/Mols.Type below contains these type numbers. + molecule_names = [ + self.readrkf("Molecules", f"Molecule name {molecule_type}") + for molecule_type in range(1, int(n_molecule_types) + 1) + ] + if any(name is None for name in molecule_names): + raise KeyError("Molecule names are incomplete in the ams.rkf file") + molecule_names = cast(List[str], molecule_names) + + try: + molecule_type_history = self.get_history_property("Mols.Type") + except KeyError as error: + raise KeyError( + "Molecule history is not present in ams.rkf. " + "Run MD with MolecularDynamics%Trajectory%WriteMolecules=True." + ) from error + if molecule_type_history is None: + raise KeyError("Molecule history is not present in the ams.rkf file") + + available_species = list(dict.fromkeys(molecule_names)) + if species is None: + selected_species = available_species + else: + selected_species = [species] if isinstance(species, str) else list(dict.fromkeys(species)) + unknown_species = set(selected_species) - set(available_species) + if unknown_species: + raise ValueError(f"Unknown molecular formula(s): {', '.join(sorted(unknown_species))}") + + try: + time = self.get_history_property("Time", "MDHistory") + except KeyError as error: + raise KeyError("MD time history is not present in the ams.rkf file") from error + if time is None: + raise KeyError("MD time history is not present in the ams.rkf file") + if len(time) != len(molecule_type_history): + raise ValueError("MD time and molecule histories have different lengths") + + counts = {name: np.zeros(len(molecule_type_history), dtype=int) for name in selected_species} + for frame_index, molecule_types in enumerate(molecule_type_history): + # Each History/Mols.Type(frame) entry lists one type number per detected molecule in + # that frame. Count them and use the Molecules lookup table to obtain populations per + # molecular formula. Multiple type numbers with the same formula are accumulated. + for molecule_type, count in Counter(molecule_types).items(): + name = molecule_names[molecule_type - 1] + if name in counts: + counts[name][frame_index] += count + frames = np.arange(1, len(molecule_type_history) + 1) + time_fs = np.asarray(time, dtype=float) + return frames, time_fs, counts + def get_property_at_step(self, step: int, varname: str, history_section: str = "History") -> Optional["TRead"]: """Return the value of *varname* in the history section *history_section at step *step*.""" if "ams" not in self.rkfs: diff --git a/src/scm/plams/tools/plot.py b/src/scm/plams/tools/plot.py index 88ea8c345..49486cc33 100644 --- a/src/scm/plams/tools/plot.py +++ b/src/scm/plams/tools/plot.py @@ -16,6 +16,7 @@ from scm.plams.core.functions import requires_optional_package from scm.plams.interfaces.adfsuite.ams import AMSJob from scm.plams.mol.molecule import Molecule +from scm.plams.tools.units import Units try: from scm.base import ChemicalSystem @@ -41,6 +42,7 @@ "plot_image_grid", "plot_correlation", "plot_msd", + "plot_molecule_counts", "plot_work_function", "plot_grid_molecules", ] @@ -806,6 +808,65 @@ def plot_msd( return ax +@requires_optional_package("matplotlib") +def plot_molecule_counts( + job: AMSJob, + species: Optional[Sequence[str]] = None, + x_axis: Literal["time", "frame"] = "time", + time_unit: str = "fs", + ax: Optional["plt.Axes"] = None, +) -> "plt.Axes": + """Plot populations of molecular species detected in a reactive MD trajectory. + + The trajectory must have been run with ``MolecularDynamics%Trajectory%WriteMolecules=True``. + + job + An AMS MD job with molecule information stored in its ``ams.rkf`` file. + + species + Molecular formulas to plot. If ``None``, plot every detected molecular formula. + + x_axis + ``"time"`` (default) plots against the saved MD time. ``"frame"`` plots against the + one-based trajectory frame number. + + time_unit + Unit for the time axis when ``x_axis="time"``. The time values stored in ``ams.rkf`` are + in fs. + + ax + Matplotlib axis. If ``None``, one is created. + + Returns + ------- + matplotlib.axes.Axes + The axis containing one curve per selected molecular formula. + """ + import matplotlib.pyplot as plt + + if x_axis not in {"time", "frame"}: + raise ValueError(f"x_axis must be 'time' or 'frame', not {x_axis!r}") + + frames, time_fs, counts = job.results.get_molecule_count_history(species=species) + if x_axis == "time": + x = Units.convert(time_fs, "fs", time_unit) + xlabel = f"Time ({time_unit})" + else: + x = frames + xlabel = "Frame" + + if ax is None: + _, ax = plt.subplots() + + for name, population in counts.items(): + ax.plot(x, population, label=name) + ax.set_xlabel(xlabel) + ax.set_ylabel("Molecule count") + ax.legend() + + return ax + + @requires_optional_package("matplotlib") def plot_work_function( coordinate: np.ndarray, diff --git a/unit_tests/test_tools_plot.py b/unit_tests/test_tools_plot.py index f6c823cc0..5483cabc2 100644 --- a/unit_tests/test_tools_plot.py +++ b/unit_tests/test_tools_plot.py @@ -25,6 +25,7 @@ plot_correlation, plot_grid_molecules, plot_molecule, + plot_molecule_counts, plot_msd, plot_work_function, ) @@ -406,6 +407,31 @@ def test_plot_correlation(run_calculations, rkf_tools_plot): assert np.allclose(y, y0, 1e-5) +# ---------------------------------------------------------- +# Testing plot_molecule_counts +# ---------------------------------------------------------- +def test_plot_molecule_counts(rkf_tools_plot): + plt.close("all") + + job = AMSJob.load_external(rkf_tools_plot / "md") + frames, time_fs, counts = job.results.get_molecule_count_history(species=["H2O"]) + assert frames[:2].tolist() == [1, 2] + assert time_fs[:2].tolist() == pytest.approx([0.0, 0.5]) + assert counts["H2O"][:2].tolist() == [16, 16] + + ax = plot_molecule_counts(job, species=["H2O"], time_unit="ps") + + assert ax.get_xlabel() == "Time (ps)" + assert ax.get_ylabel() == "Molecule count" + assert ax.lines[0].get_label() == "H2O" + assert ax.lines[0].get_xdata().tolist()[:2] == pytest.approx([0.0, 0.0005]) + assert ax.lines[0].get_ydata().tolist()[:2] == [16, 16] + + ax = plot_molecule_counts(job, x_axis="frame") + assert ax.get_xlabel() == "Frame" + assert ax.lines[0].get_xdata().tolist()[:2] == [1, 2] + + # ---------------------------------------------------------- # Testing plot_msd # ---------------------------------------------------------- From b4cc735ec283462fe3ab7519c6857c88442b4f08 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matti=20Hellstr=C3=B6m?= Date: Thu, 13 Aug 2026 11:39:11 +0200 Subject: [PATCH 2/2] fix mypy errors SO-- --- src/scm/plams/interfaces/adfsuite/ams.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/scm/plams/interfaces/adfsuite/ams.py b/src/scm/plams/interfaces/adfsuite/ams.py index 75f200daa..e95ea55a9 100644 --- a/src/scm/plams/interfaces/adfsuite/ams.py +++ b/src/scm/plams/interfaces/adfsuite/ams.py @@ -544,8 +544,8 @@ def get_molecule_count_history( *species* is ``None``, all detected formulas are returned; otherwise, only the requested formulas are included. """ - n_molecule_types = self.readrkf("Molecules", "Num molecules") - if n_molecule_types is None: + n_molecule_types_raw = self.readrkf("Molecules", "Num molecules") + if n_molecule_types_raw is None: raise KeyError( "Molecule information is not present in ams.rkf. " "Run MD with MolecularDynamics%Trajectory%WriteMolecules=True." @@ -553,13 +553,13 @@ def get_molecule_count_history( # The Molecules section defines the lookup table: molecule type N has the formula # stored in ``Molecule name N``. History/Mols.Type below contains these type numbers. + n_molecule_types = cast(int, n_molecule_types_raw) molecule_names = [ - self.readrkf("Molecules", f"Molecule name {molecule_type}") - for molecule_type in range(1, int(n_molecule_types) + 1) + cast(str, self.readrkf("Molecules", f"Molecule name {molecule_type}")) + for molecule_type in range(1, n_molecule_types + 1) ] if any(name is None for name in molecule_names): raise KeyError("Molecule names are incomplete in the ams.rkf file") - molecule_names = cast(List[str], molecule_names) try: molecule_type_history = self.get_history_property("Mols.Type")