def main(): args = parse() input = args.input n_jobs = args.n_cpu a, b, c = args.cell_vectors thresholds = args.thresholds CURRENT_PATH = path.dirname(path.realpath(__file__)) DATA_PATH = path.normpath(path.join(CURRENT_PATH, path.dirname(input))) base = path.splitext(path.basename(input))[0] # Initialize universe (time step 0.5 fs) u = mda.Universe(input, dt=5e-4) u.add_TopologyAttr('charges') u.dimensions = np.array([a, b, c, 90, 90, 90]) # Split trajectory into blocks blocks = trj2blocks.get_blocks(u, n_jobs) print('Analyzing...') results = Parallel(n_jobs=n_jobs)(delayed(hbonds)(u, thresholds, block) for block in blocks) # Average autocorrelation functions of (possibly) different length results = tolerant.mean(results) t = np.arange(len(results)) # Save results as .csv df = pd.DataFrame({'time': t, 'autocorr': results}) df.to_csv('%s/%s_%s.hblifetime' % (DATA_PATH, base, thresholds[1]), index=False) print('\nProgram executed in %.0f seconds' % (time() - t0))
def diffusion(u, t, block): '''Computes mean squared displacement of water molecules within specified layers parallel to a surface. Args: u: MDAnalysis Universe object containing trajectory. t: Thresholds specifying layer boundaries. block: Range of frames composing block. Returns: Parallel (xy) mean squared displacement. ''' # Select oxygen atoms oxygen = u.select_atoms('name O') msd = [] # Initialize 3D position array (frame, atom, dimension) position_array = np.zeros((len(block), oxygen.n_atoms, 3)) for i, ts in enumerate(u.trajectory[block.start:block.stop]): print('Processing blocks %.1f%%' % (100*i/len(block)), end='\r') # Store positions of all oxygens at current frame position_array[i, :, :] = oxygen.positions # Loop over oxygen atoms for k in range(oxygen.n_atoms): # Get atoms in region z = position_array[:, k, 2] z_bool = (z > t[0])*(z < t[1])+(z > t[2])*(z < t[3]) # Get intervals in which atom remains continuously in region contiguous_region = find_objects(label(z_bool)[0]) # Compute mean squared displacement of atom for each interval for reg in contiguous_region: msd.append(tidynamics.msd( position_array[reg[0].start:reg[0].stop, k, :2])) # Return average of all MSDs return tolerant.mean(msd)
def vacf(u, u_vel, t, species, block): '''Computes velocity-velocity autocorrelation function for given species. Args: u: MDAnalysis Universe object containing trajectory (positions). u_vel: MDAnalysis Universe object containing trajectory (velocities). t: Thresholds specifying layer boundaries. species: For which species (elements) to compute autocorrelation. block: Range of frames composing block. Returns: Velocity-velocity autocorrelation function. ''' # Which species to include in velocity autocorrelation if species == 'HH': selection = 'name H' elif species == 'OO': selection = 'name O' elif species == 'OH': selection = 'name O or name H' # Get position and velocity trajectories of selection ag_trj = u.select_atoms(selection) ag_vel = u_vel.select_atoms(selection) n_atoms = ag_vel.n_atoms # Initialize arrays for velocities and atom heights velocity_array = np.zeros((len(block), n_atoms, 3)) height_array = np.zeros((len(block), n_atoms)) corr = [] for i, ts in enumerate(u_vel.trajectory[block.start:block.stop]): print('Processing velocities %.1f%%' % (100 * i / len(block)), end='\r') # Update positions u.trajectory[ts.frame] # Loop over atoms for j, vel in enumerate(ag_vel): # Store heights and velocities of atom height_array[i, j] = ag_trj.positions[j, 2] velocity_array[i, j, :] = vel.position # Loop over atoms for k in range(n_atoms): # Get frames where atoms in specified region z = height_array[:, k] z_bool = (z > t[0]) * (z < t[1]) + (z > t[2]) * (z < t[3]) # Select intervals where remain remains continuously in region contiguous_region = find_objects(label(z_bool)[0]) # Compute dipole autocorrelation for all intervals and starting times for reg in contiguous_region: corr.append( tidynamics.acf(velocity_array[reg[0].start:reg[0].stop, k, :])) return tolerant.mean(corr)
def rotation(u, a, t, block): '''Computes water orientational (dipole) relaxation. Args: u: MDAnalysis Universe object containing trajectory. a: Lateral lattice vector assuming equal lateral dimensions t: Thresholds specifying layer boundaries. block: Range of frames composing block. Returns: Water orientational (dipole) autocorrelation function. ''' size = len(block) wor = [] # Select oxygen and hydrogen atoms oxygen = u.select_atoms('name O') hydrogen = u.select_atoms('name H') # Initialize OH distance array, dipole array and height array rOH = np.zeros((oxygen.n_atoms, hydrogen.n_atoms)) dipole_array = np.zeros((len(block), oxygen.n_atoms, 3)) height_array = np.zeros((len(block), oxygen.n_atoms)) for i, ts in enumerate(u.trajectory[block.start:block.stop]): print('Processing blocks %.1f%%' % (100 * i / size), end='\r') # Wrap atoms to box oxygen.wrap(box=u.dimensions) hydrogen.wrap(box=u.dimensions) # Compute OH distance array distance_array(oxygen.positions, hydrogen.positions, box=u.dimensions, result=rOH) # Store heights of all oxygens height_array[i, :] = oxygen.positions[:, 2] # Loop over all oxygen positions for j, pos in enumerate(oxygen.positions): # Get bound hydrogens and combine molecules broken over PBCs hbound = hydrogen.positions[(rOH < 1.2)[j]] hbound[hbound - pos < -1.2] += a hbound[hbound - pos > 1.2] -= a # Compute dipole vectors dip = np.mean(hbound, axis=0) - pos dipole_array[i, j, :] = dip / np.linalg.norm(dip) # Loop over all oxygens for k in range(oxygen.n_atoms): # Get frames where oxygen in specified region z = height_array[:, k] z_bool = (z > t[0]) * (z < t[1]) + (z > t[2]) * (z < t[3]) # Select intervals where oxygen remains continuously in region contiguous_region = find_objects(label(z_bool)[0]) # Compute dipole autocorrelation for all intervals and starting times for reg in contiguous_region: corr = tidynamics.acf(dipole_array[reg[0].start:reg[0].stop, k, :]) wor.append(lg2(corr)) # Average autocorrelation functions of (possibly) different length return tolerant.mean(wor)