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)