def test_rmsd_atom_indices_vs_ref_indices():
n_frames = 1
n_atoms_1 = 1
n_atoms_2 = 2
top_1 = md.Topology()
top_1.add_chain()
top_1.add_residue('RS2', top_1.chain(0))
top_1.add_atom('A2', 'H', top_1.residue(0))
top_2 = md.Topology()
top_2.add_chain()
top_2.add_residue('RS1', top_2.chain(0))
top_2.add_atom('A1', 'H', top_2.residue(0))
top_2.add_chain()
top_2.add_residue('RS2', top_2.chain(1))
top_2.add_atom('A2', 'H', top_2.residue(1))
# here the 2nd chain in the top_2 is rmsd-compatible to the one in the top_1 so we should be able to compute rsmd between them.
trj_1 = md.Trajectory(
np.random.RandomState(0).randn(n_frames, n_atoms_1, 3), top_1)
trj_2 = md.Trajectory(
np.random.RandomState(0).randn(n_frames, n_atoms_2, 3), top_2)
md.rmsd(trj_1, trj_2, atom_indices=[0], ref_atom_indices=[1])
md.rmsd(trj_2, trj_1, atom_indices=[1], ref_atom_indices=[0])
def create_mdtraj_traj(traj, box):
"""Convert Numpy to MDTraj trajectory.
All atoms will be argon atoms.
Keyword arguments:
traj -- Numpy array with shape (n_steps, 3, n_atoms) containing positions in nm
box -- Numpy array with shape (3) in nm
Returns:
mdtraj.Trajectory
"""
# create an argon topology
top = md.Topology()
chain = top.add_chain()
for i in range(traj.shape[2]):
residue = top.add_residue('AR', chain)
top.add_atom('Ar', md.element.argon, residue)
# create trajectory
n_frames = len(traj)
t = md.Trajectory(traj.swapaxes(1, 2),
topology=top,
unitcell_lengths=np.repeat(box[np.newaxis, ...],
n_frames,
axis=0),
unitcell_angles=np.ones((n_frames, 3)) * 90.0)
return t
def create_system(tops):
'''
tops: list
of tuple pairs of (mdtraj topology, #replicates)
Notes:
------
in future, consider creating topology from string/list specification
also can handle of tops are trajectories... then spit out a trajectory
'''
n_entries = len(tops)
sub_topologies = [None] * n_entries
new_topology = mdtraj.Topology()
for ic,chain in enumerate(tops):
tmp_top = chain[0]
if isinstance(chain[0],mdtraj.Trajectory):
tmp_top = chain[0].topology
sub_topologies[ic] = replicate_topology( tmp_top,chain[1] )
new_topology = new_topology.join(sub_topologies[ic])
if all( [isinstance(entry[0],mdtraj.Trajectory) for entry in tops] ):
xyzlist = []
# then also replicate coordinates
for ic,chain in enumerate(tops):
xyz = chain[0].xyz[0]
n_replicates = chain[1]
xyzlist.append( np.tile(xyz,[1,n_replicates,1]) )
xyzs = np.concatenate(xyzlist, axis=1)
new_topology = mdtraj.Trajectory(xyzs,new_topology)
return new_topology
def test_center_of_mass():
top = md.Topology()
chain = top.add_chain()
resi = top.add_residue("NA", chain)
top.add_atom('H1', md.element.hydrogen, resi)
top.add_atom('H2', md.element.hydrogen, resi)
top.add_atom('O', md.element.oxygen, resi)
xyz = np.array([
[
# Frame 1 - Right triangle
[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 0.0]
],
[
# Frame 2
[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.5, 0.5, 0.0]
]
])
traj = md.Trajectory(xyz, top)
com_test = mdtraj.geometry.distance.compute_center_of_mass(traj)
com_actual = (1 / 18.015324) * np.array([
[1.007947, 1.007947, 0.0],
[1.007947 + 0.5 * 15.99943, 1.007947 + 0.5 * 15.99943, 0.0],
])
eq(com_actual, com_test, decimal=4)
def test_drid_1():
n_frames = 1
n_atoms = 20
top = md.Topology()
chain = top.add_chain()
residue = top.add_residue('X', chain)
for i in range(n_atoms):
top.add_atom('X', None, residue)
t = md.Trajectory(xyz=np.random.RandomState(0).randn(n_frames, n_atoms, 3),
topology=top)
# t contains no bonds
got = compute_drid(t).reshape(n_frames, n_atoms, 3)
for i in range(n_atoms):
others = set(range(n_atoms)) - set([i])
rd = 1 / np.array(
[euclidean(t.xyz[0, i], t.xyz[0, e]) for e in others])
mean = np.mean(rd)
second = np.mean((rd - mean)**2)**(0.5)
third = scipy.special.cbrt(np.mean((rd - mean)**3))
ref = np.array([mean, second, third])
np.testing.assert_array_almost_equal(got[0, i], ref, decimal=5)
def test_drid_2():
n_frames = 3
n_atoms = 11
n_bonds = 5
top = md.Topology()
chain = top.add_chain()
residue = top.add_residue('X', chain)
for i in range(n_atoms):
top.add_atom('X', None, residue)
random = np.random.RandomState(0)
bonds = random.randint(n_atoms, size=(n_bonds, 2))
for a, b in bonds:
top.add_bond(top.atom(a), top.atom(b))
t = md.Trajectory(xyz=random.randn(n_frames, n_atoms, 3), topology=top)
got = compute_drid(t).reshape(n_frames, n_atoms, 3)
for i in range(n_frames):
recip = 1 / squareform(pdist(t.xyz[i]))
recip[np.diag_indices(n=recip.shape[0])] = np.nan
recip[bonds[:, 0], bonds[:, 1]] = np.nan
recip[bonds[:, 1], bonds[:, 0]] = np.nan
mean = nanmean(recip, axis=0)
second = nanmean((recip - mean)**2, axis=0)**(0.5)
third = scipy.special.cbrt(nanmean((recip - mean)**3, axis=0))
np.testing.assert_array_almost_equal(got[i, :, 0], mean, decimal=5)
np.testing.assert_array_almost_equal(got[i, :, 1], second, decimal=5)
np.testing.assert_array_almost_equal(got[i, :, 2], third, decimal=5)
def __init__(self,
coords: Union[np.ndarray, List[np.ndarray],
Tuple[np.ndarray]],
types: Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray]],
include_hydrogens: bool = True,
scale: float = 0.1):
self.process = None
# Setup files and mdtraj for each molecule
assert len(coords) == len(types)
self.files, self.tops, self.trajs = [], [], []
for _ in range(len(coords)):
self.files.append(tempfile.NamedTemporaryFile())
self.tops.append(mdtraj.Topology())
self.include_hydrogens = include_hydrogens
# Populate Trajectories
self.coords = [c.squeeze() for c in coords]
self.types = [t.squeeze() for t in types]
for c, typ, f, to in zip(self.coords, self.types, self.files,
self.tops):
self.add_mol_to_topology(c, typ, to)
self.trajs.append(
self._make_trajectory(coords=c, topology=to, scale=scale))
self.save_trajectory(trajectory=self.trajs[-1], filename=f.name)
self.show_vmd()
def check_topologies(map0, map1, override_topology):
"""Check if the topologies of two contact maps are ok or can be fixed"""
# Grab the two topologies
top0 = map0.topology
top1 = map1.topology
# Figure out the overlapping atoms
all_atoms0 = set(map0._all_atoms)
all_atoms1 = set(map1._all_atoms)
# Figure out overlapping residues
all_residues0 = set(map0._all_residues)
all_residues1 = set(map1._all_residues)
# This is intersect (for difference)
overlap_atoms = all_atoms0.intersection(all_atoms1)
overlap_residues = all_residues0.intersection(all_residues1)
top0, top1, topology = _get_default_topologies(top0, top1,
override_topology)
if override_topology:
override_topology = topology
all_atoms_ok = check_atoms_ok(top0, top1, overlap_atoms)
all_res_ok = check_residues_ok(top0, top1, overlap_residues,
override_topology)
if not all_res_ok and not all_atoms_ok:
topology = md.Topology()
return all_atoms_ok, all_res_ok, topology
def json_to_mdtraj_topology(json_string):
""" Copied in part from MDTraj.formats.hdf5 topology property."""
topology_dict = json.loads(json_string)
topology = mdj.Topology()
for chain_dict in sorted(topology_dict['chains'], key=operator.itemgetter('index')):
chain = topology.add_chain()
for residue_dict in sorted(chain_dict['residues'], key=operator.itemgetter('index')):
try:
resSeq = residue_dict["resSeq"]
except KeyError:
resSeq = None
warn('No resSeq information found in HDF file, defaulting to zero-based indices')
try:
segment_id = residue_dict["segmentID"]
except KeyError:
segment_id = ""
residue = topology.add_residue(residue_dict['name'], chain,
resSeq=resSeq, segment_id=segment_id)
for atom_dict in sorted(residue_dict['atoms'], key=operator.itemgetter('index')):
try:
element = elem.get_by_symbol(atom_dict['element'])
except KeyError:
element = elem.virtual
topology.add_atom(atom_dict['name'], element, residue)
atoms = list(topology.atoms)
for index1, index2 in topology_dict['bonds']:
topology.add_bond(atoms[index1], atoms[index2])
return topology
def generate_pdb_mapping(traj, mapping):
'''
take in a cg_site_of_aa mapping array, and generate a new mapping topology for the chain, annotated with cgbead site
Parameters
----------
traj : mdtraj trajectory
mapping : array
should be cg_site_of_aa array format, of a single chain
'''
mapping_top = mdtraj.Topology()
ch_previous = None
sites = []
site_to_res_map = {}
for ia, atom in enumerate(traj.top.atoms):
cg_beadname, cg_identifier = mapping[ia]
if ia == 0:
#if atom.chain != ch_previous: #techincally I envision this only working for single chain first
ch_previous = mapping_top.add_chain()
if atom.residue.chain != traj.top.atom(0).residue.chain:
raise ValueError('Multiple chains detected, not implemented yet')
if cg_identifier not in sites: #need to add a new residue
sites.append(cg_identifier)
res = mapping_top.add_residue(cg_beadname, ch_previous)
site_to_res_map[cg_identifier] = res
res = site_to_res_map[cg_identifier]
new_atom = mapping_top.add_atom(atom.name, atom.element, res)
for bond in traj.top.bonds:
a1, a2 = mapping_top.atom(bond[0].index), mapping_top.atom(
bond[1].index)
mapping_top.add_bond(a1, a2)
mapping_traj = mdtraj.Trajectory(traj.xyz, mapping_top)
return mapping_traj
def do_GRO(self, fileout, frame=None, top=None):
#embed()
if frame is not None:
mesh = self.meshpts[frame]
curr_time = np.array([self.time[frame]])
else:
mesh = self.meshpts[0]
n_atoms = mesh.shape[0]
mesh = mesh[np.newaxis, ...]
if top is None:
top = md.Topology()
c = top.add_chain()
cnt = 0
for i in range(n_atoms):
cnt += 1
r = top.add_residue('II', c)
a = top.add_atom('II', md.element.get_by_symbol('VS'), r, i)
with md.formats.GroTrajectoryFile(fileout, 'w') as f:
# Mesh pts have to be in nm
f.write(mesh / 10, top)
def setup_AVB3(multi_gpu=False):
""" Integrin-avb3-head エネルギーモデルをセットアップする
Returns
-------
top [MDTraj Topology object] : AVB3のTopologyオブジェクト
system [OpenMM System object] : AVB3のSystemオブジェクト
avb3_omm_energy [Energy model] : AVB3のEnergy model
"""
INTEGRATOR_ARGS = (300 * unit.kelvin, 1.0 / unit.picoseconds,
2.0 * unit.femtoseconds)
from simtk.openmm import app
# pdb構造をロードしてpdbオブジェクトを生成
pdb = app.PDBFile(pdb_dir + 'avb3_head.pdb')
# openMM組み込みの力場ファイルをロードしてForceFieldオブジェクトを生成 (implicit solvant[GB-obc]モデル)
forcefield = openmm.app.ForceField('amber99sbildn.xml', 'amber99_obc.xml')
# charmm力場を使う場合は以下をコメントアウトする
# forcefield = openmm.app.ForceField('charmm36.xml')
# pdbオブジェクトとForceFieldオブジェクトを合体し、計算条件を加えたsystemオブジェクトを生成
system = forcefield.createSystem(pdb.topology,
removeCMMotion=False,
nonbondedMethod=app.CutoffNonPeriodic,
nonbondedCutoff=1.0 * unit.nanometers,
constraints=None,
rigidWater=True)
# 運動方程式の積分器を定義
integrator = openmm.LangevinIntegrator(300 * unit.kelvin,
1.0 / unit.picoseconds,
2.0 * unit.femtoseconds)
# pdbファイルとsystemファイルと積分器をまとめて、simulationオブジェクトを生成
simulation = openmm.app.Simulation(pdb.topology, system, integrator)
# openMM APIを使用してエネルギー計算用のモデルを生成
avb3_omm_energy = OpenMMEnergy(openmm_system=system,
openmm_integrator=openmm.LangevinIntegrator,
length_scale=unit.nanometers,
n_atoms=md.Topology().from_openmm(
simulation.topology).n_atoms,
openmm_integrator_args=INTEGRATOR_ARGS,
multi_gpu=multi_gpu)
# MDtrajのopenMM APIを使用して、openMM用トポロジーをMDtraj用トポロジーに変換する
mdtraj_topology = md.Topology().from_openmm(pdb.topology)
return mdtraj_topology, system, avb3_omm_energy
def DoCOMMap(self):
''' Function to perform COM mapping of Loaded Trajectories
based off of residues.
'''
PBar = ProgressBar('Mapping Traj:',
Steps=(int(self.mdTraj.n_residues) - 1),
BarLen=20,
UpdateFreq=1.)
self.Write2Log('\n')
self.Write2Log('Performing Mapping:\n')
# setup topology file for COM trajectory
_top2 = md.Topology()
for rindx, residue in enumerate(self.mdTraj.topology.residues):
_top2.add_chain()
_top2.add_residue(residue.name, _top2.chain(-1))
_top2.add_atom("P", md.element.carbon, _top2.residue(-1))
start = time.time()
# setup np matrix to hold the mapped coordinates
COM_xyz = np.zeros(
(int(self.mdTraj.n_frames), int(self.mdTraj.n_residues), 3))
# loop through each residue and do mapping for ALL frames in trajectory
for rindx, residue in enumerate(self.mdTraj.topology.residues):
PBar.Update(rindx)
resTraj = self.mdTraj.atom_slice(
self.mdTraj.topology.select('resid {}'.format(rindx)))
tempmass = np.column_stack((self.ResidueMasses[residue.name],
self.ResidueMasses[residue.name],
self.ResidueMasses[residue.name]))
rsq = np.multiply(resTraj.xyz, tempmass)
rsq = np.sum(rsq, axis=1) / np.sum(
self.ResidueMasses[residue.name])
if self.WrapTraj: # WrapTrajCoordinates
rsq = np.mod(rsq, self.mdTraj.unitcell_lengths)
COM_xyz[:, rindx, :] = rsq
final = time.time()
totaltime = final - start
self.Write2Log('RunTime: {}\n'.format(totaltime))
PBar.Clear()
trajCOM = md.Trajectory(COM_xyz,
_top2,
unitcell_lengths=self.mdTraj.unitcell_lengths,
unitcell_angles=self.mdTraj.unitcell_angles)
trajCOM.save_lammpstrj(
os.path.join(self.SaveDirName, "TrajCOM.lammpstrj"))
trajCOM.save_dcd(os.path.join(self.SaveDirName, "TrajCOM.dcd"))
_PDB = md.formats.PDBTrajectoryFile(os.path.join(
self.SaveDirName, "TrajCOM.pdb"),
mode='w',
force_overwrite=True,
standard_names=True)
_PDB.write(trajCOM.xyz[0], trajCOM.topology)
def write_pdb(data, filename):
top = mdtraj.Topology()
c = top.add_chain()
for i, pos in enumerate(data):
r = top.add_residue('RHO', c)
a = top.add_atom('CUB', mdtraj.element.get_by_symbol('VS'), r, i)
with mdtraj.formats.PDBTrajectoryFile(filename, 'w') as f:
f.write(data, top)
def setRestraints(self, restrained_sets):
#=['protein and backbone'],
#forces=[100*kilojoules_per_mole/angstroms]):
"""
Set position restraints on atom set(s).
Parameters
----------
restrained_sets : list of str, optional
Selection(s) (MDTraj). The default is ['protein and backbone'].
forces : list of int, optional
The force applied in kilojoules_per_mole/angstroms. The default is 100.
Returns
-------
system : TYPE
DESCRIPTION.
"""
#('protein and name CA and not chainid 1', 'chainid 3 and resid 0')
#trajectory_out_atoms = 'protein or resname SAM or resname ZNB'
topology = md.Topology().from_openmm(self.topology)
equation="(k/2)*periodicdistance(x, y, z, x0, y0, z0)^2"
print('Applying potential: ', equation)
if len(restrained_sets['selections']) == len(restrained_sets['forces']):
for restrained_set, force_value in zip(restrained_sets['selections'], restrained_sets['forces']):
force = omm.CustomExternalForce(equation)
force.addGlobalParameter("k", force_value*kilojoules_per_mole/angstroms**2)
force.addPerParticleParameter("x0")
force.addPerParticleParameter("y0")
force.addPerParticleParameter("z0")
print(f'Restraining set ({force_value*kilojoules_per_mole/angstroms**2}): {len(topology.select(restrained_set))}')
for res_atom_index in topology.select(restrained_set):
force.addParticle(int(res_atom_index), self.positions[int(res_atom_index)].value_in_unit(unit.nanometers))
# TODO: Decide wether to spawn system reps. Streamlined now. Legacy removes later.
self.system.addForce(force)
return self.system
else:
print('Inputs for restrained sets are not the same length.')
def create_system_mapping(element_names, n_beads_TOTAL, t):
"""Create a system mapping
Parameters
----------
element_names : (???)
(???)
n_beads_TOTAL : int
Number of beads in the system
t : mdTraj.Trajectory
Initial trajectory object generated from structure and trajectory files
"""
# SLOWEST PART OF CODE IS THIS FUNCTION
# Initialize atoms with elements
## for loop to traverse element_names array for elements
## need to expand from just carbon to more/different elements
## maybe use elements from periodic package
for atom in t.top.atoms: #possible other function
atom.element = Element.getBySymbol(atom.name) # check element
#need for the xml file to have element symbol as type
# Map the beads accordingly
cg_idx = 0
start_idx = 0
propane_map = {0: [0, 1, 2]} ## mapping definition needs to be created
# from search and user files
## TEST CODE
######################################################################
######################################################################
## TEST CODE
system_mapping = {}
for n in range(n_beads_TOTAL
): # what does sections mean in this particular context
for bead, atoms in propane_map.items():
system_mapping[cg_idx] = [x + start_idx for x in atoms]
start_idx += len(atoms) # understand this part
cg_idx += 1
# Apply mapping for XYZ coordinates
cg_xyz = np.empty((t.n_frames, len(system_mapping), 3))
for cg_bead, aa_indices in system_mapping.items():
cg_xyz[:, cg_bead, :] = md.compute_center_of_mass(
t.atom_slice(aa_indices))
# Apply mapping for Topology object
cg_top = md.Topology()
for cg_bead in system_mapping.keys(): #i got the keys keys keys
cg_top.add_atom('carbon', element.virtual_site,
cg_top.add_residue('A', cg_top.add_chain()))
## Check element and name for items 'A'
## Possible interface with mbuild for better UI and aesthetics
return cg_xyz, cg_top
def _read_topology(self):
if not self._open:
raise ValueError('I/O operation on closed file')
if not self._mode == 'r':
raise ValueError('file not opened for reading')
pdb.PDBTrajectoryFile._loadNameReplacementTables()
n_atoms = None
topology = md.Topology()
chain = topology.add_chain()
residue = None
atomReplacements = {}
for ln, line in enumerate(self._file):
if ln == 1:
n_atoms = int(line.strip())
elif ln > 1 and ln < n_atoms + 2:
(thisresnum, thisresname, thisatomname, thisatomnum) = \
[line[i*5:i*5+5].strip() for i in range(4)]
thisresnum, thisatomnum = map(int, (thisresnum, thisatomnum))
if residue is None or residue.resSeq != thisresnum or residue.name != thisresname:
if residue is not None and residue.name != thisresname:
warnings.warn(
"WARNING: two consecutive residues with same number (%s, %s)"
% (thisresname, residue.name))
if thisresname in pdb.PDBTrajectoryFile._residueNameReplacements:
thisresname = pdb.PDBTrajectoryFile._residueNameReplacements[
thisresname]
residue = topology.add_residue(thisresname,
chain,
resSeq=thisresnum)
if thisresname in pdb.PDBTrajectoryFile._atomNameReplacements:
atomReplacements = pdb.PDBTrajectoryFile._atomNameReplacements[
thisresname]
else:
atomReplacements = {}
thiselem = thisatomname
if len(thiselem) > 1:
thiselem = thiselem[0] + sub('[A-Z0-9]', '', thiselem[1:])
try:
element = elem.get_by_symbol(thiselem)
except KeyError:
element = elem.virtual
if thisatomname in atomReplacements:
thisatomname = atomReplacements[thisatomname]
topology.add_atom(thisatomname,
element=element,
residue=residue,
serial=thisatomnum)
topology.create_standard_bonds()
return n_atoms, topology
def map_single(traj, mapping):
'''
take in a mapping array, and generate a new mapped topology for the chain, retaining bonds!
Parameters
----------
traj : mdtraj trajectory
mapping : array
should be aa_indices_in_cgbead array format, of a single chain
'''
cg_top = mdtraj.Topology()
#=== first add all the species ===
print('creating topology: adding species')
ch = cg_top.add_chain()
n_beads = len(mapping)
xyz = np.zeros([n_beads, 3])
cg_beads = []
for ic, cgbead_entry in enumerate(mapping):
cgbead_name, aa_indices = cgbead_entry
resname = traj.top.atom(
aa_indices[0]
).residue.name #use first atom's residue as the resname. alternative: use most common occurrence
masses = np.array(
[traj.top.atom(ia).element.mass for ia in aa_indices])
com = np.sum(traj.xyz[0, aa_indices, :] * masses[:, None],
0) / masses.sum()
#print( mdtraj.compute_center_of_mass(traj.atom_slice(aa_indices))[0] )
#print( com )
xyz[ic, :] = com
res = cg_top.add_residue(resname, ch)
element = mdtraj.element.virtual #for now... later on can create elements for CG simulation from cg bead name
bead = cg_top.add_atom(cgbead_name, element, res)
cg_beads.append(bead)
#=== figure out connectivity ===
#loop over each cg bead, then over other cg beads, then check bonding
print('creating topology: figuring out bonding')
bondgraph = traj.top.to_bondgraph()
for ii in range(n_beads):
aa_indices1 = mapping[ii][1]
for jj in range(ii + 1, n_beads):
aa_indices2 = mapping[jj][1]
bonded = any([
bondgraph.has_edge(traj.top.atom(a1), traj.top.atom(a2))
for a1 in aa_indices1 for a2 in aa_indices2
])
if bonded:
cg_top.add_bond(cg_beads[ii], cg_beads[jj])
#=== create mapped trajectory and save ===
cg_traj = mdtraj.Trajectory([xyz], cg_top)
return cg_traj
def points_on_cube():
"""Generate example data.
Returns
-------
traj : mdtraj.Trajectory
Trajectory containing points on cube.
Notes
-----
This generates an mdtraj trajectory with 8 frames and 4 atoms. Of the
8 frames, there are only 4 distinct conformations. In particular
traj.xyz[i] == traj.xyz[i + 4] for all i.
If we cluster this data into 4 clusters, we expect there to be 2
conformations assigned to each cluster. We also expect the distance
to each generator to be zero.
"""
n_frames = 8
n_atoms = 4
top = md.Topology()
chain = top.add_chain()
residue = top.add_residue("UNK", chain)
pose0 = np.zeros((n_atoms, 3))
pose1 = np.zeros((n_atoms, 3))
pose1[3, 0] = 1
pose2 = np.zeros((n_atoms, 3))
pose2[2, 1] = 1
pose2[3, 1] = 2
pose3 = np.zeros((n_atoms, 3))
pose0[0, 2] = 0.
pose0[1, 2] = 1.
pose0[2, 2] = 2.
pose0[3, 2] = 3.
for i in range(n_atoms):
atom = top.add_atom("H%d" % i, None, residue)
xyz = np.array([pose0, pose1, pose2, pose3, pose0, pose1, pose2, pose3])
traj = md.Trajectory(xyz, top)
return traj
def __init__(
self,
topology_file: typing.Optional[PathLike] = None,
trajectory_file: typing.Optional[PathLike] = None,
frame_start: int = 0,
frame_end: typing.Optional[int] = None,
):
"""Load a simulation trajectory.
:param topology_file: File containing system topology
:param trajectory_file: File containing simulation trajectory
:param frame_start: First frame of trajectory to use
:param frame_end: Last frame of trajectory to use
"""
if topology_file is not None:
try:
logging.info("Loading topology file")
self._trajectory = load_traj(topology_file)
self._topology = self._trajectory.topology
logging.info("Finished loading topology file")
if trajectory_file is not None:
try:
logging.info(
"Loading trajectory file - this may take a while")
self._trajectory = load_traj(trajectory_file,
top=self._topology)
self._trajectory = self._slice_trajectory(
frame_start, frame_end)
logging.info(
"Finished loading trajectory file - loaded %d frames",
self._trajectory.n_frames,
)
except ValueError as exc:
raise NonMatchingSystemError from exc
self._load_trajectory()
except OSError as exc:
if "no loader" in str(exc) or "format is not supported" in str(
exc):
raise UnsupportedFormatException from exc
raise
else:
# No topology - we're probably building a CG frame
self._topology = mdtraj.Topology()
def xyz_to_mdtraj(xyz, cluster_ids=None):
"""Convert a set of x,y,z coordinates to and mdtraj.Trajectory with a carbon
atom centered at each of the specified coordinates.
Each carbon will be part of a residue called POK. If cluster_ids are
specified, these will be used as the residue numbers ofr sets of carbons
in the same cluster.
Parameters
----------
xyz : np.ndarray, shape=(n_atoms, 3)
Cartesian coordinates to center carbons at.
cluster_ids : np.ndarray, shape=(n_atoms)
If specified, the numbers in this array (one corresponding to each
carbon atom to be created) will become the residue numbers for each
carbon.
Returns
-------
struct : mdtraj.Trajectory
A Trajectory with a single frame containing a carbon at each of the
specified x,y,z coordinates with residue numbers determined bye the
cluster_ids, if specified.
"""
n_xyz = xyz.shape[0]
element = md.element.carbon
top = md.Topology()
chain = top.add_chain()
if cluster_ids is None:
res = top.add_residue("POK", chain, 0)
for i in range(n_xyz):
top.add_atom('C', element, res)
sorted_xyz = xyz
else:
sorted_xyz = np.zeros((n_xyz, 3))
inds_sorted_by_cluster = np.argsort(cluster_ids)
prev_res_ind = -1
for i in range(n_xyz):
cur_res_ind = cluster_ids[inds_sorted_by_cluster[i]]
if cur_res_ind != prev_res_ind:
res = top.add_residue("POK", chain, cur_res_ind)
prev_res_ind = cur_res_ind
top.add_atom('C', element, res)
sorted_xyz[i] = xyz[inds_sorted_by_cluster[i]]
struct = md.Trajectory(sorted_xyz, top)
return struct
def get_state_as_mdtraj(simulation):
"""Construct a length-1 trajectory with unitcells from the current simulation state"""
state = simulation.context.getState(getPositions=True)
xyz = state.getPositions(asNumpy=True).value_in_unit(unit.nanometer)
box_vectors = [
np.array(v.value_in_unit(unit.nanometer))
for v in state.getPeriodicBoxVectors()
]
a_length, b_length, c_length, alpha, beta, gamma = unitcell.box_vectors_to_lengths_and_angles(
*box_vectors)
return md.Trajectory([xyz],
topology=md.Topology().from_openmm(
simulation.topology),
unitcell_lengths=(a_length, b_length, c_length),
unitcell_angles=(alpha, beta, gamma))
def replicate_topology(top,n_replicates):
'''
Try to be more efficient than sequentially adding (which I believe can cost a lot of iterations)
'''
bits = np.binary_repr(n_replicates)
power2_replicates = [top]
for ia in range( len(bits)-1 ):
power2_replicates.append( power2_replicates[-1].join(power2_replicates[-1]) )
new_top = mdtraj.Topology()
for ib,bit in enumerate(bits):
if int(bit) == 1:
new_top = new_top.join(power2_replicates[-(ib+1)])
return new_top
def make_test_topology():
t = mdtraj.Topology()
c = t.add_chain()
r1 = t.add_residue("ALA", c, resSeq=5)
r2 = t.add_residue("HOH", c, resSeq=6)
t.add_atom("CA", mdtraj.element.carbon, r1)
t.add_atom("H", mdtraj.element.hydrogen, r1)
t.add_atom("O", mdtraj.element.oxygen, r2)
t.add_atom("H1", mdtraj.element.hydrogen, r2)
t.add_atom("H2", mdtraj.element.hydrogen, r2)
return t
def _create_com_traj(xyz,
time,
unitcell_lengths,
unitcell_angles,
masses=None):
top = md.Topology()
chain = top.add_chain()
for i, com in enumerate(xyz[0]):
com_res = top.add_residue('COM{}'.format(i), chain)
top.add_atom('COM', get_by_symbol('C'), residue=com_res)
return md.Trajectory(xyz,
topology=top,
time=time,
unitcell_lengths=unitcell_lengths,
unitcell_angles=unitcell_angles)
def create_system_mapping(element_names, n_sections_TOTAL, t):
# Initialize atoms with elements
## for loop to traverse element_names array for elements
## first test - only use for carbon (one element)
## then expand later
## need to allow for different types of elements together
## maybe use element library from msibi repo
from mdtraj.core import element
list(t.top.atoms)[0].element = element.carbon # check element
list(t.top.atoms)[0].element.mass
for atom in t.top.atoms:
atom.element = element.carbon # check element
# Map the beads accordingly
cg_idx = 0
start_idx = 0
propane_map = {0: [0, 1, 2]} ## mapping definition needs to be created
# from search and user files
## TEST CODE
######################################################################
######################################################################
######################################################################
######################################################################
## TEST CODE
system_mapping = {}
for n in range(n_sections_TOTAL):
for bead, atoms in propane_map.items():
system_mapping[cg_idx] = [x + start_idx for x in atoms]
start_idx += len(atoms) # understand this part
cg_idx += 1
# Apply mapping for XYZ coordinates
cg_xyz = np.empty((t.n_frames, len(system_mapping), 3))
for cg_bead, aa_indices in system_mapping.items():
cg_xyz[:, cg_bead, :] = md.compute_center_of_mass(
t.atom_slice(aa_indices))
# Apply mapping for Topology object
cg_top = md.Topology()
for cg_bead in system_mapping.keys(): #i got the keys keys keys
cg_top.add_atom('carbon', element.virtual_site,
cg_top.add_residue('A', cg_top.add_chain()))
## Check element and name for items 'A'
## Possible interface with mbuild for better UI and aesthetics
return cg_xyz, cg_top
def butane_toy_model():
# Starting configuration is close to a cis conformation
phi0 = np.pi
phi_ini = 2*np.pi*0.8
xyz = np.zeros((1, 4, 3), float)
xyz[:,0,:] = np.array([1, -1, 0])
xyz[:,1,:] = np.array([0, -1, 0])
xyz[:,2,:] = np.array([0, 0, 0])
xyz[:,3,0] = np.cos(phi_ini)
xyz[:,3,2] = np.sin(phi_ini)
# Coordinates must be positive in gromacs
shift = np.array([5,5,5])
xyz += shift
# Create a mdtraj Topology for our butane molecule.
newtop = md.Topology()
chain = newtop.add_chain()
for i in range(4):
res = newtop.add_residue("GLY", chain, i)
new_ca = newtop.add_atom('CA', get_by_symbol('C'), res, serial=i)
#if i >= 1:
# prev_ca = chain.atom(i - 1)
# newtop.add_bond(prev_ca, new_ca)
model = mdb.models.Model(newtop, bead_repr="CA")
model.mapping.add_atoms(mass=10)
top = model.mapping.top
# Add bond interactions
model.Hamiltonian._add_bond("HARMONIC_BOND", top.atom(0), top.atom(1), 100, 1)
model.Hamiltonian._add_bond("HARMONIC_BOND", top.atom(1), top.atom(2), 100, 1)
model.Hamiltonian._add_bond("HARMONIC_BOND", top.atom(2), top.atom(3), 100, 1)
# Add angle interactions
model.Hamiltonian._add_angle("HARMONIC_ANGLE", top.atom(0), top.atom(1),
top.atom(2), 20, np.pi/2)
model.Hamiltonian._add_angle("HARMONIC_ANGLE", top.atom(1), top.atom(2),
top.atom(3), 20, np.pi/2)
# Add dihedral interaction
model.Hamiltonian._add_dihedral("COSINE_DIHEDRAL", top.atom(0), top.atom(1),
top.atom(2), top.atom(3), 0.1, phi0, 1)
return xyz, top, model
def gaussian_traj(path):
g = cc.parser.Gaussian(path)
parsed = g.parse()
top = md.Topology()
chain = top.add_chain()
residue = top.add_residue('UNK', chain)
for i, z in enumerate(parsed.atomnos):
elem = md.element.Element.getByAtomicNumber(z)
a = top.add_atom('A{}'.format(i+1), elem, residue, i)
traj = md.Trajectory(parsed.atomcoords/10., top)
traj = traj.center_coordinates()
traj = traj.superpose(traj, frame=0)
output_path = path + '.pdb'
traj.save_pdb(output_path)
print('Written {} to {}...'.format(traj, output_path))
return traj
def _random_trajs():
top = md.Topology()
c = top.add_chain()
r = top.add_residue('HET', c)
for _ in range(101):
top.add_atom('CA', md.element.carbon, r)
traj1 = md.Trajectory(xyz=np.random.uniform(size=(100, 101, 3)),
topology=top,
time=np.arange(100))
traj2 = md.Trajectory(xyz=np.random.uniform(size=(100, 101, 3)),
topology=top,
time=np.arange(100))
ref = md.Trajectory(xyz=np.random.uniform(size=(7, 101, 3)),
topology=top,
time=np.arange(7))
return traj1, traj2, ref
def __init__(self,
filename: Optional[str] = None,
include_hydrogens: bool = True):
if filename is None:
self.file = tempfile.NamedTemporaryFile()
self.filename = self.file.name
else:
self.file = None
self.filename = filename
self.include_hydrogens = include_hydrogens
self.traj_coords = []
self._max_natoms = None
self.top = mdtraj.Topology()
self.chain = self.top.add_chain()
self.process = None
