This page was generated from examples/atom-selection.ipynb.
Atom selection¶
Basics¶
In this example, we’ll go over the basics of atom and reside selection in MDTraj. First let’s load up an example trajectory.
[1]:
import mdtraj as md
traj = md.load("ala2.h5")
print(traj)
<mdtraj.Trajectory with 100 frames, 22 atoms, 3 residues, without unitcells>
/Users/singhs15/work/src/dev-projs/mdtraj/mdtraj/core/trajectory.py:441: UserWarning: top= kwargs ignored since this file parser does not support it
warnings.warn("top= kwargs ignored since this file parser does not support it")
We can also more directly find out how many atoms or residues there are by using traj.n_atoms
and traj.n_residues
.
[2]:
print("How many atoms? %s" % traj.n_atoms)
print("How many residues? %s" % traj.n_residues)
How many atoms? 22
How many residues? 3
We can also manipulate the atom positions by working with traj.xyz
, which is a NumPy array contain the xyz coordinated of each atom with dimensions (n_frames, n_atoms, 3). Let’s find the 3D coordinates of the tenth atom in frame 5.
[4]:
frame_idx = 4 # zero indexed frame number
atom_idx = 9 # zero indexed atom index
print("Where is the fifth atom at the tenth frame?")
print("x: %s\ty: %s\tz: %s" % tuple(traj.xyz[frame_idx, atom_idx, :]))
Where is the fifth atom at the tenth frame?
x: 0.69715095 y: 0.92419046 z: 0.87260383
Topology Object¶
As mentioned previously in the introduction, every Trajectory
object contains a Topology
. The Topology
of a Trajectory
contains all the connectivity information of your system and specific chain, residue, and atom information.
[5]:
topology = traj.topology
print(topology)
<mdtraj.Topology with 1 chains, 3 residues, 22 atoms, 21 bonds>
With the topology object we can select a certain atom
, or loop through them all. (Note: everything is zero-indexed)
[6]:
print("Fifth atom: %s" % topology.atom(4))
print("All atoms: %s" % [atom for atom in topology.atoms])
Fifth atom: ACE1-C
All atoms: [ACE1-H1, ACE1-CH3, ACE1-H2, ACE1-H3, ACE1-C, ACE1-O, ALA2-N, ALA2-H, ALA2-CA, ALA2-HA, ALA2-CB, ALA2-HB1, ALA2-HB2, ALA2-HB3, ALA2-C, ALA2-O, NME3-N, NME3-H, NME3-C, NME3-H1, NME3-H2, NME3-H3]
The same goes for residues.
[7]:
print("Second residue: %s" % traj.topology.residue(1))
print("All residues: %s" % [residue for residue in traj.topology.residues])
Second residue: ALA2
All residues: [ACE1, ALA2, NME3]
Additionally, every atom
and residue
is also an object, and has it’s own set of properties. Here is a simple example that showcases just a few.
[8]:
atom = topology.atom(10)
print(
"""Hi! I am the %sth atom, and my name is %s.
I am a %s atom with %s bonds.
I am part of an %s residue."""
% (atom.index, atom.name, atom.element.name, atom.n_bonds, atom.residue.name)
)
Hi! I am the 10th atom, and my name is CB.
I am a carbon atom with 4 bonds.
I am part of an ALA residue.
There are also more complex properties, like atom.is_sidechain
or residue.is_protein
, which allow for more powerful selections.
Putting Everything Together¶
Hopefully, you can see how these properties can be combined with Python’s filtered list functionality. Let’s say we want the indices of all carbon atoms in the sidechains of our molecule. We could try something like this.
[9]:
print(
[
atom.index
for atom in topology.atoms
if atom.element.symbol == "C" and atom.is_sidechain
]
)
[1, 10]
Or maybe we want all even-indexed residues in the first chain (Although this example only has the one chain….).
[10]:
print([residue for residue in topology.chain(0).residues if residue.index % 2 == 0])
[ACE1, NME3]
Atom Selection Language¶
If you’re hesistant about programming filtered lists like the ones above, MDTraj also features a rich atom selection language, similar to that of PyMol and VMD. You can access it by using topology.select
. Let’s find all atoms in the last two residues.
More information about the atom selection syntax is available in the main documentation.
[11]:
print(topology.select("resid 1 to 2"))
[ 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21]
You can also do more complex operations. Here, we’re looking for all nitrogen atoms in the backbone.
[12]:
print(topology.select("name N and backbone"))
[ 6 16]
If you ever want to see the code that generates these results you can use select_expression
, which will yield a string represention of the atom selection code.
[13]:
selection = topology.select_expression("name CA and resid 1 to 2")
print(selection)
[atom.index for atom in topology.atoms if ((atom.name == 'CA') and (1 <= atom.residue.index <= 2))
]