def test_superpose_refinds():
# make one frame far from the origin
normal = np.random.randn(1, 100, 3)
normal_xyz = normal.copy()
flipped = np.zeros_like(normal)
flipped[:, :50, :] = normal[:, 50:, :]
flipped[:, 50:, :] = normal[:, :50, :]
flipped_xyz = flipped.copy()
normal = md.Trajectory(xyz=normal, topology=None)
flipped = md.Trajectory(xyz=flipped, topology=None)
normal.superpose(flipped,
atom_indices=np.arange(0, 50),
ref_atom_indices=np.arange(50, 100))
eq(normal.xyz, normal_xyz)
flipped.superpose(normal,
atom_indices=np.arange(50, 100),
ref_atom_indices=np.arange(0, 50))
eq(flipped.xyz, flipped_xyz)
normal.superpose(flipped)
assert not np.allclose(normal.xyz, normal_xyz)
def test_angle_pbc_0():
# Check that angles are calculated correctly accross periodic boxes
random = np.random.RandomState(0)
boxLengths = [1,2,3]
N = 1000
X = np.cumsum(0.1 * random.randn(N*3, 3), axis=0).reshape(N, 3, 3)
center = X.mean(axis=1).reshape(N, 1, 3)
cell = np.floor(center / boxLengths)
X_cell = X - (cell * boxLengths)
assert not np.all(X == X_cell)
t0 = md.Trajectory(xyz=X, topology=None)
t1 = md.Trajectory(xyz=X_cell, topology=None)
t1.unitcell_vectors = np.tile(np.diag(boxLengths), [len(X), 1, 1])
r0 = md.compute_angles(t0, [[0,1,2]], opt=False).reshape(-1)
r1 = md.compute_angles(t1, [[0,1,2]], opt=False).reshape(-1)
r2 = md.compute_angles(t0, [[0,1,2]], opt=True).reshape(-1)
r3 = md.compute_angles(t1, [[0,1,2]], opt=True).reshape(-1)
np.testing.assert_array_almost_equal(r0, r1, decimal=4)
np.testing.assert_array_almost_equal(r2, r3, decimal=4)
np.testing.assert_array_almost_equal(r0, r3, decimal=4)
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 test_sasa_5():
# Test if you can change atom radii effectively without breaking things
traj_chlorine = md.Trajectory(xyz=np.zeros((1, 1, 3)), topology=topology3)
traj_sulfur = md.Trajectory(xyz=np.zeros((1, 1, 3)), topology=topology4)
probe_radius = 0.14
change_radii = {'Cl': 0.175}
calc_area_chlorine = np.sum(
md.geometry.shrake_rupley(traj_chlorine,
probe_radius=probe_radius,
change_radii=change_radii))
calc_area_sulfur = np.sum(
md.geometry.shrake_rupley(traj_sulfur,
probe_radius=probe_radius,
change_radii=change_radii))
true_area_chlorine = 4 * np.pi * (0.175 + probe_radius)**2
# This atom should remain unchanged
true_area_sulfur = 4 * np.pi * (_ATOMIC_RADII['S'] + probe_radius)**2
# And calc_area_chlorine should be different from this value
prev_area_chlorine = 4 * np.pi * (_ATOMIC_RADII['Cl'] + probe_radius)**2
assert_approx_equal(calc_area_chlorine, true_area_chlorine)
# Also ensure that rest of dict is not gettting deleted
assert_approx_equal(calc_area_sulfur, true_area_sulfur)
# Finally, ensure that we are not changing _ATOMIC_RADII
assert prev_area_chlorine != calc_area_chlorine
def lipic(traj_file, top_file, stride, toAl):
# load trajectory
traj = md.load_xtc(traj_file, top=top_file, stride=stride)
# align if needed
if toAl:
# align to the first frame
traj.superpose(traj[0])
# get average structure
avg_xyz = traj.xyz.astype(np.float64).\
mean(axis = 0, dtype=np.float64)
avg_traj = md.Trajectory([avg_xyz], traj.top)
# second step alignment to the average structure
traj.superpose(avg_traj[0])
# calculate average structure
avg_xyz = traj.xyz.astype(np.float64).\
mean(axis = 0, dtype=np.float64)
avg_traj = md.Trajectory([avg_xyz], traj.top)
# add average structure to the trajectory
traj = traj.join(avg_traj)
return traj
def get_rmsd():
root_dir = '/home/hyang/bio/erf/data/decoys/msm'
trj_dir1 = f'{root_dir}/deshaw/DESRES-Trajectory-ww_1-protein/ww_1-protein/'
trj_dir2 = f'{root_dir}/deshaw/DESRES-Trajectory-ww_2-protein/ww_2-protein/'
df = pd.read_csv(f'{root_dir}/fip35_bead.csv')
coords_native = df[['xcb', 'ycb', 'zcb']].values[None, :, :]
for i, trj_dir in enumerate([trj_dir1, trj_dir2]):
structure = md.load(f'{trj_dir}/ww-protein.pdb')
top = structure.topology
df = pd.read_csv(f'{trj_dir}/ww-protein-beads.csv')
cb_idx = df['beads_cb_index'].values
seq = df['group_name'].values
state_assignment = pd.read_csv(f'{root_dir}/traj_assigment_0',
header=None)[0].values
state_counts = pd.value_counts(state_assignment)
coords_list = []
flist = pd.read_csv(f'{trj_dir}/flist.txt')['fname']
for k, fname in enumerate(flist):
trj = md.load(f'{trj_dir}/{fname}', top=top)
coords_all = trj.xyz * 10
coords_cb_all = coords_all[:, cb_idx, :]
coords_list.append(coords_cb_all)
coords_all = np.vstack(coords_list)
score = pd.read_csv(f'{trj_dir}/energy.csv')['energy'].values
# save low energy frames to PDBs
idx = (score < 1020)
write_pdb_sample(seq, coords_all[idx],
f'{root_dir}/trj{i}_low_energy_sample.pdb')
# save 10 sample for each state
coords_sample = []
t_native = md.Trajectory(xyz=coords_native, topology=None)
for state, counts in zip(state_counts.index, state_counts.values):
if counts > 10:
coords_state = coords_all[(state_assignment == state)]
rand_idx = np.random.randint(0, coords_state.shape[0], 10)
coords_sample.append(coords_state[rand_idx])
# # compute rmsd to native and within states
t = md.Trajectory(xyz=coords_state.copy(), topology=None)
rmsd = md.rmsd(t, t_native, frame=0)
# rmsd = md.rmsd(t, t, frame=0)
print(rmsd.mean(), rmsd.std())
coords_sample = np.vstack(coords_sample)
# print(coords_sample.shape)
write_pdb_sample(seq, coords_sample,
f'{root_dir}/trj{i}_states_sample.pdb')
score = np.append(np.array([0]), score)
coords_list = [coords_native] + coords_list
coords_all = np.vstack(coords_list)
t = md.Trajectory(xyz=coords_all, topology=None)
t = t.superpose(t, frame=0)
rmsd = md.rmsd(t, t, frame=0)
df = pd.DataFrame({'energy': score, 'rmsd': rmsd})
df.to_csv(f'{trj_dir}/energy_rmsd.csv', index=False)
def test_stack_2():
t1 = md.Trajectory(xyz=np.random.rand(10,5,3), topology=None)
t2 = md.Trajectory(xyz=np.random.rand(10,6,3), topology=None)
t3 = t1.stack(t2)
eq(t3.xyz[:, :5], t1.xyz)
eq(t3.xyz[:, 5:], t2.xyz)
eq(t3.n_atoms, 11)
def test_array_vs_matrix():
top = md.load(get_fn('native.pdb')).topology
xyz = np.random.randn(1, 22, 3)
xyz_mat = np.matrix(xyz)
t1 = md.Trajectory(xyz, top)
t2 = md.Trajectory(xyz_mat, top)
eq(t1.xyz, xyz)
eq(t2.xyz, xyz)
def test_superpose_2():
t1 = md.Trajectory(xyz=np.random.randn(1, 100, 3) + 100, topology=None)
t2 = md.Trajectory(xyz=np.random.randn(1, 100, 3) + 100, topology=None)
t2_copy = t2.xyz.copy()
t1.superpose(t2)
t1.superpose(t2, atom_indices=[1, 2, 3, 4, 5, 6, 7])
# make sure that superposing doesn't alter the reference traj
eq(t2.xyz, t2_copy)
def test_minrmsd_metric(self):
# make sure impl is registered
_ = KmeansClustering(n_clusters=5)
# now we can import the impl
impl = deeptime.clustering.metrics['minRMSD']
target = np.random.uniform(size=(1, 3 * 15))
reference = np.random.uniform(size=(1, 3 * 15))
x = mdtraj.rmsd(mdtraj.Trajectory(target.reshape(1, -1, 3), None),
mdtraj.Trajectory(reference.reshape(1, -1, 3), None))
y = impl.compute_metric(target, reference)
np.testing.assert_almost_equal(x[0], y)
def test_join():
xyz = np.random.rand(10,5,3)
# overlapping frames
t1 = md.Trajectory(xyz=xyz[:5], topology=None)
t2 = md.Trajectory(xyz=xyz[4:], topology=None)
t3 = t1.join(t2, discard_overlapping_frames=True)
t4 = t1.join(t2, discard_overlapping_frames=False)
eq(t3.xyz, xyz)
eq(len(t4.xyz), 11)
eq(t4.xyz, np.vstack((xyz[:5], xyz[4:])))
def get_rmsd_native(self, walkers):
num_walkers = len(walkers)
small_lig_idxs = np.array(range(len(self.ligand_idxs)))
small_bs_idxs = np.array(
range(len(self.ligand_idxs),
len(self.ligand_idxs) + len(self.binding_site_idxs)))
keep_atoms = np.concatenate((self.ligand_idxs, self.binding_site_idxs),
axis=0)
small_pos = self._xyz_from_walkers(walkers, keep_atoms)
box_lengths = self._box_from_walkers(walkers)
newpos_small = np.zeros_like(small_pos)
for frame_idx, positions in enumerate(small_pos):
newpos_small[frame_idx, :, :] = recenter_pair(
positions, box_lengths[frame_idx], small_lig_idxs,
small_bs_idxs)
small_top = self.topology.subset(keep_atoms)
traj_rec = mdj.Trajectory(newpos_small, small_top)
traj_rec.superpose(self.comp_traj, atom_indices=small_bs_idxs)
rmsd_native = np.zeros((num_walkers))
for i in range(num_walkers):
rmsd_native[i] = calc_rmsd(traj_rec.xyz[i], self.comp_traj.xyz[0],
small_lig_idxs)
if self.alt_maps is not None:
# figure out the "small" alternative maps
small_alt_maps = deepcopy(self.alt_maps)
for i, a in enumerate(self.alt_maps):
for j, e in enumerate(a):
try:
small_alt_maps[i][j] = list(self.binding_site_idxs).index(e) +\
len(self.ligand_idxs)
except:
raise Exception(
'Alternative maps are assumed to be permutations of'
' existing binding site indices')
for alt_map in small_alt_maps:
alt_traj_rec = mdj.Trajectory(newpos_small, small_top)
alt_traj_rec.superpose(self.comp_traj,
atom_indices=small_bs_idxs,
ref_atom_indices=alt_map)
for i in range(num_walkers):
dist = calc_rmsd(alt_traj_rec.xyz[i],
self.comp_traj.xyz[0], small_lig_idxs)
if dist < rmsd_native[i]:
rmsd_native[i] = dist
return rmsd_native
def write_trajectory(positions,
topology,
outputfile='trajectory.pdb',
center=True,
offline=None):
if offline != None:
traj = md.Trajectory(positions[0::offline], topology)
else:
traj = md.Trajectory(positions, topology)
if center == True:
traj.center_coordinates()
traj.save_pdb(outputfile)
return
def create_extreme(eigvec, avg_str_xyz, max_, min_, avg_str):
extr_max_xyz = [
eigvec[i] * max_ + avg_str_xyz[i] for i in range(len(eigvec))
]
extr_max_xyz = np.array(extr_max_xyz).reshape(len(extr_max_xyz) // 3, 3)
print(extr_max_xyz)
extr_min_xyz = [
eigvec[i] * min_ + avg_str_xyz[i] for i in range(len(eigvec))
]
extr_min_xyz = np.array(extr_min_xyz).reshape(len(extr_min_xyz) // 3, 3)
print(extr_min_xyz)
extr_max_traj = md.Trajectory([extr_max_xyz], avg_str.top)
extr_min_traj = md.Trajectory([extr_min_xyz], avg_str.top)
extr_max_traj.save('extreme_1_max.pdb')
extr_min_traj.save('extreme_1_min.pdb')
def test_superpose_1():
# make one frame far from the origin
reference = md.Trajectory(xyz=np.random.randn(1, 100, 3) + 100, topology=None)
reference_xyz = reference.xyz.copy()
for indices in [None, np.arange(90)]:
# make another trajectory in a similar rotational state
query = md.Trajectory(xyz=reference.xyz + 0.01 * np.random.randn(*reference.xyz.shape), topology=None)
query.superpose(reference, 0, atom_indices=indices)
assert eq(reference.xyz, reference_xyz)
new_centers = np.mean(query.xyz[0], axis=1)
assert 80 < new_centers[0] < 120
assert 80 < new_centers[1] < 120
assert 80 < new_centers[2] < 120
def run(self):
# load and concat trjs
if self.mem_efficient:
trj_lengths, xyzs = load_as_concatenated(
filenames=self.trj_filenames,
processes=self.n_procs,
top=self.base_struct_md,
atom_indices=self.atom_indices_vals)
trjs_sub = md.Trajectory(
xyzs,
self.base_struct_md.atom_slice(
self.atom_indices_vals).topology)
else:
trj_lengths, xyzs = load_as_concatenated(
filenames=self.trj_filenames,
processes=self.n_procs,
top=self.base_struct_md)
trjs = md.Trajectory(xyzs, self.base_struct_md.topology)
trjs_sub = trjs.atom_slice(self.atom_indices_vals)
# determine if rebuilding all msm stuff
if self.build_full:
base_struct_centers = self.base_struct_md.atom_slice(
self.atom_indices_vals)
base_struct_centers.save_pdb("./centers.pdb")
self.base_struct_md.save_pdb("./prot_masses.pdb")
init_centers = None
else:
init_centers = md.load("./data/centers.xtc", top="./centers.pdb")
# fit data with base clustering object
self.base_clust_obj.fit(trjs_sub, init_centers=init_centers)
center_indices, distances, assignments, centers = \
self.base_clust_obj.result_.partition(trj_lengths)
# save data
ra.save("./data/assignments.h5", assignments)
ra.save("./data/distances.h5", distances)
trjs_sub = trjs_sub[self.base_clust_obj.center_indices_]
trjs_sub.superpose(trjs_sub[0])
trjs_sub.save_xtc("./data/centers.xtc")
if not self.mem_efficient:
full_centers = trjs[self.base_clust_obj.center_indices_]
full_centers.superpose(self.base_struct_md)
full_centers.save_xtc("./data/full_centers.xtc")
# save states
n_states = len(self.base_clust_obj.center_indices_)
unique_states = np.arange(n_states)
if init_centers is not None:
unique_states = unique_states[-(n_states - len(init_centers)):]
np.save("./data/unique_states.npy", unique_states)
def setup_module():
global DATADIR, HMM
DATADIR = tempfile.mkdtemp()
# 4 components and 3 features. Each feature is going to be the x, y, z
# coordinate of 1 atom
HMM = sklearn.hmm.GaussianHMM(n_components=4)
HMM.transmat_ = np.array([[0.9, 0.1, 0.0, 0.0], [0.1, 0.7, 0.2, 0.0],
[0.0, 0.1, 0.8, 0.1], [0.0, 0.1, 0.1, 0.8]])
HMM.means_ = np.array([[-10, -10, -10], [-5, -5, -5], [5, 5, 5],
[10, 10, 10]])
HMM.covars_ = np.array([[0.1, 0.1, 0.1], [0.5, 0.5, 0.5], [1, 1, 1],
[4, 4, 4]])
HMM.startprob_ = np.array([1, 1, 1, 1]) / 4.0
# get a 1 atom topology
topology = md.load(get_mdtraj_fn('native.pdb')).restrict_atoms([1
]).topology
# generate the trajectories and save them to disk
for i in range(10):
d, s = HMM.sample(100)
t = md.Trajectory(xyz=d.reshape(len(d), 1, 3), topology=topology)
t.save(os.path.join(DATADIR, 'Trajectory%d.h5' % i))
fetch_alanine_dipeptide()
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 onSaveTrajectory(self, target_filename=None):
if target_filename is None:
target_filename = str(QtWidgets.QFileDialog.getSaveFileName(None, 'Save H5-Model file', '', 'H5-files (*.h5)'))[0]
translation_vector = self.translation_vector
rotation_matrix = self.rotation_matrix
stride = self.stride
if self.verbose:
print("Stride: %s" % stride)
print("\nRotation Matrix")
print(rotation_matrix)
print("\nTranslation vector")
print(translation_vector)
first_frame = md.load_frame(self.trajectory_filename, 0)
traj_new = md.Trajectory(xyz=np.empty((1, first_frame.n_atoms, 3)), topology=first_frame.topology)
traj_new.save(target_filename)
chunk_size = 1000
table = tables.open_file(target_filename, 'a')
for i, chunk in enumerate(md.iterload(self.trajectory_filename, chunk=chunk_size, stride=stride)):
xyz = chunk.xyz.copy()
rotate(xyz, rotation_matrix)
translate(xyz, translation_vector)
table.root.xyz.append(xyz)
table.root.time.append(np.arange(i * chunk_size, i * chunk_size + xyz.shape[0], dtype=np.float32))
table.close()
def binding_site_atoms(mdtraj_topology, ligand_resid, coords):
# selecting ligand and protein binding site atom indices for
# resampler and boundary conditions
lig_idxs = ligand_idxs(mdtraj_topology, ligand_resid)
prot_idxs = protein_idxs(mdtraj_topology)
# select water atom indices
water_atom_idxs = mdtraj_topology.select("water")
#select protein and ligand atom indices
protein_lig_idxs = [
atom.index for atom in mdtraj_topology.atoms
if atom.index not in water_atom_idxs
]
# make a trajectory to compute the neighbors from
traj = mdj.Trajectory([coords], mdtraj_topology)
# selects protein atoms which have less than 8 A from ligand
# atoms in the crystal structure
neighbors_idxs = mdj.compute_neighbors(traj, BINDING_SITE_CUTOFF, lig_idxs)
# selects protein atoms from neighbors list
binding_selection_idxs = np.intersect1d(neighbors_idxs, prot_idxs)
return binding_selection_idxs
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 compute_neighbor_list(coords, neighbor_cutoff, max_num_neighbors,
periodic_box_size):
"""Computes a neighbor list from atom coordinates."""
N = coords.shape[0]
import mdtraj
traj = mdtraj.Trajectory(coords.reshape((1, N, 3)), None)
box_size = None
if periodic_box_size is not None:
box_size = np.array(periodic_box_size)
traj.unitcell_vectors = np.array(
[[[box_size[0], 0, 0], [0, box_size[1], 0], [0, 0, box_size[2]]]],
dtype=np.float32)
neighbors = mdtraj.geometry.compute_neighborlist(traj, neighbor_cutoff)
neighbor_list = {}
for i in range(N):
if max_num_neighbors is not None and len(
neighbors[i]) > max_num_neighbors:
delta = coords[i] - coords.take(neighbors[i], axis=0)
if box_size is not None:
delta -= np.round(delta / box_size) * box_size
dist = np.linalg.norm(delta, axis=1)
sorted_neighbors = list(zip(dist, neighbors[i]))
sorted_neighbors.sort()
neighbor_list[i] = [
sorted_neighbors[j][1] for j in range(max_num_neighbors)
]
else:
neighbor_list[i] = list(neighbors[i])
return neighbor_list
def clean_pdb(name, route=None, chain_num=None):
""" Cleans the structure to only leave the important part.
Inputs:
* name: str. route of the input .pdb file
* route: str. route of the output. will overwrite input if not provided
* chain_num: int. index of chain to select (1-indexed as pdb files)
Output: route of destin file.
"""
destin = route if route is not None else name
# read input
raw_prot = mdtraj.load_pdb(name)
# iterate over prot and select the specified chains
idxs = []
for chain in raw_prot.topology.chains:
# if arg passed, only select that chain
if chain_num is not None:
if chain_num != chain.index:
continue
# select indexes of chain
chain_idxs = raw_prot.topology.select(f"chainid == {str(chain.index)}")
idxs.extend( chain_idxs.tolist() )
# sort: topology and xyz selection are ordered
idxs = sorted(idxs)
# get new trajectory from the sleected subset of indexes and save
prot = mdtraj.Trajectory(xyz=raw_prot.xyz[:, idxs],
topology=raw_prot.topology.subset(idxs))
prot.save(destin)
return destin
def _make_nglview(topology=None,positions=None):
mdtraj_aux_topology = mdtraj.Topology.from_openmm(topology)
traj_aux = mdtraj.Trajectory(positions._value, mdtraj_aux_topology)
view = nglview.show_mdtraj(traj_aux)
view.center()
return view
def get_conf(conf_loc, ref_top, selection, debug):
ref_atoms = [a for a in ref_top.atoms]
# load conf with no H
x_top = md.load(conf_loc).top
x = md.load(conf_loc, atom_indices=x_top.select(selection))
if x.top.n_atoms != ref_top.n_atoms and debug:
print x.top.n_atoms
print ref_top.n_atoms
raise ValueError('x.top.n_atoms != ref_top.n_atoms')
x_atoms = [a for a in x.top.atoms]
# rearrange atom lines to link to reference
new_xyz = np.zeros((1, x.xyz.shape[1], 3))
for a, ref_atom in enumerate(ref_atoms):
for x_index, x_atom in enumerate(x_atoms):
if str(ref_atom) == str(x_atom):
new_xyz[0, a, :] = x.xyz[0, x_index, :]
break
new_x = md.Trajectory(new_xyz, ref_top)
return new_x
def init(self, **kwargs):
"""
Parameters
----------
**kwargs :
Returns
-------
"""
super().init(**kwargs)
# load the json topology as an mdtraj one
mdtraj_top = json_to_mdtraj_topology(self.json_main_rep_top)
# make a traj for the initial state to use as a topology for
# visualizing the walkers
init_traj = mdj.Trajectory(
[self.init_main_rep_positions],
unitcell_lengths=[self.init_unitcell_lengths],
unitcell_angles=[self.init_unitcell_angles],
topology=mdtraj_top)
# write out the init traj as a pdb
logging.info("Writing initial state to {}".format(
self.init_state_path))
init_traj.save_pdb(self.init_state_path)
def to_mdtraj(self, topology=None):
"""
Construct a mdtraj.Trajectory object from the Trajectory itself
Parameters
----------
topology : :class:`mdtraj.Topology`
If not None this topology will be used to construct the mdtraj
objects otherwise the topology object will be taken from the
configurations in the trajectory snapshots.
Returns
-------
:class:`mdtraj.Trajectory`
the trajectory
"""
if topology is None:
topology = self.topology.mdtraj
output = self.xyz
traj = md.Trajectory(output, topology)
traj.unitcell_vectors = self.box_vectors
return traj
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 _calc_min_distance(self, walker):
"""Min-min distance for a walker.
Parameters
----------
walker : object implementing the Walker interface
Returns
-------
min_distance : float
"""
cell_lengths, cell_angles = box_vectors_to_lengths_angles(walker.state['box_vectors'])
t2 = time.time()
# make a traj out of it so we can calculate distances through
# the periodic boundary conditions
walker_traj = mdj.Trajectory(walker.state['positions'],
topology=self._mdj_top,
unitcell_lengths=cell_lengths,
unitcell_angles=cell_angles)
t3 = time.time()
# calculate the distances through periodic boundary conditions
# and get hte minimum distance
min_distance = np.min(mdj.compute_distances(walker_traj,
it.product(self.ligand_idxs,
self.receptor_idxs),
periodic=self._periodic)
)
t4 = time.time()
logging.info("Make a traj: {0}; Calc dists: {1}".format(t3-t2,t4-t3))
return min_distance
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
