1+ #!/usr/bin/env python3
2+ import os
3+ import re
4+ import sys
5+ import glob
6+ import time
7+ import yaml
8+ import shutil
9+ import subprocess
10+ from ase import Atoms
11+ from ase .io import read , write
12+
13+
14+ def split_stru_with_index (input_file , selected_indices , output1 , output2 ):
15+
16+ atoms = read (input_file , verbose = True )
17+ pp = atoms .info ["pp" ]
18+ basis = atoms .info ["basis" ]
19+
20+ for i , atom in enumerate (atoms ):
21+ atom .tag = i
22+
23+ atoms1 = atoms [selected_indices ]
24+
25+
26+ remaining_indices = [i for i in range (len (atoms )) if i not in selected_indices ]
27+ atoms2 = atoms [remaining_indices ]
28+
29+
30+ write (output1 , atoms1 , format = 'abacus' , pp = pp , basis = basis )
31+ write (output2 , atoms2 , format = 'abacus' , pp = pp , basis = basis )
32+
33+ print (f" Success! \n A.stru: { len (atoms1 )} \n B.stru: { len (atoms2 )} \n " )
34+ return selected_indices , remaining_indices , pp , basis
35+
36+
37+ def merge_stru_by_index (file1 , file2 , indices1 , indices2 , pp , basis , output_file ):
38+
39+ atoms1 = read (file1 )
40+ atoms2 = read (file2 )
41+
42+
43+ total_atoms = len (atoms1 ) + len (atoms2 )
44+ merged_atoms = [None ] * total_atoms
45+
46+ for i , idx in enumerate (indices1 ):
47+ merged_atoms [idx ] = atoms1 [i ]
48+
49+ for i , idx in enumerate (indices2 ):
50+ merged_atoms [idx ] = atoms2 [i ]
51+
52+
53+ cell = atoms2 .get_cell ()
54+ pbc = atoms2 .get_pbc ()
55+
56+ positions = [atom .position for atom in merged_atoms ]
57+ symbols = [atom .symbol for atom in merged_atoms ]
58+
59+ merged = Atoms (symbols = symbols , positions = positions , cell = cell , pbc = pbc )
60+
61+ write (output_file , merged , format = 'abacus' , pp = pp , basis = basis )
62+ print (f"success! total { len (merged )} )
63+
64+
65+ def parse_forces_from_file (indices , file_path = "running_scf.log" ):
66+
67+ with open (file_path , 'r' , encoding = 'utf-8' ) as f :
68+ contents = f .read ()
69+
70+ lines = contents .split ('\n ' )
71+
72+ start_index = - 1
73+ for i , line in enumerate (lines ):
74+ if "#TOTAL-FORCE (eV/Angstrom)#" in line :
75+ start_index = i + 4
76+ break
77+
78+ results = []
79+ for i in indices :
80+ results .append (lines [start_index + i ])
81+
82+ result_text = '\n ' .join (results )
83+ return result_text
84+
85+
86+ def load_config (config_file = "config.yaml" ):
87+
88+ with open (config_file , 'r' , encoding = 'utf-8' ) as f :
89+ return yaml .safe_load (f )
90+
91+ def submit_jobs ():
92+
93+ # loading settings
94+ config = load_config ()
95+
96+ # split structure into A and B
97+ origin_structure = config ['origin_structure' ]
98+ selected_indices = config ['selected_indices' ]
99+ indices1 , indices2 , pp , basis = split_stru_with_index (origin_structure , selected_indices ,
100+ 'A.stru' , 'B.stru' )
101+
102+
103+ # setting.conf for phonopy
104+ setting_conf_content = config ['setting.conf' ]
105+ with open ("setting.conf" , "w" , encoding = "utf-8" ) as f :
106+ f .write (setting_conf_content )
107+
108+
109+ # generate perturbed structures
110+ result = subprocess .run ("phonopy setting.conf --abacus -d -c A.stru" ,
111+ shell = True , capture_output = True , text = True )
112+ print (result .stdout )
113+
114+
115+ all_files = sorted (glob .glob ("STRU-*" ))
116+ TOTAL_TASKS = len (all_files )
117+ TASKS_PER_BATCH = config ['tasks_per_batch' ]
118+ WAIT_TIME = config ['wait_time' ]
119+ TOTAL_BATCHES = (TOTAL_TASKS + TASKS_PER_BATCH - 1 ) // TASKS_PER_BATCH
120+ print (f' TOTAL_TASKS = { TOTAL_TASKS } \n TASKS_PER_BATCH = { TASKS_PER_BATCH } \n TOTAL_BATCHES = { TOTAL_BATCHES } \n ' )
121+
122+
123+ # job script
124+ job_script_content = config ['job_script' ]
125+ with open ("job.sh" , "w" , encoding = "utf-8" ) as f :
126+ f .write (job_script_content )
127+ os .chmod ("job.sh" , 0o755 )
128+
129+ # KPT
130+ kpt_content = config ['kpt' ]
131+ with open ("KPT" , "w" , encoding = "utf-8" ) as f :
132+ f .write (kpt_content )
133+
134+ input_content = config ['input' ]
135+
136+ num_digits = len (str (TOTAL_TASKS ))
137+ for batch in range (1 , TOTAL_BATCHES + 1 ):
138+ print ("-" * 30 )
139+ print (f"Current batch: { batch } { TOTAL_BATCHES } )
140+ print ("-" * 30 )
141+
142+ for task in range (1 , TASKS_PER_BATCH + 1 ):
143+ index = task - 1 + (batch - 1 ) * TASKS_PER_BATCH
144+
145+ if index >= TOTAL_TASKS :
146+ break
147+
148+
149+ new_index = f"{ index + 1 :0{num_digits + 1 }d}
150+
151+ dir_name = f"job_{ new_index }
152+ os .makedirs (dir_name , exist_ok = True )
153+ os .chdir (dir_name )
154+
155+
156+ # STRU
157+ perturbed_stru = f"../{ all_files [index ]}
158+ merge_stru_by_index (perturbed_stru , '../B.stru' , indices1 , indices2 ,
159+ pp , basis , 'STRU' )
160+
161+ # INPUT
162+ with open ("INPUT" , "w" ) as f :
163+ f .write (input_content )
164+
165+
166+ # submit job
167+ try :
168+ result = subprocess .run (["sbatch" , "../job.sh" ],
169+ check = True , capture_output = True , text = True )
170+ print (f"submit job: { new_index } )
171+ except subprocess .CalledProcessError as e :
172+ print (f"Failed - { e } )
173+
174+ os .chdir (".." )
175+
176+ # sleep if not the last batch
177+ if batch < TOTAL_BATCHES :
178+ print (f"wait { WAIT_TIME } )
179+ time .sleep (WAIT_TIME )
180+
181+ print ("-" * 30 )
182+
183+ print ("All job submitted!" )
184+ print ("-" * 30 )
185+
186+ def postprocess ():
187+
188+ # loading settings
189+ config = load_config ()
190+
191+ selected_indices = config ['selected_indices' ]
192+ natom = len (selected_indices )
193+
194+ all_files = sorted (glob .glob ("STRU-*" ))
195+ TOTAL_TASKS = len (all_files )
196+ num_digits = len (str (TOTAL_TASKS ))
197+
198+ for task in range (1 , TOTAL_TASKS + 1 ):
199+
200+ new_index = f"{ task :0{num_digits + 1 }d}
201+ dir_name = f"job_{ new_index }
202+ os .chdir (dir_name )
203+
204+ out_folders = glob .glob ("OUT.*" )
205+ out_folder = out_folders [0 ]
206+ os .chdir (out_folder )
207+
208+ atom_forces = parse_forces_from_file (selected_indices )
209+
210+ # phonon.log
211+ log = f"""TOTAL ATOM NUMBER = { natom }
212+
213+ #TOTAL-FORCE (eV/Angstrom)#
214+ -------------------------------------------------------------------------
215+ Atoms Force_x Force_y Force_z
216+ -------------------------------------------------------------------------\n """
217+
218+ with open ("phonon.log" , "w" ) as f :
219+ f .write (log )
220+ f .write (str (atom_forces ))
221+ f .write ("\n -------------------------------------------------------------------------" )
222+
223+
224+ os .chdir ("../.." )
225+
226+ result = subprocess .run ("phonopy -f job_*/OUT*/phonon.log" ,
227+ shell = True , check = True , capture_output = True , text = True )
228+
229+
230+ mesh_content = config ['mesh.conf' ]
231+ with open ("mesh.conf" , "w" ) as f :
232+ f .write (mesh_content )
233+
234+ result = subprocess .run ("phonopy -t mesh.conf" ,
235+ shell = True , check = True , capture_output = True , text = True )
236+
237+
238+
239+
240+ def main ():
241+ if len (sys .argv ) != 2 :
242+ print ("usage: python3 selective_dynamics.py --submit" )
243+ print ("usage: python3 selective_dynamics.py --post" )
244+ sys .exit (1 )
245+
246+ if sys .argv [1 ] == "--submit" :
247+ submit_jobs ()
248+ elif sys .argv [1 ] == "--post" :
249+ postprocess ()
250+ else :
251+ print (f"No such parameter: { sys .argv [1 ]} )
252+
253+ if __name__ == "__main__" :
254+ main ()
0 commit comments