def number_1st(thickness=0.5):
'''
The number of cations and anions in the first layer at anode.
'''
id_cations = [
atom.id for atom in top.atoms if atom.type in ('CR', 'NA', 'CW')
]
id_anions = [
atom.id for atom in top.atoms if atom.type in ('N3A', 'CZA', 'NZA')
]
frame_list = []
N_cation_list = []
N_anion_list = []
for i in range(args.begin, args.end, args.skip):
frame_list.append(i)
frame = trj.read_frame(i)
sys.stdout.write('\r frame %i' % i)
z_cations = frame.positions[id_cations][:, 2]
z_anions = frame.positions[id_anions][:, 2]
N_cation = 0.5 * erfc(
(anode - thickness - z_cations) / 2**0.5 / args.sigma).sum() * 0.2
N_anion = 0.5 * erfc(
(anode - thickness - z_anions) / 2**0.5 / args.sigma).sum() * 0.2
N_cation_list.append(N_cation)
N_anion_list.append(N_anion)
name_column_dict = {
'frame': frame_list,
'N_cation': N_cation_list,
'N_anion': N_anion_list
}
print_data_to_file(name_column_dict, f'{args.output}-n1st.txt')
def charge_petersen():
_conv = ELEMENTARY_CHARGE / area / NANO**2 * 1000 # convert from charge (e) to charge density (mC/m^2)
top_atom_charges = np.array([atom.charge for atom in top.atoms],
dtype=np.float32)
frame_list = []
q_left_list = []
for i in range(args.begin, args.end, args.skip):
frame_list.append(i)
frame = trj.read_frame(i)
sys.stdout.write('\r frame %i' % i)
z_to_left = frame.positions[:, 2] - elecd_l
ids_ils = np.where((z_to_left > 0.1) & (z_to_left < lz - 0.1))[0]
z_to_left_ils = z_to_left[ids_ils]
q_ils = frame.charges[
ids_ils] if frame.has_charge else top_atom_charges[ids_ils]
q_left = np.sum(q_ils * z_to_left_ils) / lz + \
args.voltage * area / lz * VACUUM_PERMITTIVITY / ELEMENTARY_CHARGE * NANO
q_left_list.append(q_left)
charge_densities = np.array(q_left_list) * _conv
voltages_surface = charge_densities / 1000 * lz * NANO / VACUUM_PERMITTIVITY
print('\n%-8s %10s %10s %10s %10s' %
('n_frame', 'rho_q', 'var_rho_q', 'q_atom', 'V_surface'))
print('%-8i %10.4f %10.4f %10.6f %10.4f' %
(len(q_left_list), charge_densities.mean(), charge_densities.var(),
charge_densities.mean() / _conv / len(ids_cathode),
voltages_surface.mean()))
name_column_dict = {
'frame': frame_list,
'rho_q': charge_densities,
'V_surface': voltages_surface
}
print_data_to_file(name_column_dict, f'{args.output}-surface_charge.txt')
def permittivity():
'''
Calculate the static relative permittivity from the fluctuation of dipole
\eps_r = 1 + (<M^2> - <M>^2) / (3VkT 4\pi\eps_0)
'''
top_atom_charges = np.array([atom.charge for atom in top.atoms],
dtype=np.float32)
frame_list = []
dipoles_scalar: [float] = []
for i in range(args.begin, args.end, args.skip):
frame = trj.read_frame(i)
sys.stdout.write('\r frame %i' % i)
frame_list.append(i)
charges = frame.charges if frame.has_charge else top_atom_charges
dipole = np.sum(frame.positions * np.transpose([charges] * 3), axis=0)
dipoles_scalar.append(np.sqrt(dipole.dot(dipole)))
eps_list = [1 + np.var(dipoles_scalar[: i + 1]) * (ELEMENTARY_CHARGE * NANO) ** 2 \
/ (3 * frame0.cell.volume * NANO ** 3) \
/ (8.314 * 298 / AVOGADRO) \
/ (4 * PI * VACUUM_PERMITTIVITY)
for i in range(len(frame_list))]
print('\nRelative permittivity = ', eps_list[-1])
name_column_dict = {
'frame': frame_list,
'dipole': dipoles_scalar,
'permittivity': eps_list
}
print_data_to_file(name_column_dict, f'{args.output}-permittivity.txt')
def charge_1st(thickness=0.5):
top_atom_charges = np.array([atom.charge for atom in top.atoms],
dtype=np.float32)
frame_list = []
q1st_list = []
for i in range(args.begin, args.end, args.skip):
frame_list.append(i)
frame = trj.read_frame(i)
sys.stdout.write('\r frame %i' % i)
z_to_left = frame.positions[:, 2] - elecd_l
ids_ils = np.where((z_to_left > 0.1) & (z_to_left < lz - 0.1))[0]
z_ils = frame.positions[ids_ils, 2]
q_ils = frame.charges[
ids_ils] if frame.has_charge else top_atom_charges[ids_ils]
q = (0.5 * erfc(
(anode - thickness - z_ils) / 2**0.5 / args.sigma) * q_ils).sum()
q1st_list.append(q)
name_column_dict = {'frame': frame_list, 'q1st': q1st_list}
print_data_to_file(name_column_dict, f'{args.output}-q1st.txt')
def voltage():
charges = np.zeros(n_bin)
top_atom_charges = np.array([atom.charge for atom in top.atoms],
dtype=np.float32)
n_frame = 0
for i in range(args.begin, args.end, args.skip):
n_frame += 1
frame = trj.read_frame(i)
sys.stdout.write('\r frame %i' % i)
positions = frame.positions
if args.reverse:
positions[:, 2] = elecd_l + elecd_r - positions[:, 2]
z_to_left = positions[:, 2] - elecd_l
ids_ils = np.where((z_to_left > 0.1) & (z_to_left < lz - 0.1))[0]
z_ils = positions[ids_ils, 2]
q_ils = frame.charges[
ids_ils] if frame.has_charge else top_atom_charges[ids_ils]
i_bin_ils = ((z_ils - edges[0]) / dz).astype(int)
for i, i_bin in enumerate(i_bin_ils):
charges[i_bin] += q_ils[i]
charges /= n_frame
### add back charges on electrodes
if args.charge != 0:
q_electrode = args.charge * MILLI * area * NANO * NANO / ELEMENTARY_CHARGE
idx_cat = int((cathode - edges[0]) / dz)
idx_ano = int((anode - edges[0]) / dz)
charges[idx_cat] += q_electrode
charges[idx_ano] -= q_electrode
charges_cumulative = np.cumsum(charges)
charges /= area * dz # e/nm^3
e_field = itg.cumtrapz(charges, dx=dz, initial=0) \
* ELEMENTARY_CHARGE / NANO ** 2 / VACUUM_PERMITTIVITY
voltage = -itg.cumtrapz(e_field, dx=dz, initial=0) * NANO
name_column_dict = {
'z': z_array,
'rho_q': charges,
'cum_q': charges_cumulative,
'EF': e_field,
'V': voltage
}
print_data_to_file(name_column_dict, f'{args.output}-voltage.txt')
fig, ax = plt.subplots()
ax.set(xlabel='z (nm)', ylabel='charge density (e/nm$^3$)')
ax.plot(z_array, charges)
ax.plot(z_array, [0] * n_bin, '--')
fig.tight_layout()
fig.savefig(f'{args.output}-charge.png')
fig, ax = plt.subplots()
ax.set(xlabel='z (nm)', ylabel='cumulative charges (e)')
ax.plot(z_array, charges_cumulative)
ax.plot(z_array, [0] * n_bin, '--')
fig.tight_layout()
fig.savefig(f'{args.output}-charge_cumulative.png')
fig, ax = plt.subplots()
ax.set(xlabel='z (nm)', ylabel='electric field (V/m)')
ax.plot(z_array, e_field)
ax.plot(z_array, [0] * n_bin, '--')
fig.tight_layout()
fig.savefig(f'{args.output}-efield.png')
fig, ax = plt.subplots()
ax.set(xlabel='z (nm)', ylabel='voltage (V)')
ax.plot(z_array, voltage)
ax.plot(z_array, [0] * n_bin, '--')
fig.tight_layout()
fig.savefig(f'{args.output}-voltage.png')
def dipole():
'''
Calculate the distribution of molecular dipoles
'''
top_atom_charges = np.array([atom.charge for atom in top.atoms],
dtype=np.float32)
name_z_dipoles_dict = {
} # {'ring': [[array, array, ...], [array, array, ...] , ...]}
name_z_dipole_dict = {} # {'ring': array([array, array , ...])}
n_frame = 0
for i in range(args.begin, args.end, args.skip):
n_frame += 1
frame = trj.read_frame(i)
sys.stdout.write('\r frame %i' % i)
positions = frame.positions
if args.reverse:
positions[:, 2] = elecd_l + elecd_r - positions[:, 2]
for mol in top.molecules:
if mol.name not in mol_names:
continue
section = ini['molecule.%s' % (mol.name)]
dipoles = section['dipoles'].split(';')
for d in dipoles:
name, atoms = [x.strip() for x in d.split(':')]
if name not in name_z_dipoles_dict:
name_z_dipoles_dict[name] = [[] for z in z_array]
atoms = _get_atoms(mol, atoms.split())
ids = [atom.id for atom in atoms]
charges = frame.charges[
ids] if frame.has_charge else top_atom_charges[ids]
rel_charges = charges - charges.mean()
poss = positions[ids]
dipole = np.sum(poss * np.transpose([rel_charges] * 3), axis=0)
com = _get_weighted_center(poss, [atom.mass for atom in atoms])
idx = int((com[2] - edges[0]) / dz)
name_z_dipoles_dict[name][idx].append(dipole)
for name, z_dipoles in name_z_dipoles_dict.items():
name_z_dipole_dict[name] = np.array([
np.mean(dipoles, axis=0)
if len(dipoles) > n_frame else np.array([0, 0, 0])
for dipoles in z_dipoles
])
name_column_dict = {'z': z_array}
fig, ax = plt.subplots()
ax.set(xlabel='z (nm)', ylabel='mean dipole (e*nm)')
for name, z_dipole in name_z_dipole_dict.items():
ax.plot(z_array,
z_dipole[:, 2],
label='dipoleZ - ' + name,
color='darkred' if name.startswith('dca') else None)
name_column_dict.update({
'dipX-' + name: z_dipole[:, 0],
'dipY-' + name: z_dipole[:, 1],
'dipZ-' + name: z_dipole[:, 2],
})
ax.legend()
fig.tight_layout()
fig.savefig(f'{args.output}-dipole.png')
print_data_to_file(name_column_dict, f'{args.output}-dipole.txt')
def diffusion():
name_atoms_dict = {
} # {'ring': [[atom1, atom2, ...], [atom11, atom12, ...], ...]}
t_list = []
z_dict = {} # {'ring': [[], [], ...]}
residence_zrange_dict = {}
residence_dict = {} # {'ring': [[], [], ...]}
acf_dict = {} # {'ring': []}
for mol in top.molecules:
if mol.name not in mol_names:
continue
diffusions = ini['molecule.%s' % (mol.name)]['diffusions'].split(';')
for diffusion in diffusions:
name, atoms = [x.strip() for x in diffusion.split(':')]
if name not in name_atoms_dict:
name_atoms_dict[name] = []
z_dict[name] = []
residence_dict[name] = []
z_range_str = ini['molecule.%s' %
(mol.name)]['diffusion.%s.residence_zrange' %
name].split()
residence_zrange_dict[name] = [
elecd_l + float(x) for x in z_range_str
]
name_atoms_dict[name].append(_get_atoms(mol, atoms.split()))
z_dict[name].append([])
residence_dict[name].append([])
for i in range(args.begin, args.end, args.skip):
frame = trj.read_frame(i)
sys.stdout.write('\r frame %i' % i)
if frame.time != -1:
t_list.append(frame.time / 1E3) # convert ps to ns
else:
t_list.append(i * args.dt / 1E3)
for name, atoms_list in name_atoms_dict.items():
for k, atoms in enumerate(atoms_list):
com_position = _get_com_position(frame.positions, atoms)
z_dict[name][k].append(com_position[2])
residence = int(
com_position[2] >= residence_zrange_dict[name][0]
and com_position[2] <= residence_zrange_dict[name][1])
residence_dict[name][k].append(residence)
for name, residence_series_list in residence_dict.items():
acf_series_list = []
# for series in residence_series_list:
# acf_series_list.append(_calc_acf(series - np.mean(residence_series_list)))
# acf_dict[name] = np.mean(acf_series_list, axis=0) / np.var(residence_series_list)
for series in residence_series_list:
acf_series_list.append(_calc_acf(series))
acf_dict[name] = np.mean(acf_series_list,
axis=0) / np.mean(residence_series_list)
print('')
t_array = np.array(t_list) - t_list[0]
name_column_dict = {'time': t_array}
for name, z_series_list in z_dict.items():
fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(8, 8))
ax1.set(ylim=[elecd_r - 1.0, elecd_r], ylabel='z (nm)')
ax2.set(ylim=[elecd_l, elecd_l + 1.0],
xlabel='time (ns)',
ylabel='z (nm)')
ax1.spines['bottom'].set_visible(False)
ax2.spines['top'].set_visible(False)
ax1.xaxis.tick_top()
ax2.xaxis.tick_bottom()
for z_series in z_series_list:
ax1.plot(t_array, z_series, linewidth=1)
ax2.plot(t_array, z_series, linewidth=1)
plt.subplots_adjust(hspace=0.1)
fig.tight_layout()
fig.savefig(f'{args.output}-diffusion-{name}.png')
fig, ax = plt.subplots()
ax.set(xlim=[0, t_array[-1] * 0.75],
ylim=[0, 1.2],
xlabel='time (ns)',
ylabel='residence auto correlation')
for name, z_series_list in z_dict.items():
_len = len(acf_dict[name])
ax.plot(t_array[:_len],
acf_dict[name],
label=name,
color='darkred' if name.startswith('dca') else None)
name_column_dict['time'] = t_array[:_len]
name_column_dict['acf-' + name] = acf_dict[name]
ax.plot([0, t_array[-1] * 0.75], [np.exp(-1), np.exp(-1)],
'--',
label='$e^{-1}$')
ax.legend()
fig.tight_layout()
fig.savefig(f'{args.output}-residence.png')
print_data_to_file(name_column_dict, f'{args.output}-residence.txt')
def distribution():
mol_dists_atom = {}
mol_dists_com = {}
mol_dists_angle = {}
z_atom_name_atoms = {}
z_com_name_atoms = {}
theta_name_atoms_com_zrange = {}
z_atom_dict = {}
z_com_dict = {}
theta_dict = {}
for mol_name in mol_names:
section = ini['molecule.%s' % (mol_name)]
dists = section['distributions'].split(';')
mol_dists_atom[mol_name] = []
for dist in dists:
name, atoms_str = [x.strip() for x in dist.split(':')]
mol_dists_atom[mol_name].append(name)
z_atom_name_atoms[name] = atoms_str.split()
z_atom_dict[name] = []
com_dists = section['com_distributions'].split(';')
mol_dists_com[mol_name] = []
for dist in com_dists:
name, atoms_str = [x.strip() for x in dist.split(':')]
mol_dists_com[mol_name].append(name)
z_com_name_atoms[name] = atoms_str.split()
z_com_dict[name] = []
try:
angles = section['angles'].split(';')
except:
angles = []
mol_dists_angle[mol_name] = []
for angle in angles:
name, atoms_str = [x.strip() for x in angle.split(':')]
mol_dists_angle[mol_name].append(name)
com_atoms_str, z_range_str = [
x.strip()
for x in section['angle.%s.com_zrange' % name].split(':')
]
if args.reverse:
com_z_range = [
elecd_r - float(x) for x in reversed(z_range_str.split())
]
else:
com_z_range = [elecd_l + float(x) for x in z_range_str.split()]
theta_name_atoms_com_zrange[name] = (atoms_str.split(),
com_atoms_str.split(),
com_z_range)
theta_dict[name] = []
all_frames = list(range(args.begin, args.end, args.skip))
n_group = math.ceil(len(all_frames) / args.nproc)
for i_group in range(n_group):
i_frames = all_frames[i_group * args.nproc:(i_group + 1) * args.nproc]
frames = trj.read_frames(i_frames)
sys.stdout.write('\r frames %s' % ' '.join(map(str, i_frames)))
queue = multiprocessing.Queue()
def worker(frame):
_tmp_z_atom_dict = {name: [] for name in z_atom_dict}
_tmp_z_com_dict = {name: [] for name in z_com_dict}
_tmp_theta_dict = {name: [] for name in theta_dict}
positions = frame.positions
if args.reverse:
positions[:, 2] = elecd_l + elecd_r - positions[:, 2]
for mol in top.molecules:
if mol.name in mol_dists_atom:
for name in mol_dists_atom[mol.name]:
atom_names = z_atom_name_atoms[name]
atoms = _get_atoms(mol, atom_names)
_tmp_z_atom_dict[name] += [
positions[atom.id][2] for atom in atoms
]
if mol.name in mol_dists_com:
for name in mol_dists_com[mol.name]:
atom_names = z_com_name_atoms[name]
atoms = _get_atoms(mol, atom_names)
_tmp_z_com_dict[name].append(
_get_com_position(positions, atoms)[2])
if mol.name in mol_dists_angle:
for name in mol_dists_angle[mol.name]:
atom_names, com_atom_names, com_z_range = theta_name_atoms_com_zrange[
name]
com_atoms = _get_atoms(mol, com_atom_names)
com_pos = _get_com_position(positions, com_atoms)
if com_pos[2] < com_z_range[0] or com_pos[
2] > com_z_range[1]:
continue
if len(atom_names) == 2:
a1, a2 = _get_atoms(mol, atom_names)
vec = positions[a2.id] - positions[a1.id]
_tmp_theta_dict[name].append(_calc_angle_xy(vec))
elif len(atom_names) == 3:
a1, a2, a3 = _get_atoms(mol, atom_names)
vec_12 = positions[a2.id] - positions[a1.id]
vec_13 = positions[a3.id] - positions[a1.id]
vec_normal = np.cross(vec_12, vec_13)
_tmp_theta_dict[name].append(
90 - _calc_angle_xy(vec_normal))
else:
raise Exception('Invalid angle definition', name,
atom_names)
queue.put((_tmp_z_atom_dict, _tmp_z_com_dict, _tmp_theta_dict))
jobs = []
for frame in frames:
job = multiprocessing.Process(target=worker, args=(frame, ))
jobs.append(job)
job.start()
for job in jobs:
_tmp_z_atom_dict, _tmp_z_com_dict, _tmp_theta_dict = queue.get()
for k, v in z_atom_dict.items():
z_atom_dict[k] += _tmp_z_atom_dict[k]
for k, v in z_com_dict.items():
z_com_dict[k] += _tmp_z_com_dict[k]
for k, v in theta_dict.items():
theta_dict[k] += _tmp_theta_dict[k]
for job in jobs:
job.join()
print('')
n_frame = len(all_frames)
name_column_dict = {'z': z_array}
fig, ax = plt.subplots()
ax.set(xlim=[edges[0], edges[-1]],
xlabel='z (nm)',
ylabel='particle density (/$nm^3$)')
for name, z_list in z_atom_dict.items():
x, y = histogram(z_list, bins=edges)
ax.plot(x,
y / area / dz / n_frame,
label=name,
color='darkred' if name.startswith('dca') else None)
name_column_dict['rho-' + name] = y / area / dz / n_frame
ax.legend()
fig.tight_layout()
fig.savefig(f'{args.output}-dist.png')
print_data_to_file(name_column_dict, f'{args.output}-dist.txt')
name_column_dict = {'z': z_array}
fig, ax = plt.subplots()
ax.set(xlim=[edges[0], edges[-1]],
xlabel='z (nm)',
ylabel='molecule density (/$nm^3$)')
if args.cn:
ax2 = ax.twinx()
ax2.set_ylabel('cumulative molecule number')
for name, z_list in z_com_dict.items():
x, y_com = histogram(z_list, bins=edges)
ax.plot(x,
y_com / area / dz / n_frame,
label=name,
color='darkred' if name.startswith('dca') else None)
if args.cn:
ax2.plot(x,
np.cumsum(y_com) / n_frame,
'--',
label=name,
color='darkred' if name.startswith('dca') else None)
name_column_dict['rho-' + name] = y_com / area / dz / n_frame
name_column_dict['cum-' + name] = np.cumsum(y_com) / n_frame
ax.legend()
fig.tight_layout()
fig.savefig(f'{args.output}-dist-com.png')
print_data_to_file(name_column_dict, f'{args.output}-dist-com.txt')
name_column_dict = {}
fig, ax = plt.subplots()
ax.set(xlim=[0, 90], ylim=[0, 0.1], xlabel='theta', ylabel='probability')
for name, t_list in theta_dict.items():
x, y = histogram(t_list, bins=np.linspace(0, 90, 91), normed=True)
ax.plot(x,
y,
label=name,
linestyle='--' if name.endswith('-') else None,
color='darkred' if name.startswith('dca') else None)
name_column_dict.update({'theta': x, 'prob-' + name: y})
ax.legend()
fig.tight_layout()
fig.savefig(f'{args.output}-angle.png')
print_data_to_file(name_column_dict, f'{args.output}-angle.txt')
com_group1 = _get_com_group(frame.positions, ids_group1, masses_group1)
com_group2 = _get_com_group(frame.positions, ids_group2, masses_group2)
vol = frame.cell.volume
density_pair = len(group1) * len(group2) / vol
density_array = np.zeros(len(r_array), dtype=np.float32)
for com1 in com_group1:
for com2 in com_group2:
distance = _get_distance_periodic(com1, com2, frame.cell.size,
args.maxr + dr / 2)
if 0 < distance < args.maxr + dr / 2:
idx_r = int((distance - edges[0]) / dr)
density_array[idx_r] += 1
rdf_array[1:] += density_array[1:] / (4 * np.pi * r_array[1:]**2 *
dr) / density_pair
rdf_array = rdf_array / n_frame
name_column_dict = {'r': r_array, 'rdf': rdf_array}
print_data_to_file(name_column_dict, f'{args.output}-rdf.txt')
fig, ax = plt.subplots()
ax.set(xlabel='r (nm)', ylabel='RDF')
ax.plot(r_array, rdf_array, label='RDF')
ax.legend()
fig.tight_layout()
fig.savefig(f'{args.output}-rdf.png')