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)