diff --git a/ipsuite/data_loading/add_data_gromacs.py b/ipsuite/data_loading/add_data_gromacs.py index d3a47444..6355e0c1 100644 --- a/ipsuite/data_loading/add_data_gromacs.py +++ b/ipsuite/data_loading/add_data_gromacs.py @@ -34,7 +34,15 @@ def _get_symbols(u: mda.Universe) -> list[str]: - """Extract element symbols from a Universe, trying multiple strategies.""" + """ + Produce a list of element symbols for the atoms in an MDAnalysis Universe by using available per-atom metadata and sensible fallbacks. + + Parameters: + u (mda.Universe): MDAnalysis Universe containing atoms to derive symbols for. + + Returns: + list[str]: Element symbols (e.g., "C", "Cl", "Na") for each atom in the Universe in atom order. + """ # 1. Use elements attribute if available try: return list(u.atoms.elements) @@ -228,6 +236,14 @@ class Gmx2Frames(zntrack.Node): frames_path: Path = zntrack.outs_path(zntrack.nwd / "frames.h5") def run(self) -> None: + """ + Convert the configured GROMACS inputs into ASE Atoms frames and persist them + to the node's HDF5 output at self.frames_path. + + The node's topology, optional trajectory, optional EDR file, and slicing + parameters (start, stop, step) are used to produce the frames which are + written to the frames_path via znh5md. + """ data = gmx_to_ase( topology=str(self.topology), trajectory=str(self.trajectory) if self.trajectory else None, @@ -241,6 +257,12 @@ def run(self) -> None: @property def frames(self) -> typing.List[Atoms]: + """ + Return all ASE `Atoms` frames stored in the node's HDF5 frames file. + + Returns: + typing.List[Atoms]: A list of ASE `Atoms` objects read from the HDF5 file at `self.frames_path`. + """ with self.state.fs.open(self.frames_path, "rb") as f: with h5py.File(f) as file: return znh5md.IO(file_handle=file)[:]