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], :]