def distmatrix(uni,idxs,chosenselections,co,mindist=False, mindist_mode=None, type1char='p',type2char='n'):
numframes = uni.trajectory.numframes
final_percmat = np.zeros((len(chosenselections),len(chosenselections)))
log.info("Distance matrix will be %dx%d (%d elements)" % (len(idxs),len(idxs),len(idxs)**2))
a=1
distmats=[]
if mindist:
P = []
N = []
Pidxs = []
Nidxs = []
Psizes = []
Nsizes = []
for i in range(len(idxs)):
if idxs[i][3][-1] == type1char:
P.append(chosenselections[i])
Pidxs.append(idxs[i])
Psizes.append(len(chosenselections[i]))
elif idxs[i][3][-1] == type2char:
N.append(chosenselections[i])
Nidxs.append(idxs[i])
Nsizes.append(len(chosenselections[i]))
else:
raise
Nsizes = np.array(Nsizes, dtype=np.int)
Psizes = np.array(Psizes, dtype=np.int)
if mindist_mode == "diff":
sets = [(P,N)]
sets_idxs = [(Pidxs,Nidxs)]
sets_sizes = [(Psizes,Nsizes)]
elif mindist_mode == "same":
sets = [(P,P),(N,N)]
sets_idxs = [(Pidxs,Pidxs),(Nidxs,Nidxs)]
sets_sizes = [(Psizes,Psizes),(Nsizes,Nsizes)]
elif mindist_mode == "both":
sets = [(chosenselections, chosenselections)]
sets_idxs = [(idxs, idxs)]
sizes = [len(s) for s in chosenselections]
sets_sizes = [(sizes,sizes)]
else: raise
percmats = []
coords = []
for s in sets:
coords.append([[],[]])
for ts in uni.trajectory:
sys.stdout.write( "Caching coordinates: frame %d / %d (%3.1f%%)\r" % (a,numframes,float(a)/float(numframes)*100.0) )
sys.stdout.flush()
a+=1
for si,s in enumerate(sets):
if s[0] == s[1]: # triangular case
log.info("Caching coordinates...")
for group in s[0]:
coords[si][0].append(group.coordinates())
coords[si][1].append(group.coordinates())
else: # square case
log.info("Caching coordinates...")
for group in s[0]:
coords[si][0].append(group.coordinates())
for group in s[1]:
coords[si][1].append(group.coordinates())
for si,s in enumerate(sets): # recover the final matrix
if s[0] == s[1]:
this_coords = np.array(np.concatenate(coords[si][0]),dtype=np.float64)
inner_loop = LoopDistances(this_coords, this_coords, co)
percmats.append(inner_loop.run_triangular_mindist(sets_sizes[si][0]))
else:
this_coords1 = np.array(np.concatenate(coords[si][0]),dtype=np.float64)
this_coords2 = np.array(np.concatenate(coords[si][1]),dtype=np.float64)
inner_loop = LoopDistances(this_coords1, this_coords2, co)
percmats.append( inner_loop.run_square_mindist(sets_sizes[si][0], sets_sizes[si][1]))
for si,s in enumerate(sets): # recover the final matrix
Pidxs = sets_idxs[si][0]
Nidxs = sets_idxs[si][1]
if s[0] == s[1]: # triangular case
for j in range(len(s[0])):
for k in range(0,j):
final_percmat[idxs.index(Pidxs[j]), idxs.index(Pidxs[k])] = percmats[si][j,k]
final_percmat[idxs.index(Pidxs[k]), idxs.index(Pidxs[j])] = percmats[si][j,k]
else: # square case
for j in range(len(s[0])):
for k in range(len(s[1])):
final_percmat[idxs.index(Pidxs[j]), idxs.index(Nidxs[k])] = percmats[si][j,k]
final_percmat[idxs.index(Nidxs[k]), idxs.index(Pidxs[j])] = percmats[si][j,k]
final_percmat = np.array(final_percmat, dtype=np.float)/numframes*100.0
else:
all_coms = []
for ts in uni.trajectory:
sys.stdout.write( "now analyzing: frame %d / %d (%3.1f%%)\r" % (a,numframes,float(a)/float(numframes)*100.0) )
sys.stdout.flush()
a+=1
distmat = np.zeros((len(chosenselections),len(chosenselections)))
coms = np.zeros([len(chosenselections),3])
for j in range(len(chosenselections)):
coms[j,:] = chosenselections[j].centerOfMass()
all_coms.append(coms)
all_coms = np.concatenate(all_coms)
inner_loop = LoopDistances(all_coms, all_coms, co)
percmat = inner_loop.run_triangular_distmatrix(coms.shape[0])
distmats = []
final_percmat = np.array(percmat, dtype=np.float)/numframes*100.0
return (final_percmat,distmats)
