Feeding atomate Fireworks with structures from the fw_spec

Hi everybody,

I am planning a quite complex workflow that should generate structures on the fly (e.g. make slabs), relax them, and use them for other stuff (e.g. interfaces between different materials).

My question is: Is there a clever way of using predefined fireworks of atomate (e.g. OptimizeFW) and feed them an input structure that is defined in the fw_spec?
Right now I have a structure stored in a ‘structures’ sub-dictionary of the fw_spec under the key ‘to_optimize’ and I use the custom FireTask below to add the OptimizeFW to my workflow, but this seems unnecessarily complicated and it is difficult to go on from there. Is there a way to access the fw_spec and pass the structure directly on the workflow level? Or am I maybe thinking in the wrong way entirely? Many thanks, Michael

@explicit_serialize
class SpawnOptimizeFW(FiretaskBase):
“”"Spawns a new optimization workflow using a structure saved in the spec.

The structure has to be saved in the 'structures' dictionary within the
fw_spec and has to have the 'to_optimize' key.
"""

_fw_name = 'Spawn Optimize_Workflow'

def run_task(self, fw_spec):
    from atomate.vasp.fireworks.core import OptimizeFW
    
    structure = fw_spec['structures']['to_optimize']
    new_optimize_fw = OptimizeFW(structure)
    
    return FWAction(additions = [new_optimize_fw])