Skip to content

Commit

Permalink
modified contacts to check error
Browse files Browse the repository at this point in the history
  • Loading branch information
@rsexton2
rsexton2 committed Feb 14, 2024
1 parent b03fcfb commit 6b078d1
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 37 deletions.
13 changes: 7 additions & 6 deletions basicrta/contacts.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,8 @@ def run(self):
i, aslice in enumerate(sliced_frames)]

lens = []
with Pool(nproc, initializer=tqdm.set_lock, initargs=(Lock(),)) as p:
with (Pool(self.nproc, initializer=tqdm.set_lock, initargs=(Lock(),))
as p):
for alen in tqdm(p.istarmap(self._run_contacts, input_list),
total=self.nslices, position=0,
desc='overall progress'):
Expand Down Expand Up @@ -123,14 +124,14 @@ def run(self):
'contacts file using the "map_name" '
'argument')

self.ts = siground(np.unique(memmap[1:, 4]-memmap[:-1, 4])[1], 1)
self.ts = dtype.metadata['u'].trajectory.dt/1000 # convert to ns
lresids = np.unique(memmap[:, 2])
params = [[res, memmap[memmap[:, 2] == res], i] for i, res in
enumerate(lresids)]
pool = Pool(self.nproc, initializer=tqdm.set_lock, initargs=(Lock(),))

try:
lens = pool.starmap(self._lipswap, params)
lens = pool.istarmap(self._lipswap, params)
except KeyboardInterrupt:
pool.terminate()
pool.close()
Expand All @@ -146,9 +147,9 @@ def run(self):
contact_map.flush()

contact_map.dump(f'contacts_{self.cutoff}.pkl', protocol=5)
os.remove('.tmpmap')
cfiles = glob.glob('.contacts*')
[os.remove(f) for f in cfiles]
# os.remove('.tmpmap')
# cfiles = glob.glob('.contacts*')
# [os.remove(f) for f in cfiles]
print(f'\nSaved contacts to "contacts_{self.cutoff}.npy"')


Expand Down
86 changes: 55 additions & 31 deletions basicrta/weighted_density.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@
from MDAnalysis.lib.util import realpath


class WeightedDensity(object):
def __init__(self, gibbs, contacts, step=1, n=None):
self.gibbs, self.N, self.step = gibbs, n, step
class MapKinetics(object):
def __init__(self, gibbs, contacts, n=None):
self.gibbs, self.N = gibbs, n
self.cutoff = float(contacts.split('/')[-1].strip('.pkl').
split('_')[-1])
self.write_sel = None
Expand All @@ -22,39 +22,55 @@ def __init__(self, gibbs, contacts, step=1, n=None):
traj = realpath(metadata['traj'])
self.ag1, self.ag2 = metadata['ag1'], metadata['ag2']
self.u = mda.Universe(top, traj)
self.dataname = (f'{self.gibbs.residue}/den_write_data_'
f'step{self.step}.npy')

self.topname = (f'{self.gibbs.residue}/reduced.pdb')
self.trajname = (f'{self.gibbs.residue}/chol_traj_step{self.step}.xtc')
if n is not None:
self.dataname = (f'{self.gibbs.residue}/den_write_data_'
f'top{self.N}.npy')
self.trajname = (f'{self.gibbs.residue}/chol_traj_top{self.N}.xtc')
else:
self.dataname = (f'{self.gibbs.residue}/den_write_data_all.npy')
self.trajname = (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
if os.path.exists(f'{self.gibbs.residue}/den_write_data_all.npy') and \
self.N is not None:
tmp = np.load(f'{self.gibbs.residue}/den_write_data_all.npy')
wf, wl, wi = (tmp[:, 0].astype(int), tmp[:, 1].astype(int),
tmp[:, 2:])
sortinds = [wi[:, i].argsort()[::-1][:self.N] for i in
range(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
else:
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 #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):
Expand All @@ -67,15 +83,23 @@ def _create_traj(self):
tmp = np.load(self.dataname)
wf, wl, wi = tmp[:, 0].astype(int), tmp[:, 1].astype(int), tmp[:, 2:]

if self.N is not None:
sortinds = [wi[:, i].argsort()[::-1][:self.N] for i in
range(self.gibbs.processed_results.ncomp)]
else:
sortinds = None

if not os.path.exists(self.trajname):
with mda.Writer(self.trajname, len(write_ag.atoms)) as W:
for i, ts in tqdm(enumerate(self.u.trajectory[wf]),
total=len(wf), desc='writing trajectory'):
for i, ts in tqdm(enumerate(self.u.trajectory[wf[sortinds]]),
total=self.N, desc='writing trajectory'):
W.write(self.ag1.atoms +
self.ag2.residues[wl[i]].atoms)
self.ag2.residues[wl[sortinds][i]].atoms)
else:



def run(self):
def compute_weighted_densities(self):
if not os.path.exists(self.trajname):
self._create_traj()

Expand Down

0 comments on commit 6b078d1

Please sign in to comment.