# 定义辅助函数
def rearrange_atoms(old_gro, new_gro):
traj = md.load(old_gro)
new_index = []
Li_index = traj.topology.select("name Li")
new_index += Li_index.tolist()
P_index = traj.topology.select("name P")
for i in range(len(P_index)):
F_index = compute_neighbors(traj, 0.17, [P_index[i]])
if len(F_index[0]) != 6:
raise ValueError("P atom is not surrounded by 6 F atoms, maybe the cutoff is too small")
new_index.append(P_index[i])
new_index += F_index[0].tolist()
new_atom_names = [traj.topology.atom(i).name for i in new_index]
new_coordinates = traj.xyz[:, new_index, :]
traj.xyz = new_coordinates
for i, atom in enumerate(traj.topology.atoms):
atom.name = new_atom_names[i]
traj.save_gro(new_gro)
def gen_unwrap_top(filename,dataset_path=dataset_path):
traj = md.load(filename)
n_LIPF6 = len(traj.topology.select("element P"))
unwrap_top_str =f"""; Include forcefield parameters
#include "amber99sb.ff/forcefield.itp"
#include "{dataset_path}/topol/Atomtypes.itp"
#include "amber99sb.ff/ions.itp"
#include "{dataset_path}/topol/PF6.itp"
[ system ]
GAFF force field
[ molecules ]
; Compound nmols
LI {n_LIPF6}
PF6 {n_LIPF6}
"""
with open("unwrap.top", "w") as f:
f.write(unwrap_top_str)
def find_closest_pairs(array1, array2):
min_diffs = [np.inf, np.inf]
closest_pairs = [None, None]
for i in range(len(array2)):
for j in range(i+1, len(array2)):
avg = (array2[i] + array2[j]) / 2
diff = np.linalg.norm(avg - array1)
if diff < max(min_diffs):
index = min_diffs.index(max(min_diffs))
min_diffs[index] = diff
closest_pairs[index] = [i, j]
if max(min_diffs) == 0:
break
if max(min_diffs) == 0:
break
if max(min_diffs) > 0.1:
raise ValueError("Something went wrong, the closest pairs are too far apart")
remaining_indices = list(set(range(len(array2))) - set(np.array(closest_pairs).flatten()))
closest_pairs.append(remaining_indices)
return closest_pairs
def ajust_F_order(old_gro, new_gro):
traj = md.load(old_gro)
new_index = []
Li_index = traj.topology.select("name LI")
new_index += Li_index.tolist()
P_index = traj.topology.select("name P")
for i in P_index:
new_index.append(i)
F_index = np.linspace(i+1, i+6, 6, dtype=int)
P_atom = traj.xyz[0, i, :]
F_atoms = traj.xyz[0, F_index, :]
closest_pairs = find_closest_pairs(P_atom, F_atoms)
F_array = closest_pairs[0]
for pair in closest_pairs[1:]:
mid = len(F_array) // 2
F_array = F_array[:mid] + pair + F_array[mid:]
if len(F_array) != 6:
raise ValueError("Something went wrong, the new array is not of length 6")
new_index += F_index[F_array].tolist()
traj.xyz = traj.xyz[:, new_index, :]
traj.save_pdb(new_gro)
def get_system_info(filename):
traj = md.load(filename)
com = md.compute_center_of_mass(traj)[0]
n_LIPF6 = len(traj.topology.select("element P"))
box_len = np.array([1.6*(n_LIPF6)**(1/3)]*3)
n_EC = 7*n_LIPF6
n_DMC = 7*n_LIPF6
LIPF6_pos = box_len/2 - com
return box_len, n_LIPF6, n_EC, n_DMC, LIPF6_pos
def gen_packmol_input(filename,dataset_path=dataset_path):
box_len, n_LIPF6, n_EC, n_DMC, LIPF6_pos = get_system_info(filename)
box_len_str = " ".join([str(i) for i in box_len*10])
LIPF6_pos_str = " ".join([str(i) for i in LIPF6_pos*10])
packmol_input =f"""tolerance 2.0
filetype pdb
output system.pdb
structure {filename}
number 1
fixed {LIPF6_pos_str} 0.0 0.0 0.0
end structure
structure {dataset_path}/EC.pdb
number {n_EC}
inside box 0. 0. 0. {box_len_str} # Box dimensions
end structure
structure {dataset_path}/DMC.pdb
number {n_DMC}
inside box 0. 0. 0. {box_len_str} # Box dimensions
end structure"""
with open("packmol.inp", "w") as f:
f.write(packmol_input)
return packmol_input
def gen_md_top(filename,dataset_path=dataset_path):
box_len, n_LIPF6, n_EC, n_DMC, LIPF6_pos = get_system_info(filename)
top_str =f"""; Include forcefield parameters
#include "amber99sb.ff/forcefield.itp"
#include "{dataset_path}/topol/Atomtypes.itp"
#include "amber99sb.ff/ions.itp"
#include "{dataset_path}/topol/PF6.itp"
#include "{dataset_path}/topol/EC.itp"
#include "{dataset_path}/topol/DMC.itp"
[ system ]
GAFF force field
[ molecules ]
; Compound nmols
LI {n_LIPF6}
PF6 {n_LIPF6}
EC {n_EC}
DMC {n_DMC}
"""
with open("topol.top", "w") as f:
f.write(top_str)
return top_str