add plot_molecule_counts and AMSResults.get_molecule_count_history() …

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.
 
-

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_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."
+            )
+
+        # 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 = [
+            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")
+
+        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:

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,

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
 # ----------------------------------------------------------