def _process_frame(self, trj, energy, hbonds, entropy):
"""
Frame wise calculation of GIST quantities.
Parameters
----------
trj :
Molecular dynamic trajectory representing current frame.
energy :
If True, solute-water and water-water energies are calculated for each water in each voxel in current
frame.
hbonds : bool
If True, solute-water and water-water hydrogen bonds are calculated for each water in each voxel in current
frame.
entropy : bool
If True, water coordinates and quaternions are stored for each water in each voxel in current frame.
"""
nbr_cutoff_sq = 3.5 ** 2
trj.xyz *= 10.0
coords = trj.xyz
uc = trj.unitcell_vectors[0]*10.
waters = []
calc.assign_voxels(trj.xyz, self.dims, self.gridmax, self.origin, waters, self.wat_oxygen_atom_ids)
for wat in waters:
self.voxeldata[wat[0], 4] += 1
if energy or hbonds:
e_lj_array, e_elec_array = np.copy(self.acoeff), np.copy(self.chg_product)
distance_matrix = np.zeros((self.water_sites, self.all_atom_ids.shape[0]))
calc.get_pairwise_distances(wat, self.all_atom_ids, coords, uc, distance_matrix)
wat_nbrs = self.wat_oxygen_atom_ids[np.where(
(distance_matrix[0, :][self.wat_oxygen_atom_ids] <= nbr_cutoff_sq) & (
distance_matrix[0, :][self.wat_oxygen_atom_ids] > 0.0))]
self.voxeldata[wat[0], 19] += wat_nbrs.shape[0]
calc.calculate_energy(wat[1], distance_matrix, e_elec_array, e_lj_array, self.bcoeff)
if self.prot_atom_ids.shape[0] != 0:
#self.voxeldata[wat[0], 13] += np.sum(e_lj_array[:, self.prot_atom_ids])
#self.voxeldata[wat[0], 13] += np.sum(e_elec_array[:, self.prot_atom_ids])
self.voxeldata[wat[0], 13] += np.sum(e_lj_array[:, self.non_water_atom_ids])
self.voxeldata[wat[0], 13] += np.sum(e_elec_array[:, self.non_water_atom_ids])
self.voxeldata[wat[0], 15] += np.sum(
e_lj_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(e_lj_array[:, wat[1] + self.water_sites:])
self.voxeldata[wat[0], 15] += np.sum(
e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(
e_elec_array[:, wat[1] + self.water_sites:])
e_nbr_list = [np.sum(e_lj_array[:, wat_nbrs + i] + e_elec_array[:, wat_nbrs + i]) for i in
range(self.water_sites)]
self.voxeldata[wat[0], 17] += np.sum(e_nbr_list)
# H-bond calculations
if hbonds:
prot_nbrs_all = self.prot_atom_ids[
np.where(distance_matrix[0, :][self.prot_atom_ids] <= nbr_cutoff_sq)]
prot_nbrs_hb = prot_nbrs_all[np.where(self.prot_hb_types[prot_nbrs_all] != 0)]
if wat_nbrs.shape[0] > 0:
hb_ww = self.calculate_hydrogen_bonds(trj, wat[1], wat_nbrs)
acc_ww = hb_ww[:, 0][np.where(hb_ww[:, 0] == wat[1])].shape[0]
don_ww = hb_ww.shape[0] - acc_ww
self.voxeldata[wat[0], 25] += hb_ww.shape[0]
self.voxeldata[wat[0], 31] += don_ww
self.voxeldata[wat[0], 33] += acc_ww
if wat_nbrs.shape[0] != 0 and hb_ww.shape[0] != 0:
self.voxeldata[wat[0], 21] += wat_nbrs.shape[0] / hb_ww.shape[0]
if prot_nbrs_hb.shape[0] > 0:
hb_sw = self.calculate_hydrogen_bonds(trj, wat[1], prot_nbrs_hb, water_water=False)
acc_sw = hb_sw[:, 0][np.where(hb_sw[:, 0] == wat[1])].shape[0]
don_sw = hb_sw.shape[0] - acc_sw
self.voxeldata[wat[0], 23] += hb_sw.shape[0]
self.voxeldata[wat[0], 27] += don_sw
self.voxeldata[wat[0], 29] += acc_sw
if entropy:
self.calculate_euler_angles(wat, coords[0, :, :])
def _process_frame(self, trj, energy, hbonds, entropy):
"""
Frame wise calculation of GIST quantities.
Parameters
----------
trj :
Molecular dynamic trajectory representing current frame.
energy :
If True, solute-water and water-water energies are calculated for each water in each voxel in current
frame.
hbonds : bool
If True, solute-water and water-water hydrogen bonds are calculated for each water in each voxel in current
frame.
entropy : bool
If True, water coordinates and quaternions are stored for each water in each voxel in current frame.
"""
nbr_cutoff_sq = 3.5**2
trj.xyz *= 10.0
coords = trj.xyz
uc = trj.unitcell_vectors[0] * 10.
waters = []
calc.assign_voxels(trj.xyz, self.dims, self.gridmax, self.origin,
waters, self.wat_oxygen_atom_ids)
distance_matrix = np.zeros(
(self.water_sites, self.all_atom_ids.shape[0]))
for wat in waters:
self.voxeldata[wat[0], 4] += 1
### wat[0]: Voxel index
### wat[1]: water molecule oxygen atom index
if energy or hbonds:
e_lj_array, e_elec_array = np.copy(self.acoeff), np.copy(
self.chg_product)
### Here it is assumed that oxygen atom is always the first atom. Is that always the case?
### We should probably check that somewhere during init?
###
### The self.neighbor_ids array contains atom indices of all atoms that should
### be considered as potential neighbors. Valid_neighbors has same shape as
### self.neighbor_ids and is False at the position where index wat occures in
### self.neighbor_ids, otherwise it is True. neighbor_ids stores the indices of
### the actual neighbor candidates that will be commited to get_pairwise_distances
### routine and has length of self.neighbor_ids-1. wat_nbrs_shell is of length neighbor_ids
### and holds the shell_index of each neighbor candidate atom (0:first shell, 1: beyond first
### shell)
valid_neighbors = np.ones(self.neighbor_ids.shape[0],
dtype=bool)
valid_neighbors[np.where(self.neighbor_ids == wat)] = False
neighbor_ids = self.neighbor_ids[valid_neighbors]
wat_nbrs_shell = self.wat_nbrs_shell[valid_neighbors]
calc.get_pairwise_distances(wat, self.all_atom_ids,
np.array([nbr_cutoff_sq]),
neighbor_ids, wat_nbrs_shell,
coords, uc, distance_matrix, 0)
wat_nbrs = self.all_atom_ids[np.where(wat_nbrs_shell == 0)]
self.voxeldata[wat[0], 19] += wat_nbrs.shape[0]
calc.calculate_energy(wat[1], distance_matrix, e_elec_array,
e_lj_array, self.bcoeff)
if self.prot_atom_ids.shape[0] != 0:
#self.voxeldata[wat[0], 13] += np.sum(e_lj_array[:, self.prot_atom_ids])
#self.voxeldata[wat[0], 13] += np.sum(e_elec_array[:, self.prot_atom_ids])
self.voxeldata[wat[0], 13] += np.sum(
e_lj_array[:, self.non_water_atom_ids])
self.voxeldata[wat[0], 13] += np.sum(
e_elec_array[:, self.non_water_atom_ids])
self.voxeldata[wat[0], 15] += np.sum(
e_lj_array[:,
self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(
e_lj_array[:, wat[1] + self.water_sites:])
self.voxeldata[wat[0], 15] += np.sum(
e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]
) + np.sum(e_elec_array[:, wat[1] + self.water_sites:])
e_nbr_list = [
np.sum(e_lj_array[:, wat_nbrs + i] +
e_elec_array[:, wat_nbrs + i])
for i in range(self.water_sites)
]
self.voxeldata[wat[0], 17] += np.sum(e_nbr_list)
### Might be usefull for API to have the neighbors and shell
### indices available.
self.wat_nbrs_shell[valid_neighbors] = wat_nbrs_shell
self.neighbor_ids[valid_neighbors] = neighbor_ids
"""
###DEBUG START###
elj_sw = np.sum(e_lj_array[:, :self.wat_oxygen_atom_ids[0]])
eelec_sw = np.sum(e_elec_array[:, :self.wat_oxygen_atom_ids[0]])
elj_ww = np.sum(e_lj_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(e_lj_array[:, wat[1] + 1:])
eelec_ww = np.sum(e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(e_elec_array[:, wat[1] + self.water_sites:])
e_nbr_list = [np.sum(e_lj_array[:, wat_nbrs + i] + e_elec_array[:, wat_nbrs + i]) for i in xrange(self.water_sites)]
enbr = np.sum(e_nbr_list)
print "Calc: ", elj_sw, eelec_sw, elj_ww, eelec_ww, enbr
distance_matrix = np.sqrt(distance_matrix)
energy_lj, energy_elec = self.calculate_energy(distance_matrix)
test_1 = np.sum(energy_lj[:self.wat_oxygen_atom_ids[0]:])
test_2 = np.sum(energy_elec[:, self.non_water_atom_ids])
test_3 = np.nansum(energy_lj[self.wat_oxygen_atom_ids[0]:])
test_4 = np.sum(energy_elec[:, self.wat_atom_ids[0]:wat[1]]) + np.sum(energy_elec[:, wat[1] + self.water_sites:])
test_5 = 0.0
test_5 += np.sum(energy_lj[self.wat_oxygen_atom_ids[0]:][(wat_nbrs - self.wat_oxygen_atom_ids[0]) / self.water_sites])
for i in range(self.water_sites):
test_5 += np.sum(energy_elec[:, wat_nbrs + i])
print "Ref: ", test_1, test_2, test_3, test_4, test_5
###DEBUG END###
"""
# H-bond calculations
if hbonds:
prot_nbrs_all = self.prot_atom_ids[np.where(
distance_matrix[0, :][
self.prot_atom_ids] <= nbr_cutoff_sq)]
prot_nbrs_hb = prot_nbrs_all[np.where(
self.prot_hb_types[prot_nbrs_all] != 0)]
if wat_nbrs.shape[0] > 0:
hb_ww = self.calculate_hydrogen_bonds(
trj, wat[1], wat_nbrs)
acc_ww = hb_ww[:, 0][np.where(
hb_ww[:, 0] == wat[1])].shape[0]
don_ww = hb_ww.shape[0] - acc_ww
self.voxeldata[wat[0], 25] += hb_ww.shape[0]
self.voxeldata[wat[0], 31] += don_ww
self.voxeldata[wat[0], 33] += acc_ww
if wat_nbrs.shape[0] != 0 and hb_ww.shape[0] != 0:
self.voxeldata[
wat[0],
21] += wat_nbrs.shape[0] / hb_ww.shape[0]
if prot_nbrs_hb.shape[0] > 0:
hb_sw = self.calculate_hydrogen_bonds(
trj, wat[1], prot_nbrs_hb, water_water=False)
acc_sw = hb_sw[:, 0][np.where(
hb_sw[:, 0] == wat[1])].shape[0]
don_sw = hb_sw.shape[0] - acc_sw
self.voxeldata[wat[0], 23] += hb_sw.shape[0]
self.voxeldata[wat[0], 27] += don_sw
self.voxeldata[wat[0], 29] += acc_sw
if entropy:
self.calculate_euler_angles(wat, coords[0, :, :])
def process_chunk(self, begin_chunk, chunk_size, topology, energy, hbonds,
entropy):
nbr_cutoff_sq = 3.5**2
with md.open(self.trajectory) as f:
f.seek(begin_chunk)
trj = f.read_as_traj(topology, n_frames=chunk_size, stride=1)
trj.xyz *= 10.0
pbc = md.utils.in_units_of(trj.unitcell_lengths, "nanometers",
"angstroms")
frame_data = [[] for i in range(trj.n_frames)]
calc.assign_voxels(trj.xyz, self.dims, self.gridmax, self.origin,
frame_data, self.wat_oxygen_atom_ids)
for frame in range(trj.n_frames):
coords = trj.xyz[frame, :, :].reshape(1, trj.xyz.shape[1],
trj.xyz.shape[2])
periodic_box = pbc[frame].reshape(1, pbc.shape[1])
waters = frame_data[frame]
for wat in waters:
self.voxeldata[wat[0], 4] += 1
if energy or hbonds:
e_lj_array, e_elec_array = np.copy(
self.acoeff), np.copy(self.chg_product)
distance_matrix = np.zeros(
(self.water_sites, self.all_atom_ids.shape[0]))
calc.get_pairwise_distances(wat, self.all_atom_ids,
coords, pbc,
distance_matrix)
wat_nbrs = self.wat_oxygen_atom_ids[np.where(
(distance_matrix[0, :][self.wat_oxygen_atom_ids] <=
nbr_cutoff_sq) & (distance_matrix[0, :][
self.wat_oxygen_atom_ids] > 0.0))]
self.voxeldata[wat[0], 17] += wat_nbrs.shape[0]
calc.calculate_energy(wat[1], distance_matrix,
e_elec_array, e_lj_array,
self.bcoeff)
self.voxeldata[wat[0], 11] += np.sum(
e_lj_array[:, :self.wat_oxygen_atom_ids[0]])
self.voxeldata[wat[0], 11] += np.sum(
e_elec_array[:, :self.wat_oxygen_atom_ids[0]])
self.voxeldata[wat[0], 13] += np.sum(
e_lj_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]
) + np.sum(e_lj_array[:, wat[1] + self.water_sites:])
self.voxeldata[wat[0], 13] += np.sum(
e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]
) + np.sum(e_elec_array[:, wat[1] + self.water_sites:])
e_nbr_list = [
np.sum(e_lj_array[:, wat_nbrs + i] +
e_elec_array[:, wat_nbrs + i])
for i in xrange(self.water_sites)
]
self.voxeldata[wat[0], 15] += np.sum(e_nbr_list)
"""
###DEBUG START###
elj_sw = np.sum(e_lj_array[:, :self.wat_oxygen_atom_ids[0]])
eelec_sw = np.sum(e_elec_array[:, :self.wat_oxygen_atom_ids[0]])
elj_ww = np.sum(e_lj_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(e_lj_array[:, wat[1] + 1:])
eelec_ww = np.sum(e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(e_elec_array[:, wat[1] + self.water_sites:])
e_nbr_list = [np.sum(e_lj_array[:, wat_nbrs + i] + e_elec_array[:, wat_nbrs + i]) for i in xrange(self.water_sites)]
enbr = np.sum(e_nbr_list)
print "Calc: ", elj_sw, eelec_sw, elj_ww, eelec_ww, enbr
distance_matrix = np.sqrt(distance_matrix)
energy_lj, energy_elec = self.calculate_energy(distance_matrix)
test_1 = np.sum(energy_lj[:self.wat_oxygen_atom_ids[0]:])
test_2 = np.sum(energy_elec[:, self.non_water_atom_ids])
test_3 = np.nansum(energy_lj[self.wat_oxygen_atom_ids[0]:])
test_4 = np.sum(energy_elec[:, self.wat_atom_ids[0]:wat[1]]) + np.sum(energy_elec[:, wat[1] + self.water_sites:])
test_5 = 0.0
test_5 += np.sum(energy_lj[self.wat_oxygen_atom_ids[0]:][(wat_nbrs - self.wat_oxygen_atom_ids[0]) / self.water_sites])
for i in range(self.water_sites):
test_5 += np.sum(energy_elec[:, wat_nbrs + i])
print "Ref: ", test_1, test_2, test_3, test_4, test_5
###DEBUG END###
"""
# H-bond calculations
if hbonds:
prot_nbrs_all = self.non_water_atom_ids[np.where(
distance_matrix[0, :][
self.non_water_atom_ids] <= nbr_cutoff_sq)]
prot_nbrs_hb = prot_nbrs_all[np.where(
self.prot_hb_types[prot_nbrs_all] != 0)]
if wat_nbrs.shape[0] != 0 and prot_nbrs_hb.shape[
0] != 0:
# hb_ww, hb_sw = self.calculate_hydrogen_bonds2(coords, wat[1], wat_nbrs, prot_nbrs_hb)
hb_ww, hb_sw = self.calculate_hydrogen_bonds(
trj, wat[1], wat_nbrs, prot_nbrs_hb)
acc_ww = hb_ww[:, 0][np.where(
hb_ww[:, 0] == wat[1])].shape[0]
don_ww = hb_ww.shape[0] - acc_ww
acc_sw = hb_sw[:, 0][np.where(
hb_sw[:, 0] == wat[1])].shape[0]
don_sw = hb_sw.shape[0] - acc_sw
self.voxeldata[wat[0], 23] += hb_sw.shape[0]
self.voxeldata[wat[0], 25] += hb_ww.shape[0]
self.voxeldata[wat[0], 27] += don_sw
self.voxeldata[wat[0], 29] += acc_sw
self.voxeldata[wat[0], 31] += don_ww
self.voxeldata[wat[0], 33] += acc_ww
if wat_nbrs.shape[0] != 0 and hb_ww.shape[
0] != 0:
self.voxeldata[wat[0],
19] += wat_nbrs.shape[
0] / hb_ww.shape[0]
# f_enc = 1.0 - (wat_nbrs.shape[0] / 5.25)
# if f_enc < 0.0:
# f_enc = 0.0
# self.voxeldata[wat[0], 21] += f_enc
if entropy:
self.calculate_euler_angles(wat, coords[0, :, :])
def _process_frame(self, trj, frame_i, energy, hbonds, entropy,
energy_lr_breakdown, angular_structure, shell_radii,
r_theta_cutoff):
"""Calculates hydration site properties for a given frame.
Parameters
----------
trj : mdtraj.trajectory
A trajectory object containing only one frame.
frame_i : int
Index of the frame to be processed
energy : bool
Flag for energy calculations
hbonds : bool
Flag for hydrogen bond calculations
entropy :bool
Flag for entropy calculations
Returns
-------
None : NoneType
"""
site_waters_copy = list(self.site_waters)
nbr_cutoff_sq = 3.5**2
# TODO: Use a robust unit conversion approach
trj.xyz *= 10.0
coords = trj.xyz
trj.unitcell_lengths *= 10.0
uc = trj.unitcell_vectors[0] * 10.
# Iterate over each site in the current frame if it has a water present
for site_i in range(self.hsa_data.shape[0]):
wat_O = None
if self.is_site_waters_populated:
if len(site_waters_copy[site_i]) != 0:
if site_waters_copy[site_i][0][0] == frame_i:
wat_O = site_waters_copy[site_i].pop(0)[1]
index = int(self.hsa_data[site_i, 4]) * 3
index_pairs = list(
zip(list(range(wat_O, wat_O + 3)),
list(range(index, index + 3))))
for index_pair in index_pairs:
self.hsa_dict[site_i][-1][index_pair[1]] += coords[
0, index_pair[0], :]
self.hsa_data[site_i, 4] += 1
if wat_O is not None and (energy or hbonds):
distance_matrix = np.zeros(
(self.water_sites, self.all_atom_ids.shape[0]), np.float_)
calc.get_pairwise_distances(np.asarray([site_i, wat_O]),
self.all_atom_ids, coords, uc,
distance_matrix)
wat_nbrs = self.wat_oxygen_atom_ids[np.where(
(distance_matrix[0, :][self.wat_oxygen_atom_ids] <=
nbr_cutoff_sq)
& (distance_matrix[0, :][self.wat_oxygen_atom_ids] > 0.0))]
self.hsa_dict[site_i][17].append(wat_nbrs.shape[0])
if energy:
e_lj_array, e_elec_array = np.copy(self.acoeff), np.copy(
self.chg_product)
calc.calculate_energy(wat_O, distance_matrix, e_elec_array,
e_lj_array, self.bcoeff)
e_lj_sw = np.sum(e_lj_array[:, self.non_water_atom_ids])
e_elec_sw = np.sum(e_elec_array[:,
self.non_water_atom_ids])
e_lj_ww_left = e_lj_array[:, self.
wat_oxygen_atom_ids[0]:wat_O]
e_lj_ww_right = e_lj_array[:, wat_O + self.water_sites:]
e_lj_ww = np.sum(e_lj_ww_left) + np.sum(e_lj_ww_right)
e_elec_ww_left = e_elec_array[:, self.
wat_oxygen_atom_ids[0]:wat_O]
e_elec_ww_right = e_elec_array[:,
wat_O + self.water_sites:]
e_elec_ww = np.sum(e_elec_ww_left) + np.sum(
e_elec_ww_right)
e_nbr_list = [
np.sum(e_lj_array[:, nbr:nbr + self.water_sites] +
e_elec_array[:, nbr:nbr + self.water_sites])
for nbr in wat_nbrs
]
e_lj_ww = np.sum(
e_lj_array[:, self.wat_oxygen_atom_ids[0]:wat_O]
) + np.sum(e_lj_array[:, wat_O + self.water_sites:])
e_elec_ww = np.sum(
e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat_O]
) + np.sum(e_elec_array[:, wat_O + self.water_sites:])
self.hsa_dict[site_i][7].append(e_lj_sw)
self.hsa_dict[site_i][8].append(e_elec_sw)
self.hsa_dict[site_i][10].append(e_lj_ww)
self.hsa_dict[site_i][11].append(e_elec_ww)
self.hsa_dict[site_i][6].append(e_lj_sw + e_elec_sw)
self.hsa_dict[site_i][9].append(e_lj_ww + e_elec_ww)
self.hsa_dict[site_i][12].append(e_lj_sw + e_elec_sw +
e_lj_ww + e_elec_ww)
self.hsa_dict[site_i][13].extend(
e_nbr_list) # print(e_lj_sw/2.0)
if energy_lr_breakdown:
for s in range(1, len(shell_radii)):
wat_nbrs = self.wat_oxygen_atom_ids[np.where((
distance_matrix[0, :][self.wat_oxygen_atom_ids]
<= shell_radii[s]) & (distance_matrix[0, :][
self.wat_oxygen_atom_ids] > shell_radii[
s - 1]))]
e_nbr_list = [
np.sum(e_lj_array[:,
nbr:nbr + self.water_sites] +
e_elec_array[:, nbr:nbr +
self.water_sites])
for nbr in wat_nbrs
]
self.energy_ww_lr_breakdown[site_i][s - 1] += sum(
e_nbr_list)
if hbonds:
hbtot = 0
prot_nbrs_all = self.prot_atom_ids[np.where(
distance_matrix[0, :][
self.prot_atom_ids] <= nbr_cutoff_sq)]
prot_nbrs_hb = prot_nbrs_all[np.where(
self.prot_hb_types[prot_nbrs_all] != 0)]
if wat_nbrs.shape[0] > 0:
hb_ww = self.calculate_hydrogen_bonds(
trj, wat_O, wat_nbrs)
acc_ww = hb_ww[:,
0][np.where(hb_ww[:,
0] == wat_O)].shape[0]
don_ww = hb_ww.shape[0] - acc_ww
self.hsa_dict[site_i][18].append(hb_ww.shape[0])
self.hsa_dict[site_i][23].append(acc_ww)
self.hsa_dict[site_i][24].append(don_ww)
hbtot += hb_ww.shape[0]
if wat_nbrs.shape[0] != 0 and hb_ww.shape[0] != 0:
self.hsa_dict[site_i][21].append(hb_ww.shape[0] /
wat_nbrs.shape[0])
if prot_nbrs_hb.shape[0] > 0:
hb_sw = self.calculate_hydrogen_bonds(
trj, wat_O, prot_nbrs_hb, water_water=False)
acc_sw = hb_sw[:,
0][np.where(hb_sw[:,
0] == wat_O)].shape[0]
don_sw = hb_sw.shape[0] - acc_sw
don_sw_ids = hb_sw[:, 1][np.where(hb_sw[:,
0] == wat_O)]
acc_sw_ids = hb_sw[:,
0][np.where(hb_sw[:, 0] != wat_O)]
self.hsa_dict[site_i][19].append(hb_sw.shape[0])
self.hsa_dict[site_i][25].append(acc_sw)
self.hsa_dict[site_i][26].append(don_sw)
self.hsa_dict[site_i][27].extend(acc_sw_ids)
self.hsa_dict[site_i][28].extend(don_sw_ids)
hbtot += hb_sw.shape[0]
self.hsa_dict[site_i][20].append(hbtot)
if angular_structure:
wat_nbrs = self.wat_oxygen_atom_ids[np.where(
(distance_matrix[0, :][self.wat_oxygen_atom_ids] <=
r_theta_cutoff**2) & (distance_matrix[0, :][
self.wat_oxygen_atom_ids] > 0.0))]
angles = self.water_nbr_orientations(
trj, wat_O, wat_nbrs)
dist = np.sqrt(distance_matrix[0, wat_nbrs])
self.angular_st_distribution[site_i].extend(
zip(dist, angles))
if entropy:
# save coordinates of hsa region waters in current frame
for index, wat_O in enumerate(
self.hsa_region_O_ids[frame_i - self.start_frame]):
flat_id = self.hsa_region_flat_ids[frame_i -
self.start_frame][index]
index_pairs = list(
zip(list(range(wat_O, wat_O + 3)),
list(range(flat_id, flat_id + 3))))
for index_pair in index_pairs:
self.hsa_region_water_coords[index_pair[1], :] += coords[
0, index_pair[0], :]
def _process_frame(self, trj, frame_i, energy, hbonds, entropy):
nbr_cutoff_sq = 3.5**2
trj.xyz *= 10.0
coords = trj.xyz
periodic_box = md.utils.in_units_of(trj.unitcell_lengths, "nanometers",
"angstroms")
waters = []
calc.assign_voxels(trj.xyz, self.dims, self.gridmax, self.origin,
waters, self.wat_oxygen_atom_ids)
for wat in waters:
self.voxeldata[wat[0], 4] += 1
if energy or hbonds:
e_lj_array, e_elec_array = np.copy(self.acoeff), np.copy(
self.chg_product)
distance_matrix = np.zeros(
(self.water_sites, self.all_atom_ids.shape[0]))
calc.get_pairwise_distances(wat, self.all_atom_ids, coords,
periodic_box, distance_matrix)
wat_nbrs = self.wat_oxygen_atom_ids[np.where(
(distance_matrix[0, :][self.wat_oxygen_atom_ids] <=
nbr_cutoff_sq)
& (distance_matrix[0, :][self.wat_oxygen_atom_ids] > 0.0))]
self.voxeldata[wat[0], 19] += wat_nbrs.shape[0]
calc.calculate_energy(wat[1], distance_matrix, e_elec_array,
e_lj_array, self.bcoeff)
if self.prot_atom_ids.shape[0] != 0:
self.voxeldata[wat[0], 13] += np.sum(
e_lj_array[:, self.prot_atom_ids])
self.voxeldata[wat[0], 13] += np.sum(
e_elec_array[:, self.prot_atom_ids])
self.voxeldata[wat[0], 15] += np.sum(
e_lj_array[:,
self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(
e_lj_array[:, wat[1] + self.water_sites:])
self.voxeldata[wat[0], 15] += np.sum(
e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]
) + np.sum(e_elec_array[:, wat[1] + self.water_sites:])
e_nbr_list = [
np.sum(e_lj_array[:, wat_nbrs + i] +
e_elec_array[:, wat_nbrs + i])
for i in xrange(self.water_sites)
]
self.voxeldata[wat[0], 17] += np.sum(e_nbr_list)
"""
###DEBUG START###
elj_sw = np.sum(e_lj_array[:, :self.wat_oxygen_atom_ids[0]])
eelec_sw = np.sum(e_elec_array[:, :self.wat_oxygen_atom_ids[0]])
elj_ww = np.sum(e_lj_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(e_lj_array[:, wat[1] + 1:])
eelec_ww = np.sum(e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat[1]]) + np.sum(e_elec_array[:, wat[1] + self.water_sites:])
e_nbr_list = [np.sum(e_lj_array[:, wat_nbrs + i] + e_elec_array[:, wat_nbrs + i]) for i in xrange(self.water_sites)]
enbr = np.sum(e_nbr_list)
print "Calc: ", elj_sw, eelec_sw, elj_ww, eelec_ww, enbr
distance_matrix = np.sqrt(distance_matrix)
energy_lj, energy_elec = self.calculate_energy(distance_matrix)
test_1 = np.sum(energy_lj[:self.wat_oxygen_atom_ids[0]:])
test_2 = np.sum(energy_elec[:, self.non_water_atom_ids])
test_3 = np.nansum(energy_lj[self.wat_oxygen_atom_ids[0]:])
test_4 = np.sum(energy_elec[:, self.wat_atom_ids[0]:wat[1]]) + np.sum(energy_elec[:, wat[1] + self.water_sites:])
test_5 = 0.0
test_5 += np.sum(energy_lj[self.wat_oxygen_atom_ids[0]:][(wat_nbrs - self.wat_oxygen_atom_ids[0]) / self.water_sites])
for i in range(self.water_sites):
test_5 += np.sum(energy_elec[:, wat_nbrs + i])
print "Ref: ", test_1, test_2, test_3, test_4, test_5
###DEBUG END###
"""
# H-bond calculations
if hbonds:
prot_nbrs_all = self.non_water_atom_ids[np.where(
distance_matrix[0, :][
self.non_water_atom_ids] <= nbr_cutoff_sq)]
prot_nbrs_hb = prot_nbrs_all[np.where(
self.prot_hb_types[prot_nbrs_all] != 0)]
if wat_nbrs.shape[0] > 0:
hb_ww = self.calculate_hydrogen_bonds(
trj, wat[1], wat_nbrs)
acc_ww = hb_ww[:, 0][np.where(
hb_ww[:, 0] == wat[1])].shape[0]
don_ww = hb_ww.shape[0] - acc_ww
self.voxeldata[wat[0], 25] += hb_ww.shape[0]
self.voxeldata[wat[0], 31] += don_ww
self.voxeldata[wat[0], 33] += acc_ww
if wat_nbrs.shape[0] != 0 and hb_ww.shape[0] != 0:
self.voxeldata[
wat[0],
21] += wat_nbrs.shape[0] / hb_ww.shape[0]
if prot_nbrs_hb.shape[0] > 0:
hb_sw = self.calculate_hydrogen_bonds(
trj, wat[1], prot_nbrs_hb, water_water=False)
acc_sw = hb_sw[:, 0][np.where(
hb_sw[:, 0] == wat[1])].shape[0]
don_sw = hb_sw.shape[0] - acc_sw
self.voxeldata[wat[0], 23] += hb_sw.shape[0]
self.voxeldata[wat[0], 27] += don_sw
self.voxeldata[wat[0], 29] += acc_sw
if entropy:
self.calculate_euler_angles(wat, coords[0, :, :])
def process_chunk(self, begin_chunk, chunk_size, topology, energy, hbonds, entropy):
site_waters_copy = list(self.site_waters)
nbr_cutoff_sq = 3.5 ** 2
with md.open(self.trajectory) as f:
f.seek(begin_chunk)
trj = f.read_as_traj(topology, n_frames=chunk_size, stride=1)
#md.utils.in_units_of(trj.xyz, "nanometers", "angstroms", inplace=True)
trj.xyz *= 10.0
trj.unitcell_lengths *= 10.0
pbc = trj.unitcell_lengths
for frame in range(trj.n_frames):
frame_i = frame + (begin_chunk * chunk_size)
coords = trj.xyz[frame, :, :].reshape(1, trj.xyz.shape[1], trj.xyz.shape[2])
oxygen_pos = coords[0, self.wat_oxygen_atom_ids, :]
for site_i in range(self.hsa_data.shape[0]):
wat_O = None
if self.is_site_waters_populated:
if len(site_waters_copy[site_i]) != 0:
if site_waters_copy[site_i][0][0] == frame_i:
wat_O = site_waters_copy[site_i].pop(0)[1]
index = int(self.hsa_data[site_i, 4]) * 3
index_pairs = zip(xrange(wat_O, wat_O + 3), xrange(index, index + 3))
for index_pair in index_pairs:
self.hsa_dict[site_i][-1][index_pair[1]] += coords[0, index_pair[0], :]
self.hsa_data[site_i, 4] += 1
else:
cluster_search_space = NeighborSearch(oxygen_pos, 1.0)
cluster_center_coords = (self.hsa_data[site_i, 1], self.hsa_data[site_i, 2], self.hsa_data[site_i, 3])
nbr_indices = cluster_search_space.query_nbrs_single_point(cluster_center_coords)
cluster_wat_oxygens = [self.wat_oxygen_atom_ids[nbr_index] for nbr_index in nbr_indices]
if len(cluster_wat_oxygens) != 0:
wat_O = cluster_wat_oxygens[0]
self.site_waters[site_i].append((frame_i, wat_O))
index = int(self.hsa_data[site_i, 4]) * 3
index_pairs = zip(xrange(wat_O, wat_O + 3), xrange(index, index + 3))
for index_pair in index_pairs:
self.hsa_dict[site_i][-1][index_pair[1]] += coords[0, index_pair[0], :]
self.hsa_data[site_i, 4] += 1
if wat_O is not None and (energy or hbonds):
distance_matrix = np.zeros((self.water_sites, self.all_atom_ids.shape[0]), np.float_)
calc.get_pairwise_distances(np.asarray([site_i, wat_O]), self.all_atom_ids, coords, pbc, distance_matrix)
wat_nbrs = self.wat_oxygen_atom_ids[np.where(
(distance_matrix[0, :][self.wat_oxygen_atom_ids] <= nbr_cutoff_sq) & (
distance_matrix[0, :][self.wat_oxygen_atom_ids] > 0.0))]
self.hsa_dict[site_i][17].append(wat_nbrs.shape[0])
if energy:
e_lj_array, e_elec_array = np.copy(self.acoeff), np.copy(self.chg_product)
calc.calculate_energy(wat_O, distance_matrix, e_elec_array, e_lj_array, self.bcoeff)
e_lj_sw = np.sum(e_lj_array[:, :self.wat_oxygen_atom_ids[0]])
e_elec_sw = np.sum(e_elec_array[:, :self.wat_oxygen_atom_ids[0]])
e_lj_ww = np.sum(
e_lj_array[:, self.wat_oxygen_atom_ids[0]:wat_O]) + np.sum(e_lj_array[:, wat_O + self.water_sites:])
e_elec_ww = np.sum(
e_elec_array[:, self.wat_oxygen_atom_ids[0]:wat_O]) + np.sum(
e_elec_array[:, wat_O + self.water_sites:])
e_nbr_list = [np.sum(e_lj_array[:, nbr:nbr + self.water_sites] + e_elec_array[:, nbr:nbr + self.water_sites]) for nbr in wat_nbrs]
#e_nbr_list = [np.sum(e_lj_array[:, wat_nbrs + i] + e_elec_array[:, wat_nbrs + i]) for i in
# xrange(self.water_sites)]
self.hsa_dict[site_i][7].append(e_lj_sw)
self.hsa_dict[site_i][8].append(e_elec_sw)
self.hsa_dict[site_i][10].append(e_lj_ww)
self.hsa_dict[site_i][11].append(e_elec_ww)
self.hsa_dict[site_i][6].append(e_lj_sw + e_elec_sw)
self.hsa_dict[site_i][9].append(e_lj_ww + e_elec_ww)
self.hsa_dict[site_i][12].append(e_lj_sw + e_elec_sw + e_lj_ww + e_elec_ww)
self.hsa_dict[site_i][13].extend(e_nbr_list) # print(e_lj_sw/2.0)
if hbonds:
prot_nbrs_all = self.non_water_atom_ids[
np.where(distance_matrix[0, :][self.non_water_atom_ids] <= nbr_cutoff_sq)]
prot_nbrs_hb = prot_nbrs_all[np.where(self.prot_hb_types[prot_nbrs_all] != 0)]
#print wat_O, wat_nbrs
if wat_nbrs.shape[0] + prot_nbrs_hb.shape[0] > 0:
hb_ww, hb_sw = self.calculate_hydrogen_bonds(trj, wat_O, wat_nbrs, prot_nbrs_hb)
#print wat_nbrs, hb_ww
acc_ww = hb_ww[:, 0][np.where(hb_ww[:, 0] == wat_O)].shape[0]
don_ww = hb_ww.shape[0] - acc_ww
acc_sw = hb_sw[:, 0][np.where(hb_sw[:, 0] == wat_O)].shape[0]
don_sw = hb_sw.shape[0] - acc_sw
don_sw_ids = hb_sw[:, 1][np.where(hb_sw[:, 0] == wat_O)]
acc_sw_ids = hb_sw[:, 0][np.where(hb_sw[:, 0] != wat_O)]
self.hsa_dict[site_i][18].append(hb_ww.shape[0])
self.hsa_dict[site_i][19].append(hb_sw.shape[0])
self.hsa_dict[site_i][20].append(hb_ww.shape[0] + hb_sw.shape[0])
self.hsa_dict[site_i][23].append(acc_ww)
self.hsa_dict[site_i][24].append(don_ww)
self.hsa_dict[site_i][25].append(acc_sw)
self.hsa_dict[site_i][26].append(don_sw)
self.hsa_dict[site_i][27].extend(acc_sw_ids)
self.hsa_dict[site_i][28].extend(don_sw_ids)
if wat_nbrs.shape[0] != 0 and hb_ww.shape[0] != 0:
self.hsa_dict[site_i][21].append(
hb_ww.shape[0] / wat_nbrs.shape[0])
if entropy:
# save coordinates of hsa region waters in current frame
for index, wat_O in enumerate(self.hsa_region_O_ids[frame_i - self.start_frame]):
flat_id = self.hsa_region_flat_ids[frame_i - self.start_frame][index]
index_pairs = zip(xrange(wat_O, wat_O + 3), xrange(flat_id, flat_id + 3))
for index_pair in index_pairs:
self.hsa_region_water_coords[index_pair[1], :] += coords[0, index_pair[0], :]