Passing structure object from Firework to atomate2 flow

FireWorks
1

Hi y’all,

I am trying to create a Firework using MatGL to relax a structure before passing it to VASP for additional dft calculations. I have created a Firework using PyTask, and I am able to use this Firework to perform quick relaxations. This is the function I’m using:

def relax_structure(struct):
    # To suppress warnings for clearer output
    warnings.simplefilter("ignore")
    
    # Load M3GNet potential and relax structure
    pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")
    relaxer = Relaxer(potential=pot)
    relax_results = relaxer.relax(struct, fmax=0.01)
    
    # Return the relaxed structure
    relaxed_structure = relax_results["final_structure"]
    
    return relaxed_structure

And I’m creating a Firework with PyTask using a quick package I created that contains this function

# Create a pytask from a python function relax_structure
fw1 = Firework(PyTask(func='matglrelaxer.relaxer.relax_structure',args=[struct],outputs=['relaxed_structure']),name="matgl relax")

This seems to work fine. However, the issue I encounter is that to pass this structure to an atomate2 workflow, it appears that you need to have the structure available immediately to define the flow object. I am trying to combine these workflows as follows:

# Create a vasp relax workflow
relax_flow = DoubleRelaxMaker(name="firework_test").make(structure = struct) # this is the issue
workflow = flow_to_workflow(relax_flow)
fws = workflow.fws

# Create a new workflow that includes both the FireTask and the existing workflow
new_workflow = Workflow([fw1] + fws,{fw1:fws[0]},name="firework_test")

The new workflow structure seems to be what I want, but I do not know how to pass the relaxed structure object to the DoubleRelaxMaker before it is created Can I pass a reference through spec?. Any ideas?

Thanks,

Kat