def _find_interface_lipid(traj, headgroup_indices, return_variance=False):
""" Find the interface based on lipid head groups"""
# Sort into top and bottom leaflet
#midplane = np.mean(traj.xyz[:,headgroup_indices,2])
midplane = np.mean(traj.unitcell_lengths[:, 2]) / 2
bot_leaflet = [
a for a in headgroup_indices
if traj.xyz[0, a,
2] < midplane and abs(traj.xyz[0, a, 2] - midplane) > 1
]
top_leaflet = [
a for a in headgroup_indices
if traj.xyz[0, a,
2] > midplane and abs(traj.xyz[0, a, 2] - midplane) > 1
]
com_bot = mdtraj.compute_center_of_mass(traj.atom_slice(bot_leaflet))
com_top = mdtraj.compute_center_of_mass(traj.atom_slice(top_leaflet))
z_interface_bot = np.mean(com_bot, axis=0)[2]
z_interface_top = np.mean(com_top, axis=0)[2]
if len(bot_leaflet) == 0:
z_interface_bot = None
if len(top_leaflet) == 0:
z_interface_top = None
if return_variance:
bot_variance = np.std(traj.xyz[:, bot_leaflet, 2])
top_variance = np.std(traj.xyz[:, top_leaflet, 2])
z_variance = np.mean([bot_variance, top_variance])
return z_interface_bot, z_interface_top, z_variance
return z_interface_bot, z_interface_top
def check_mdtraj_close_to_boundary(self, traj, milestone_variables,
verbose=False, max_avg=0.03, max_std=0.05):
"""
"""
traj1 = traj.atom_slice(self.group1)
traj2 = traj.atom_slice(self.group2)
com1_array = mdtraj.compute_center_of_mass(traj1)
com2_array = mdtraj.compute_center_of_mass(traj2)
distances = []
for frame_index in range(traj.n_frames):
com1 = com1_array[frame_index,:]
com2 = com2_array[frame_index,:]
radius = np.linalg.norm(com2-com1)
milestone_radius = milestone_variables["radius"]
distances.append(radius - milestone_radius)
avg_distance = np.mean(distances)
std_distance = np.std(distances)
if abs(avg_distance) > max_avg or std_distance > max_std:
if verbose:
warnstr = """The distance between the system and central
milestone were found on average to be {:.4f} nm apart.
The standard deviation was {:.4f} nm.""".format(avg_distance, std_distance)
print(warnstr)
return False
return True
def check_mdtraj_within_boundary(self,
traj,
milestone_variables,
verbose=False):
"""
Check if an mdtraj Trajectory describes a system that remains
within the expected anchor. Return True if passed, return
False if failed.
"""
traj1 = traj.atom_slice(self.group1)
traj2 = traj.atom_slice(self.group2)
com1_array = mdtraj.compute_center_of_mass(traj1)
com2_array = mdtraj.compute_center_of_mass(traj2)
for frame_index in range(traj.n_frames):
com1 = com1_array[frame_index, :]
com2 = com2_array[frame_index, :]
radius = np.linalg.norm(com2 - com1)
milestone_k = milestone_variables["k"]
milestone_radius = milestone_variables["radius"]
if milestone_k * (radius - milestone_radius) > 0.0:
if verbose:
warnstr = """The center of masses of atom group1 and atom
group2 were found to be {:.4f} nm apart.
This distance falls outside of the milestone
boundary at {:.4f} nm.""".format(radius, milestone_radius)
print(warnstr)
return False
return True
def compute_COM_distance(query1, query2, simu_traj):
com1 = md.compute_center_of_mass(
simu_traj.atom_slice(simu_traj.top.select(query1)))
com2 = md.compute_center_of_mass(
simu_traj.atom_slice(simu_traj.top.select(query2)))
#TODO fix so that we handle periodic BC correctly here
return np.linalg.norm(com1 - com2, axis=1)
def get_com_distances(proteinfile, ligandfile, genfile, topology):
lig_indices=loadtxt(ligandfile, ndmin=1, dtype=int)
prot_indices=loadtxt(proteinfile, ndmin=1, dtype=int)
prot_gens=md.load(genfile, top=topology, atom_indices=prot_indices)
lig_gens=md.load(genfile, top=topology, atom_indices=lig_indices)
prot_coms=md.compute_center_of_mass(prot_gens)
lig_coms=md.compute_center_of_mass(lig_gens)
com_distances=[10*get_displacement(i,j) for (i,j) in zip(prot_coms, lig_coms)]
return array(com_distances), lig_gens
def _centerBilayer(self):
print("*"*20)
print("Centering bilayer...")
print("*"*20)
non_water = self._traj.topology.select('not water')
sub_traj = self._traj.atom_slice(non_water)
# Get center of mass of the bilayer
com = mdtraj.compute_center_of_mass(sub_traj)
box = [self._traj.unitcell_lengths[0,0],
self._traj.unitcell_lengths[0,1],
self._traj.unitcell_lengths[0,2]]
# Shift all coordinates so bilayer is at center of box
for i, val in enumerate(self._traj.xyz):
self._traj.xyz[i] = [[x, y, z-com[0,2]+(box[2]/2)] for x,y,z
in self._traj.xyz[i][:]]
# Go back and fix for pbc
for i, val in enumerate(self._traj.xyz):
for j,atom in enumerate(self._traj.xyz[i]):
for k, coord in enumerate(self._traj.xyz[i,j]):
if coord < 0:
self._traj.xyz[i,j,k] += box[k]
elif coord > box[k]:
self._traj.xyz[i,j,k] -= box[k]
self._traj.save(os.path.join(self._baseDir,'centered.gro'))
#self.grofile = 'centered.gro'
self._grofile = (os.path.join(self._baseDir,'centered.gro'))
def results(self):
p_slice = self.trj.atom_slice(self._principal_axis_atoms)
paxis = principal_axis(p_slice)
paxis_com = md.compute_center_of_mass(p_slice)
# Handle the 180 degree flip based on quadrant so there's no uncertainty
quad_dot_paxis = np.dot(paxis, self._principal_axis_quadrant)
paxis[quad_dot_paxis < 0] *= -1
temp_results = []
for ref_atom, helix_atoms in zip(self._reference_atoms, self._helix_atoms):
if len(ref_atom) > 1:
ref_pos = self.trj.xyz[:,ref_atom,:].mean(axis=1)
else:
ref_pos = self.trj.xyz[:,ref_atom[0],:]
helix_com = self.trj.xyz[:, helix_atoms, :].mean(axis=1)
# Principal axis defines a plane for projection
proj_of_com = vector_projection_onto_plane(paxis_com - helix_com, paxis)
proj_of_ref = vector_projection_onto_plane(ref_pos - helix_com, paxis)
a = angle_between_vectors(proj_of_com, proj_of_ref)
# Compute the signed angle
cross = np.cross(proj_of_com, proj_of_ref)
cross_dot_paxis = np.einsum('ij,ij->i', paxis, cross)
a[cross_dot_paxis < 0] *= -1
temp_results.append(a)
return pd.DataFrame(temp_results)
def get_com_ligand(filename):
"""get center of mass from a trajectory"""
print("Trajectory: ", filename)
topol_name = 'topol.prmtop'
pdb_ref = 'start_conf.pdb'
T1 = md.load(filename, top=topol_name)
topol = md.load(pdb_ref)
topology = topol.topology
selection = 'resname GLA'
atom_to_keep = topology.select(selection)
print("Number of frames in trajectory:", T1.n_frames)
T1.restrict_atoms(atom_to_keep)
print("Number of atoms in selection:", T1.n_atoms)
coord = md.compute_center_of_mass(T1)
coord_cog = get_cog(T1)
print("Dimensions of COM coordinates:", coord.shape, "\n")
basename = os.path.basename(filename)
out = basename.split('.')
output_file = out[0] + '_COM.dat'
np.savetxt(output_file, coord, fmt='%3.4f', delimiter='\t') # Save array
output_file = out[0] + '_COG.dat'
np.savetxt(output_file, coord_cog, fmt='%3.4f',
delimiter='\t') # Save array
def main():
traj = md.load_pdb(sys.argv[1])
com = md.compute_center_of_mass(traj)
with open(
os.path.splitext(os.path.basename(sys.argv[1]))[0] + '.CoMz.dat',
'w') as f:
for e in com[:, 2] * 10:
print >> f, '%.3f' % e
def get_current_com(simulation, top, peptide_particles):
peptide_top = top
positions = simulation.context.getState(getPositions=True).getPositions(
asNumpy=True)
new_positions = np.zeros((1, len(peptide_particles), 3))
new_positions[0, :, :] = positions[peptide_particles, :]
#print positions.shape, peptide_particles.shape, np.array([ positions[peptide_particles, :] ])
peptide_traj = md.Trajectory(new_positions, peptide_top)
return md.compute_center_of_mass(peptide_traj)[0]
def calc_hexagonal_order_freud(traj_filename, top_filename, output_filename,
ndx_filename, n_chains):
from freud.box import Box
from freud.order import HexOrderParameter
topology = md.load(top_filename).topology
atoms = np.array(list(topology.atoms))
atom_names = [atom.name for atom in atoms]
groups = read_ndx(ndx_filename)
bottom_chains = np.array(groups['bottom_chains']) - 1
top_chains = np.array(groups['top_chains']) - 1
bottom_chains = np.array_split(bottom_chains, n_chains)
top_chains = np.array_split(top_chains, n_chains)
bottom_termini = np.array(groups['bottom_termini']) - 1
top_termini = np.array(groups['top_termini']) - 1
bottom_termini = np.array_split(bottom_termini, n_chains)
top_termini = np.array_split(top_termini, n_chains)
traj = md.load(traj_filename, top=top_filename)
for i, group in enumerate(
[bottom_chains, top_chains, bottom_termini, top_termini]):
group_COM = []
for chain in group:
chain_slice = traj.atom_slice(chain)
chain_com = md.compute_center_of_mass(chain_slice)
for com in chain_com:
for j in range(2):
if com[j] < 0:
com[j] += traj.unitcell_lengths[0, j]
elif com[j] > traj.unitcell_lengths[0, j]:
com[j] -= traj.unitcell_lengths[0, j]
group_COM.append(chain_com)
group_COM = np.array(group_COM).transpose(1, 0, 2)
rmax = 0.75
box = Box(Lx=traj.unitcell_lengths[0, 0],
Ly=traj.unitcell_lengths[0, 1],
Lz=traj.unitcell_lengths[0, 2])
order_parameter = []
for xyz in group_COM:
median = np.median(xyz[:, 2])
if i == 2:
xyz = np.array(
[point for point in xyz if point[2] > median - 0.5])
elif i == 3:
xyz = np.array(
[point for point in xyz if point[2] < median + 0.5])
hex_order = HexOrderParameter(rmax, 6)
hex_order.compute(box, xyz.astype(np.float32))
order_parameter.append(abs(hex_order.getPsi()).mean())
print(np.mean(order_parameter[int(len(order_parameter) / 2):]))
import pdb
pdb.set_trace()
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 calc_number_density(coord_file, trj_file, bin_width, area, dim, box_range,
data_path):
"""
Return a 1-D number density profiles for each residue across a channel.
"""
first_frame = md.load_frame(trj_file, top=coord_file, index=0)
resnames = np.unique([x.name for x in first_frame.topology.residues])
open('{0}/resnames.txt'.format(data_path), 'w').close()
for resname in resnames:
traj = md.load(trj_file,
top=coord_file,
atom_indices=first_frame.topology.select(
'resname {}'.format(resname)))
indices = [[at.index for at in compound.atoms]
for compound in list(traj.topology.residues)]
if 0 in [
x.mass for x in [atom.element for atom in traj.topology.atoms]
]:
warnings.warn(
"mdtraj found zero mass, setting element to hydrogen",
UserWarning)
for atom in traj.topology.atoms:
if atom.element in [Element.virtual, Element.virtual_site]:
atom.element = Element.hydrogen
com = list()
for i, ids in enumerate(indices):
sub_traj = traj.atom_slice(ids)
com.append(md.compute_center_of_mass(sub_traj))
com = np.array(com)
x = np.histogram(com[:, 1:, dim].reshape((-1, 1)),
bins=np.linspace(
box_range[0],
box_range[1],
num=1 + round(
(box_range[1] - box_range[0]) / bin_width)))
np.savetxt(
'{0}/{1}-number-density.txt'.format(data_path, resname),
np.vstack([
x[1][:-1] + np.mean(x[1][:2]) - box_range[0],
x[0] / (area * bin_width * (len(traj) - 1))
]).transpose())
with open('{0}/resnames.txt'.format(data_path), "a") as myfile:
myfile.write(resname + '\n')
def extractCoordinatesXTCFile(file_name, params, topology, selected_indices):
trajectory = md.load(file_name, top=topology)
if params.com:
# getCOM case
# convert nm to A
coordinates = 10 * md.compute_center_of_mass(
trajectory.atom_slice(selected_indices))
else:
coordinates = 10 * trajectory.xyz[:, selected_indices, :].reshape(
(trajectory.n_frames, -1))
return coordinates
def find_interface_lipid(traj, headgroup_indices):
""" Find the interface based on lipid head groups"""
# Sort into top and bottom leaflet
#midplane = np.mean(traj.xyz[:,headgroup_indices,2])
midplane = np.mean(traj.unitcell_lengths[:, 2]) / 2
bot_leaflet = [
a for a in headgroup_indices
if traj.xyz[0, a,
2] < midplane and abs(traj.xyz[0, a, 2] - midplane) > 1
]
top_leaflet = [
a for a in headgroup_indices
if traj.xyz[0, a,
2] > midplane and abs(traj.xyz[0, a, 2] - midplane) > 1
]
com_bot = mdtraj.compute_center_of_mass(traj.atom_slice(bot_leaflet))
com_top = mdtraj.compute_center_of_mass(traj.atom_slice(top_leaflet))
return com_bot[:, 2], com_top[:, 2]
def lipic(traj_file, top_file, stride):
traj = md.load_xtc(traj_file, top=top_file, stride=stride)
traj = traj.superpose(traj[0])
traj.unitcell_vectors = None
avg_xyz = traj.xyz.astype(np.float64)
avg_xyz = avg_xyz.mean(axis=0, dtype=np.float64)
avg_traj = md.Trajectory([avg_xyz], traj.top)
new_traj = md.Trajectory(
traj[0].xyz.astype(np.float64) -
md.compute_center_of_mass(traj[0]).astype(np.float64), traj[0].top)
for i in range(1, len(traj)):
new_traj = new_traj.join(
md.Trajectory(
traj[i].xyz.astype(np.float64) -
md.compute_center_of_mass(traj[i]).astype(np.float64),
traj[i].top))
new_traj = new_traj.superpose(avg_traj)
new_traj = new_traj.join(avg_traj)
print(new_traj.xyz.shape)
return new_traj
def interfacial_fraction(traj_filename, top_filename, output_filename,
ndx_filename, n_chains):
"""Estimate the fraction of chains where termini are at the interface
Some description of what is returned
Parameters
----------
traj_filename : str
Name of trajectory file
top_filename : str
Name of topology file
output_filename : str
Name of output file
ndx_filename : str
Name of Gromacs .ndx file which specifies atom groups
n_chains : int
Number of chains per monolayer
"""
traj = md.load(traj_filename, top=top_filename)
interfacial_frac_bottom = []
interfacial_frac_top = []
groups = read_ndx(ndx_filename)
for group in ['bottom_termini', 'top_termini']:
termini = np.array(ndx[group]) - 1
termini = np.array_split(termini, n_chains)
xyz = []
for chain in termini:
chain_slice = traj.atom_slice(chain)
chain_com = md.compute_center_of_mass(chain_slice)
xyz.append(chain_com[0])
termini_z = np.array([pos[2] for pos in bottom_xyz])
if group == 'bottom_termini':
cut = np.median(termini_z) - 0.5
termini_xy = np.array([
coords[:2] for i, coords in enumerate(xyz)
if termini_z[i] > cut
])
interfacial_frac_bottom.append(len(termini_xy) / 100)
else:
cut = np.median(termini_z) + 0.5
termini_xy = np.array([
coords[:2] for i, coords in enumerate(xyz)
if termini_z[i] < cut
])
interfacial_frac_top.append(len(termini_xy) / 100)
print('Under construction!')
'''
def calc_com_2Dmsd(traj_filename, top_filename, output_filetag, ndx_filename,
chainlength):
topology = md.load(top_filename).topology
atoms = np.array(list(topology.atoms))
atom_names = [atom.name for atom in atoms]
groups = read_ndx(ndx_filename)
bottom_chains = [
id - 1 for id in groups['bottom_chains'] if atom_names[id - 1] == 'C'
]
n_chains = int(len(bottom_chains) / chainlength)
bottom_chains = [
chain.tolist() for chain in np.array_split(bottom_chains, n_chains)
]
top_chains = [
id - 1 for id in groups['top_chains'] if atom_names[id - 1] == 'C'
]
top_chains = [
chain.tolist() for chain in np.array_split(top_chains, n_chains)
]
traj = md.load(traj_filename, top=top_filename)
locations = []
for i, chain in enumerate(bottom_chains):
chain_slice = traj.atom_slice(chain)
com = md.compute_center_of_mass(chain_slice)
for com_frame in com:
for j in range(2):
if com_frame[j] < 0:
com_frame[j] += traj.unitcell_lengths[0, j]
elif com_frame[j] > traj.unitcell_lengths[0, j]:
com_frame[j] -= traj.unitcell_lengths[0, j]
locations.append(com)
locations = np.array(locations).transpose(1, 0, 2)
com_traj = _create_com_traj(locations, traj.time, traj.unitcell_lengths,
traj.unitcell_angles)
msd = []
for frame in com_traj:
msd_frame = 0
for i, atom in enumerate(com_traj.top.atoms):
disp = frame.xyz[0, i, :2] - com_traj.xyz[0, i, :2]
for j, dim in enumerate(disp):
if dim > com_traj.unitcell_lengths[0, j] / 2:
dim -= traj.unitcell_lengths[0, j]
elif dim < -com_traj.unitcell_lengths[0, j] / 2:
dim += traj.unitcell_lengths[0, j]
msd_frame += np.linalg.norm(dim)**2
msd_frame /= com_traj.top.n_atoms
msd.append(msd_frame)
np.savetxt(output_filetag + '.txt', np.column_stack((traj.time, msd)))
def extractCoordinatesXTCFile(file_name, ligand, atom_Ids, writeCA, topology,
selected_indices, sidechains):
trajectory = md.load(file_name, top=topology)
if not writeCA and (atom_Ids is None
or len(atom_Ids) == 0) and not sidechains:
# getCOM case
# convert nm to A
coordinates = 10 * md.compute_center_of_mass(
trajectory.atom_slice(selected_indices))
else:
coordinates = 10 * trajectory.xyz[:, selected_indices, :].reshape(
(trajectory.n_frames, -1))
return coordinates
def SUMO_ligand_dist(tr):
#coordinates for the Cgamma of SUMO1_F36, SUMO2_F31, or SUMO3_F31:
select_str = '(resname==PHE and (resid==15 or resid==30 or resid==17)) and (name==CG)'
atom_ix = tr.topology.select(select_str)[0]
a = tr.xyz[:, atom_ix]
# ligand all atom coordinatess:
lig = tr.atom_slice(tr.topology.select('chainid==1'))
# ligand center of mass:
b = md.compute_center_of_mass(lig)
# distance between K37/K32_CA and ligand center of mass:
return (((a - b)**2).sum(1))**0.5
def _get_z_variation(sweep, sim, resname=''):
os.chdir(os.path.join(curr_dir, sweep, sim))
xtc_file = glob.glob("Stage4_Eq*.xtc")[0]
gro_file = glob.glob("Stage4_Eq*.gro")[0]
traj = mdtraj.load(xtc_file, top=gro_file)
topol = traj.topology
lipid_dict, headgroup_dict = bilayer_analysis_functions.get_lipids(topol)
lipid_tails, lipid_heads = bilayer_analysis_functions.get_lipid_tails(
topol, lipid_dict)
bot_leaf, top_leaf = bilayer_analysis_functions.identify_leaflets(
traj, topol, lipid_dict)
n_lipid = len(lipid_dict.keys())
n_lipid_tails = len(lipid_tails.keys())
n_tails_per_lipid = n_lipid_tails / n_lipid
# Iterate through the bot leaflet or top leaflet
# Get an atom, its residue, and its residue index
# If the residue index hasn't been tabulated, add the average z coordinate
already_examined_resids = set()
top_avgs = []
top_stds = []
bot_avgs = []
bot_stds = []
for i, leaflet in enumerate((bot_leaf, top_leaf)):
all_avgs = []
all_stds = []
for atomid in leaflet:
resindex = topol.atom(atomid).residue.index
resname_i = topol.atom(atomid).residue.name
if str(resindex) in lipid_heads.keys() and resname in resname_i:
if resindex not in already_examined_resids:
#z_stats = np.array([[np.mean(xyz), np.std(xyz)] for xyz in \
# traj.xyz[:,lipid_heads[str(resindex)],2]])
#pdb.set_trace()
#z_stats = [mdtraj.compute_center_of_mass(frame)[0,2] for frame \
#in traj.atom_slice(lipid_heads[str(resindex)])]
z_stats = mdtraj.compute_center_of_mass(
traj.atom_slice(lipid_heads[str(resindex)]))[:, 2]
already_examined_resids.add(resindex)
all_avgs.append(z_stats)
#all_avgs.append(z_stats[:,0])
#all_stds.append(z_stats[:,1])
if i == 0:
bot_avgs = np.mean(np.asarray(all_avgs), axis=0)
bot_stds = np.std(np.asarray(all_avgs), axis=0)
else:
top_avgs = np.mean(np.asarray(all_avgs), axis=0)
top_stds = np.std(np.asarray(all_avgs), axis=0)
return bot_stds, top_stds
def get_atoms_nearest_centroid(kinase,run_index=0):
t=mdt.load("/home/msultan/research/kinase/drug_binding/2016/drug_binding/%s_models/%d/prot.pdb"%(kinase,run_index))
all_protein_slice = t.atom_slice([i.index for i in t.top.atoms if i.residue.is_protein])
protein_slice = t.atom_slice([i.index for i in t.top.atoms if i.residue.is_protein and i.element.name!="hydrogen"])
drug_slice = t.atom_slice([i.index for i in t.top.atoms if i.residue.name=='LIG'])
prot_com = mdt.compute_center_of_mass(protein_slice)
drug_com = mdt.compute_center_of_mass(drug_slice)
prot_atom_index = np.argmin([np.linalg.norm(t.xyz[:,a.index,:]-prot_com)
for a in t.top.atoms])
#since we need to limit to drug, i slice and then add back
drug_atom_index = np.argmin([np.linalg.norm(t.xyz[:,a.index,:] - drug_com) \
for a in t.top.atoms if a.residue.name=='LIG']) + all_protein_slice.n_atoms
assert(t.top.atom(prot_atom_index).residue.is_protein)
assert(t.top.atom(drug_atom_index).residue.name=="LIG")
distance_between_atoms = mdt.compute_distances(t,[[prot_atom_index,drug_atom_index]])
return prot_atom_index,drug_atom_index,distance_between_atoms
def _map_waters(traj, water_start, frame_index):
""" Worker function to parallelize mapping waters via kmeans
Parameters
----------
traj : mdtraj trajectory
full atomstic trajectory
frame index : int
parallelizing calculation frame by frame
water_start : int
counter denoting which index in the CG coordiantes is water
"""
from sklearn import cluster
start = time.time()
frame = traj[frame_index]
# Get atom indices of all water oxygens
waters = frame.topology.select('resname tip3p and name O1')
# Get coordinates and number of all water oxygens
n_aa_water = len(waters)
aa_water_xyz = frame.atom_slice(waters).xyz[0, :, :]
# Number of CG water molecules based on mapping scheme
water_bead_mapping = 4
n_cg_water = int(n_aa_water / water_bead_mapping)
# Water clusters are a list (n_cg_water) of empty lists
water_clusters = [[] for i in range(n_cg_water)]
# Perform the k-means clustering based on the AA water xyz
k_means = cluster.KMeans(n_clusters=n_cg_water)
k_means.fit(aa_water_xyz)
# Each cluster index says which cluster an atom belongs to
for atom_index, cluster_index in enumerate(k_means.labels_):
# Sort each water atom into the corresponding cluster
# The item being added should be an atom index
water_clusters[cluster_index].append(waters[atom_index])
single_frame_coms = []
# For each cluster, compute enter of mass
for cg_index, water_cluster in enumerate(water_clusters):
com = mdtraj.compute_center_of_mass(frame.atom_slice(water_cluster))
single_frame_coms.append((frame_index, cg_index + water_start, com))
#CG_xyz[frame_index, cg_index + water_start,:] = com
end = time.time()
print("K-means for frame {}: {}".format(frame_index, end - start))
return single_frame_coms
def run(self):
# determine if file already exists
if os.path.exists(self.output_name):
pass
else:
# load centers
centers = md.load("./data/full_centers.xtc",
top="./prot_masses.pdb")
# optionall calculate center of mass between pairs
if self.center_of_mass:
# slice domains
iis_domain0 = self.atom_pairs[:, 0]
iis_domain0 = np.array(iis_domain0[np.where(iis_domain0)],
dtype=int)
iis_domain1 = self.atom_pairs[:, 1]
iis_domain1 = np.array(iis_domain1[np.where(iis_domain1)],
dtype=int)
domain0 = centers.atom_slice(iis_domain0)
domain1 = centers.atom_slice(iis_domain1)
# obtain masses
center_of_mass_domain0 = md.compute_center_of_mass(domain0)
center_of_mass_domain1 = md.compute_center_of_mass(domain1)
# obtain distances
diffs = np.abs(center_of_mass_domain0 - center_of_mass_domain1)
distances = np.sqrt(np.einsum('ij,ij->i', diffs, diffs))[:,
None]
else:
# get distances of atom pairs
distances = md.compute_distances(centers,
atom_pairs=self.atom_pairs)
if self.set_points is not None:
diffs = np.abs(distances - self.set_points)
else:
diffs = np.abs(distances)
norm_dists = np.sum(diffs**self.p_norm, axis=1)**(1 / self.p_norm)
np.save(self.output_name, norm_dists)
def SUMO_ligand_dist(traj):
'''
given a SUMO1-compound trajectory,
compute distance between center of compound and F36CG
'''
#coordinates for the Cgamma of SUMO1_F36, SUMO2/3_F31
atom_ix = traj.topology.select('residue==36 and name==CG')[0]
a = traj.xyz[:, atom_ix]
# ligand all atom coordinatess:
lig = traj.atom_slice(traj.topology.select('chainid==1'))
# ligand center of mass:
b = md.compute_center_of_mass(lig)
return (((a - b) ** 2).sum(1)) ** 0.5
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 protsize(path,name,wd):
if not os.path.isfile(path+"/"+name):
print("File "+path+"/"+name+" does not exist! Please check.")
quit
else:
prot=md.load(path+"/"+name)
minx=np.min(prot.xyz[0][:,0])
miny=np.min(prot.xyz[0][:,1])
minz=np.min(prot.xyz[0][:,2])
maxx=np.max(prot.xyz[0][:,0])
maxy=np.max(prot.xyz[0][:,1])
maxz=np.max(prot.xyz[0][:,2])
protcom=md.compute_center_of_mass(prot)
#process pdb with AMBERTool
os.system("reduce -Trim "+path+"/"+name+" > "+wd+"/prot-reduce.pdb ")
return minx,maxx,miny,maxy,minz,maxz,protcom[0]
def calc_hexagonal_order_freud(traj_filename, top_filename, output_filename,
ndx_filename, chainlength):
from freud.box import Box
from freud.order import HexOrderParameter
topology = md.load(top_filename).topology
atoms = np.array(list(topology.atoms))
atom_names = [atom.name for atom in atoms]
groups = read_ndx(ndx_filename)
bottom_chains = [
id - 1 for id in groups['bottom_chains'] if atom_names[id - 1] == 'C'
]
n_chains = int(len(bottom_chains) / chainlength)
bottom_chains = np.array_split(bottom_chains, n_chains)
top_chains = [
id - 1 for id in groups['top_chains'] if atom_names[id - 1] == 'C'
]
top_chains = np.array_split(top_chains, n_chains)
traj = md.load(traj_filename, top=top_filename)
for i, group in enumerate([bottom_chains, top_chains]):
group_COM = []
for chain in group:
chain_slice = traj.atom_slice(chain)
chain_com = md.compute_center_of_mass(chain_slice)
for com in chain_com:
for j in range(2):
if com[j] < 0:
com[j] += traj.unitcell_lengths[0, j]
elif com[j] > traj.unitcell_lengths[0, j]:
com[j] -= traj.unitcell_lengths[0, j]
group_COM.append(chain_com)
group_COM = np.array(group_COM).transpose(1, 0, 2)
rmax = 0.75
box = Box(Lx=traj.unitcell_lengths[0, 0],
Ly=traj.unitcell_lengths[0, 1],
Lz=traj.unitcell_lengths[0, 2])
order_parameter = []
for xyz in group_COM:
hex_order = HexOrderParameter(rmax, 6)
hex_order.compute(box, xyz.astype(np.float32))
order_parameter.append(abs(hex_order.getPsi()).mean())
print(np.mean(order_parameter[int(len(order_parameter) / 2):]))
def batch_coordinate_loader(trajectory,
surface_parameters,
topology=None,
chunk=500):
"""Generates molecular positions and centre of mass for each frame
Parameters
----------
trajectory: str
Path to trajectory file
surface_parameters: instance of SurfaceParameters
Parameters for intrinsic surface
topology: str, optional
Path to topology file
chunk int, optional
Maximum chunk size for mdtraj batch loading
"""
mol_traj = np.empty((0, surface_parameters.n_mols, 3))
mol_vec = np.empty((0, surface_parameters.n_mols, 3))
com_traj = np.empty((0, 3))
cell_dim = np.zeros((0, 3))
masses = np.repeat(surface_parameters.masses, surface_parameters.n_mols)
for index, traj in enumerate(
md.iterload(trajectory, chunk=chunk, top=topology)):
cell_dim_chunk = traj.unitcell_lengths * 10
com_chunk = md.compute_center_of_mass(traj) * 10
traj = traj.atom_slice(surface_parameters.atom_indices)
mol_chunk = coordinate_arrays(traj,
surface_parameters.atoms,
masses,
mode=surface_parameters.com_mode,
com_sites=surface_parameters.com_sites)
mol_traj = np.concatenate([mol_traj, mol_chunk])
com_traj = np.concatenate([com_traj, com_chunk])
cell_dim = np.concatenate([cell_dim, cell_dim_chunk])
vec_chunk = orientation(traj, surface_parameters.center_atom,
surface_parameters.vector_atoms)
mol_vec = np.concatenate([mol_vec, vec_chunk])
return mol_traj, com_traj, cell_dim, mol_vec
def make_comtrj(trj):
"""Takes a trj and returns a trj with COM positions as atoms"""
comtop = md.Topology()
coords = np.ndarray(shape=(trj.n_frames, trj.n_residues, 3))
for j, res in enumerate(trj.topology.residues):
comtop.add_atom(res.name, virtual_site, comtop.add_residue(res.name, comtop.add_chain()))
res_frame = trj.atom_slice([at.index for at in res.atoms])
coords[:, j, :] = md.compute_center_of_mass(res_frame)
comtrj = md.Trajectory(xyz=coords,
topology=comtop,
time=trj.time,
unitcell_angles=trj.unitcell_angles,
unitcell_lengths=trj.unitcell_lengths)
return comtrj
def setUp(self):
super(Test_COM_utils, self).setUp()
self.pdb_file = test_filenames.top_pdb
self.file_xtc = test_filenames.traj_xtc_stride_20
self.top = md.load(self.pdb_file).top
self.traj = md.load(self.file_xtc, top=self.top)
self.traj_5_frames = self.traj[:5]
# Very slow, but what other way of directly computing the COM is there?
COMS_mdtraj = [
md.compute_center_of_mass(
self.traj_5_frames.atom_slice([aa.index for aa in rr.atoms]))
for rr in self.top.residues
]
# re order along the time axis
self.COMS_mdtraj = _np.zeros((5, self.top.n_residues, 3))
for ii in range(5):
self.COMS_mdtraj[ii, :, :] = [jcom[ii] for jcom in COMS_mdtraj]
def _generateWindows(self):
""" Build simWindows out from the center of mass"""
print("*"*20)
print("Generating windows from center of mass...")
print("*"*20)
non_water = self._traj.topology.select('not water')
sub_traj = self._traj.atom_slice(non_water)
com = mdtraj.compute_center_of_mass(sub_traj)
center_z = round(com[0,2], 2)
self._windows = np.zeros(self._n_windows)
midpoint = int(self._n_windows/2)
self._windows[midpoint] = center_z
# Fill in lower half from centerpoint outward
for i in range(1, midpoint+1):
self._windows[midpoint-i] = center_z - i * self._dz
# Fill in top half from centerpoint outward
for i in range(midpoint+1, self._n_windows):
self._windows[i] = self._windows[0] + i * self._dz
def _getTracerCoordinates(self, tracer):
"""
Parameters
----------
tracer : int
Molecule number
Notes
-----
Gromacs is 1-indexed, while mdtraj is 0-indexed
Getting the coordinates of gmx resid 3 means
getting the coordinates of mdtraj resid 2
"""
tracer_atoms = self._traj.topology.select('resid {}'.format(tracer-1))
sub_traj = self._traj.atom_slice(tracer_atoms)
com = mdtraj.compute_center_of_mass(sub_traj)
return [com[0,0], com[0,1], com[0,2]]