asyncmd.trajectory.convert.FrameExtractor#

The asyncmd.trajectory.convert module contains various predefined classes to extract single frames from Trajectories, possibly after applying a modification. It is also very easy to write your own FrameExtractor class with a custom modification. This is as easy as subclassing from asyncmd.trajectory.convert.FrameExtractor and implementing the apply_modification method (see below).

The asyncmd.trajectory.convert module also contains a TrajectoryConcatenator, which can be used to concatenate/write out trajectories from a list of trajectories and slices (it is explained in more detail in the notebook ConditionalTrajectoryPropagator.ipynp).

Imports and some basic checks that everything is available#

%%bash
# if using the module system to make gromacs and friends available:
# check that they are loaded!
#module list
%%bash
# unix only, check that gmx is available
which gmx
/usr/local/gromacs-2022.4/bin/gmx
%matplotlib inline
import asyncio
import os
import numpy as np
from asyncmd import gromacs as asyncgmx
Could not initialize SLURM cluster handling. If you are sure SLURM (sinfo/sacct/etc) is available try calling `asyncmd.config.set_slurm_settings()` with the appropriate arguments.

Setup working directory#

We will write the trajectory output to it.

workdir = "."

Create a short MD trajectory to extract frames from#

Load and modify the parameter file (mdp file) for the molecular dynamics simulations#

# Pcoupl = C-rescale needs gromacs version >= 2021
mdp = asyncgmx.MDP("../resources/gromacs/capped_alanine_dipeptide/md.mdp")
print(mdp["Pcoupl"])
# set nstxout-compressed, such that the engines will produce XTC trajectories
mdp["nstxout-compressed"] = 20
# and deactivate trr trajectory output
mdp["nstxout"] = mdp["nstvout"] = 0
C-rescale
# if your gmx version is >= 2021 you should comment the next line since C-rescale give the correct ensemble (and Berendsen doesnt!)
#mdp["Pcoupl"] = "Berendsen"

Initialize and run a GmxEngine to create a short trajectory#

We will use this trajectory to extract frames from using the various FrameExtractor classes.

engine = asyncgmx.GmxEngine(mdconfig=mdp,
                            gro_file="../resources/gromacs/capped_alanine_dipeptide/conf.gro",
                            top_file="../resources/gromacs/capped_alanine_dipeptide/topol_amber99sbildn.top",
                            mdrun_extra_args="-nt 2",
                           )

await engine.prepare(starting_configuration=None, workdir=workdir, deffnm="traj_to_extract_from")
traj = await engine.run_steps(nsteps=1e4)

Extract Frames using the predefined FrameExtractor classes#

Each FrameExtractor takes (one of) the arguments mda_transformations and mda_transformations_setup_func which allow you to pass/setup MDAnalysis on-the-fly transformations to e.g. center on a given molecule and wrap all molecules/atoms back into the simulation box while extracting and writing out the frame. See the FrameExtractor docstring for when to use mda_transformations and when mda_transformations_setup_func and see https://docs.mdanalysis.org/stable/documentation_pages/trajectory_transformations.html for more on MDAnalysis transformations.

# these are all the FrameExtractors,
# note that FrameExtractor is an abstract base class, i.e. you can not instantiate it (or its subclasses without implementing the apply_modification method)
from asyncmd.trajectory.convert import (FrameExtractor, NoModificationFrameExtractor,
                                        InvertedVelocitiesFrameExtractor, RandomVelocitiesFrameExtractor)
# extract a number of frames, each FrameExtractor works the same, so we will only use the RandomVelocitiesFrameExtractor
extractor = RandomVelocitiesFrameExtractor(T=303, # temperature for Maxwell-Boltzmann velocities in Kelvin
                                           )
n_frames = 10
for i in range(n_frames):
    # the extract method returns the frame as an asyncmd.Trajectory
    print(extractor.extract(outfile=os.path.join(workdir, f"frame_{i}.trr"),  # where to write the frame to
                            traj_in=traj,  # the trajectory from which we take the original frame
                            idx=np.random.randint(len(traj)),  # the index of the frame in traj_in
                            )
          )
Trajectory(trajectory_files=frame_0.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_1.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_2.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_3.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_4.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_5.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_6.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_7.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_8.trr, structure_file=traj_to_extract_from.tpr)
Trajectory(trajectory_files=frame_9.trr, structure_file=traj_to_extract_from.tpr)

The extract method also has an async counterpart#

# it has exactly the same arguments as the extract method
await asyncio.gather(*(extractor.extract_async(outfile=os.path.join(workdir, f"frame_{i}.trr"),  # where to write the frame to
                                               traj_in=traj,  # the trajectory from which we take the original frame
                                               idx=np.random.randint(len(traj)),  # the index of the frame in traj_in
                                               overwrite=True,  # overwrite=True makes sure we overwrite the existing outfiles (from the cell above) without asking/error
                                               )
                       for i in range(n_frames)
                       )
                      )
[Trajectory(trajectory_files=frame_0.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_1.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_2.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_3.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_4.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_5.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_6.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_7.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_8.trr, structure_file=traj_to_extract_from.tpr),
 Trajectory(trajectory_files=frame_9.trr, structure_file=traj_to_extract_from.tpr)]

Writing your own FrameExtractor subclass#

# it is as easy as this:
class CustomFrameExtractor(FrameExtractor):
    def apply_modification(self, universe, ts):
        # universe is the mdanalysis universe of the Trajectory/Frame that is being extracted
        # ts is the timestep of the Frame that is being extracted
        
        # Here you can no apply you desired modifications to the timestep
        ts.positions *= 100 ## dont do this in real live ;)

        # the function does not (need to) return anything. Any return will be ignored
        # But the changes to the timestep and universe will naturally be written out with the extracted frame


# see also the implementations of the InvertedVelocitiesFrameExtractor and the RandomVelocitiesFrameExtractor in asyncmd/trajectory/convert.py
custom_extract = CustomFrameExtractor()
custom_extract.extract(outfile=os.path.join(workdir, "frame_custom.trr"),
                       traj_in=traj,
                       idx=0,
                      )
Trajectory(trajectory_files=frame_custom.trr, structure_file=traj_to_extract_from.tpr)