def getFirstShellWater(TrajName, LDUpperCuts, FirstShellCut = 6.0, Prefix = 'LD_wm', target_dir = os.getcwd(), frame_start = 0, frame_stop = 100, stepfreq = 100): ''' Calculates the local densities and first shell waters for each monomer for several different cutoffs. Also generates the monomer local density distribution. ''' # Get un-mapped trajectory Trj = getPickle(TrajName) ## Parse trajectory and make frame stepper init = Trj[0] BoxL = Trj.FrameData['BoxL']; BoxL = BoxL[0] FrameRange = range(frame_start, frame_stop, stepfreq) NFrames = len(FrameRange) AtomTypes = Trj.AtomTypes ## Make Local density count array NCuts = len(LDUpperCuts) N_mon = len(np.where(AtomTypes == 3)[0]) N_water = len(np.where(AtomTypes == 2)[0]) LD_mon = np.zeros([N_mon*NFrames,NCuts], np.float64) FirstShellWaters = np.zeros(N_mon*NFrames, np.float64) ## Precompute LD coefficients coeff = np.zeros([NCuts, 4], np.float64) for i, cut in enumerate(LDUpperCuts): coeff[i] = calcCoeff(cut) pb = sim.utility.ProgressBar(Text = 'Processing frame by frame...', Steps = NFrames) # Frame stepping count = 0 for frame in FrameRange: Pos = Trj[frame] # Call fortran subroutine for per-frame computation fsw = np.zeros(N_mon, np.float64) ld = np.zeros([N_mon, NCuts], np.float64) (ld, fsw) = fortlib.find_firstshellwaters(ld = ld, fsw = fsw, pos = Pos, atomtypes = AtomTypes, n_water = N_water, lduppercuts = LDUpperCuts, firstshellcut = FirstShellCut, coeff = coeff, boxl = BoxL, atomtypedefs = atomtypedefs) ## Average #LD_mon += ld #FirstShellWaters += fsw ## Per-frame FirstShellWaters[count:count+N_mon] = fsw for j in range(NCuts): LD_mon[count:count+N_mon,j] = ld[:,j] count += N_mon pb.Update(frame/stepfreq) #LD_mon /= NFrames #FirstShellWaters /= NFrames # pickling data pickleName = os.path.join(target_dir, Prefix + '.pickle') pickle.dump((LDUpperCuts, LD_mon, FirstShellWaters), open(pickleName, 'w'))
