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)