def get_final_structure(cls, wf):
assert wf.metadata['workflow_class'] == cls.workflow_class
assert wf.metadata['workflow_module'] == cls.workflow_module
step_index = -1
final_fw_id = None
for fw_id, fw in wf.id_fw.items():
if 'SRC_task_index' in fw.spec:
if fw.tasks[-1].src_type != 'run':
continue
task_index = fw.spec['SRC_task_index']
if task_index.task_type == 'MPRelaxVasp':
if task_index.index > step_index:
step_index = task_index.index
final_fw_id = fw_id
if final_fw_id is None:
raise RuntimeError('Final structure not found ...')
myfw = wf.id_fw[final_fw_id]
mytask = myfw.tasks[-1]
last_launch = (myfw.archived_launches + myfw.launches)[-1]
# myfw.tasks[-1].set_workdir(workdir=last_launch.launch_dir)
# mytask.setup_rundir(last_launch.launch_dir, create_dirs=False)
helper = MPRelaxTaskHelper()
helper.set_task(mytask)
helper.task.setup_rundir(last_launch.launch_dir)
structure = helper.get_final_structure()
return {'structure': structure.as_dict()}
def get_final_structure(cls, wf):
assert wf.metadata['workflow_class'] == cls.workflow_class
assert wf.metadata['workflow_module'] == cls.workflow_module
step_index = -1
final_fw_id = None
nrelaxations = 1
for fw_id, fw in wf.id_fw.items():
if 'SRC_task_index' in fw.spec:
if fw.tasks[-1].src_type != 'run':
continue
task_index = fw.spec['SRC_task_index']
if 'MPRelaxVasp' in task_index.task_type:
if task_index.task_type == 'MPRelaxVasp':
if task_index.index > step_index:
step_index = task_index.index
final_fw_id = fw_id
else:
match = re.match("MPRelaxVasp-\d+", task_index.task_type)
if match and match.string == task_index.task_type:
nn = int(task_index.task_type[12:])
if nn == nrelaxations:
if task_index.index > step_index:
step_index = task_index.index
final_fw_id = fw_id
elif nn > nrelaxations:
step_index = task_index.index
final_fw_id = fw_id
nrelaxations = nn
if final_fw_id is None:
raise RuntimeError('Final structure not found ...')
myfw = wf.id_fw[final_fw_id]
mytask = myfw.tasks[-1]
last_launch = (myfw.archived_launches + myfw.launches)[-1]
# myfw.tasks[-1].set_workdir(workdir=last_launch.launch_dir)
# mytask.setup_rundir(last_launch.launch_dir, create_dirs=False)
helper = MPRelaxTaskHelper()
helper.set_task(mytask)
helper.task.setup_rundir(last_launch.launch_dir)
structure = helper.get_final_structure()
return {'structure': structure.as_dict()}
def __init__(self, vasp_input_set, spec):
# Initializes fws list and links_dict
fws = []
links_dict = {}
if 'additional_controllers' in spec:
additional_controllers = spec['additional_controllers']
spec.pop('additional_controllers')
else:
additional_controllers = [
WalltimeController(),
MemoryController(),
VaspXMLValidatorController()
]
control_procedure = ControlProcedure(
controllers=additional_controllers)
task_helper = MPRelaxTaskHelper()
task_type = task_helper.task_type
src_fws = createVaspSRCFireworks(vasp_input_set=vasp_input_set,
task_helper=task_helper,
task_type=task_type,
control_procedure=control_procedure,
custodian_handlers=[],
max_restarts=10,
src_cleaning=None,
task_index=None,
spec=None,
setup_spec_update=None,
run_spec_update=None)
fws.extend(src_fws['fws'])
links_dict_update(links_dict=links_dict,
links_update=src_fws['links_dict'])
self.wf = Workflow(fireworks=fws,
links_dict=links_dict,
metadata={
'workflow_class': self.workflow_class,
'workflow_module': self.workflow_module
},
name='MPRelaxFWWorkflowSRC')
def __init__(self,
neb_vasp_input_set,
spec,
neb_terminals,
relax_terminals=True,
n_insert=1,
n_nebs=3,
relax_vasp_input_set=None,
initial_neb_structures=None,
climbing_image=True):
user_incar_settings = {'NPAR': 4, 'ISIF': 0, 'SIGMA': 0.2, 'ISMEAR': 0}
if n_nebs < 1:
raise ValueError('Minimum one NEB ...')
if relax_terminals and initial_neb_structures is not None:
raise ValueError(
'Cannot relax terminals and start from initial NEB structures')
# Initializes fws list and links_dict
fws = []
links_dict = {}
if 'additional_controllers' in spec:
additional_controllers = spec['additional_controllers']
spec.pop('additional_controllers')
else:
additional_controllers = [
WalltimeController(),
MemoryController(),
VaspXMLValidatorController()
]
# Control procedure
control_procedure = ControlProcedure(
controllers=additional_controllers)
# First NEB
gen_neb_spec = spec.copy()
if relax_terminals:
gen_neb_spec['terminal_start'] = None
gen_neb_spec['terminal_end'] = None
else:
gen_neb_spec['terminal_start'] = neb_terminals[0]
gen_neb_spec['terminal_end'] = neb_terminals[1]
# gen_neb_spec['structures'] = neb_terminals
gen_neb_spec = set_short_single_core_to_spec(gen_neb_spec)
terminal_start_relax_task_type = 'MPRelaxVasp-start'
terminal_end_relax_task_type = 'MPRelaxVasp-end'
terminal_start_task_type = None
terminal_end_task_type = None
if relax_terminals:
terminal_start_task_type = terminal_start_relax_task_type
terminal_end_task_type = terminal_end_relax_task_type
gen_neb_task = GenerateNEBRelaxationTask(
n_insert=n_insert,
user_incar_settings=user_incar_settings,
climbing_image=climbing_image,
task_index='neb1',
terminal_start_task_type=terminal_start_task_type,
terminal_end_task_type=terminal_end_task_type)
gen_neb_fw = Firework([gen_neb_task],
spec=gen_neb_spec,
name='gen-neb1')
fws.append(gen_neb_fw)
if relax_terminals:
# Start terminal
relax_task_helper = MPRelaxTaskHelper()
vis_start = relax_vasp_input_set(
neb_terminals[0], user_incar_settings=user_incar_settings)
start_src_fws = createVaspSRCFireworks(
vasp_input_set=vis_start,
task_helper=relax_task_helper,
task_type=terminal_start_relax_task_type,
control_procedure=control_procedure,
custodian_handlers=[],
max_restarts=10,
src_cleaning=None,
task_index=None,
spec=None,
setup_spec_update=None,
run_spec_update=None)
fws.extend(start_src_fws['fws'])
links_dict_update(links_dict=links_dict,
links_update=start_src_fws['links_dict'])
linkupdate = {start_src_fws['control_fw'].fw_id: gen_neb_fw.fw_id}
links_dict_update(links_dict=links_dict, links_update=linkupdate)
# End terminal
vis_end = relax_vasp_input_set(
neb_terminals[1], user_incar_settings=user_incar_settings)
end_src_fws = createVaspSRCFireworks(
vasp_input_set=vis_end,
task_helper=relax_task_helper,
task_type=terminal_end_relax_task_type,
control_procedure=control_procedure,
custodian_handlers=[],
max_restarts=10,
src_cleaning=None,
task_index=None,
spec=None,
setup_spec_update=None,
run_spec_update=None)
fws.extend(end_src_fws['fws'])
links_dict_update(links_dict=links_dict,
links_update=end_src_fws['links_dict'])
linkupdate = {end_src_fws['control_fw'].fw_id: gen_neb_fw.fw_id}
links_dict_update(links_dict=links_dict, links_update=linkupdate)
if n_nebs > 1:
for ineb in range(2, n_nebs + 1):
prev_gen_neb_fw = gen_neb_fw
gen_neb_spec = spec.copy()
gen_neb_spec['structures'] = None
gen_neb_spec = set_short_single_core_to_spec(gen_neb_spec)
gen_neb_task = GenerateNEBRelaxationTask(
n_insert=n_insert,
user_incar_settings=user_incar_settings,
climbing_image=climbing_image,
task_index='neb{:d}'.format(ineb),
prev_neb_task_type='neb{:d}'.format(ineb - 1),
terminal_start_task_type=terminal_start_task_type,
terminal_end_task_type=terminal_end_task_type)
gen_neb_fw = Firework([gen_neb_task],
spec=gen_neb_spec,
name='gen-neb{:d}'.format(ineb))
fws.append(gen_neb_fw)
linkupdate = {prev_gen_neb_fw.fw_id: gen_neb_fw.fw_id}
links_dict_update(links_dict=links_dict,
links_update=linkupdate)
if relax_terminals:
linkupdate = {
start_src_fws['control_fw'].fw_id: gen_neb_fw.fw_id
}
links_dict_update(links_dict=links_dict,
links_update=linkupdate)
linkupdate = {
end_src_fws['control_fw'].fw_id: gen_neb_fw.fw_id
}
links_dict_update(links_dict=links_dict,
links_update=linkupdate)
if climbing_image:
wfname = "MPcNEBRelaxFWWorkflowSRC"
else:
wfname = "MPcNEBRelaxFWWorkflowSRC"
self.wf = Workflow(fireworks=fws,
links_dict=links_dict,
metadata={
'workflow_class': self.workflow_class,
'workflow_module': self.workflow_module
},
name=wfname)
def get_terminal_computed_entries(cls, wf):
assert wf.metadata['workflow_class'] == cls.workflow_class
assert wf.metadata['workflow_module'] == cls.workflow_module
terminal_start_step_index = -1
terminal_start_final_fw_id = None
terminal_start_setup_step_index = -1
terminal_start_setup_fw_id = None
terminal_end_step_index = -1
terminal_end_final_fw_id = None
terminal_end_setup_step_index = -1
terminal_end_setup_fw_id = None
for fw_id, fw in wf.id_fw.items():
if 'SRC_task_index' in fw.spec:
src_type = fw.tasks[-1].src_type
if src_type in ['run', 'setup']:
task_index = fw.spec['SRC_task_index']
if task_index.task_type == 'MPRelaxVasp-start':
if src_type == 'run':
if task_index.index > terminal_start_step_index:
terminal_start_step_index = task_index.index
terminal_start_final_fw_id = fw_id
elif src_type == 'setup':
if task_index.index > terminal_start_setup_step_index:
terminal_start_setup_step_index = task_index.index
terminal_start_setup_fw_id = fw_id
elif task_index.task_type == 'MPRelaxVasp-end':
if src_type == 'run':
if task_index.index > terminal_end_step_index:
terminal_end_step_index = task_index.index
terminal_end_final_fw_id = fw_id
elif src_type == 'setup':
if task_index.index > terminal_end_setup_step_index:
terminal_end_setup_step_index = task_index.index
terminal_end_setup_fw_id = fw_id
if terminal_start_final_fw_id is None:
raise RuntimeError('No terminal-start relaxation found ...')
if terminal_end_final_fw_id is None:
raise RuntimeError('No terminal-end relaxation found ...')
# Terminal start dir
terminal_start_fw = wf.id_fw[terminal_start_final_fw_id]
terminal_start_last_launch = (terminal_start_fw.archived_launches + terminal_start_fw.launches)[-1]
terminal_start_run_dir = terminal_start_last_launch.launch_dir
# Terminal end dir
terminal_end_fw = wf.id_fw[terminal_end_final_fw_id]
terminal_end_last_launch = (terminal_end_fw.archived_launches + terminal_end_fw.launches)[-1]
terminal_end_run_dir = terminal_end_last_launch.launch_dir
helper = MPRelaxTaskHelper()
# Get the information for the terminal start
terminal_start_task = terminal_start_fw.tasks[-1]
helper.set_task(terminal_start_task)
helper.task.setup_rundir(terminal_start_run_dir)
terminal_start_all_info = helper.get_all_info(raise_error=True)
terminal_start_setup_fw = wf.id_fw[terminal_start_setup_fw_id]
terminal_start_setup_task = terminal_start_setup_fw.tasks[-1]
terminal_start_vasp_input_set = terminal_start_setup_task.vasp_input_set
terminal_start_kpts_generation_description = get_kpoints_generation_description(terminal_start_vasp_input_set)
terminal_start_computed_entry = terminal_start_all_info['computed_entry']
terminal_start_vasprun = terminal_start_all_info['vasprun']
terminal_start_calculation_parameters = {'INCAR': terminal_start_vasprun.incar,
'INCAR_ALL': terminal_start_vasprun.parameters,
'potcar_symbols': terminal_start_vasprun.potcar_symbols,
'ENCUT': terminal_start_vasprun.incar['ENCUT'] or terminal_start_vasprun.parameters['ENCUT'],
'EDIFFG': terminal_start_vasprun.incar.get('EDIFFG') or terminal_start_vasprun.parameters.get('EDIFFG'),
'ISMEAR': terminal_start_vasprun.incar.get('ISMEAR') or terminal_start_vasprun.parameters.get('ISMEAR'),
'SIGMA': terminal_start_vasprun.incar.get('SIGMA') or terminal_start_vasprun.parameters.get('SIGMA'),
'kpoints_generation_description': terminal_start_kpts_generation_description}
if terminal_start_computed_entry.data is None:
terminal_start_computed_entry.data = {}
terminal_start_computed_entry.data['calculation_parameters'] = terminal_start_calculation_parameters
# Get the information for the terminal end
terminal_end_task = terminal_end_fw.tasks[-1]
helper.set_task(terminal_end_task)
helper.task.setup_rundir(terminal_end_run_dir)
terminal_end_all_info = helper.get_all_info(raise_error=True)
terminal_end_setup_fw = wf.id_fw[terminal_end_setup_fw_id]
terminal_end_setup_task = terminal_end_setup_fw.tasks[-1]
terminal_end_vasp_input_set = terminal_end_setup_task.vasp_input_set
terminal_end_kpts_generation_description = get_kpoints_generation_description(
terminal_end_vasp_input_set)
terminal_end_computed_entry = terminal_end_all_info['computed_entry']
terminal_end_vasprun = terminal_end_all_info['vasprun']
terminal_end_calculation_parameters = {'INCAR': terminal_end_vasprun.incar,
'INCAR_ALL': terminal_end_vasprun.parameters,
'potcar_symbols': terminal_end_vasprun.potcar_symbols,
'ENCUT': terminal_end_vasprun.incar['ENCUT'] or
terminal_end_vasprun.parameters['ENCUT'],
'EDIFFG': terminal_end_vasprun.incar.get(
'EDIFFG') or terminal_end_vasprun.parameters.get('EDIFFG'),
'ISMEAR': terminal_end_vasprun.incar.get(
'ISMEAR') or terminal_end_vasprun.parameters.get('ISMEAR'),
'SIGMA': terminal_end_vasprun.incar.get(
'SIGMA') or terminal_end_vasprun.parameters.get('SIGMA'),
'kpoints_generation_description': terminal_end_kpts_generation_description}
if terminal_end_computed_entry.data is None:
terminal_end_computed_entry.data = {}
terminal_end_computed_entry.data['calculation_parameters'] = terminal_end_calculation_parameters
return {'computed_entry_terminal_start': terminal_start_computed_entry.as_dict(),
'computed_entry_terminal_end': terminal_end_computed_entry.as_dict()}
def get_final_results(cls, wf):
assert wf.metadata['workflow_class'] == cls.workflow_class
assert wf.metadata['workflow_module'] == cls.workflow_module
step_index = -1
final_fw_id = None
nrelaxations = 1
fw_ids = {}
final_setup_fw_id = None
final_setup_step_index = -1
final_setup_nrelaxations = 1
for fw_id, fw in wf.id_fw.items():
if 'SRC_task_index' in fw.spec:
src_type = fw.tasks[-1].src_type
if src_type in ['run', 'setup']:
task_index = fw.spec['SRC_task_index']
if 'MPRelaxVasp' in task_index.task_type:
if task_index.task_type == 'MPRelaxVasp':
if src_type == 'run':
fw_ids[(1, task_index.index)] = fw_id
if task_index.index > step_index:
step_index = task_index.index
final_fw_id = fw_id
elif src_type == 'setup':
if task_index.index > final_setup_step_index:
final_setup_step_index = task_index.index
final_setup_fw_id = fw_id
else:
match = re.match("MPRelaxVasp-\d+", task_index.task_type)
if match and match.string == task_index.task_type:
nn = int(task_index.task_type[12:])
if src_type == 'run':
fw_ids[(nn, task_index.index)] = fw_id
if nn == nrelaxations:
if task_index.index > step_index:
step_index = task_index.index
final_fw_id = fw_id
elif nn > nrelaxations:
step_index = task_index.index
final_fw_id = fw_id
nrelaxations = nn
elif src_type == 'setup':
if nn == final_setup_nrelaxations:
if task_index.index > final_setup_step_index:
final_setup_step_index = task_index.index
final_setup_fw_id = fw_id
elif nn > final_setup_nrelaxations:
final_setup_step_index = task_index.index
final_setup_fw_id = fw_id
final_setup_nrelaxations = nn
if final_fw_id is None:
raise RuntimeError('Final structure not found ...')
rlx_and_task_indices = sorted(fw_ids.keys())
ionic_steps = []
helper = MPRelaxTaskHelper()
for rlx_and_task_index in rlx_and_task_indices:
fw_id = fw_ids[rlx_and_task_index]
myfw = wf.id_fw[fw_id]
mytask = myfw.tasks[-1]
last_launch = (myfw.archived_launches + myfw.launches)[-1]
helper.set_task(mytask)
helper.task.setup_rundir(last_launch.launch_dir)
all_info = helper.get_all_info(raise_error=False)
if all_info:
task_ionic_steps = all_info['ionic_steps']
for ionstep in task_ionic_steps:
ionstep.pop('structure')
ionic_steps.append(task_ionic_steps)
final_setup_fw = wf.id_fw[final_setup_fw_id]
final_setup_task = final_setup_fw.tasks[-1]
final_vasp_input_set = final_setup_task.vasp_input_set
kpoints_generation_description = get_kpoints_generation_description(final_vasp_input_set)
all_info_final = helper.get_all_info(raise_error=True)
computed_entry = all_info_final['computed_entry']
vasprun = all_info_final['vasprun']
calculation_parameters = {'INCAR': vasprun.incar,
'INCAR_ALL': vasprun.parameters,
'potcar_symbols': vasprun.potcar_symbols,
'ENCUT': vasprun.incar['ENCUT'] or vasprun.parameters['ENCUT'],
'EDIFFG': vasprun.incar.get('EDIFFG') or vasprun.parameters.get('EDIFFG'),
'ISMEAR': vasprun.incar.get('ISMEAR') or vasprun.parameters.get('ISMEAR'),
'SIGMA': vasprun.incar.get('SIGMA') or vasprun.parameters.get('SIGMA'),
'kpoints_generation_description': kpoints_generation_description}
if computed_entry.data is None:
computed_entry.data = {}
computed_entry.data['calculation_parameters'] = calculation_parameters
return {'computed_entry': computed_entry.as_dict(),
'final_structure': all_info_final['final_structure'].as_dict(),
'ionic_steps': ionic_steps}
