simple_workflow.py

"""Very simple example file for setting up lots of calculations.

IT IS NOT INTENDED TO BE RUN AS IS.

It serves only as an example. We can see that we read in a given input structure and set up many calculations for
that structure with varied volume. All calculations are executed via jobflow-remote.

Could simply add a loop over more starting structures.
"""

from typing import Any

from excitingtools import ExcitingInputXML
from excitingtools.structure.pymatgen_utilities import pymatgen_to_exciting_structure
from jobflow import Flow
from jobflow_remote import submit_flow
from pymatgen.core import Structure
from qtoolkit import QResources

from excitingworkflow import ExcitingCalculation
from excitingworkflow.remote_runner import RemoteRunner, SettingsBinary
from excitingworkflow.workflows.calculation_maker import ExcitingCalculationMaker


def vary_volume(struct: Any) -> list[Structure]:
    """Need to vary the volume. Needs to be implemented.

    :param struct: the original structure
    :return: list of varied structures.
    """


def main() -> None:
    """Create and submit the workflow."""
    pymatgen_structure = Structure.from_file("/path/to/struct.civ")

    input_structures = vary_volume(pymatgen_structure)

    job_list = []
    for i, struct in enumerate(input_structures):
        name = f"run_{i}"
        exciting_structure = pymatgen_to_exciting_structure(struct)
        groundstate = {"rgkmax": 7, "nempty": 5, "ngridk": [4, 4, 4]}
        inp = ExcitingInputXML(groundstate=groundstate, structure=exciting_structure, title=name)

        runner = RemoteRunner(SettingsBinary, run_cmd="srun")

        calc = ExcitingCalculation(name, "", inp, "path/to/speciesfiles", runner=runner)

        maker = ExcitingCalculationMaker(name)
        my_job = maker.make(calc)
        my_job.name = name
        job_list.append(my_job)

    flow = Flow(job_list)
    flow.name = "my_example_flow"
    resources = QResources(nodes=1, processes_per_node=24, threads_per_process=4, time_limit=3600)
    print(submit_flow(flow, "my_worker", exec_config="my_exec_config", resources=resources))


if __name__ == "__main__":
    main()