def main():
arg_parser = argparse.ArgumentParser(description='通过给定残基名称，残基内原子数目，两个原子在残基内的索引(从0开始)，计算所有残基内这两个原子之间的直线距离。')
arg_parser.add_argument('resname', action='store', help='残基名称')
arg_parser.add_argument('atoms_num', type=int, action='store', help='残基内原子数目')
arg_parser.add_argument('index1', type=int, action='store', help='第一个原子的索引，索引从0开始')
arg_parser.add_argument('index2', type=int, action='store', help='第二个原子的索引，索引从0开始')
arg_parser.add_argument('topology_file', action='store', help='拓扑文件，例如gro, pdb')
args = arg_parser.parse_args()
resname, atoms_num, index1, index2 = args.resname, args.atoms_num, args.index1, args.index2
universe = Universe(args.topology_file)
atom_groups = universe.selectAtoms("resname " + resname)
if len(atom_groups) % atoms_num != 0:
print("拓扑文件内对应残基原子总数不是所给原子数目的整数倍，请给予正确的原子数目。")
exit(1)
atoms1 = []
atoms2 = []
for i in range(0, len(atom_groups), atoms_num):
atoms1.append(atom_groups[i:i + atoms_num][index1])
atoms2.append(atom_groups[i:i + atoms_num][index2])
dists = dist(AtomGroup(atoms1), AtomGroup(atoms2))
print("The distance between atoms %s and %s is:" % (index1, index2))
for i in dists[2]:
print(i)
print("The average distance between atoms %s and %s is:" % (index1, index2))
print(np.average(dists[2]))
def __init__(self,
otherUniverse,
selections,
names=None,
collapse_hydrogens=False,
lammps_force_dump=None,
residue_reduction_map=None):
if (names is None):
names = ['%dX' % x for x in range(len(selections))]
if (len(names) != len(selections)):
raise ValueError(
'Length of slections (%d) must match lenght of names (%d)' %
(len(selections), len(names)))
self.ref_u = otherUniverse
self.atoms = AtomGroup([])
self.selections = selections
self.names = names
self.chydrogens = collapse_hydrogens
self.universe = self
self.residue_reduction_map = residue_reduction_map
if (lammps_force_dump):
self.lfdump = open(lammps_force_dump, 'r')
else:
self.lfdump = None
self._build_structure()
print('Topology mapping by center of mass, forces by sum')
print('This is %s a periodic trajectory. %d Frames' %
('' if self.trajectory.periodic else 'not',
self.trajectory.numframes))
def select_dihedrals(univer):
CA = univer.selectAtoms('name CA and not resname ACE and not resname NH2')
C = univer.selectAtoms('name C and not resname ACE and not resname NH2')
N = univer.selectAtoms('name N and not resname ACE and not resname NH2')
##########ILLUSTRATION: 2 DIHEDRAL ANGLE FORMD BY 3 RESIDUES############
# dihedral: // or \\
# O
# "
# Ca N C Ca O-H
# / \ //\ / \\ / \ /
# N C Ca N C
# " "
# O O
tets = []
for ca1, c, n, ca2 in zip(CA[:-1], C[:-1], N[1:], CA[1:]):
verify_atom_order(ca1, c, n, ca2)
tet = [ca1, c, n, ca2]
tets.append(AtomGroup(tet))
# TEST CASE: angle around 180
# ca1 = univer.selectAtoms("resid 18 and name CA")[0]
# c = univer.selectAtoms("resid 18 and name C")[0]
# n = univer.selectAtoms("resid 19 and name N")[0]
# ca2 = univer.selectAtoms("resid 19 and name CA")[0]
# tets = [AtomGroup([ca1, c, n, ca2])]
return tets
def search_all(self, radius, level="A"):
"""All neighbor search.
Search all entities that have atoms pairs within
*radius*.
:Arguments:
*radius*
float
*level*
char (A, R, S); what entitity level is returned
(e.g. atoms or residues) is determined by level (A=atoms,
R=residues, S=Segments).
"""
if not level in entity_levels:
raise ValueError("%s: Unknown level" % level)
self.kdt.all_search(radius)
indices = self.kdt.all_get_indices()
atom_list = self.atom_list
atom_pair_list = []
for i1, i2 in indices:
a1 = atom_list[i1]
a2 = atom_list[i2]
atom_pair_list.append(AtomGroup([a1, a2]))
if level == "A":
return atom_pair_list # return atoms as list of AtomGroup pairs
elif level == "R":
return self._get_unique_pairs('residue', atom_pair_list)
elif level == "S":
return self._get_unique_pairs('segment', atom_pair_list)
else:
raise NotImplementedError("level=%s not implemented" % level)
def _index2level(self, indices, level):
"""Convert list of atom_indices in a AtomGroup to either the
Atoms or segments/residues containing these atoms.
Parameters
----------
indices
list of atom indices
level : str
char (A, R, S). Return atoms(A), residues(R) or segments(S) within
*radius* of *atoms*.
"""
n_atom_list = [self.atom_group[i] for i in indices]
if level == 'A':
if len(n_atom_list) == 0:
return []
else:
return AtomGroup(n_atom_list)
elif level == 'R':
return list({a.residue for a in n_atom_list})
elif level == 'S':
return list(set([a.segment for a in n_atom_list]))
else:
raise NotImplementedError(
'{0}: level not implemented'.format(level))
def __init__(self, otherUniverse, selections, names = None, collapse_hydrogens = False, lammps_force_dump = None, residue_reduction_map = None):
#this line of code here should work,
super(Universe, self).__init__()
#but I don't know why it doesn't. So instead I also do this:
self._cache = dict()
self._topology = dict()
self.atoms = AtomGroup([])
if(names is None):
names = [ '%dX' % x for x in range(len(selections))]
if(len(names) != len(selections)):
raise ValueError('Length of slections (%d) must match lenght of names (%d)' % (len(selections), len(names)))
self.fgref_u = otherUniverse
self.atoms = AtomGroup([])
self.selections = selections
self.names = names
self.chydrogens = collapse_hydrogens
self.universe = self
self.residue_reduction_map = residue_reduction_map
if(lammps_force_dump):
self.lfdump = open(lammps_force_dump, 'r')
#find which column is id and which are forces
line = ""
while(not line.startswith('ITEM: ATOMS')):
line = self.lfdump.readline()
self.lfdump_map = {'id':-1, 'force':-1}
for i, item in enumerate(line.split()):
if(item == 'id'):
self.lfdump_map['id'] = i-2
elif(item == 'fx'):
self.lfdump_map['force'] = i-2
if(self.lfdump_map['id'] == -1 or self.lfdump_map['force'] == -1):
raise IOError('The lammps force dump, {}, does not contain a force or id column'.format(lammps_force_dump))
else:
self.lfdump = None
self.lfdump_map = None
self._build_structure()
print('Topology mapping by center of mass, forces by sum')
print('This is %s a periodic trajectory. %d Frames' % ('' if self.trajectory.periodic else 'not', self.trajectory.numframes))
def _apply_map(self, mapping):
"""Apply the mapping scheme to the beads"""
beads = []
for r, m in mapping.items():
for res in self.atu.selectAtoms('resname {}'.format(r)).residues:
beads.extend([AtomGroup(res.atoms[idx]) for idx in m])
return beads
def test_verticalTG(self):
b1 = self.universe.atoms[0].torsions[0]
b2 = self.universe.atoms[20].torsions[0]
TG = TopologyGroup([b1, b2])
forwards = [AtomGroup([b1[i], b2[i]]) for i in range(4)]
backwards = [AtomGroup([b2[i], b1[i]]) for i in range(4)]
verts = [TG.atom1, TG.atom2, TG.atom3, TG.atom4]
# the lists might be in one of two formats, but always in a strict order
# ie any(1234 or 4321) but not (1324)
assert_equal(
any([
all([list(x) == list(y) for x, y in zip(forwards, verts)]),
all([list(x) == list(y) for x, y in zip(backwards, verts)])
]), True)
def __init__(self, otherUniverse, selections = None, names = None, collapse_hydrogens = False, lammps_force_dump = None, residue_reduction_map = None):
#this line of code here should work,
super(Universe, self).__init__()
#but I don't know why it doesn't. So instead I also do this:
self._cache = dict()
self._topology = dict()
self.atoms = AtomGroup([])
self.fgref_u = otherUniverse
self.atoms = AtomGroup([])
if(lammps_force_dump):
self.lfdump = open(lammps_force_dump, 'r')
#find which column is id and which are forces
line = ""
while(not line.startswith('ITEM: ATOMS')):
line = self.lfdump.readline()
self.lfdump_map = {'id':-1, 'force':-1}
for i, item in enumerate(line.split()):
if(item == 'id'):
self.lfdump_map['id'] = i-2
elif(item == 'fx'):
self.lfdump_map['force'] = i-2
if(self.lfdump_map['id'] == -1 or self.lfdump_map['force'] == -1):
raise IOError('The lammps force dump, {}, does not contain a force or id column'.format(lammps_force_dump))
else:
self.lfdump = None
self.lfdump_map = None
if selections is not None:
self.map_to_selection(selections, names, collapse_hydrogens, residue_reduction_map)
self._build_cg_maps()
def __init__(self, *args, **kwargs):
"""Initialise like a normal MDAnalysis Universe but give the mapping keyword.
Mapping must be a dictionary with resnames as keys.
Each resname must then correspond to a list of list of indices,
signifying how to split up a single residue into many beads.
eg:
mapping = {'A':[[0, 1, 2], [3, 4, 5]],
'B':[[0, 1], [2, 3], [4, 5]]}
Would split 'A' residues into 2 beads of 3 atoms,
and 'B' residues into 3 beads of 2 atoms.
Note that the indices inside the mapping refer to the atom's position
inside the residue, not their absolute position in the Universe.
"""
try:
mapping = kwargs.pop('mapping')
except KeyError:
raise ValueError("CGUniverse requires the mapping keyword")
# Atomistic Universe
# TODO: wrap this in try/except and process errors in constructing
# atomistic universe properly.
self.atu = mda.Universe(*args, **kwargs)
# Coarse grained Universe
# Make a blank Universe for myself.
super(CGUniverse, self).__init__()
# Fake up some beads
# TODO: Names, better topology etc
beads = self._apply_map(mapping)
# beads = [AtomGroup(self.atu.atoms[m]) for m in mapping]
beadgroup = [
Bead(b, i, 'BEAD', 'BEAD', b[0].resname, b[0].resid, b[0].segid,
1.0, 1.0) for i, b in enumerate(beads)
]
for b in beadgroup:
b.universe = self
self.atoms = AtomGroup(beadgroup)
self.beads = beadgroup
# TODO: Wrap this in try except to catch errors in making fake Reader
# This replaces load_new in a traditional Universe
self.trajectory = CGTraj(self.atu.trajectory, self.beads)
def _index2level(self, indices, level):
if not level in entity_levels:
raise ValueError("%s: Unknown level" % level)
n_atom_list = [self.atom_list[i] for i in indices]
if level == "A":
try:
return AtomGroup(n_atom_list)
except:
return [
] # empty n_atom_list (AtomGroup throws exception, can't be easily fixed...)
elif level == "R":
residues = set([a.residue for a in n_atom_list])
return list(residues)
elif level == "S":
segments = set([a.segment for a in n_atom_list])
return list(segments)
else:
raise NotImplementedError("level=%s not implemented" % level)
def between(group, A, B, distance):
"""Return sub group of *group* that is within *distance* of both *A* and *B*.
*group*, *A*, and *B* must be
:class:`~MDAnalysis.core.AtomGroup.AtomGroup` instances. Works best
if *group* is bigger than either *A* or *B*. This function is not
aware of periodic boundary conditions.
Can be used to find bridging waters or molecules in an interface.
Similar to "*group* and (AROUND *A* *distance* and AROUND *B* *distance*)".
.. SeeAlso:: Makes use of :mod:`MDAnalysis.lib.NeighborSearch`.
.. versionadded: 0.7.5
"""
from MDAnalysis.core.AtomGroup import AtomGroup
ns_group = AtomNeighborSearch(group)
resA = set(ns_group.search_list(A, distance))
resB = set(ns_group.search_list(B, distance))
return AtomGroup(resB.intersection(resA))
def __init__(self, otherUniverse, selections, names = None, collapse_hydrogens = False, lammps_force_dump = None, residue_reduction_map = None):
if(names is None):
names = [ '%dX' % x for x in range(len(selections))]
if(len(names) != len(selections)):
raise ValueError('Length of slections (%d) must match lenght of names (%d)' % (len(selections), len(names)))
self.ref_u = otherUniverse
self.atoms = AtomGroup([])
self.selections = selections
self.names = names
self.chydrogens = collapse_hydrogens
self.universe = self
self.residue_reduction_map = residue_reduction_map
if(lammps_force_dump):
self.lfdump = open(lammps_force_dump, 'r')
else:
self.lfdump = None
self._build_structure()
print('Topology mapping by center of mass, forces by sum')
print('This is %s a periodic trajectory. %d Frames' % ('' if self.trajectory.periodic else 'not', self.trajectory.numframes))
def between(group, A, B, distance):
"""Return sub group of `group` that is within `distance` of both `A` and `B`
This function is not aware of periodic boundary conditions.
Can be used to find bridging waters or molecules in an interface.
Similar to "*group* and (AROUND *A* *distance* and AROUND *B* *distance*)".
Parameters
----------
group : AtomGroup
Find members of `group` that are between `A` and `B`
A, B : AtomGroups
`A` and `B` are :class:`~MDAnalysis.core.AtomGroup.AtomGroup`
instances. Works best if `group` is bigger than either `A` or
`B`.
distance : float
maximum distance for an atom to be counted as in the vicinity of
`A` or `B`
Returns
-------
AtomGroup
:class:`~MDAnalysis.core.AtomGroup.AtomGroup` of atoms that
fulfill the criterion
.. versionadded: 0.7.5
"""
from MDAnalysis.core.AtomGroup import AtomGroup
ns_group = AtomNeighborSearch(group)
resA = set(ns_group.search_list(A, distance))
resB = set(ns_group.search_list(B, distance))
return AtomGroup(resB.intersection(resA))
[[0, 0, 0], [1, 1, 0], [-1, 1, 0], [-1, -1, 0], [1, -1, 0]],
dtype=numpy.float32)
CNS = CoordinateNeighborSearch(coords)
center = numpy.array([x, y, z])
found_indices = CNS.search(center, R)
# check manually
diff = coords[found_indices] - center[numpy.newaxis, :]
return found_indices, numpy.sqrt(numpy.sum(diff * diff, axis=1))
if __name__ == "__main__":
import numpy
class Atom:
def __init__(self):
self.coord = 100 * numpy.random.random(3)
def coordinates(self):
return self.coord
class AtomGroup(list):
def coordinates(self):
return numpy.array([atom.coordinates() for atom in self])
for i in range(0, 20):
al = AtomGroup([Atom() for i in range(0, 1000)])
ns = AtomNeighborSearch(al)
print "Found ", len(ns.search_all(5.0))
class CGUniverse(Universe):
''' Class which uses center of mass mappings to reduce the number
of degrees of freedom in a given trajectory/structure
file. The selections define how atoms are grouped PER
residue. Thus, if there is only one residue defined, then all
atoms in that selection will be one bead in the CG system.
A residue_reduction_map may be used so that multiple residues
are combined. It should be an array with a length equal to the
total number of desired output residues. Each element should
be an array containing indices that point to the fine-grain
universe residue indices.
'''
def __init__(self,
otherUniverse,
selections,
names=None,
collapse_hydrogens=False,
lammps_force_dump=None,
residue_reduction_map=None):
if (names is None):
names = ['%dX' % x for x in range(len(selections))]
if (len(names) != len(selections)):
raise ValueError(
'Length of slections (%d) must match lenght of names (%d)' %
(len(selections), len(names)))
self.ref_u = otherUniverse
self.atoms = AtomGroup([])
self.selections = selections
self.names = names
self.chydrogens = collapse_hydrogens
self.universe = self
self.residue_reduction_map = residue_reduction_map
if (lammps_force_dump):
self.lfdump = open(lammps_force_dump, 'r')
else:
self.lfdump = None
self._build_structure()
print('Topology mapping by center of mass, forces by sum')
print('This is %s a periodic trajectory. %d Frames' %
('' if self.trajectory.periodic else 'not',
self.trajectory.numframes))
def _build_structure(self):
#atoms cannot be written out of index order, so we
#need to iterate residue by residue
index = 0
reverse_map = {}
residues = {}
#keep track of selections, so we can throw a useful error if we don't end up selecting anything
selection_count = {}
segments = {}
for s in self.selections:
selection_count[s] = 0
#if we're reducing the residues, we'll need to take care of that
ref_residues = self.ref_u.residues
if (self.residue_reduction_map):
#reduce them
ref_residues = []
for i, ri in enumerate(self.residue_reduction_map):
ref_residues.append(
Residue(name='CMB',
id=i + 1,
atoms=reduce(lambda x, y: x + y,
[self.ref_u.residues[j] for j in ri]),
resnum=i + 1))
for r in ref_residues:
residue_atoms = []
for s, n in zip(self.selections, self.names):
group = r.selectAtoms(s)
#check if there were any selected atoms
if (len(group) == 0):
continue
selection_count[s] += len(group)
#make new atom
new_mass = sum([x.mass if x in group else 0 for x in r])
if (sum([1 if x in group else 0 for x in r]) > 0
and new_mass == 0):
raise ValueError(
'Zero mass CG particle found! Please check all-atom masses and/or set them manually via \"fine_grain_universe.selectAtoms(...).set_mass(...)\"'
)
a = Atom(index, n, n, r.name, r.id, r.atoms[0].segid, new_mass,
0)
index += 1
for ra in group:
if (ra in reverse_map):
raise ValueError(
'Attemtping to map {} to {} and {}'.format(
ra, a, reverse_map[ra]))
reverse_map[ra] = a
#append atom to new residue atom group
residue_atoms.append(a)
#add the atom to Universe
self.atoms += a
#now actually create new residue and give atoms a reference to it
residues[r.id] = Residue(r.name,
r.id,
residue_atoms,
resnum=r.resnum)
for a in residue_atoms:
a.residue = residues[r.id]
#take care of putting residue into segment
segid = None if len(residue_atoms) == 0 else residue_atoms[0].segid
if (segid in segments):
segments[segid].append(residues[r.id])
elif (segid):
segments[segid] = [residues[r.id]]
#check to make sure we selected something
total_selected = 0
for s in self.selections:
count = selection_count[s]
total_selected += count
if (count == 0):
raise ValueError('Selection "%s" matched no atoms' % s)
#check counting
if (len(self.ref_u.atoms) < total_selected):
print 'Warining: some atoms placed into more than 1 CG Site'
elif (len(self.ref_u.atoms) > total_selected):
print 'Warning: some atoms not placed into CG site'
#find hydrogens and collapse them into beads
if (self.chydrogens):
for b in self.ref_u.bonds:
#my hack for inferring a hydrogen
for a1, a2 in [(b.atom1, b.atom2), (b.atom2, b.atom1)]:
if (a1.type.startswith('H') and a1.mass < 4.):
reverse_map[a1] = reverse_map[a2]
#add the mass
reverse_map[a2].mass += a1.mass
#generate matrix mappings for center of mass and sum of forces
# A row is a mass normalized cg site defition. or unormalized 1s for forces
self.top_map = npsp.lil_matrix(
(self.atoms.numberOfAtoms(), self.ref_u.atoms.numberOfAtoms()),
dtype=np.float32)
self.force_map = npsp.lil_matrix(
(self.atoms.numberOfAtoms(), self.ref_u.atoms.numberOfAtoms()),
dtype=np.float32)
for a in self.ref_u.atoms:
try:
self.top_map[reverse_map[a].number,
a.number] = a.mass / reverse_map[a].mass
self.force_map[reverse_map[a].number, a.number] = 1.
except KeyError:
#was not selected
pass
#Put them into efficient sparse matrix.
self.top_map = self.top_map.tobsr()
self.force_map = self.force_map.tobsr()
#add bonds using the reverse map
self.bonds = []
for b in self.ref_u.bonds:
try:
cgatom1 = reverse_map[b.atom1]
cgatom2 = reverse_map[b.atom2]
for cbg in self.bonds:
if (not (cbg.atom1 in [cgatom1, cgatom2])
and not (cbg.atom2 in [cgatom1, cgatom2])):
#OK, no bond exists yet
self.bonds.append(Bond(cgatom1, cgatom2))
except KeyError:
#was not in selection
pass
self.__trajectory = CGReader(self, self.ref_u.trajectory, self.top_map,
self.force_map, self.lfdump)
for a in self.atoms:
a.universe = self
#take care of segments now
segment_groups = {}
for k, v in segments.iteritems():
segment_groups[k] = Segment(k, v)
for a in self.atoms:
a.segment = segment_groups[a.segid]
self.atoms._rebuild_caches()
@property
def trajectory(self):
return self.__trajectory
def cache(self, directory='cg_cache'):
'''This precomputes the trajectory and structure so that it doesn't need to be
recalculated at each timestep. Especially useful for random access.
Returns a Universe object corresponding to the cached trajectory
'''
if (not os.path.exists(directory)):
os.mkdir(directory)
structure = os.path.join(directory, 'cg.pdb')
trajectory = os.path.join(directory, 'cg.trr')
write_structure(self, structure, bonds='all')
write_trajectory(self, trajectory)
u = Universe(structure, trajectory)
u.trajectory.periodic = self.trajectory.periodic
apply_mass_map(u, create_mass_map(self))
return u
class CGUniverse(Universe):
''' Class which uses center of mass mappings to reduce the number
of degrees of freedom in a given trajectory/structure
file. The selections define how atoms are grouped PER
residue. Thus, if there is only one residue defined, then all
atoms in that selection will be one bead in the CG system.
A residue_reduction_map may be used so that multiple residues
are combined. It should be an array with a length equal to the
total number of desired output residues. Each element should
be an array containing indices that point to the fine-grain
universe residue indices.
'''
def __init__(self, otherUniverse, selections, names = None, collapse_hydrogens = False, lammps_force_dump = None, residue_reduction_map = None):
if(names is None):
names = [ '%dX' % x for x in range(len(selections))]
if(len(names) != len(selections)):
raise ValueError('Length of slections (%d) must match lenght of names (%d)' % (len(selections), len(names)))
self.ref_u = otherUniverse
self.atoms = AtomGroup([])
self.selections = selections
self.names = names
self.chydrogens = collapse_hydrogens
self.universe = self
self.residue_reduction_map = residue_reduction_map
if(lammps_force_dump):
self.lfdump = open(lammps_force_dump, 'r')
else:
self.lfdump = None
self._build_structure()
print('Topology mapping by center of mass, forces by sum')
print('This is %s a periodic trajectory. %d Frames' % ('' if self.trajectory.periodic else 'not', self.trajectory.numframes))
def _build_structure(self):
#atoms cannot be written out of index order, so we
#need to iterate residue by residue
index = 0
reverse_map = {}
residues = {}
#keep track of selections, so we can throw a useful error if we don't end up selecting anything
selection_count = {}
segments = {}
for s in self.selections:
selection_count[s] = 0
#if we're reducing the residues, we'll need to take care of that
ref_residues = self.ref_u.residues
if(self.residue_reduction_map):
#reduce them
ref_residues = []
for i,ri in enumerate(self.residue_reduction_map):
ref_residues.append(Residue(name='CMB', id=i+1,
atoms=reduce(lambda x,y: x+y, [self.ref_u.residues[j] for j in ri]),
resnum=i+1))
for r in ref_residues:
residue_atoms = []
for s,n in zip(self.selections, self.names):
group = r.selectAtoms(s)
#check if there were any selected atoms
if(len(group) == 0):
continue
selection_count[s] += len(group)
#make new atom
new_mass = sum([x.mass if x in group else 0 for x in r])
if(sum([1 if x in group else 0 for x in r]) > 0 and new_mass == 0):
raise ValueError('Zero mass CG particle found! Please check all-atom masses and/or set them manually via \"fine_grain_universe.selectAtoms(...).set_mass(...)\"')
a = Atom(index, n, n, r.name, r.id, r.atoms[0].segid, new_mass, 0)
index += 1
for ra in group:
if(ra in reverse_map):
raise ValueError('Attemtping to map {} to {} and {}'.format(ra, a, reverse_map[ra]))
reverse_map[ra] = a
#append atom to new residue atom group
residue_atoms.append(a)
#add the atom to Universe
self.atoms += a
#now actually create new residue and give atoms a reference to it
residues[r.id] = Residue(r.name, r.id, residue_atoms, resnum=r.resnum)
for a in residue_atoms:
a.residue = residues[r.id]
#take care of putting residue into segment
segid = None if len(residue_atoms) == 0 else residue_atoms[0].segid
if(segid in segments):
segments[segid].append(residues[r.id])
elif(segid):
segments[segid] = [residues[r.id]]
#check to make sure we selected something
total_selected = 0
for s in self.selections:
count = selection_count[s]
total_selected += count
if(count == 0):
raise ValueError('Selection "%s" matched no atoms' % s)
#check counting
if(len(self.ref_u.atoms) < total_selected):
print 'Warining: some atoms placed into more than 1 CG Site'
elif(len(self.ref_u.atoms) > total_selected):
print 'Warning: some atoms not placed into CG site'
#find hydrogens and collapse them into beads
if(self.chydrogens):
for b in self.ref_u.bonds:
#my hack for inferring a hydrogen
for a1,a2 in [(b.atom1, b.atom2), (b.atom2, b.atom1)]:
if(a1.type.startswith('H') and a1.mass < 4.):
reverse_map[a1] = reverse_map[a2]
#add the mass
reverse_map[a2].mass += a1.mass
#generate matrix mappings for center of mass and sum of forces
# A row is a mass normalized cg site defition. or unormalized 1s for forces
self.top_map = npsp.lil_matrix( (self.atoms.numberOfAtoms(), self.ref_u.atoms.numberOfAtoms()) , dtype=np.float32)
self.force_map = npsp.lil_matrix( (self.atoms.numberOfAtoms(), self.ref_u.atoms.numberOfAtoms()) , dtype=np.float32)
for a in self.ref_u.atoms:
try:
self.top_map[reverse_map[a].number, a.number] = a.mass / reverse_map[a].mass
self.force_map[reverse_map[a].number, a.number] = 1.
except KeyError:
#was not selected
pass
#Put them into efficient sparse matrix.
self.top_map = self.top_map.tobsr()
self.force_map = self.force_map.tobsr()
#add bonds using the reverse map
self.bonds = []
for b in self.ref_u.bonds:
try:
cgatom1 = reverse_map[b.atom1]
cgatom2 = reverse_map[b.atom2]
for cbg in self.bonds:
if(not (cbg.atom1 in [cgatom1, cgatom2]) and not( cbg.atom2 in [cgatom1, cgatom2])):
#OK, no bond exists yet
self.bonds.append( Bond(cgatom1, cgatom2) )
except KeyError:
#was not in selection
pass
self.__trajectory = CGReader(self, self.ref_u.trajectory, self.top_map, self.force_map, self.lfdump)
for a in self.atoms:
a.universe = self
#take care of segments now
segment_groups = {}
for k,v in segments.iteritems():
segment_groups[k] = Segment(k, v)
for a in self.atoms:
a.segment = segment_groups[a.segid]
self.atoms._rebuild_caches()
@property
def trajectory(self):
return self.__trajectory
def cache(self, directory='cg_cache'):
'''This precomputes the trajectory and structure so that it doesn't need to be
recalculated at each timestep. Especially useful for random access.
Returns a Universe object corresponding to the cached trajectory
'''
if(not os.path.exists(directory)):
os.mkdir(directory)
structure = os.path.join(directory, 'cg.pdb')
trajectory = os.path.join(directory, 'cg.trr')
write_structure(self, structure, bonds='all')
write_trajectory(self, trajectory)
u = Universe(structure, trajectory)
u.trajectory.periodic = self.trajectory.periodic
apply_mass_map(u, create_mass_map(self))
return u
def __init__(self, atoms, *args, **kwargs):
"""atoms should be an AtomGroup representing the bead"""
self.atoms = AtomGroup(atoms)
super(Bead, self).__init__(*args, **kwargs)
def run(self, **kwargs):
"""Analyze trajectory and produce timeseries.
Stores the hydrogen bond data per frame
as :attr:`HydrogenBondAnalysis.timeseries` (see there for
output format).
.. SeeAlso:: :meth:`HydrogenBondAnalysis.generate_table` for processing
the data into a different format.
.. versionchanged:: 0.7.6
Results are not returned, only stored in
:attr:`~HydrogenBondAnalysis.timeseries` and duplicate hydrogen bonds
are removed from output (can be suppressed with *remove_duplicates* =
``False``)
"""
logger.info("HBond analysis: starting")
logger.debug("HBond analysis: donors %r", self.donors)
logger.debug("HBond analysis: acceptors %r", self.acceptors)
remove_duplicates = kwargs.pop('remove_duplicates',
True) # False: old behaviour
if not remove_duplicates:
logger.warn(
"Hidden feature remove_duplicates = True activated: you will probably get duplicate H-bonds."
)
verbose = kwargs.pop('verbose', False)
if verbose != self.verbose:
self.verbose = verbose
logger.debug("Toggling verbose to %r", self.verbose)
if not self.verbose:
logger.debug(
"HBond analysis: For full step-by-step debugging output use verbose=True"
)
self.timeseries = []
self.timesteps = []
logger.info("checking trajectory...") # numframes can take a while!
try:
frames = numpy.arange(self.u.trajectory.numframes)[self.traj_slice]
except:
logger.error(
"Problem reading trajectory or trajectory slice incompatible.")
logger.exception()
raise
pm = ProgressMeter(
len(frames),
format="HBonds frame %(step)5d/%(numsteps)d [%(percentage)5.1f%%]\r"
)
try:
self.u.trajectory.time
def _get_timestep():
return self.u.trajectory.time
logger.debug("HBond analysis is recording time step")
except NotImplementedError:
# chained reader or xyz(?) cannot do time yet
def _get_timestep():
return self.u.trajectory.frame
logger.warn(
"HBond analysis is recording frame number instead of time step"
)
logger.info(
"Starting analysis (frame index start=%d stop=%d, step=%d)",
(self.traj_slice.start or 0),
(self.traj_slice.stop or self.u.trajectory.numframes),
self.traj_slice.step or 1)
for ts in self.u.trajectory[self.traj_slice]:
# all bonds for this timestep
frame_results = []
# dict of tuples (atomid, atomid) for quick check if
# we already have the bond (to avoid duplicates)
already_found = {}
frame = ts.frame
timestep = _get_timestep()
self.timesteps.append(timestep)
pm.echo(ts.frame)
self.logger_debug(
"Analyzing frame %(frame)d, timestep %(timestep)f ps", vars())
if self.update_selection1:
self._update_selection_1()
if self.update_selection2:
self._update_selection_2()
if self.selection1_type in ('donor', 'both'):
self.logger_debug("Selection 1 Donors <-> Acceptors")
ns_acceptors = NS.AtomNeighborSearch(self._s2_acceptors)
for i, donor_h_set in self._s1_donors_h.items():
d = self._s1_donors[i]
for h in donor_h_set:
res = ns_acceptors.search_list(AtomGroup([h]),
self.distance)
for a in res:
angle = self.calc_angle(d, h, a)
if self.distance_type in ('hydrogen'):
dist = self.calc_eucl_distance(h, a)
if self.distance_type in ('heavy'):
dist = self.calc_eucl_distance(d, a)
#print "{0}\n{1}\n{2}\n{3}\n{4}\n{5}\n".format(d,h,a,dist,dist2,angle)
if angle >= self.angle and dist <= self.distance:
self.logger_debug(
"S1-D: %s <-> S2-A: %s %f A, %f DEG" %
(h.number + 1, a.number + 1, dist, angle))
#self.logger_debug("S1-D: %r <-> S2-A: %r %f A, %f DEG" % (h, a, dist, angle))
frame_results.append([
h.number + 1, a.number + 1,
'%s%s:%s' %
(h.resname, repr(h.resid), h.name),
'%s%s:%s' %
(a.resname, repr(a.resid), a.name), dist,
angle
])
already_found[(h.number + 1,
a.number + 1)] = True
if self.selection1_type in ('acceptor', 'both'):
self.logger_debug("Selection 1 Acceptors <-> Donors")
ns_acceptors = NS.AtomNeighborSearch(self._s1_acceptors)
for i, donor_h_set in self._s2_donors_h.items():
d = self._s2_donors[i]
for h in donor_h_set:
res = ns_acceptors.search_list(AtomGroup([h]),
self.distance)
for a in res:
if remove_duplicates and \
((h.number+1, a.number+1) in already_found or \
(a.number+1, h.number+1) in already_found):
continue
angle = self.calc_angle(d, h, a)
if self.distance_type in ('hydrogen'):
dist = self.calc_eucl_distance(h, a)
if self.distance_type in ('heavy'):
dist = self.calc_eucl_distance(d, a)
if angle >= self.angle and dist <= self.distance:
self.logger_debug(
"S1-A: %s <-> S2-D: %s %f A, %f DEG" %
(a.number + 1, h.number + 1, dist, angle))
#self.logger_debug("S1-A: %r <-> S2-D: %r %f A, %f DEG" % (a, h, dist, angle))
frame_results.append([
h.number + 1, a.number + 1,
'%s%s:%s' %
(h.resname, repr(h.resid), h.name),
'%s%s:%s' %
(a.resname, repr(a.resid), a.name), dist,
angle
])
self.timeseries.append(frame_results)
logger.info(
"HBond analysis: complete; timeseries with %d hbonds in %s.timeseries",
self.count_by_time().count.sum(), self.__class__.__name__)
class CGUniverse(Universe):
''' Class which uses center of mass mappings to reduce the number
of degrees of freedom in a given trajectory/structure
file. The selections define how atoms are grouped PER
residue. Thus, if there is only one residue defined, then all
atoms in that selection will be one bead in the CG system.
A residue_reduction_map may be used so that multiple residues
are combined. It should be an array with a length equal to the
total number of desired output residues. Each element should
be an array containing indices that point to the fine-grain
universe residue indices.
'''
def __init__(self, otherUniverse, selections, names = None, collapse_hydrogens = False, lammps_force_dump = None, residue_reduction_map = None):
#this line of code here should work,
super(Universe, self).__init__()
#but I don't know why it doesn't. So instead I also do this:
self._cache = dict()
self._topology = dict()
self.atoms = AtomGroup([])
if(names is None):
names = [ '%dX' % x for x in range(len(selections))]
if(len(names) != len(selections)):
raise ValueError('Length of slections (%d) must match lenght of names (%d)' % (len(selections), len(names)))
self.fgref_u = otherUniverse
self.atoms = AtomGroup([])
self.selections = selections
self.names = names
self.chydrogens = collapse_hydrogens
self.universe = self
self.residue_reduction_map = residue_reduction_map
if(lammps_force_dump):
self.lfdump = open(lammps_force_dump, 'r')
#find which column is id and which are forces
line = ""
while(not line.startswith('ITEM: ATOMS')):
line = self.lfdump.readline()
self.lfdump_map = {'id':-1, 'force':-1}
for i, item in enumerate(line.split()):
if(item == 'id'):
self.lfdump_map['id'] = i-2
elif(item == 'fx'):
self.lfdump_map['force'] = i-2
if(self.lfdump_map['id'] == -1 or self.lfdump_map['force'] == -1):
raise IOError('The lammps force dump, {}, does not contain a force or id column'.format(lammps_force_dump))
else:
self.lfdump = None
self.lfdump_map = None
self._build_structure()
print('Topology mapping by center of mass, forces by sum')
print('This is %s a periodic trajectory. %d Frames' % ('' if self.trajectory.periodic else 'not', self.trajectory.numframes))
def _build_structure(self):
#atoms cannot be written out of index order, so we
#need to iterate residue by residue
index = 0
fgtocg_incl_map = {} #keys = fine atoms, values = coarse beads
residues = {}
#keep track of selections, so we can throw a useful error if we don't end up selecting anything
selection_count = {}
segments = {}
for s in self.selections:
selection_count[s] = 0
#if we're reducing the residues, we'll need to take care of that
ref_residues = self.fgref_u.residues
if(self.residue_reduction_map):
#reduce them
ref_residues = []
for i,ri in enumerate(self.residue_reduction_map):
ref_residues.append(Residue(name='CMB', id=i+1,
atoms=reduce(lambda x,y: x+y, [self.fgref_u.residues[j] for j in ri]),
resnum=i+1))
total_masses = []
for r in ref_residues:
residue_atoms = []
for s,n in zip(self.selections, self.names):
group = r.selectAtoms(s)
#check if there were any selected atoms
if(len(group) == 0):
continue
selection_count[s] += len(group)
#calculate the total mass lumped into the new CG atom
total_mass = sum([x.mass if x in group else 0 for x in r])
if(sum([1 if x in group else 0 for x in r]) > 0 and total_mass == 0):
raise ValueError('Zero mass CG particle found! Please check all-atom masses and/or set them manually via \"fine_grain_universe.selectAtoms(...).set_mass(...)\"')
#make new atom, using total mass as mass
#masses are corrected after hydrogens are collapsed in
#and the fg to cg maps are defined
a = Atom(index, n, n, r.name, r.id, r.atoms[0].segid, total_mass, 0)
index += 1
for ra in group:
if(ra in fgtocg_incl_map):
raise ValueError('Attemtping to map {} to {} and {}'.format(ra, a, fgtocg_incl_map[ra]))
fgtocg_incl_map[ra] = a
#append atom to new residue atom group
residue_atoms.append(a)
#add the atom to Universe
self.atoms += a
#now actually create new residue and give atoms a reference to it
residues[r.id] = Residue(r.name, r.id, residue_atoms, resnum=r.resnum)
for a in residue_atoms:
a.residue = residues[r.id]
#take care of putting residue into segment
segid = None if len(residue_atoms) == 0 else residue_atoms[0].segid
if(segid in segments):
segments[segid].append(residues[r.id])
elif(segid):
segments[segid] = [residues[r.id]]
#check to make sure we selected something
total_selected = 0
for s in self.selections:
count = selection_count[s]
total_selected += count
if(count == 0):
raise ValueError('Selection "%s" matched no atoms' % s)
#check counting
if(len(self.fgref_u.atoms) < total_selected):
print 'Warining: some atoms placed into more than 1 CG Site'
elif(len(self.fgref_u.atoms) > total_selected):
print 'Warning: some atoms not placed into CG site'
#find hydrogens and collapse them into beads
if(self.chydrogens):
for b in self.fgref_u.bonds:
#my hack for inferring a hydrogen
for a1,a2 in [(b[0], b[1]), (b[1], b[0])]:
if(a1.type.startswith('H') and a1.mass < 4.):
fgtocg_incl_map[a1] = fgtocg_incl_map[a2]
#add the mass
fgtocg_incl_map[a2].mass += a1.mass
#generate matrix mappings for center of mass and sum of forces
# A row is a mass normalized cg site defition. or unormalized 1s for forces
self.pos_map = npsp.lil_matrix( (self.atoms.numberOfAtoms(), self.fgref_u.atoms.numberOfAtoms()) , dtype=np.float32)
self.force_map = npsp.lil_matrix( (self.atoms.numberOfAtoms(), self.fgref_u.atoms.numberOfAtoms()) , dtype=np.float32)
#keep a forward map for use in state-masks
#The key is a CG atom number and the value is an atom group of the fine-grain atoms
self.cgtofg_fiber_map = [[] for x in range(self.atoms.numberOfAtoms())]
for a in self.fgref_u.atoms:
try:
self.pos_map[fgtocg_incl_map[a].number, a.number] = a.mass / fgtocg_incl_map[a].mass
self.force_map[fgtocg_incl_map[a].number, a.number] = 1.
self.cgtofg_fiber_map[fgtocg_incl_map[a].number] += [a.number]
except KeyError:
#was not selected
pass
#set the masses correctly now that the total is used and
#a cg to fg map is available
#see Noid et al 2008 (MS-CG I), M_I = (\sum_i c_Ii^2 / m_i)^-1 for nonzero c_Ii
#c_Ii = fg atom mass / cg atom total mass for all fg atoms included
#in the definition of cg atom I
for cg_idx in range(self.atoms.numberOfAtoms()):
new_mass = 1.0 / sum([(self.fgref_u.atoms[fg_idx].mass / self.atoms[cg_idx].mass ** 2) for fg_idx in self.cgtofg_fiber_map[cg_idx]])
self.atoms[cg_idx].mass = new_mass
for i in range(self.atoms.numberOfAtoms()):
#convert the list of atom indices into an AtomGroup object
#it has to be wraped so we can manipulate it later
self.cgtofg_fiber_map[i] = AtomGroupWrapper(reduce(lambda x,y: x + y, [self.fgref_u.atoms[x] for x in self.cgtofg_fiber_map[i]]))
#Put them into efficient sparse matrix.
self.pos_map = self.pos_map.tobsr()
self.force_map = self.force_map.tobsr()
#add bonds using the reverse map
#There is new syntax in 0.92 that uses topolgy groups instead of lists.
#delete using attribute
del self.bonds
for b in self.fgref_u.bonds:
try:
cgatom1 = fgtocg_incl_map[b[0]]
cgatom2 = fgtocg_incl_map[b[1]]
for cbg in self.bonds:
if(not (cbg[0] in [cgatom1, cgatom2]) and not( cbg[1] in [cgatom1, cgatom2])):
#OK, no bond exists yet
self.bonds += Bond( (cgatom1, cgatom2) )
except KeyError:
#was not in selection
pass
self.__trajectory = CGReader(self, self.fgref_u.trajectory, self.pos_map, self.force_map, self.lfdump, self.lfdump_map)
for a in self.atoms:
a.universe = self
#take care of segments now
segment_groups = {}
for k,v in segments.iteritems():
segment_groups[k] = Segment(k, v)
for a in self.atoms:
a.segment = segment_groups[a.segid]
self.atoms._rebuild_caches()
@property
def trajectory(self):
return self.__trajectory
def cache(self, directory='cg_cache', maximum_frames=0):
'''This precomputes the trajectory and structure so that it doesn't need to be
recalculated at each timestep. Especially useful for random access.
Returns a Universe object corresponding to the cached trajectory
'''
write_files = True
if(mpi_support):
#check if we're running with MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
write_files = rank == 0
if(not os.path.exists(directory) and write_files):
os.mkdir(directory)
structure = os.path.join(directory, 'cg.pdb')
trajectory = os.path.join(directory, 'cg.trr')
if(write_files):
write_structure(self, structure, bonds='all')
write_trajectory(self, trajectory, maximum_frames)
#sync area
if(mpi_support):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
#synchronize by using a broadcast
comm.bcast(0)
u = Universe(structure, trajectory)
u.trajectory.periodic = self.trajectory.periodic
apply_mass_map(u, create_mass_map(self))
return u
def make_state_mask(self, state_function, state_number):
"""Creates a State_Mask object. The sate_function should take
an MDAnalysis atom group from the fine-grain universe that
made up the CG site. The function should return a vector of
membership percentages for each state. The number of states
is state_number. This returns an array of masks for each
possible state"""
return [State_Mask(self.cgtofg_fiber_map, state_function, x) for x in range(state_number)]
