Extending the set of abstract variables that can be referenced as part of built in atomate firetasks

Atomate
1

Hi All,

I have been wondering how to seamlessly extend the number of abstract variables that are referenced in built-in Firetasks like RunVaspCustodian. In particular, I am trying to abstract the ‘wall_time’ parameter that is extracted from the optional parameters of the RunVaspCustodian firetask. One way that I was thinking of doing it is to create an inital firetask that sets a set of abstract variables based on where the FW is running and then mutates the fw_spec such that all subsequent firetasks have access to the values of those abstract variables. See the small prototype below:

from atomate.vasp.fireworks.core import OptimizeFW
from fireworks import FiretaskBase, FWAction, explicit_serialize, LaunchPad, Firework, Workflow
from pymatgen.core import Structure
from atomate.utils.utils import env_chk
from CatFlows.dft_settings.settings import MOSurfaceSet
from CatFlows.firetasks.handlers import ContinueOptimizeFW

@explicit_serialize
class SetAbstractVariables(FiretaskBase):

def run_task(self, fw_spec):
    wall_time = env_chk(self.get('wall_time'), fw_spec)
    # Mutate the wall_time variable in the Firework spec for all subsequent firetasks, particularly RunVaspCustodian
    fw_spec['wall_time'] = fw_spec['_fw_env']['wall_time'] # This will take different values based on where the FW runs

Create a new class that inherits from OptimizeFW but which can access user-defined abstract variables

class Opt_FW(OptimizeFW):

def __init__(self, structure, vasp_input_set,  **kwargs):
    super(Opt_FW, self).__init__(structure=structure, vasp_input_set=vasp_input_set)
    # Insert the SetAbstractVariables Firetask into Custom Opt_FW
    self.tasks.insert(0, SetAbstractVariables())

if name == ‘main’:
launchpad = LaunchPad() # Use local MongoDB
launchpad.reset("", require_password=False)
struct = Structure.from_file(‘POSCAR_center_zero’)
vis = MOSurfaceSet(struct, bulk=False)
fw = Opt_FW(structure=struct, vasp_input_set=vis, spec={“wall_time”:">>wall_time<<",})
wf = Workflow([fw], name=“Test_Abstract_Vars”)
launchpad.add_wf(wf)

However, because the RunVaspCustodian file uses self[‘wall_time’] and self.get(‘wall_time’) to retrieve the value of the parameter, it is not feasible for me to do so. Is there any reason we do not do fw_spec[‘wall_time’] in order to get the variable in a dynamic way: None on a certain server that is not constrained by wall_time but 172800 s on NERSC?

Any help is much appreciated.

Thanks.

Best,

Yuri.