Update aiida.py

Author: jan-janssen

python_workflow_definition/src/python_workflow_definition/aiida.py

@@ -123,215 +123,3 @@ def write_workflow_json(wg, file_name):
         json.dump(data, f, indent=2)
 
     return data
-
-
-def construct_wg_qe(
-    get_bulk_structure,
-    calculate_qe,
-    generate_structures,
-    plot_energy_volume_curve,
-    strain_lst,
-):
-
-    from .shared import get_dict
-    from .shared import get_list
-
-    # NOTE: `get_dict` is `get_input_dict`, to compile the input values for the calc tasks
-    # NOTE: `add_link` must be from outputs to inputs
-    wg = WorkGraph("wg-qe")
-
-    get_bulk_structure_task = wg.add_task(
-        get_bulk_structure,
-        name="get_bulk_structure",
-        register_pickle_by_value=True,
-    )
-
-    relax_task = wg.add_task(
-        calculate_qe,
-        # ! I don't like the `mini` name...
-        name="mini",
-        register_pickle_by_value=True,
-    )
-
-    generate_structures_task = wg.add_task(
-        generate_structures,
-        name="generate_structures",
-        register_pickle_by_value=True,
-    )
-
-    # here we add the structure outputs based on the number of strains
-    del wg.tasks.generate_structures.outputs["result"]
-
-    scf_qe_tasks = []
-    for i, strain in enumerate(strain_lst):
-        generate_structures_task.add_output("workgraph.any", f"s_{i}")
-
-        scf_qe_task = wg.add_task(
-            calculate_qe,
-            name=f"qe_{i}",
-            register_pickle_by_value=True,
-        )
-        scf_qe_tasks.append(scf_qe_task)
-
-    plot_energy_volume_curve_task = wg.add_task(
-        plot_energy_volume_curve,
-        name="plot_energy_volume_curve",
-        register_pickle_by_value=True,
-    )
-
-    pickle_element_task = wg.add_task(
-        pickle_node,
-        name="pickle_element",
-        value="Al",
-    )
-
-    pickle_a_task = wg.add_task(
-        pickle_node, name="pickle_a", value=4.05
-    )
-
-    pickle_cubic_task = wg.add_task(
-        pickle_node, name="pickle_cubic", value=True
-    )
-
-    pickle_relax_workdir_task = wg.add_task(
-        pickle_node,
-        name="pickle_relax_workdir",
-        value="mini",
-    )
-
-    # ? relax or SCF, or general? -> Should be relax
-    relax_get_dict_task = wg.add_task(
-        task.pythonjob(
-            # outputs=["structure", "calculation", "kpts", "pseudopotentials", "smearing"]
-            # outputs=["dict"]
-        )(get_dict),
-        name="relax_get_dict",
-        register_pickle_by_value=True,
-    )
-
-    pickle_pp_task = wg.add_task(
-        pickle_node,
-        name="pseudopotentials",
-        value={"Al": "Al.pbe-n-kjpaw_psl.1.0.0.UPF"},
-    )
-
-    pickle_kpts_task = wg.add_task(
-        pickle_node, name="kpts_task", value=[3, 3, 3]  # FIXME: Back to [3, 3, 3]
-    )
-
-    pickle_calc_type_relax_task = wg.add_task(
-        pickle_node,
-        name="calc_type_relax",
-        value="vc-relax",
-    )
-
-    pickle_smearing_task = wg.add_task(
-        pickle_node, name="smearing", value=0.02
-    )
-
-    strain_lst_task = wg.add_task(
-        pickle_node,
-        name="pickle_strain_lst",
-        value=strain_lst,
-    )
-
-    strain_dir_tasks, scf_get_dict_tasks = [], []
-    for i, strain in enumerate(strain_lst):
-        strain_dir = f"strain_{i}"
-
-        strain_dir_task = wg.add_task(
-            pickle_node,
-            name=f"pickle_{strain_dir}_dir",
-            value=strain_dir,
-            register_pickle_by_value=True,
-        )
-        strain_dir_tasks.append(strain_dir_task)
-
-        scf_get_dict_task = wg.add_task(
-            task.pythonjob()(get_dict),
-            name=f"get_dict_{i}",
-            register_pickle_by_value=True,
-        )
-        scf_get_dict_tasks.append(scf_get_dict_task)
-
-        if i == 0:
-            pickle_calc_type_scf_task = wg.add_task(
-                pickle_node,
-                name="calc_type_scf",
-                value="scf",
-            )
-
-    get_volumes_task = wg.add_task(
-        task.pythonjob()(get_list),
-        name="get_volumes",
-        register_pickle_by_value=True,
-    )
-    
-    get_energies_task = wg.add_task(
-        task.pythonjob()(get_list),
-        name="get_energies",
-        register_pickle_by_value=True,
-    )
-
-    # Add remaining links
-    wg.add_link(
-        pickle_element_task.outputs.result, get_bulk_structure_task.inputs.element
-    )
-    wg.add_link(pickle_a_task.outputs.result, get_bulk_structure_task.inputs.a)
-    wg.add_link(pickle_cubic_task.outputs.result, get_bulk_structure_task.inputs.cubic)
-
-    # `.set` rather than `.add_link`, as get_dict takes `**kwargs` as input
-    relax_get_dict_task.set(
-        {
-            "structure": get_bulk_structure_task.outputs.result,
-            "calculation": pickle_calc_type_relax_task.outputs.result,
-            "kpts": pickle_kpts_task.outputs.result,
-            "pseudopotentials": pickle_pp_task.outputs.result,
-            "smearing": pickle_smearing_task.outputs.result,
-        }
-    )
-
-    wg.add_link(relax_get_dict_task.outputs.result, relax_task.inputs.input_dict)
-    wg.add_link(
-        pickle_relax_workdir_task.outputs.result,
-        relax_task.inputs.working_directory,
-    )
-
-    wg.add_link(relax_task.outputs.structure, generate_structures_task.inputs.structure)
-    wg.add_link(
-        strain_lst_task.outputs.result, generate_structures_task.inputs.strain_lst
-    )
-
-    for i, (scf_get_dict_task, scf_qe_task, strain_dir_task) in enumerate(
-        list(zip(scf_get_dict_tasks, scf_qe_tasks, strain_dir_tasks))
-    ):
-        scf_get_dict_task.set(
-            {
-                "structure": generate_structures_task.outputs[f"s_{i}"],
-                "calculation": pickle_calc_type_scf_task.outputs.result,
-                "kpts": pickle_kpts_task.outputs.result,
-                "pseudopotentials": pickle_pp_task.outputs.result,
-                "smearing": pickle_smearing_task.outputs.result,
-            }
-        )
-        wg.add_link(scf_get_dict_task.outputs.result, scf_qe_task.inputs.input_dict)
-        wg.add_link(
-            strain_dir_task.outputs.result, scf_qe_task.inputs.working_directory
-        )
-
-        # collect energy and volume
-        # wg.add_link(scf_qe_task.outputs.energy, get_energies_task.inputs.kwargs)
-        get_energies_task.set({f"{i}": scf_qe_task.outputs.energy})
-        # wg.add_link(scf_qe_task.outputs.volume, get_volumes_task.inputs.kwargs)
-        get_volumes_task.set({f"{i}": scf_qe_task.outputs.volume})
-
-    wg.add_link(
-        get_volumes_task.outputs.result,
-        plot_energy_volume_curve_task.inputs.volume_lst,
-    )
-    wg.add_link(
-        get_energies_task.outputs.result,
-        plot_energy_volume_curve_task.inputs.energy_lst,
-    )
-
-    return wg