def run(self, selecttext, windowsize):
if self.input_type == None:
raise LoadError(1)
# setup a temporary dataset
tmpDS = self._custom_dataset()
tmpDS.copy_metadata(self.primaryDS)
# create feature descriptors
self.descriptor_list = []
atoms = self.u.select_atoms(selecttext)
for atom in atoms:
tmpDS.add_feature(
('%s_%s_%s' % (atom.segid, atom.resid, atom.name),
'rmsf_over_%d_frame_window' % windowsize,
'segid %s and resid %s and name %s' %
(atom.segid, atom.resid, atom.name)),
width=3)
position_data = AnalysisFromFunction(lambda x: x.positions,
atoms).run().results
# hold the coordinates in the temporary dataset
tmpDS.format_data(position_data.reshape(position_data.shape[0], -1).T)
# use array convolution to get rolling averages - set up convolution array
convolution = np.empty((windowsize, ), dtype=np.float64)
convolution.fill(1.0)
convolution = convolution / float(windowsize)
# average time values
time_conv = sp_conv(tmpDS['time'], convolution, mode='valid')
# average coordinates
running_avgs = []
for i in range(tmpDS.data.shape[0]):
i_conv = sp_conv(tmpDS.data[i, :], convolution, mode='valid')
running_avgs.append(i_conv)
running_avg_array = np.stack(running_avgs, axis=0)
# setup array to hold final values
final_data = np.empty(
(running_avg_array.shape[0] / 3, running_avg_array.shape[1]))
# iterate over time windows
for i in range(running_avg_array.shape[1]):
# broadcasting doesn't work here, so make a stretched array of avg coords manually
i_avg_coord = np.repeat(running_avg_array[:, i].reshape(1, -1).T,
windowsize,
axis=1)
window_coords = tmpDS.data[:, i:i + windowsize]
diff_coords = window_coords - i_avg_coord
# iterate over atoms and calculate rmsf
for j in range(running_avg_array.shape[0] / 3):
j_diff_coords = diff_coords[j * 3:j * 3 + 3, :]
disp = np.linalg.norm(j_diff_coords, axis=0)
disp_sq = np.square(disp)
ij_rmsf = np.sqrt(np.mean(disp_sq))
final_data[j, i] = ij_rmsf
# load up the dataset
for elem in tmpDS.feature_list:
self.primaryDS.add_feature(elem.get_descriptor(), width=1)
self.primaryDS.format_data(final_data)
self.primaryDS.time = time_conv
self.primaryDS.feature_dict['time'] = self.primaryDS.time
def __init__(self, top, traj, ncut=25):
self.top = top
self.traj = traj
univ1 = mda.Universe(self.top, self.traj, verbose=True)
protein = univ1.select_atoms("protein")
coords_protein = AnalysisFromFunction(lambda ag: ag.positions.copy(),
protein).run().results
univ2 = mda.Merge(protein) # create the protein-only Universe
univ2.load_new(coords_protein, format=MemoryReader)
cut = []
for i in np.arange(0, len(coords_protein[:, 0, 0]), ncut):
cut.append(coords_protein[i, :, :])
coord_cut = np.zeros((len(cut), 304, 3), dtype=np.float64)
coord_cut[:, :, :] = [cut[i] for i in np.arange(0, len(cut))]
univ_cut = mda.Universe("folded.pdb", coord_cut)
self.univ_cut = univ_cut
self.coords_protein = coords_protein
self.coord_cut = coord_cut
self.univ2 = univ2
self.ncut = ncut
ind_end = (rank + 1) * local_n
else:
ind_end = N
if rank == 0:
inds = np.arange(N)
else:
inds = None
local_inds = np.empty(local_n, dtype=np.int)
comm.Scatter(inds, local_inds, root=0)
#print(chains)
for i in local_inds:
filename = folder + 'chain_{}.xtc'.format(i + 1)
print(filename)
sys.stdout.flush()
protein = chains[i]
if not os.path.exists(filename):
u2 = md.Merge(protein).load_new(AnalysisFromFunction(
lambda ag: ag.positions.copy(), protein).run().results,
format=MemoryReader)
with md.Writer(filename, protein.n_atoms) as W:
for ts in u2.trajectory:
W.write(u2)
if rank == 0:
chains[0].write(folder + 'chain.gro')
chains[0].write(folder + 'chain.pdb')
def extract_combined_grid(struc_a, xtc_a, struc_b, xtc_b, atomgroup,
write_grid_as, out_name):
"""
Writes out combined atomgroup density for both input simulations.
Parameters
----------
struc_a : str
File name for the reference file (PDB or GRO format).
xtc_a : str
File name for the trajectory (xtc format).
struc_b : str
File name for the reference file (PDB or GRO format).
xtc_b : str
File name for the trajectory (xtc format).
atomgroup : str
Atomgroup selection to calculate the density for (atom name in structure_input).
write_grid_as : str
The water model to convert the density into.
Options are: SPC, TIP3P, TIP4P, water
out_name : str
Prefix for written filename.
"""
if not os.path.exists('dens/'):
os.makedirs('dens/')
condition_a = mda.Universe(struc_a, xtc_a)
condition_b = mda.Universe(struc_b, xtc_b)
# # # Combine both ensembles' atoms into one universe
Combined_conditions = mda.Merge(condition_a.atoms, condition_b.atoms)
# # # Extract trajectory coordinates
aligned_coords_a = AnalysisFromFunction(_copy_coords,
condition_a.atoms).run().results
aligned_coords_b = AnalysisFromFunction(_copy_coords,
condition_b.atoms).run().results
# # # The density needs to be formed from an even contribution of both conditions
# # # otherwise it will be unevely biased towards one condition.
# # # So we match the simulation lengths first
sim1_coords, sim2_coords = _match_sim_lengths(aligned_coords_a,
aligned_coords_b)
# # # Then we merge the coordinates into one system
merged_coords = np.hstack([sim1_coords, sim2_coords])
# # # We load in the merged coordinated into our new universe that contains
# # # the receptor in both conditions
Combined_conditions.load_new(merged_coords, format=MemoryReader)
# # # We extract the density grid from the combined condition universe
density_atomgroup = Combined_conditions.select_atoms("name " + atomgroup)
# a resolution of delta=1.0 ensures the coordinates of the maxima match the coordinates of the simulation box
D = DensityAnalysis(density_atomgroup, delta=1.0)
D.run()
D.density.convert_density(write_grid_as)
D.density.export('dens/' + out_name + atomgroup + "_density.dx",
type="double")
def iter_from_trajectory(
nneighbor_cutoff,
universe,
selection='all',
r_cut=10.,
period=1):
''' This generator will process information from a trajectory and
yield a tuple of ``[nlist, positions, box, sample_weight]`` and ``MDAnalysis.TimeStep`` object.
The first list can be directly used to call a :py:class:`.SimModel` (e.g., ``model(inputs)``).
See :py:meth:`.SimModel.compute` for details of these terms.
Here's an example:
.. code:: python
model = MyModel(16)
for inputs, ts in iter_from_trajectory(16, universe):
result = model(inputs)
:param nneighbor_cutoff: The maximum size of neighbor list
:type nneighbor_cutoff: int
:param universe: The MDAnalysis universe
:param selection: The atom groups to extract from universe
:type selection: string
:param r_cut: The cutoff raduis to use in neighbor list
calculations
:type r_cut: float
:param period: Period of reading the trajectory frames
:type period: int
'''
import MDAnalysis
# Modifying the universe for none 'all' atom selections.
if selection != 'all':
from MDAnalysis.analysis.base import AnalysisFromFunction
p = universe.select_atoms(selection)
dt = universe.trajectory[0].dt
dimensions = universe.trajectory[0].dimensions
if universe.trajectory[0].has_forces is False:
# Only include positions if traj does not have forces
x = AnalysisFromFunction(lambda ag: [ag.positions.copy()], p).run().results
# Construct new_trajectory from the MemoryReader explicitly:
new_traj = MDAnalysis.coordinates.memory.MemoryReader(
x[:, 0], dimensions=dimensions, dt=dt)
else:
# Include positions, velocities and forces:
xvf = AnalysisFromFunction(lambda ag: [ag.positions.copy(
), ag.velocities.copy(), ag.forces.copy()], p).run().results
new_traj = MDAnalysis.coordinates.memory.MemoryReader(
xvf[:, 0], velocities=xvf[:, 1], forces=xvf[:, 2], dimensions=dimensions, dt=dt)
universe.trajectory = new_traj
print('The universe was redefined based on the atom group selection input.')
# read trajectory
box = universe.dimensions
# define the system
hoomd_box = np.array([[box[0], 0, 0], [0, box[1], 0], [0, 0, box[2]]])
# make type array
# Select atom group to use in the system
atom_group = universe.select_atoms(selection)
# get unique atom types in the selected atom group
try:
types = list(np.unique(atom_group.atoms.types))
# associate atoms types with individual atoms
type_array = np.array([types.index(i)
for i in atom_group.atoms.types]).reshape(-1, 1)
except MDAnalysis.exceptions.NoDataError:
type_array = np.zeros(len(atom_group)).reshape(-1, 1)
# define nlist operation
# box_size = [box[0], box[1], box[2]]
nlist = compute_nlist(
atom_group.positions,
r_cut=r_cut,
NN=nneighbor_cutoff,
box_size=[
box[0],
box[1],
box[2]])
# Run the model at every nth frame, where n = period
for i, ts in enumerate(universe.trajectory):
if i % period == 0:
yield [nlist, np.concatenate(
(atom_group.positions,
type_array),
axis=1), hoomd_box, 1.0], ts
def coarse_grain(universe,
residue_list,
simulation_name='simulation_name',
export=False):
# ============== Misc Initiation ============== #
with open('src/mapping_dict.json', "r") as f:
mapping_dict = load(f)
with open('src/abrev_dict.json', "r") as f:
abrev_dict = load(f)
u = universe
# ================= Execution ================= #
print('Calculating Bond connections...')
resnames = ' '.join(residue_list)
original_bond_count = len(u.bonds)
u.select_atoms(f'resname {resnames}').guess_bonds(vdwradii=config.vdw_radi)
print(
f'Original file contained {original_bond_count} bonds. {len(u.bonds) - original_bond_count} additional bonds infered.'
)
print(f'Begining Coarse-Graining process...')
bead_data = []
cg_beads = []
dummy_parents = {}
non_water_atoms = u.select_atoms('not resname WAT')
for residue in non_water_atoms.residues: # loops thu each matching residue id
resid = residue.resid # store int id
resname = residue.resname
if resname in residue_list: # if resname == "PHOSPHATE" or resname == "RIBOSE":
# resname_atoms = u.atoms.select_atoms('resname DA DT DG DC DU')
# else:
# resname_atoms = u.atoms.select_atoms(f'resname {resname}') # selects all resname-specific atoms
if len(resname) == 4 and resname[0] == 'D': # for D-varants
resname_key = resname[1:]
else:
resname_key = resname
try:
segments = mapping_dict[resname_key].keys()
for segment in segments: # loops thru each segment of each residue
params = 'name ' + ' '.join(
mapping_dict[resname_key][segment]
['atoms']) # generates param
# selects all atoms in a given residue segment
atms = residue.atoms.select_atoms(params)
dummy = atms[0]
# names dummy atom in propper format
dummy.name = str(abrev_dict[resname_key]) + str(
segment[0]) + str(resid)
dummy.type = mapping_dict[resname_key][segment]['name']
dummy.charge = mapping_dict[resname_key][segment]['charge']
bead_data.append((dummy, atms))
cg_beads.append(dummy)
for atm in atms:
dummy_parents[atm.ix] = dummy
except KeyError:
print(
f'{resname_key} was not found in mapping/abrev_dict, skipping coarse grain. Please add its parameters to the dictionary. (See README section A3. for help)'
)
new_bonds = []
# for residue in residue_list:
# for mapping in mapping_dict[residue]["Bonds"]:
# first_code = mapping[0] # segment, resid offset, resname
# seccond_code = mapping[1] if isinstance(mapping[1], list) else [mapping[1], 0, residue]
# type_params = list(mapping_dict[residue]["Mapping"].keys())[first_code]
# first_atoms = cg_beads.select_atoms(f'resname {residue} and type {type_params}')
# for first_atom in first_atoms:
# type_params = list(mapping_dict[residue]["Mapping"].keys())[seccond_code[0]] # segment
# seccond_atom_resid = int(first_atom.resid) + int(seccond_code[1])
# try:
# seccond_atom = cg_beads.atoms.select_atoms(f'resname {seccond_code[2]} and type {type_params} and resid {seccond_atom_resid}')
# except IndexError:
# pass
# if isinstance(seccond_atom, mda.core.groups.AtomGroup):
# closest = seccond_atom[0]
# closest_dist = mda.AtomGroup([first_atom, seccond_atom[0]]).bond.length()
# for atom in seccond_atom:
# dist = mda.AtomGroup([first_atom, atom]).bond.length()
# if dist < closest_dist:
# closest = atom
# closest_dist = dist
# seccond_atom = closest
# new_bonds.append([first_atom.index, seccond_atom.index])
# new_bonds = []
for dummy, atms in bead_data:
# connect all parents with connected children
for atom in atms:
for bond in atom.bonds:
for bonded_atom in bond.atoms:
if bonded_atom not in atms: # make more efficent if atms were a set
# by the end of all these loops and ifs, every bonded_atom that gets to this point is an atom connected to the edge of the cluster of atoms assigned to the coarse grain dummy bead in question
try:
new_bonds.append([
cg_beads.index(dummy),
cg_beads.index(dummy_parents[bonded_atom.ix])
]) # type is used to store the cluster dummy
except KeyError: # raises if atom does not belong to a coarse grain bead
pass
# try:
# new_bonds.append([cg_beads.index(dummy), cg_beads.index(bonded_atom)]) # adds the bond between the dummies
# except ValueError: # if the other atom is just an atom withouot a coarse grain bead parent, ignore it
# pass
cg_beads = mda.AtomGroup(cg_beads)
# TODO: EXPORT NEW_U INSTEAD OF OLD U
# TODO: EXPORT NEW_U TO HAVE APPROPRIATE FRAMES
# TODO: SHIFT THE DEFINITION OF CENTERS IN THE UNIVERSE EVEN IF NOT EXPORTING
# TODO: AUTOTUNE THE CURVE TO FIND THE RIGHT STEP
progress(0)
number_of_frames = len(u.trajectory)
for frame in u.trajectory: # loops tru each frame
f = frame.frame
# positions a dummy atoms at cluster center of mass
for dummy, atms in bead_data:
dummy.position = AtomGroup(atms).center_of_mass()
progress(f / number_of_frames)
progress(1)
print()
for dummy, atms in bead_data:
dummy.mass = AtomGroup(atms).masses.sum()
# purge existing reminant bonds
u.delete_bonds(u.bonds)
u.delete_angles(u.angles)
u.delete_dihedrals(u.dihedrals)
print(f'Building new coarse-grained universe...')
coordinates = AnalysisFromFunction(lambda ag: ag.positions.copy(),
cg_beads).run().results
new_u = mda.Merge(cg_beads)
new_u.load_new(coordinates, format=MemoryReader)
new_u.add_TopologyAttr('bonds', new_bonds)
new_u.add_TopologyAttr('angles', guess_angles(new_u.bonds))
new_u.add_TopologyAttr('dihedrals', guess_dihedrals(new_u.angles))
print(
f'Built universe with {len(new_u.atoms)} coarse-grained beads, {len(new_u.bonds)} bonds, {len(new_u.angles)} angles, and {len(new_u.dihedrals)} dihedrals'
)
if export:
print('Writing Output Files...')
out_file = f'outputs/CoarseGrain/{simulation_name}_CG.pdb'
with open(out_file, 'w+') as _:
new_u.atoms.write(out_file, bonds='all')
print(f'Topology written to {simulation_name}_CG.pdb!')
is_multiframe = number_of_frames > 1
with mda.Writer(f'outputs/CoarseGrain/{simulation_name}_CG.dcd',
new_u.atoms.n_atoms,
multiframe=is_multiframe,
bonds='all') as w:
for frame in new_u.trajectory[1:]: # loops tru each frame
w.write(new_u.atoms)
print('Generated All Coarse Grained Molecules!')
print(f'Trajectory written to {simulation_name}_CG.dcd!')
# for dummy, atms in bead_data:
# dummy.type = ''
print(f'Reduced {len(u.atoms)} atoms to {len(new_u.atoms)} beads!')
print('Coarse Graining Task complete!')
return new_u
def _expand_universe(universe, length):
coordinates = AnalysisFromFunction(lambda ag: ag.positions.copy(),
universe.atoms).run().results
coordinates = np.tile(coordinates, (length, 1, 1))
universe.load_new(coordinates, format=MemoryReader)