Skip to content

Commit

Permalink
modified weighted density to inclute top_n arg
Browse files Browse the repository at this point in the history
  • Loading branch information
@rsexton2
rsexton2 committed Feb 19, 2024
1 parent e5128c4 commit c0ad2cd
Showing 1 changed file with 63 additions and 42 deletions.
105 changes: 63 additions & 42 deletions basicrta/weighted_density.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,54 +9,53 @@


class MapKinetics(object):
def __init__(self, gibbs, contacts, n=None):
self.gibbs, self.N = gibbs, n
def __init__(self, gibbs, contacts):
self.gibbs= gibbs
self.cutoff = float(contacts.split('/')[-1].strip('.pkl').
split('_')[-1])
self.write_sel = None
with open(contacts, 'rb') as f:
self.contacts = pickle.load(f)

metadata = self.contacts.dtype.metadata
self.ag1, self.ag2 = metadata['ag1'], metadata['ag2']
self.ag1 = metadata['ag1']
self.ag2 = metadata['ag2']
self.u = metadata['u']

self.dataname = f'{self.gibbs.residue}/den_write_data_all.npy'
self.topname = f'{self.gibbs.residue}/reduced.pdb'
self.fulltraj = f'{self.gibbs.residue}/chol_traj_all.xtc'


def _create_data(self):
contacts = self.contacts
resid = int(self.gibbs.residue[1:])
ncomp = self.gibbs.processed_results.ncomp

times = np.array(contacts[contacts[:, 0] == resid][:, 3])
trajtimes = np.array(contacts[contacts[:, 0] == resid][:, 2])
lipinds = np.array(contacts[contacts[:, 0] == resid][:, 1])
dt = self.u.trajectory.ts.dt/1000 #nanoseconds
contacts = self.contacts
resid = int(self.gibbs.residue[1:])
ncomp = self.gibbs.processed_results.ncomp

indicators = self.gibbs.processed_results.indicator
times = np.array(contacts[contacts[:, 0] == resid][:, 3])
trajtimes = np.array(contacts[contacts[:, 0] == resid][:, 2])
lipinds = np.array(contacts[contacts[:, 0] == resid][:, 1])
dt = self.u.trajectory.ts.dt/1000 # covert to nanoseconds

bframes, eframes = get_start_stop_frames(trajtimes, times, dt)
tmp = [np.arange(b, e) for b, e in zip(bframes, eframes)]
tmpL = [np.ones_like(np.arange(b, e)) * l for b, e, l in
zip(bframes, eframes, lipinds)]
tmpI = [indic * np.ones((len(np.arange(b, e)), ncomp))
for b, e, indic in zip(bframes, eframes, indicators)]
indicators = self.gibbs.processed_results.indicator

write_frames = np.concatenate([*tmp]).astype(int)
write_linds = np.concatenate([*tmpL]).astype(int)
write_indics = np.concatenate([*tmpI])
bframes, eframes = get_start_stop_frames(trajtimes, times, dt)
tmp = [np.arange(b, e) for b, e in zip(bframes, eframes)]
tmpL = [np.ones_like(np.arange(b, e)) * l for b, e, l in
zip(bframes, eframes, lipinds)]
tmpI = [indic * np.ones((len(np.arange(b, e)), ncomp))
for b, e, indic in zip(bframes, eframes, indicators)]

darray = np.zeros((len(write_frames), ncomp + 2))
darray[:, 0], darray[:, 1], darray[:, 2:] = (write_frames,
write_linds,
write_indics)
np.save(self.dataname, darray)
write_frames = np.concatenate([*tmp]).astype(int)
write_linds = np.concatenate([*tmpL]).astype(int)
write_indics = np.concatenate([*tmpI])

darray = np.zeros((len(write_frames), ncomp + 2))
darray[:, 0], darray[:, 1], darray[:, 2:] = (write_frames,
write_linds,
write_indics)
np.save(self.dataname, darray)

def create_traj(self):
def create_traj(self, top_n=None):
write_ag = self.ag1.atoms + self.ag2.residues[0].atoms
write_ag.atoms.write(self.topname)

Expand Down Expand Up @@ -86,8 +85,7 @@ def create_traj(self):
W.write(self.ag1 +
self.ag2.select_atoms(f'resid {wl[i]}').atoms)


def weighted_densities(self, step=1):
def weighted_densities(self, step=1, top_n=None):
if not os.path.exists(self.fulltraj):
self.create_traj()

Expand All @@ -103,25 +101,48 @@ def weighted_densities(self, step=1):
u_red = mda.Universe(self.topname, self.fulltraj)
chol_red = u_red.select_atoms('not protein')

sortinds = [wi[:, i].argsort()[::-1] for i in
range(self.gibbs.processed_results.ncomp)]

for k in range(self.gibbs.processed_results.ncomp):
frames = np.where(wi[sortinds[k], k] > 0)[0][::step]
tmpwi = wi[frames, k]
d = WDensityAnalysis(chol_red, tmpwi,
if top_n is None:
from basicrta.pwdensity import WDensityAnalysis
u_red = mda.Universe(self.topname, self.fulltraj)
chol_red = u_red.select_atoms('not protein')
d = WDensityAnalysis(chol_red, wi,
gridcenter=u_red.select_atoms(f'protein and '
f'resid {resid}')
.center_of_geometry(), xdim=40, ydim=40,
zdim=40)
d.run(verbose=True, frames=sortinds[k][frames])
if self.N is not None:
outname = f'{self.gibbs.residue}/wcomp{k}_top{self.N}.dx'
d.run(verbose=True, step=step)
if step > 1:
outnames = [f'{self.gibbs.residue}/wcomp{k}_all_step{step}.dx'
for k in range(self.gibbs.processed_results.ncomp)]
else:
outname = f'{self.gibbs.residue}/wcomp{k}_all.dx'
outnames = [f'{self.gibbs.residue}/wcomp{k}_all.dx' for k in
range(self.gibbs.processed_results.ncomp)]

d.results.density.export(outname)
[den.export(outnames[k]) for k, den in
enumerate(d.results.densities)]

else:
from basicrta.wdensity import WDensityAnalysis
sortinds = [wi[:, i].argsort()[::-1] for i in
range(self.gibbs.processed_results.ncomp)]

for k in range(self.gibbs.processed_results.ncomp):
frames = np.where(wi[sortinds[k], k] > 0)[0][::step]
tmpwi = wi[frames, k]
d = WDensityAnalysis(chol_red, tmpwi,
gridcenter=u_red.select_atoms(f'protein '
f'and resid '
f'{resid}')
.center_of_geometry(), xdim=40, ydim=40,
zdim=40)
d.run(verbose=True, frames=sortinds[k][frames][::step])
if step > 1:
outname = (f'{self.gibbs.residue}/wcomp{k}_top{top_n}_step'
f'{step}.dx')
else:
outname = f'{self.gibbs.residue}/wcomp{k}_top{top_n}.dx'

d.results.density.export(outname)


if __name__ == "__main__":
Expand Down

0 comments on commit c0ad2cd

Please sign in to comment.