pwact 0.2.0__py3-none-any.whl → 0.2.2.dev0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -30,8 +30,10 @@ from pwact.utils.draw.hist_model_devi import draw_hist_list
30
30
  from pwact.utils.app_lib.lammps import make_lammps_input
31
31
  from pwact.data_format.configop import save_config, get_atom_type
32
32
 
33
+ from pwdata import Config
33
34
  import os
34
35
  import glob
36
+ import numpy as np
35
37
  import pandas as pd
36
38
  """
37
39
  md_dir:
@@ -66,11 +68,13 @@ class Explore(object):
66
68
  self.explore_dir = os.path.join(self.input_param.root_dir, itername, TEMP_STRUCTURE.tmp_run_iter_dir, AL_STRUCTURE.explore)
67
69
  self.md_dir = os.path.join(self.explore_dir, EXPLORE_FILE_STRUCTURE.md)
68
70
  self.select_dir = os.path.join(self.explore_dir, EXPLORE_FILE_STRUCTURE.select)
71
+ self.direct_dir = os.path.join(self.explore_dir, EXPLORE_FILE_STRUCTURE.direct)
69
72
  self.kpu_dir = os.path.join(self.explore_dir, EXPLORE_FILE_STRUCTURE.kpu) # for kpu calculate
70
73
 
71
74
  self.real_explore_dir = os.path.join(self.input_param.root_dir, itername, AL_STRUCTURE.explore)
72
75
  self.real_md_dir = os.path.join(self.real_explore_dir, EXPLORE_FILE_STRUCTURE.md)
73
76
  self.real_select_dir = os.path.join(self.real_explore_dir, EXPLORE_FILE_STRUCTURE.select)
77
+ self.real_direct_dir = os.path.join(self.real_explore_dir, EXPLORE_FILE_STRUCTURE.direct)
74
78
  self.real_kpu_dir = os.path.join(self.real_explore_dir, EXPLORE_FILE_STRUCTURE.kpu) # for kpu calculate
75
79
 
76
80
  def back_explore(self):
@@ -334,6 +338,14 @@ class Explore(object):
334
338
  os.path.join(self.select_dir, "model_devi_distribution-{}.png".format(os.path.basename(md_sys_dir)))
335
339
  )
336
340
  copy_dir(self.select_dir, self.real_select_dir)
341
+ if self.input_param.strategy.direct:
342
+ # delete nouse files
343
+ del_file_list_by_patten(self.direct_dir, "slurm-*")
344
+ del_file_list_by_patten(self.direct_dir, "*tag*")
345
+ # copy direct dir
346
+ copy_dir(source_dir=self.direct_dir,
347
+ target_dir=self.real_direct_dir)
348
+
337
349
 
338
350
  '''
339
351
  description:
@@ -369,4 +381,102 @@ class Explore(object):
369
381
  print("Image select result:\n {}\n\n".format(summary_info))
370
382
  return summary
371
383
 
384
+
385
+ def make_drct_work(self):
386
+ md_work_list = []
387
+ # read candidate.csv
388
+ candidate = pd.read_csv(os.path.join(self.select_dir, EXPLORE_FILE_STRUCTURE.candidate))
389
+ # make scf work dir
390
+ image_list = None
391
+ for index, row in candidate.iterrows():
392
+ config_index = int(row["config_index"])
393
+ sub_md_sys_path = row["file_path"]
394
+ atom_names = None
395
+ with open(os.path.join(sub_md_sys_path, LAMMPS.atom_type_file), 'r') as rf:
396
+ line = rf.readline()
397
+ atom_names = line.split()
398
+ if image_list is None:
399
+ image_list = Config(data_path=os.path.join(sub_md_sys_path, EXPLORE_FILE_STRUCTURE.traj, "{}{}".format(config_index, LAMMPS.traj_postfix)),
400
+ format=PWDATA.lammps_dump, atom_names=atom_names)
401
+ else:
402
+ image_list.append(Config(data_path=os.path.join(sub_md_sys_path, EXPLORE_FILE_STRUCTURE.traj, "{}{}".format(config_index, LAMMPS.traj_postfix)),
403
+ format=PWDATA.lammps_dump, atom_names=atom_names))
404
+ # cvt_lammps.dump to extxyz
405
+ image_list.to(data_path=self.direct_dir, format=PWDATA.extxyz, data_name="{}".format(EXPLORE_FILE_STRUCTURE.candidate_xyz))
406
+ # copy direct script
407
+ copy_file(self.input_param.strategy.direct_script, os.path.join(self.direct_dir, os.path.basename(self.input_param.strategy.direct_script)))
408
+ # make direct work slurm file
409
+
410
+ self.make_drct_slurm_jobs([self.direct_dir])
411
+
412
+ def make_drct_slurm_jobs(self, md_work_list:list[str]):
413
+ # delete old job file
414
+ del_file_list_by_patten(self.direct_dir, "*{}".format(EXPLORE_FILE_STRUCTURE.direct_job))
415
+ group_list = split_job_for_group(1, md_work_list, 1)
416
+ for g_index, group in enumerate(group_list):
417
+ if group[0] == "NONE":
418
+ continue
419
+ jobname = "direct{}".format(g_index)
420
+ tag_name = "{}-{}".format(g_index, EXPLORE_FILE_STRUCTURE.direct_tag)
421
+ tag = os.path.join(self.direct_dir, tag_name)
422
+
423
+ run_cmd = self.resource.direct_resource.command
424
+
425
+ group_slurm_script = set_slurm_script_content(
426
+ gpu_per_node=self.resource.direct_resource.gpu_per_node,
427
+ number_node = self.resource.direct_resource.number_node, #1
428
+ cpu_per_node = self.resource.direct_resource.cpu_per_node,
429
+ queue_name = self.resource.direct_resource.queue_name,
430
+ custom_flags = self.resource.direct_resource.custom_flags,
431
+ env_script = self.resource.direct_resource.env_script,
432
+ job_name = jobname,
433
+ run_cmd_template = run_cmd,
434
+ group = group,
435
+ job_tag = tag,
436
+ task_tag = EXPLORE_FILE_STRUCTURE.direct_tag,
437
+ task_tag_faild = EXPLORE_FILE_STRUCTURE.direct_tag_faild,
438
+ parallel_num=self.resource.direct_resource.parallel_num,
439
+ check_type=None
440
+ )
441
+ slurm_script_name = "{}-{}".format(g_index, EXPLORE_FILE_STRUCTURE.direct_job)
442
+ slurm_job_file = os.path.join(self.direct_dir, slurm_script_name)
443
+ write_to_file(slurm_job_file, group_slurm_script, "w")
444
+
445
+ def do_drct_jobs(self):
446
+ mission = Mission()
447
+ slurm_remain, slurm_success = get_slurm_job_run_info(self.direct_dir, \
448
+ job_patten="*-{}".format(EXPLORE_FILE_STRUCTURE.direct_job), \
449
+ tag_patten="*-{}".format(EXPLORE_FILE_STRUCTURE.direct_tag))
450
+ # for slurm remain, check if tags done
451
+ slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False
452
+ if slurm_done is False:
453
+ slurm_remain = recheck_slurm_by_jobtag(slurm_remain, EXPLORE_FILE_STRUCTURE.md_tag)
454
+ if len(slurm_remain) > 0:
455
+ #recover slurm jobs
456
+ if len(slurm_remain) > 0:
457
+ print("Run these MD Direct Jobs:\n")
458
+ print(slurm_remain)
459
+ for i, script_path in enumerate(slurm_remain):
460
+ slurm_job = SlurmJob()
461
+ tag_name = "{}-{}".format(os.path.basename(script_path).split('-')[0].strip(), EXPLORE_FILE_STRUCTURE.direct_tag)
462
+ tag = os.path.join(os.path.dirname(script_path),tag_name)
463
+ slurm_job.set_tag(tag, job_type=SLURM_JOB_TYPE.direct)
464
+ slurm_job.set_cmd(script_path)
465
+ mission.add_job(slurm_job)
466
+
467
+ if len(mission.job_list) > 0:
468
+ mission.commit_jobs()
469
+ mission.check_running_job()
470
+ mission.all_job_finished()
471
+
472
+ def post_drct(self):
473
+ select_idx = np.loadtxt(os.path.join(self.direct_dir, "{}".format(EXPLORE_FILE_STRUCTURE.select_idx)))
474
+ count = 0
475
+ with open(os.path.join(self.direct_dir, EXPLORE_FILE_STRUCTURE.candidate_xyz), 'r') as file:
476
+ for line in file:
477
+ count += line.count("Lattice")
372
478
 
479
+ summary = "Direct selection: Before : {} configs, After : {} selected.\n\n".format(select_idx.shape[0], count)
480
+ write_to_file(os.path.join(self.input_param.root_dir, EXPLORE_FILE_STRUCTURE.iter_select_file), summary, mode='a')
481
+ print(summary)
482
+ return summary
@@ -245,7 +245,7 @@ def count_info(save_dir, error_pd, accurate_pd, rand_candi, remove_candi, md_run
245
245
  format(error_pd.shape[0], EXPLORE_FILE_STRUCTURE.failed)
246
246
  if md_run_info is not None:
247
247
  summary_info += md_run_info
248
- summary_info += "\n\n"
248
+ summary_info += "\n"
249
249
 
250
250
  summary += md_run_info
251
251
  summary += "\n"
@@ -0,0 +1,182 @@
1
+ """
2
+ do md with bigmodel
3
+ --/work_path/
4
+ ------------/relax
5
+ -----------------/*_init_config/*_init_.config, etot.input, pseudo files, final.config, mvms_relax
6
+ ------------/super_cell_scale
7
+ -----------------/*_init_config/scale.config, from final.config or *_init_.config
8
+ ------------/perturb
9
+ --------------------/*_init_config/scale.config or final.config
10
+ ----------------------------------/structures/*.config_perturb
11
+
12
+ ------------/AIMD
13
+ -----------------/*_init_config/*_config_dir the config file from pertub or *_init_.config
14
+ -------------------------------/etot.input, atom.config, ...
15
+ or
16
+ ------------/BIGMODEL
17
+ -----------------/*_init_config/*_config_dir the config file from pertub or *_init_.config
18
+ -------------------------------/bigmodel_md_script.py, POSCAR, ...
19
+ ------------/DIRECT
20
+ -----------------/*_init_config/*_config_dir the config file from pertub or *_init_.config
21
+ -------------------------------/direct_script.py, POSCAR, ...
22
+
23
+ -----------/result
24
+ -----------------/*_init_config/relax_mvm_*_init_config if need; aimd_mvm_*_init_config if need
25
+
26
+ """
27
+ import os
28
+
29
+ from pwact.active_learning.user_input.resource import Resource
30
+ from pwact.active_learning.user_input.init_bulk_input import InitBulkParam
31
+ from pwact.active_learning.init_bulk.duplicate_scale import get_config_files_with_order
32
+
33
+ from pwact.utils.constant import PWMAT, INIT_BULK, TEMP_STRUCTURE, SLURM_OUT, DFT_STYLE
34
+ from pwact.active_learning.slurm.slurm import SlurmJob, Mission
35
+ from pwact.utils.slurm_script import get_slurm_job_run_info, split_job_for_group, set_slurm_script_content
36
+
37
+ from pwact.utils.file_operation import write_to_file, link_file, del_file_list_by_patten, copy_file
38
+ from pwact.utils.app_lib.common import link_pseudo_by_atom, set_input_script
39
+ from pwact.data_format.configop import save_config, get_atom_type
40
+
41
+ class DIRECT(object):
42
+ def __init__(self, resource: Resource, input_param:InitBulkParam):
43
+ self.resource = resource
44
+ self.input_param = input_param
45
+ self.init_configs = self.input_param.sys_config
46
+ self.relax_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.relax)
47
+ self.super_cell_scale_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.super_cell_scale)
48
+ self.pertub_dir = os.path.join(self.input_param.root_dir,TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.pertub)
49
+ self.bigmodel_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.bigmodel)
50
+ self.real_bigmodel_dir = os.path.join(self.input_param.root_dir, INIT_BULK.bigmodel)
51
+ self.direct_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.direct)
52
+ self.real_direct_dir = os.path.join(self.input_param.root_dir, INIT_BULK.direct)
53
+
54
+ def make_bigmodel_work(self):
55
+ bigmodel_paths = []
56
+ for init_config in self.init_configs:
57
+ if init_config.bigmodel is False:
58
+ continue
59
+ init_config_name = "init_config_{}".format(init_config.config_index)
60
+ # config_list, config_type = get_config_files_with_order(self.super_cell_scale_dir, self.relax_dir, init_config_name, init_config.config, self.pertub_dir)
61
+ config_list, config_type = get_config_files_with_order(
62
+ super_cell_scale_dir=self.super_cell_scale_dir,
63
+ relax_dir=self.relax_dir,
64
+ init_config_dirname=init_config_name,
65
+ init_config_path=init_config.config_file,
66
+ pertub_dir=self.pertub_dir,
67
+ dft_style=self.resource.dft_style
68
+ )
69
+
70
+ for index, config in enumerate(config_list):
71
+ if config_type == INIT_BULK.pertub:
72
+ tmp_config_dir = os.path.basename(os.path.basename(os.path.dirname(config)))
73
+ elif config_type == INIT_BULK.super_cell or config_type == INIT_BULK.scale:
74
+ tmp_config_dir = os.path.basename(config).replace(DFT_STYLE.get_postfix(self.resource.dft_style), "")
75
+ elif config_type == INIT_BULK.relax:
76
+ tmp_config_dir = INIT_BULK.relax
77
+ else:
78
+ tmp_config_dir = INIT_BULK.init
79
+ bigmodel_dir = os.path.join(self.bigmodel_dir, init_config_name, tmp_config_dir, "{}_{}".format(index, INIT_BULK.bigmodel))
80
+ if not os.path.exists(bigmodel_dir):
81
+ os.makedirs(bigmodel_dir)
82
+
83
+ self.make_bigmodel_file(
84
+ bigmodel_dir=bigmodel_dir,
85
+ config_file=config,
86
+ conifg_format=DFT_STYLE.get_format_by_postfix(os.path.basename(config)),
87
+ input_file=init_config.bigmodel_input_file)
88
+
89
+ bigmodel_paths.append(bigmodel_dir)
90
+ # make slurm script and slurm job
91
+ self.make_bigmodel_slurm_job_files(bigmodel_paths)
92
+
93
+ def check_work_done(self):
94
+ slurm_remain, slurm_success = get_slurm_job_run_info(self.bigmodel_dir, \
95
+ job_patten="*-{}".format(INIT_BULK.bigmodel_job), \
96
+ tag_patten="*-{}".format(INIT_BULK.bigmodel_tag))
97
+ slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False # len(slurm_remain) > 0 exist slurm jobs need to do
98
+ return slurm_done
99
+
100
+ def do_bigmodel_jobs(self):
101
+ mission = Mission()
102
+ slurm_remain, slurm_success = get_slurm_job_run_info(self.bigmodel_dir, \
103
+ job_patten="*-{}".format(INIT_BULK.bigmodel_job), \
104
+ tag_patten="*-{}".format(INIT_BULK.bigmodel_tag))
105
+ slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False # len(slurm_remain) > 0 exist slurm jobs need to do
106
+ if slurm_done is False:
107
+ #recover slurm jobs
108
+ if len(slurm_remain) > 0:
109
+ print("Run these bigmodel Jobs:\n")
110
+ print(slurm_remain)
111
+ for i, script_path in enumerate(slurm_remain):
112
+ slurm_job = SlurmJob()
113
+ tag_name = "{}-{}".format(os.path.basename(script_path).split('-')[0].strip(), INIT_BULK.bigmodel_tag)
114
+ tag = os.path.join(os.path.dirname(script_path),tag_name)
115
+ slurm_job.set_tag(tag)
116
+ slurm_job.set_cmd(script_path)
117
+ mission.add_job(slurm_job)
118
+
119
+ if len(mission.job_list) > 0:
120
+ mission.commit_jobs()
121
+ mission.check_running_job()
122
+ mission.all_job_finished()
123
+ # mission.move_slurm_log_to_slurm_work_dir()
124
+
125
+ '''
126
+ description:
127
+ input_file is aimd input control file, for vasp is incar, for pwmat is etot.input
128
+ return {*}
129
+ author: wuxingxing
130
+ '''
131
+ def make_bigmodel_file(
132
+ bigmodel_dir=str,
133
+ config_file=str,
134
+ conifg_format=str,
135
+ input_file=str):
136
+ #1. set config file, cvt to poscar format
137
+ target_config = save_config(config=config_file,
138
+ input_format=conifg_format,
139
+ wrap = False,
140
+ direct = True,
141
+ sort = True,
142
+ save_format=DFT_STYLE.poscar,
143
+ save_path=bigmodel_dir,
144
+ save_name=DFT_STYLE.get_normal_config(DFT_STYLE.vasp)
145
+ )
146
+
147
+ atom_type_list, _ = get_atom_type(config_file, conifg_format)
148
+ #2. copy script file
149
+ copy_file(input_file, os.path.join(bigmodel_dir, os.path.basename(input_file)))
150
+
151
+ def make_bigmodel_slurm_job_files(self, bigmodel_dir_list:list[str]):
152
+ del_file_list_by_patten(self.bigmodel_dir, "*{}".format(INIT_BULK.bigmodel_job))
153
+ group_list = split_job_for_group(self.resource.explore_resource.group_size, bigmodel_dir_list, self.resource.explore_resource.parallel_num)
154
+ for group_index, group in enumerate(group_list):
155
+ if group[0] == "NONE":
156
+ continue
157
+ jobname = "bigmodel{}".format(group_index)
158
+ tag_name = "{}-{}".format(group_index, INIT_BULK.bigmodel_tag)
159
+ tag = os.path.join(self.bigmodel_dir, tag_name)
160
+ run_cmd = self.resource.explore_resource.command
161
+ group_slurm_script = set_slurm_script_content(gpu_per_node=self.resource.explore_resource.gpu_per_node,
162
+ number_node = self.resource.explore_resource.number_node,
163
+ cpu_per_node = self.resource.explore_resource.cpu_per_node,
164
+ queue_name = self.resource.explore_resource.queue_name,
165
+ custom_flags = self.resource.explore_resource.custom_flags,
166
+ env_script = self.resource.explore_resource.env_script,
167
+ job_name = jobname,
168
+ run_cmd_template = run_cmd,
169
+ group = group,
170
+ job_tag = tag,
171
+ task_tag = INIT_BULK.bigmodel_tag,
172
+ task_tag_faild = INIT_BULK.bigmodel_tag_failed,
173
+ parallel_num=self.resource.dft_resource.parallel_num,
174
+ check_type=self.resource.dft_style
175
+ )
176
+ slurm_script_name = "{}-{}".format(group_index, INIT_BULK.bigmodel_job)
177
+ slurm_job_file = os.path.join(self.bigmodel_dir, slurm_script_name)
178
+ write_to_file(slurm_job_file, group_slurm_script, "w")
179
+
180
+ def do_post_process(self):
181
+ if os.path.exists(self.bigmodel_dir):
182
+ link_file(self.bigmodel_dir, self.real_bigmodel_dir)
@@ -1,5 +1,5 @@
1
1
  import os
2
-
2
+ import bisect
3
3
  from pwact.utils.file_operation import search_files, link_file, copy_file, write_to_file
4
4
  from pwact.utils.constant import INIT_BULK, DFT_STYLE, TEMP_STRUCTURE
5
5
 
@@ -0,0 +1,301 @@
1
+ """
2
+ do md with bigmodel
3
+ --/work_path/
4
+ ------------/relax
5
+ -----------------/*_init_config/*_init_.config, etot.input, pseudo files, final.config, mvms_relax
6
+ ------------/super_cell_scale
7
+ -----------------/*_init_config/scale.config, from final.config or *_init_.config
8
+ ------------/perturb
9
+ --------------------/*_init_config/scale.config or final.config
10
+ ----------------------------------/structures/*.config_perturb
11
+
12
+ ------------/AIMD
13
+ -----------------/*_init_config/*_config_dir the config file from pertub or *_init_.config
14
+ -------------------------------/etot.input, atom.config, ...
15
+ or
16
+ ------------/BIGMODEL
17
+ -----------------/*_init_config/*_config_dir the config file from pertub or *_init_.config
18
+ -------------------------------/bigmodel_md_script.py, atom.config, ...
19
+
20
+ -----------/result
21
+ -----------------/*_init_config/relax_mvm_*_init_config if need; aimd_mvm_*_init_config if need
22
+
23
+ """
24
+ import os
25
+
26
+ from pwact.active_learning.user_input.resource import Resource
27
+ from pwact.active_learning.user_input.init_bulk_input import InitBulkParam
28
+ from pwact.active_learning.init_bulk.duplicate_scale import get_config_files_with_order
29
+
30
+ from pwact.utils.constant import PWMAT, INIT_BULK, TEMP_STRUCTURE, SLURM_OUT, DFT_STYLE, PWDATA
31
+ from pwact.active_learning.slurm.slurm import SlurmJob, Mission
32
+ from pwact.utils.slurm_script import get_slurm_job_run_info, split_job_for_group, set_slurm_script_content
33
+
34
+ from pwact.utils.file_operation import write_to_file, link_file, del_file_list_by_patten, copy_file, merge_files_to_one, save_json_file
35
+ from pwact.utils.app_lib.common import link_pseudo_by_atom, set_input_script
36
+ from pwact.data_format.configop import save_config, get_atom_type
37
+ from pwdata import Config
38
+
39
+ class BIGMODEL(object):
40
+ def __init__(self, resource: Resource, input_param:InitBulkParam):
41
+ self.resource = resource
42
+ self.input_param = input_param
43
+ self.init_configs = self.input_param.sys_config
44
+ self.relax_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.relax)
45
+ self.super_cell_scale_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.super_cell_scale)
46
+ self.pertub_dir = os.path.join(self.input_param.root_dir,TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.pertub)
47
+ self.bigmodel_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.bigmodel)
48
+ self.real_bigmodel_dir = os.path.join(self.input_param.root_dir, INIT_BULK.bigmodel)
49
+
50
+ self.direct_dir = os.path.join(self.bigmodel_dir, INIT_BULK.direct)
51
+ self.real_direct_dir = os.path.join(self.real_bigmodel_dir, INIT_BULK.direct)
52
+
53
+ self.model_traj_files = []
54
+
55
+ def make_bigmodel_work(self):
56
+ bigmodel_paths = []
57
+
58
+ for init_config in self.init_configs:
59
+ if init_config.bigmodel is False:
60
+ continue
61
+ init_config_name = "init_config_{}".format(init_config.config_index)
62
+ # config_list, config_type = get_config_files_with_order(self.super_cell_scale_dir, self.relax_dir, init_config_name, init_config.config, self.pertub_dir)
63
+ config_list, config_type = get_config_files_with_order(
64
+ super_cell_scale_dir=self.super_cell_scale_dir,
65
+ relax_dir=self.relax_dir,
66
+ init_config_dirname=init_config_name,
67
+ init_config_path=init_config.config_file,
68
+ pertub_dir=self.pertub_dir,
69
+ dft_style=self.resource.dft_style
70
+ )
71
+
72
+ for index, config in enumerate(config_list):
73
+ if config_type == INIT_BULK.pertub:
74
+ tmp_config_dir = os.path.basename(os.path.basename(os.path.dirname(config)))
75
+ elif config_type == INIT_BULK.super_cell or config_type == INIT_BULK.scale:
76
+ tmp_config_dir = os.path.basename(config).replace(DFT_STYLE.get_postfix(self.resource.dft_style), "")
77
+ elif config_type == INIT_BULK.relax:
78
+ tmp_config_dir = INIT_BULK.relax
79
+ else:
80
+ tmp_config_dir = INIT_BULK.init
81
+ bigmodel_dir = os.path.join(self.bigmodel_dir, init_config_name, tmp_config_dir, "{}_{}".format(index, INIT_BULK.bigmodel))
82
+ if not os.path.exists(bigmodel_dir):
83
+ os.makedirs(bigmodel_dir)
84
+
85
+ self.make_bigmodel_file(
86
+ bigmodel_dir,
87
+ config,
88
+ DFT_STYLE.get_format_by_postfix(os.path.basename(config)),
89
+ init_config.bigmodel_input_file)
90
+
91
+ self.model_traj_files.append(os.path.join(bigmodel_dir, INIT_BULK.bigmodel_traj))
92
+
93
+ bigmodel_paths.append(bigmodel_dir)
94
+ # make slurm script and slurm job
95
+ self.make_bigmodel_slurm_job_files(bigmodel_paths)
96
+
97
+ def check_work_done(self):
98
+ slurm_remain, slurm_success = get_slurm_job_run_info(self.bigmodel_dir, \
99
+ job_patten="*-{}".format(INIT_BULK.bigmodel_job), \
100
+ tag_patten="*-{}".format(INIT_BULK.bigmodel_tag))
101
+ slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False # len(slurm_remain) > 0 exist slurm jobs need to do
102
+ return slurm_done
103
+
104
+ def do_bigmodel_jobs(self):
105
+ mission = Mission()
106
+ slurm_remain, slurm_success = get_slurm_job_run_info(self.bigmodel_dir, \
107
+ job_patten="*-{}".format(INIT_BULK.bigmodel_job), \
108
+ tag_patten="*-{}".format(INIT_BULK.bigmodel_tag))
109
+ slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False # len(slurm_remain) > 0 exist slurm jobs need to do
110
+ if slurm_done is False:
111
+ #recover slurm jobs
112
+ if len(slurm_remain) > 0:
113
+ print("Run these bigmodel Jobs:\n")
114
+ print(slurm_remain)
115
+ for i, script_path in enumerate(slurm_remain):
116
+ slurm_job = SlurmJob()
117
+ tag_name = "{}-{}".format(os.path.basename(script_path).split('-')[0].strip(), INIT_BULK.bigmodel_tag)
118
+ tag = os.path.join(os.path.dirname(script_path),tag_name)
119
+ slurm_job.set_tag(tag)
120
+ slurm_job.set_cmd(script_path)
121
+ mission.add_job(slurm_job)
122
+
123
+ if len(mission.job_list) > 0:
124
+ mission.commit_jobs()
125
+ mission.check_running_job()
126
+ mission.all_job_finished()
127
+ # mission.move_slurm_log_to_slurm_work_dir()
128
+
129
+ '''
130
+ description:
131
+ input_file is aimd input control file, for vasp is incar, for pwmat is etot.input
132
+ return {*}
133
+ author: wuxingxing
134
+ '''
135
+ def make_bigmodel_file(self,
136
+ bigmodel_dir:str,
137
+ config_file:str,
138
+ conifg_format:str,
139
+ input_file:str):
140
+ #1. set config file, cvt to poscar format
141
+ target_config = save_config(config=config_file,
142
+ input_format=conifg_format,
143
+ wrap = False,
144
+ direct = True,
145
+ sort = True,
146
+ save_format=PWDATA.vasp_poscar,
147
+ save_path=bigmodel_dir,
148
+ save_name=DFT_STYLE.get_normal_config(DFT_STYLE.vasp)
149
+ )
150
+
151
+ atom_type_list, _ = get_atom_type(config_file, conifg_format)
152
+ #2. copy script file
153
+ copy_file(input_file, os.path.join(bigmodel_dir, os.path.basename(input_file)))
154
+
155
+ def make_bigmodel_slurm_job_files(self, bigmodel_dir_list:list[str]):
156
+ del_file_list_by_patten(self.bigmodel_dir, "*{}".format(INIT_BULK.bigmodel_job))
157
+ group_list = split_job_for_group(self.resource.explore_resource.group_size, bigmodel_dir_list, self.resource.explore_resource.parallel_num)
158
+ for group_index, group in enumerate(group_list):
159
+ if group[0] == "NONE":
160
+ continue
161
+ jobname = "bigmodel{}".format(group_index)
162
+ tag_name = "{}-{}".format(group_index, INIT_BULK.bigmodel_tag)
163
+ tag = os.path.join(self.bigmodel_dir, tag_name)
164
+ run_cmd = self.resource.explore_resource.command
165
+ group_slurm_script = set_slurm_script_content(gpu_per_node=self.resource.explore_resource.gpu_per_node,
166
+ number_node = self.resource.explore_resource.number_node,
167
+ cpu_per_node = self.resource.explore_resource.cpu_per_node,
168
+ queue_name = self.resource.explore_resource.queue_name,
169
+ custom_flags = self.resource.explore_resource.custom_flags,
170
+ env_script = self.resource.explore_resource.env_script,
171
+ job_name = jobname,
172
+ run_cmd_template = run_cmd,
173
+ group = group,
174
+ job_tag = tag,
175
+ task_tag = INIT_BULK.bigmodel_tag,
176
+ task_tag_faild = INIT_BULK.bigmodel_tag_failed,
177
+ parallel_num=self.resource.dft_resource.parallel_num,
178
+ check_type=None
179
+ )
180
+ slurm_script_name = "{}-{}".format(group_index, INIT_BULK.bigmodel_job)
181
+ slurm_job_file = os.path.join(self.bigmodel_dir, slurm_script_name)
182
+ write_to_file(slurm_job_file, group_slurm_script, "w")
183
+
184
+ ###### for direct
185
+ def get_traj_files(self):
186
+ model_traj_files = []
187
+ config_dix = []
188
+ for init_dix, init_config in enumerate(self.init_configs):
189
+ if init_config.bigmodel is False:
190
+ continue
191
+ init_config_name = "init_config_{}".format(init_config.config_index)
192
+ # config_list, config_type = get_config_files_with_order(self.super_cell_scale_dir, self.relax_dir, init_config_name, init_config.config, self.pertub_dir)
193
+ config_list, config_type = get_config_files_with_order(
194
+ super_cell_scale_dir=self.super_cell_scale_dir,
195
+ relax_dir=self.relax_dir,
196
+ init_config_dirname=init_config_name,
197
+ init_config_path=init_config.config_file,
198
+ pertub_dir=self.pertub_dir,
199
+ dft_style=self.resource.dft_style
200
+ )
201
+
202
+ for index, config in enumerate(config_list):
203
+ if config_type == INIT_BULK.pertub:
204
+ tmp_config_dir = os.path.basename(os.path.basename(os.path.dirname(config)))
205
+ elif config_type == INIT_BULK.super_cell or config_type == INIT_BULK.scale:
206
+ tmp_config_dir = os.path.basename(config).replace(DFT_STYLE.get_postfix(self.resource.dft_style), "")
207
+ elif config_type == INIT_BULK.relax:
208
+ tmp_config_dir = INIT_BULK.relax
209
+ else:
210
+ tmp_config_dir = INIT_BULK.init
211
+ bigmodel_dir = os.path.join(self.bigmodel_dir, init_config_name, tmp_config_dir, "{}_{}".format(index, INIT_BULK.bigmodel))
212
+ model_traj_files.append(os.path.join(bigmodel_dir, INIT_BULK.bigmodel_traj))
213
+ config_dix.append(init_dix)
214
+ return model_traj_files, config_dix
215
+
216
+ def make_direct_work(self):
217
+ if not os.path.exists(self.direct_dir):
218
+ os.makedirs(self.direct_dir)
219
+ # convert configs to xyz format
220
+ model_traj_files, config_dix = self.get_traj_files()
221
+ candidate_dict = {}
222
+ merge_files_to_one(model_traj_files, os.path.join(self.direct_dir, INIT_BULK.candidate_xyz))
223
+ idx = 0
224
+ for traj_idx, traj in enumerate(model_traj_files):
225
+ image = Config(data_path=traj, format=PWDATA.extxyz)
226
+ candidate_dict[idx] = {}
227
+ candidate_dict[idx]["idx"] = config_dix[traj_idx]
228
+ candidate_dict[idx]["num"] = len(image.images)
229
+ idx += 1
230
+ save_json_file(candidate_dict, os.path.join(self.direct_dir, INIT_BULK.candidate_idx))
231
+ # copy script
232
+ copy_file(self.input_param.dft_input.direct_input_list[0].input_file,
233
+ os.path.join(self.direct_dir, os.path.basename(self.input_param.dft_input.direct_input_list[0].input_file)))
234
+ # make slurm script and slurm job
235
+ self.make_direct_slurm_job_files([self.direct_dir])
236
+
237
+ def check_direct_done(self):
238
+ slurm_remain, slurm_success = get_slurm_job_run_info(self.direct_dir, \
239
+ job_patten="*-{}".format(INIT_BULK.direct_job), \
240
+ tag_patten="*-{}".format(INIT_BULK.direct_tag))
241
+ slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False # len(slurm_remain) > 0 exist slurm jobs need to do
242
+ return slurm_done
243
+
244
+ def do_direct_jobs(self):
245
+ mission = Mission()
246
+ slurm_remain, slurm_success = get_slurm_job_run_info(self.direct_dir, \
247
+ job_patten="*-{}".format(INIT_BULK.direct_job), \
248
+ tag_patten="*-{}".format(INIT_BULK.direct_tag))
249
+ slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False # len(slurm_remain) > 0 exist slurm jobs need to do
250
+ if slurm_done is False:
251
+ #recover slurm jobs
252
+ if len(slurm_remain) > 0:
253
+ print("Run these direct Jobs:\n")
254
+ print(slurm_remain)
255
+ for i, script_path in enumerate(slurm_remain):
256
+ slurm_job = SlurmJob()
257
+ tag_name = "{}-{}".format(os.path.basename(script_path).split('-')[0].strip(), INIT_BULK.direct_tag)
258
+ tag = os.path.join(os.path.dirname(script_path),tag_name)
259
+ slurm_job.set_tag(tag)
260
+ slurm_job.set_cmd(script_path)
261
+ mission.add_job(slurm_job)
262
+
263
+ if len(mission.job_list) > 0:
264
+ mission.commit_jobs()
265
+ mission.check_running_job()
266
+ mission.all_job_finished()
267
+ # mission.move_slurm_log_to_slurm_work_dir()
268
+
269
+ def make_direct_slurm_job_files(self, direct_dir_list:list[str]):
270
+ del_file_list_by_patten(self.direct_dir, "*{}".format(INIT_BULK.direct_job))
271
+ group_list = split_job_for_group(self.resource.direct_resource.group_size, direct_dir_list, self.resource.direct_resource.parallel_num)
272
+ for group_index, group in enumerate(group_list):
273
+ if group[0] == "NONE":
274
+ continue
275
+ jobname = "direct{}".format(group_index)
276
+ tag_name = "{}-{}".format(group_index, INIT_BULK.direct_tag)
277
+ tag = os.path.join(self.direct_dir, tag_name)
278
+ run_cmd = self.resource.direct_resource.command
279
+ group_slurm_script = set_slurm_script_content(gpu_per_node=self.resource.direct_resource.gpu_per_node,
280
+ number_node = self.resource.direct_resource.number_node,
281
+ cpu_per_node = self.resource.direct_resource.cpu_per_node,
282
+ queue_name = self.resource.direct_resource.queue_name,
283
+ custom_flags = self.resource.direct_resource.custom_flags,
284
+ env_script = self.resource.direct_resource.env_script,
285
+ job_name = jobname,
286
+ run_cmd_template = run_cmd,
287
+ group = group,
288
+ job_tag = tag,
289
+ task_tag = INIT_BULK.direct_tag,
290
+ task_tag_faild = INIT_BULK.direct_tag_failed,
291
+ parallel_num=self.resource.dft_resource.parallel_num,
292
+ check_type=None
293
+ )
294
+ slurm_script_name = "{}-{}".format(group_index, INIT_BULK.direct_job)
295
+ slurm_job_file = os.path.join(self.direct_dir, slurm_script_name)
296
+ write_to_file(slurm_job_file, group_slurm_script, "w")
297
+
298
+ def do_post_process(self):
299
+ if os.path.exists(self.direct_dir):
300
+ link_file(self.direct_dir, self.real_direct_dir)
301
+