diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 75516dd8..39acc0b7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,6 +1,11 @@ name: continuous-integration -on: [push, pull_request] +on: + push: + # only pushes to main trigger + branches: [main] + pull_request: + # always triggered jobs: @@ -11,29 +16,8 @@ jobs: steps: - uses: actions/checkout@v2 - - name: Set up Python 3.12 + - name: Set up Python 3.11 uses: actions/setup-python@v2 with: - python-version: 3.12 + python-version: 3.11 - uses: pre-commit/action@v2.0.0 - - build: - - runs-on: ubuntu-latest - timeout-minutes: 30 - - steps: - - uses: actions/checkout@v2 - - name: Set up Python 3.12 - uses: actions/setup-python@v2 - with: - python-version: 3.12 - - name: Install python dependencies - run: | - pip install --upgrade pip - pip install -r requirements.txt - - name: "Build HTML docs" - run: | - make -C docs html linkcheck - env: - SPHINXOPTS: -nW --keep-going diff --git a/.readthedocs.yml b/.readthedocs.yml index 0fd6f0e3..d7c9d85f 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -7,7 +7,21 @@ version: 2 build: os: "ubuntu-22.04" tools: - python: "3.12" + python: "miniconda3-3.12-24.1" # note that libmamba-solver is available since 22.1 + jobs: + post_create_environment: + # - python -m pip install --exists-action=w --no-cache-dir -r requirements.txt + - rabbitmq-server -detached + - sleep 10 + - rabbitmq-diagnostics status + - verdi presto + - verdi daemon start + - verdi status + - aiida-pseudo install sssp -x PBEsol + - verdi group list + - cat /proc/cpuinfo | grep processor | wc -l + - python -m sphinx -T --keep-going -b linkcheck -d docs/_build/doctrees -D language=en docs $READTHEDOCS_OUTPUT/html + # Build documentation in the docs/ directory with Sphinx sphinx: @@ -19,7 +33,7 @@ sphinx: # See https://docs.readthedocs.io/en/latest/yaml-config.html#formats formats: [] +conda: + environment: environment.yml + # Optionally set the version of Python and requirements required to build your docs -python: - install: - - requirements: requirements.txt diff --git a/README.md b/README.md index 60d0bb65..6e1da2fc 100644 --- a/README.md +++ b/README.md @@ -26,17 +26,19 @@ If you have a question, feel free to just [open an issue](https://github.com/aii ### Prerequisites -* python 3.12 or greater +* python 3.11 ### Build instructions ```bash git clone https://github.com/aiidateam/aiida-tutorials.git cd aiida-tutorials -pip install -r requirements.txt +conda env create --quiet --name aiida-tutorials --file environment.yml +conda activate aiida-tutorials pre-commit install # enable pre-commit hooks (optional) -cd docs/ -make +make -C docs html # to build docs +make -C docs linkcheck # to run link checks (only for dev) + # open build/html/index.html ``` diff --git a/docs/conf.py b/docs/conf.py index 3cf19061..a4f30e55 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -42,6 +42,7 @@ "sphinx_copybutton", "sphinx_panels", "sphinx_tabs.tabs", + "sphinx_gallery.gen_gallery", ] myst_enable_extensions = [ @@ -68,7 +69,7 @@ ipython_mplbackend = "" copybutton_selector = "div:not(.no-copy)>div.highlight pre" -copybutton_prompt_text = ">>> |\.\.\. |\$ |In \[\d*\]: | {2,5}\.\.\.: | {5,8}: " +copybutton_prompt_text = ">>> |... |$ |In [d*]: | {2,5}...: | {5,8}: " copybutton_prompt_is_regexp = True todo_include_todos = True @@ -391,13 +392,12 @@ suppress_warnings = ["misc.highlighting_failure"] # Links we ignore, because they do not work temporary and we cannot fix it -linkcheck_ignore = ["https://www.big-map.eu/"] - - -def setup(app): - """Setup function called by sphinx.""" - app.add_css_file("css/custom.css") - +linkcheck_ignore = [ + "https://www.big-map.eu/", + r".*concept/index.html", + r".*howto/index.html", + "http://127.0.0.1:8000/workgraph", +] # We are not installing a full aiida environment nb_execution_mode = "off" @@ -408,6 +408,29 @@ def setup(app): "plumpy": ("https://plumpy.readthedocs.io/en/latest/", None), } +gallery_src_relative_dir = "gallery" # relative path of the gallery src wrt. sphinx src + +# we mimik the structure in the sphinx src directory in the gallery src directory +sphinx_src_autogen_dirs = ["sections/writing_workflows_with_workgraph/autogen"] + +gallery_src_dirs = [ + os.path.join(gallery_src_relative_dir, autogen_dir) + for autogen_dir in sphinx_src_autogen_dirs +] # path of the python scripts that should be executed +sphinx_gallery_conf = { + "filename_pattern": "/*", + "examples_dirs": gallery_src_dirs, # in sphinx-gallery doc referred as gallery source + "gallery_dirs": sphinx_src_autogen_dirs, # path to where to gallery puts generated files +} + +# ignore in the autogenerated ipynb files to surpress warning +exclude_patterns.extend( + [ + os.path.join(sphinx_src_autogen_dir, "*ipynb") + for sphinx_src_autogen_dir in sphinx_src_autogen_dirs + ] +) + # Compile all things needed before building the docs # For instance, convert the notebook templates to actual tutorial and solution versions print( @@ -416,3 +439,57 @@ def setup(app): universal_newlines=True, ) ) + +import shutil +from pathlib import Path + + +def copy_html_files(app, exception): + """ + Copy all .html files from source to build directory, maintaining the directory structure. + """ + print("Copying HTML files to build directory") + copy_print_info = "Copying HTML files to build directory" + print() + print(copy_print_info) + print(len(copy_print_info) * "=") + if exception is not None: # Only copy files if the build succeeded + print( + "Build failed, but we still try to copy the HTML files to the build directory" + ) + try: + src_path = Path(app.builder.srcdir) + build_path = Path(app.builder.outdir) + + copy_print_info = f"Copying html files from sphinx src directory {src_path}" + print() + print(copy_print_info) + print(len(copy_print_info) * "-") + for html_file in src_path.rglob("*.html"): + relative_path = html_file.relative_to(src_path) + destination_file = build_path / relative_path + destination_file.parent.mkdir(parents=True, exist_ok=True) + shutil.copy(html_file, destination_file) + print(f"Copy {html_file} to {destination_file}") + + gallery_src_path = Path(app.builder.srcdir / Path(gallery_src_relative_dir)) + + copy_print_info = ( + f"Copying html files from gallery src directory {gallery_src_path} to build" + ) + print() + print(copy_print_info) + print(len(copy_print_info) * "-") + for html_file in gallery_src_path.rglob("*.html"): + relative_path = html_file.relative_to(gallery_src_path) + destination_file = build_path / relative_path + destination_file.parent.mkdir(parents=True, exist_ok=True) + shutil.copy(html_file, destination_file) + print(f"Copy {html_file} to {destination_file}") + except Exception as e: + print(f"Failed to copy HTML files: {e}") + + +def setup(app): + app.add_css_file("css/custom.css") + app.connect("build-finished", copy_html_files) diff --git a/docs/gallery/sections/writing_workflows_with_workgraph/autogen/GALLERY_HEADER.rst b/docs/gallery/sections/writing_workflows_with_workgraph/autogen/GALLERY_HEADER.rst new file mode 100644 index 00000000..71f830dd --- /dev/null +++ b/docs/gallery/sections/writing_workflows_with_workgraph/autogen/GALLERY_HEADER.rst @@ -0,0 +1,5 @@ +:orphan: + +================================ +Running workflows with WorkGraph +================================ diff --git a/docs/gallery/sections/writing_workflows_with_workgraph/autogen/eos.py b/docs/gallery/sections/writing_workflows_with_workgraph/autogen/eos.py new file mode 100644 index 00000000..8d88a42a --- /dev/null +++ b/docs/gallery/sections/writing_workflows_with_workgraph/autogen/eos.py @@ -0,0 +1,283 @@ +""" +================================== +Equation of state (EOS) WorkGraph +================================== + +""" + +# %% +# To run this tutorial, you need to install aiida-workgraph and restart the daemon. Open a terminal and run: +# +# .. code-block:: console +# +# pip install aiida-workgraph[widget] aiida-quantumespresso +# +# Restart (or start) the AiiDA daemon if needed: +# +# .. code-block:: console +# +# verdi daemon restart +# +# Create the calcfunction task +# ============================ +# + + +from aiida import orm +from aiida_workgraph import task + +# +# explicitly define the output socket name to match the return value of the function +@task.calcfunction(outputs=[{"name": "structures"}]) +def scale_structure(structure, scales): + """Scale the structure by the given scales.""" + atoms = structure.get_ase() + structures = {} + for i in range(len(scales)): + atoms1 = atoms.copy() + atoms1.set_cell(atoms.cell * scales[i], scale_atoms=True) + structure = orm.StructureData(ase=atoms1) + structures[f"s_{i}"] = structure + return {"structures": structures} + + +# +# Output result from context to the output socket +@task.graph_builder(outputs=[{"name": "result", "from": "context.result"}]) +def all_scf(structures, scf_inputs): + """Run the scf calculation for each structure.""" + from aiida_workgraph import WorkGraph + from aiida_quantumespresso.calculations.pw import PwCalculation + + wg = WorkGraph() + for key, structure in structures.items(): + pw1 = wg.add_task(PwCalculation, name=f"pw1_{key}", structure=structure) + pw1.set(scf_inputs) + # save the output parameters to the context + pw1.set_context({"output_parameters": f"result.{key}"}) + return wg + + +# + + +@task.calcfunction() +# because this is a calcfunction, and the input datas are dynamic, we need use **datas. +def eos(**datas): + """Fit the EOS of the data.""" + from ase.eos import EquationOfState + + # + volumes = [] + energies = [] + for _, data in datas.items(): + volumes.append(data.dict.volume) + energies.append(data.dict.energy) + unit = data.dict.energy_units + # + eos = EquationOfState(volumes, energies) + v0, e0, B = eos.fit() + eos = orm.Dict({"unit": unit, "v0": v0, "e0": e0, "B": B}) + return eos + + +# %% +# Build the workgraph +# =================== +# Three steps: +# +# - create an empty WorkGraph +# - add tasks: scale_structure, all_scf and eos. +# - link the output and input sockets for the tasks. +# +# Visualize the workgraph +# ----------------------- +# If you are running in a jupyter notebook, you can visualize the workgraph directly. +# + +from aiida_workgraph import WorkGraph + +# +wg = WorkGraph("eos") +scale_structure1 = wg.add_task(scale_structure, name="scale_structure1") +all_scf1 = wg.add_task(all_scf, name="all_scf1") +eos1 = wg.add_task(eos, name="eos1") +wg.add_link(scale_structure1.outputs["structures"], all_scf1.inputs["structures"]) +wg.add_link(all_scf1.outputs["result"], eos1.inputs["datas"]) +wg.to_html() +# visualize the workgraph in jupyter-notebook +# wg + + +# %% +# Prepare inputs and run +# ---------------------- +# + + +from aiida import load_profile +from aiida.common.exceptions import NotExistent +from aiida.orm import ( + Dict, + KpointsData, + StructureData, + load_code, + load_group, + InstalledCode, + load_computer, +) +from ase.build import bulk + +# +load_profile() +# create pw code +try: + pw_code = load_code( + "qe-7.2-pw@localhost" + ) # The computer label can also be omitted here +except NotExistent: + pw_code = InstalledCode( + computer=load_computer("localhost"), + filepath_executable="pw.x", + label="qe-7.2-pw", + default_calc_job_plugin="quantumespresso.pw", + ).store() +# +si = orm.StructureData(ase=bulk("Si")) +pw_paras = Dict( + { + "CONTROL": { + "calculation": "scf", + }, + "SYSTEM": { + "ecutwfc": 30, + "ecutrho": 240, + "occupations": "smearing", + "smearing": "gaussian", + "degauss": 0.1, + }, + } +) +# Load the pseudopotential family. +pseudo_family = load_group("SSSP/1.3/PBEsol/efficiency") +pseudos = pseudo_family.get_pseudos(structure=si) +# +metadata = { + "options": { + "resources": { + "num_machines": 1, + "num_mpiprocs_per_machine": 1, + }, + } +} +# +kpoints = orm.KpointsData() +kpoints.set_kpoints_mesh([3, 3, 3]) +pseudos = pseudo_family.get_pseudos(structure=si) +scf_inputs = { + "code": pw_code, + "parameters": pw_paras, + "kpoints": kpoints, + "pseudos": pseudos, + "metadata": metadata, +} +# ------------------------------------------------------- +# set the input parameters for each task +wg.tasks["scale_structure1"].set({"structure": si, "scales": [0.95, 1.0, 1.05]}) +wg.tasks["all_scf1"].set({"scf_inputs": scf_inputs}) +print("Waiting for the workgraph to finish...") +wg.submit(wait=True, timeout=300) +# one can also run the workgraph directly +# wg.run() + + +# %% +# Print out the results: +# + + +data = wg.tasks["eos1"].outputs["result"].value.get_dict() +print("B: {B}\nv0: {v0}\ne0: {e0}\nv0: {v0}".format(**data)) + +# %% +# Use graph builder +# ================= +# The Graph Builder allow user to create a dynamic workflow based on the input value, as well as nested workflows. +# + +from aiida_workgraph import WorkGraph, task + +# +@task.graph_builder(outputs=[{"name": "result", "from": "eos1.result"}]) +def eos_workgraph(structure=None, scales=None, scf_inputs=None): + wg = WorkGraph("eos") + scale_structure1 = wg.add_task( + scale_structure, name="scale_structure1", structure=structure, scales=scales + ) + all_scf1 = wg.add_task(all_scf, name="all_scf1", scf_inputs=scf_inputs) + eos1 = wg.add_task(eos, name="eos1") + wg.add_link(scale_structure1.outputs["structures"], all_scf1.inputs["structures"]) + wg.add_link(all_scf1.outputs["result"], eos1.inputs["datas"]) + return wg + + +# %% +# Then we can use the `eos_workgraph` in two ways: +# +# - Direct run the function and generate the workgraph, then submit +# - Use it as a task inside another workgraph to create nested workflow. +# +# Use the graph builder directly +# ------------------------------ +# + +wg = eos_workgraph(structure=si, scales=[0.95, 1.0, 1.05], scf_inputs=scf_inputs) +# One can submit the workgraph directly +# wg.submit(wait=True, timeout=300) +wg.to_html() +# visualize the workgraph in jupyter-notebook +# wg + +# %% +# Use it inside another workgraph +# ------------------------------- +# For example, we want to combine relax with eos. +# + + +from aiida_workgraph import WorkGraph +from copy import deepcopy +from aiida_quantumespresso.calculations.pw import PwCalculation + +# +# ------------------------------------------------------- +relax_pw_paras = deepcopy(pw_paras) +relax_pw_paras["CONTROL"]["calculation"] = "vc-relax" +relax_inputs = { + "structure": si, + "code": pw_code, + "parameters": relax_pw_paras, + "kpoints": kpoints, + "pseudos": pseudos, + "metadata": metadata, +} +# ------------------------------------------------------- +wg = WorkGraph("relax_eos") +relax_task = wg.add_task(PwCalculation, name="relax1") +relax_task.set(relax_inputs) +eos_wg_task = wg.add_task( + eos_workgraph, name="eos1", scales=[0.95, 1.0, 1.05], scf_inputs=scf_inputs +) +wg.add_link(relax_task.outputs["output_structure"], eos_wg_task.inputs["structure"]) +# ------------------------------------------------------- +# One can submit the workgraph directly +# wg.submit(wait=True, timeout=300) + +wg.to_html() +# visualize the workgraph in jupyter-notebook +# wg + +# %% +# Summary +# ======= +# There are many ways to create the workflow using graph builder. For example, one can add the relax step inside the `eos_workgraph`, and add a `run_relax` argument to control the logic. diff --git a/docs/gallery/sections/writing_workflows_with_workgraph/autogen/qe.py b/docs/gallery/sections/writing_workflows_with_workgraph/autogen/qe.py new file mode 100644 index 00000000..3cc2861f --- /dev/null +++ b/docs/gallery/sections/writing_workflows_with_workgraph/autogen/qe.py @@ -0,0 +1,357 @@ +""" +=============================== +Computational materials science +=============================== + +""" + +# %% +# Introduction +# ============ +# In this tutorial, you will use `AiiDA-WorkGraph` to carry out a DFT calculation using Quantum ESPRESSO. +# +# Requirements +# ------------ +# To run this tutorial, you need to install `aiida-workgraph`, `aiida-quantumespresso` and `aiida-pseudo`. Open a terminal and run: +# +# .. code-block:: console +# +# pip install aiida-workgraph aiida-quantumespresso aiida-pseudo +# aiida-pseudo install sssp -x PBEsol +# +# Start the AiiDA daemon if needed: +# +# .. code-block:: console +# +# verdi daemon start +# +# Start the web server +# -------------------- +# +# Open a terminal, and run: +# +# .. code-block:: console +# +# workgraph web start +# +# Then visit the page `http://127.0.0.1:8000/workgraph`, where you can view the workgraph later. +# +# Load the AiiDA profile. +# + + +from aiida import load_profile + +load_profile() +# +# %% +# First workflow: calculate the energy of N2 molecule +# =================================================== +# Define a workgraph +# ------------------ +# aiida-quantumespresso provides a CalcJob: `PwCalculation` to run a PW calculation. we can use it directly in the WorkGraph. The inputs and outputs of the task is automatically generated based on the `PwCalculation` CalcJob. +# + +from aiida_quantumespresso.calculations.pw import PwCalculation +from aiida_workgraph import WorkGraph + +# +wg = WorkGraph("energy_n2") +pw1 = wg.add_task(PwCalculation, name="pw1") +pw1.to_html() +# +# visualize the task in jupyter-notebook +# pw1 +# + +# %% +# Prepare the inputs and submit the workflow +# ------------------------------------------ +# +# + +from aiida import load_profile +from aiida.common.exceptions import NotExistent +from aiida.orm import ( + Dict, + KpointsData, + StructureData, + load_code, + load_group, + InstalledCode, + load_computer, +) +from ase.build import molecule + +# +load_profile() +# create pw code +try: + pw_code = load_code( + "qe-7.2-pw@localhost" + ) # The computer label can also be omitted here +except NotExistent: + pw_code = InstalledCode( + computer=load_computer("localhost"), + filepath_executable="pw.x", + label="qe-7.2-pw", + default_calc_job_plugin="quantumespresso.pw", + ).store() +# create input structure +mol = molecule("N2") +mol.center(vacuum=1.5) +mol.pbc = True +structure_n2 = StructureData(ase=mol) +paras = Dict( + { + "CONTROL": { + "calculation": "scf", + }, + "SYSTEM": { + "ecutwfc": 30, + "ecutrho": 240, + "occupations": "smearing", + "smearing": "gaussian", + "degauss": 0.1, + }, + } +) +kpoints = KpointsData() +kpoints.set_kpoints_mesh([1, 1, 1]) +# Load the pseudopotential family. +pseudo_family = load_group("SSSP/1.3/PBEsol/efficiency") +pseudos = pseudo_family.get_pseudos(structure=structure_n2) +# +metadata = { + "options": { + "resources": { + "num_machines": 1, + "num_mpiprocs_per_machine": 1, + }, + } +} +# +# ------------------------- Set the inputs ------------------------- +pw1.set( + { + "code": pw_code, + "structure": structure_n2, + "parameters": paras, + "kpoints": kpoints, + "pseudos": pseudos, + "metadata": metadata, + } +) +# ------------------------- Submit the calculation ------------------------- +wg.submit(wait=True, timeout=200) +# ------------------------- Print the output ------------------------- +print( + "Energy of an un-relaxed N2 molecule: {:0.3f}".format( + pw1.outputs["output_parameters"].value.get_dict()["energy"] + ) +) +# + +# %% +# Generate node graph from the AiiDA process: +# + +from aiida_workgraph.utils import generate_node_graph + +generate_node_graph(wg.pk) + + +# %% +# Second workflow: atomization energy of N2 molecule +# ================================================== +# +# The atomization energy of :math:`N_2` is defined as the energy difference between the :math:`N_2` molecule and two isolated N atoms. +# +# .. code-block:: python +# +# e_atomization = 2 * e_atom - e_molecule + +# Define a calcfunction to calculate the atomization energy +# --------------------------------------------------------- +# + +from aiida_workgraph import task + +# +@task.calcfunction() +def atomization_energy(output_atom, output_mol): + from aiida.orm import Float + + e = output_atom["energy"] * output_mol["number_of_atoms"] - output_mol["energy"] + return Float(e) + + +# %% +# Create the structure of nitrogen Atom. +# + +from ase import Atoms +from aiida.orm import StructureData + +# +atoms = Atoms("N") +atoms.center(vacuum=1.5) +atoms.pbc = True +structure_n = StructureData(ase=atoms) + +# %% +# Create a workgraph +# ------------------ + + +from aiida_workgraph import WorkGraph +from aiida.orm import load_code + +# +# load the PW code +pw_code = load_code("qe-7.2-pw@localhost") +# +wg = WorkGraph("atomization_energy") +# +# create the PW task +pw_n = wg.add_task(PwCalculation, name="pw_n") +pw_n.set( + { + "code": pw_code, + "structure": structure_n, + "parameters": paras, + "kpoints": kpoints, + "pseudos": pseudos, + "metadata": metadata, + } +) +pw_n2 = wg.add_task(PwCalculation, name="pw_n2") +pw_n2.set( + { + "code": pw_code, + "structure": structure_n2, + "parameters": paras, + "kpoints": kpoints, + "pseudos": pseudos, + "metadata": metadata, + } +) +# create the task to calculate the atomization energy +atomization = wg.add_task(atomization_energy, name="atomization_energy") +wg.add_link(pw_n.outputs["output_parameters"], atomization.inputs["output_atom"]) +wg.add_link(pw_n2.outputs["output_parameters"], atomization.inputs["output_mol"]) +wg.to_html() + + +# %% +# Submit the workgraph and print the atomization energy. +# + + +wg.submit(wait=True, timeout=300) +print( + "Atomization energy: {:0.3f} eV".format(atomization.outputs["result"].value.value) +) + + +# %% +# If you start the web app (`workgraph web start`), you can visit the page http://127.0.0.1:8000/workgraph to view the tasks. +# +# You can also generate node graph from the AiiDA process: +# + + +from aiida_workgraph.utils import generate_node_graph + +generate_node_graph(wg.pk) + +# %% +# Use already existing workchain +# ============================== +# Can we register a task from a workchain? Can we set the a input item of a namespace? Yes, we can! +# +# In the `PwRelaxWorkChain`, one can set the relax type (`calculation` key) in the input namespace `base.pw.parameters`. Now we create a new task to update the pw parameters. +# + +from aiida_workgraph import task + + +@task.calcfunction() +def pw_parameters(paras, relax_type): + paras1 = paras.clone() + paras1["CONTROL"]["calculation"] = relax_type + return paras1 + + +# %% +# Now, we create the workgraph to relax the structure of N2 molecule. +# + +from aiida_quantumespresso.workflows.pw.relax import PwRelaxWorkChain + +# +wg = WorkGraph("test_pw_relax") +# pw task +pw_relax1 = wg.add_task(PwRelaxWorkChain, name="pw_relax1") +# Load the pseudopotential family. +pseudos = pseudo_family.get_pseudos(structure=structure_n2) +pw_relax1.set( + { + "base": { + "pw": {"code": pw_code, "pseudos": pseudos, "metadata": metadata}, + "kpoints": kpoints, + }, + "structure": structure_n2, + }, +) +paras_task = wg.add_task(pw_parameters, "parameters", paras=paras, relax_type="relax") +wg.add_link(paras_task.outputs[0], pw_relax1.inputs["base.pw.parameters"]) +# One can submit the workgraph directly +# wg.submit(wait=True, timeout=200) +# print( +# "\nEnergy of a relaxed N2 molecule: {:0.3f}".format( +# pw_relax1.node.outputs.output_parameters.get_dict()["energy"] +# ) +# ) + + +# %% +# Use `protocol` to set input parameters (Experimental) +# ===================================================== +# The aiida-quantumespresso package supports setting input parameters from protocol. For example, the PwRelaxWorkChain has a `get_builder_from_protocol` method. In this tutorial, we will show how to use the `protocol` to set the input parameters inside the WorkGraph. +# + +from aiida_workgraph import build_task, WorkGraph +from aiida_quantumespresso.workflows.pw.relax import PwRelaxWorkChain +from ase.build import bulk +from aiida import orm +from pprint import pprint + +# +pw_code = orm.load_code("qe-7.2-pw@localhost") +wg = WorkGraph("test_pw_relax") +structure_si = orm.StructureData(ase=bulk("Si")) +pw_relax1 = wg.add_task(PwRelaxWorkChain, name="pw_relax1") +# set the inputs from the protocol +# this will call the `PwRelaxWorkChain.get_builder_from_protocol` method +# to set the inputs of the workchain +pw_relax1.set_from_protocol( + pw_code, structure_si, protocol="fast", pseudo_family="SSSP/1.2/PBEsol/efficiency" +) +# we can now inspect the inputs of the workchain +print("The inputs for the PwBaseWorkchain are:") +print("-" * 80) +pprint(pw_relax1.inputs["base"].value) +print("\nThe input parameters for pw are:") +print("-" * 80) +pprint(pw_relax1.inputs["base"].value["pw"]["parameters"].get_dict()) + + +# %% +# One can also adjust the parameters of the `PwRelaxWorkChain` to from protocol. +# + +# For example, we want to remove the `base_final_scf` from the inputs, so that the `PwRelaxWorkChain` will not run the `base_final_scf` step. +pw_relax1.inputs["base_final_scf"].value = None +# submit the workgraph +# wg.submit(wait=True, timeout=200) diff --git a/docs/gallery/sections/writing_workflows_with_workgraph/autogen/zero_to_hero.py b/docs/gallery/sections/writing_workflows_with_workgraph/autogen/zero_to_hero.py new file mode 100644 index 00000000..664ca9a2 --- /dev/null +++ b/docs/gallery/sections/writing_workflows_with_workgraph/autogen/zero_to_hero.py @@ -0,0 +1,527 @@ +""" +====================================== +AiiDA-WorkGraph: From Zero To Hero +====================================== + +""" + +# %% +# In this tutorial, you will learn `AiiDA-WorkGraph` to build your workflow to carry out DFT calculation. It's recommended to run this tutorial inside a Jupyter notebook. +# +# Requirements +# =============== +# To run this tutorial, you need to install `aiida-workgraph`, `aiida-quantumespresso`. Open a terminal and run: +# +# .. code-block:: console +# +# pip install aiida-workgraph[widget] aiida-quantumespresso +# +# Restart (or start) the AiiDA daemon if needed: +# +# .. code-block:: console +# +# verdi daemon restart +# +# Load the AiiDA profile. +# + + +from aiida import load_profile + +load_profile() +# + +# %% +# First workflow +# =============== +# Suppose we want to calculate ```(x + y) * z ``` in two steps. First, add `x` and `y`, then multiply the result with `z`. +# +# In AiiDA, we can define two `calcfunction` to do the `add` and `mutiply`: +# + +from aiida_workgraph import task + + +@task.calcfunction() +def add(x, y): + return x + y + + +@task.calcfunction() +def multiply(x, y): + return x * y + + +# %% +# Create the workflow +# -------------------- +# Three steps: +# +# - create a empty WorkGraph +# - add tasks: `add` and `multiply`. +# - link the output of the `add` task to the `x` input of the `multiply` task. +# +# +# In a jupyter notebook, you can visualize the workgraph directly. +# + +from aiida_workgraph import WorkGraph + +# +wg = WorkGraph("add_multiply_workflow") +wg.add_task(add, name="add1") +wg.add_task(multiply, name="multiply1", x=wg.tasks["add1"].outputs["result"]) +# export the workgraph to html file so that it can be visualized in a browser +wg.to_html() +# visualize the workgraph in jupyter-notebook +# wg +# + +# %% +# Submit the workgraph +# ------------------------- +# + + +from aiida_workgraph.utils import generate_node_graph +from aiida.orm import Int + +# +# ------------------------- Submit the calculation ------------------- +wg.submit( + inputs={"add1": {"x": Int(2), "y": Int(3)}, "multiply1": {"y": Int(4)}}, wait=True +) +# ------------------------- Print the output ------------------------- +assert wg.tasks["multiply1"].outputs["result"].value == 20 +print( + "\nResult of multiply1 is {} \n\n".format( + wg.tasks["multiply1"].outputs["result"].value + ) +) +# ------------------------- Generate node graph ------------------- +generate_node_graph(wg.pk) +# + +# %% +# CalcJob and WorkChain +# ======================= +# AiiDA uses `CalcJob` to run a calculation on a remote computer. AiiDA community also provides a lot of well-written `calcfunction` and `WorkChain`. One can use these AiiDA component direclty in the WorkGraph. The inputs and outputs of the task is automatically generated based on the input and output port of the AiiDA component. +# +# Here is an example of using the `ArithmeticAddCalculation` Calcjob inside the workgraph. Suppose we want to calculate ```(x + y) + z ``` in two steps. +# + + +from aiida_workgraph import WorkGraph +from aiida.calculations.arithmetic.add import ArithmeticAddCalculation + +# +wg = WorkGraph("test_calcjob") +new = wg.add_task +new(ArithmeticAddCalculation, name="add1") +wg.add_task(ArithmeticAddCalculation, name="add2", x=wg.tasks["add1"].outputs["sum"]) +wg.to_html() +# + +# %% +# Inspect the node +# ---------------- +# How do I know which input and output to connect? +# +# The inputs and outputs of a task are generated automatically based on the inputs/outputs of the AiiDA component. WorkGraph also has some built-in ports, like `_wait` and `_outputs`. One can inpsect a task's inputs and outputs. +# +# Note: special case for `calcfunction`, the default name of its output is `result`. +# + + +# visualize the task +wg.tasks["add1"].to_html() +# + +# %% +# First Real-world Workflow: atomization energy of molecule +# ========================================================== +# +# The atomization energy, $\Delta E$, of a molecule can be expressed as: +# +# .. math:: +# +# \Delta E = n_{\text{atom}} \times E_{\text{atom}} - E_{\text{molecule}} +# +# Where: +# +# - :math:`\Delta E` is the atomization energy of the molecule. +# - :math:`n_{\text{atom}}` is the number of atoms. +# - :math:`E_{\text{atom}}` is the energy of an isolated atom. +# - :math:`E_{\text{molecule}}` is the energy of the molecule. +# +# Define a workgraph +# ------------------- +# aiida-quantumespresso provides `PwCalculation` CalcJob and `PwBaseWorkChain` to run a PW calculation. we can use it directly in the WorkGraph. Here we use the `PwCalculation` CalcJob. +# + + +from aiida_workgraph import WorkGraph +from aiida.engine import calcfunction +from aiida_quantumespresso.calculations.pw import PwCalculation + +# + + +@calcfunction +def atomization_energy(output_atom, output_mol): + from aiida.orm import Float + + e = output_atom["energy"] * output_mol["number_of_atoms"] - output_mol["energy"] + return Float(e) + + +# +wg = WorkGraph("atomization_energy") +pw_atom = wg.add_task(PwCalculation, name="pw_atom") +pw_mol = wg.add_task(PwCalculation, name="pw_mol") +# create the task to calculate the atomization energy +wg.add_task( + atomization_energy, + name="atomization_energy", + output_atom=pw_atom.outputs["output_parameters"], + output_mol=pw_mol.outputs["output_parameters"], +) +# export the workgraph to html file so that it can be visualized in a browser +wg.to_html() +# visualize the workgraph in jupyter-notebook +# wg +# + +# %% +# Prepare the inputs and submit the workflow +# ------------------------------------------- +# You need to set up the code, computer, and pseudo potential for the calculation. Please refer to the this [documentation](https://aiida-quantumespresso.readthedocs.io/en/latest/installation/index.html) for more details. +# +# You can also stip this step. +# + +from aiida import load_profile +from aiida.common.exceptions import NotExistent +from aiida.orm import ( + Dict, + KpointsData, + StructureData, + load_code, + load_group, + InstalledCode, + load_computer, +) +from ase.build import molecule +from ase import Atoms + +# +load_profile() +# create pw code +try: + pw_code = load_code( + "qe-7.2-pw@localhost" + ) # The computer label can also be omitted here +except NotExistent: + pw_code = InstalledCode( + computer=load_computer("localhost"), + filepath_executable="pw.x", + label="qe-7.2-pw", + default_calc_job_plugin="quantumespresso.pw", + ).store() +# create structure +n_atom = Atoms("N") +n_atom.center(vacuum=1.5) +n_atom.pbc = True +structure_n = StructureData(ase=n_atom) +structure_n2 = StructureData(ase=molecule("N2", vacuum=1.5, pbc=True)) +# create the PW task +paras = Dict( + { + "CONTROL": { + "calculation": "scf", + }, + "SYSTEM": { + "ecutwfc": 30, + "ecutrho": 240, + "occupations": "smearing", + "smearing": "gaussian", + "degauss": 0.1, + }, + } +) +kpoints = KpointsData() +kpoints.set_kpoints_mesh([1, 1, 1]) +# Load the pseudopotential family. +pseudo_family = load_group("SSSP/1.3/PBEsol/efficiency") +pseudos = pseudo_family.get_pseudos(structure=structure_n2) +# +metadata = { + "options": { + "resources": { + "num_machines": 1, + "num_mpiprocs_per_machine": 1, + }, + } +} +# +# ------------------------- Set the inputs ------------------------- +wg.tasks["pw_atom"].set( + { + "code": pw_code, + "structure": structure_n, + "parameters": paras, + "kpoints": kpoints, + "pseudos": pseudos, + "metadata": metadata, + } +) +wg.tasks["pw_mol"].set( + { + "code": pw_code, + "structure": structure_n2, + "parameters": paras, + "kpoints": kpoints, + "pseudos": pseudos, + "metadata": metadata, + } +) +# ------------------------- Submit the calculation ------------------- +wg.submit(wait=True, timeout=200) +# ------------------------- Print the output ------------------------- +print( + "Energy of a N atom: {:0.3f}".format( + wg.tasks["pw_atom"].outputs["output_parameters"].value.get_dict()["energy"] + ) +) +print( + "Energy of an un-relaxed N2 molecule: {:0.3f}".format( + wg.tasks["pw_mol"].outputs["output_parameters"].value.get_dict()["energy"] + ) +) +print( + "Atomization energy: {:0.3f} eV".format( + wg.tasks["atomization_energy"].outputs["result"].value.value + ) +) +# + +# %% +# Generate node graph from the AiiDA process: +# + + +from aiida_workgraph.utils import generate_node_graph + +generate_node_graph(wg.pk) +# + +# %% +# Advanced Topic: Dynamic Workgraph +# ================================== +# +# Graph builder +# -------------- +# If we want to generate the workgraph on-the-fly, for example, if you want to use `if` to create the tasks, or repeat a calculation until it converges, you can use Graph Builder. +# +# Suppose we want to calculate: +# +# .. code-block:: python +# +# # step 1 +# result = add(x, y) +# # step 2 +# if result > 0: +# result = add(result, y) +# else: +# result = multiply(result, y) +# # step 3 +# result = add(result, y) +# + + +# Create a WorkGraph which is dynamically generated based on the input +# then we output the result of from the context +from aiida_workgraph import task + +# +@task.graph_builder(outputs=[{"name": "result", "from": "context.result"}]) +def add_multiply_if_generator(x, y): + wg = WorkGraph() + if x.value > 0: + add1 = wg.add_task(add, name="add1", x=x, y=y) + # export the result of add1 to the context, so that context.result = add1.results + add1.set_context({"result": "result"}) + else: + multiply1 = wg.add_task(multiply, name="multiply1", x=x, y=y) + # export the result of multiply1 to the context + multiply1.set_context({"result": "result"}) + return wg + + +# +wg = WorkGraph("if_task") +wg.add_task(add, name="add1") +wg.add_task( + add_multiply_if_generator, + name="add_multiply_if1", + x=wg.tasks["add1"].outputs["result"], +) +wg.add_task(add, name="add2", x=wg.tasks["add_multiply_if1"].outputs["result"]) +wg.to_html() +# + +# %% +# Submit the WorkGraph + + +wg.submit( + inputs={ + "add1": {"x": 1, "y": 2}, + "add_multiply_if1": {"y": 2}, + "add2": {"y": 2}, + }, + wait=True, +) +# ------------------------- Print the output ------------------------- +assert wg.tasks["add2"].outputs["result"].value == 7 +print("\nResult of add2 is {} \n\n".format(wg.tasks["add2"].outputs["result"].value)) +# +# %% +# Note: one can not see the detail of the `add_multiply_if1` before you running it. +# +# Second Real-world Workflow: Equation of state (EOS) WorkGraph +# ============================================================= +# +# First, create the calcfunction for the job. +# + +from aiida import orm +from aiida_workgraph import task + +# +# explicitly define the output socket name to match the return value of the function +@task.calcfunction(outputs=[{"name": "structures"}]) +def scale_structure(structure, scales): + """Scale the structure by the given scales.""" + atoms = structure.get_ase() + structures = {} + for i in range(len(scales)): + atoms1 = atoms.copy() + atoms1.set_cell(atoms.cell * scales[i], scale_atoms=True) + structure = orm.StructureData(ase=atoms1) + structures[f"s_{i}"] = structure + return {"structures": structures} + + +# +# Output result from context to the output socket +@task.graph_builder(outputs=[{"name": "result", "from": "context.result"}]) +def all_scf(structures, scf_inputs): + """Run the scf calculation for each structure.""" + from aiida_workgraph import WorkGraph + from aiida_quantumespresso.calculations.pw import PwCalculation + + wg = WorkGraph() + for key, structure in structures.items(): + pw1 = wg.add_task(PwCalculation, name=f"pw1_{key}", structure=structure) + pw1.set(scf_inputs) + # save the output parameters to the context + pw1.set_context({"output_parameters": f"result.{key}"}) + return wg + + +# + + +@task.calcfunction() +# because this is a calcfunction, and the input datas are dynamic, we need use **datas. +def eos(**datas): + """Fit the EOS of the data.""" + from ase.eos import EquationOfState + + # + volumes = [] + energies = [] + for _, data in datas.items(): + volumes.append(data.dict.volume) + energies.append(data.dict.energy) + unit = data.dict.energy_units + # + eos = EquationOfState(volumes, energies) + v0, e0, B = eos.fit() + eos = orm.Dict({"unit": unit, "v0": v0, "e0": e0, "B": B}) + return eos + + +# %% +# Define the WorkGraph +# ---------------------- +# + + +from aiida_workgraph import WorkGraph + +# +wg = WorkGraph("eos") +scale_structure1 = wg.add_task(scale_structure, name="scale_structure1") +all_scf1 = wg.add_task( + all_scf, name="all_scf1", structures=scale_structure1.outputs["structures"] +) +eos1 = wg.add_task(eos, name="eos1", datas=all_scf1.outputs["result"]) +wg.to_html() +# + +# %% +# Combine with a relax task +# -------------------------- +# + + +from aiida_workgraph import WorkGraph, task +from aiida_quantumespresso.calculations.pw import PwCalculation + +# +@task.graph_builder(outputs=[{"name": "result", "from": "eos1.result"}]) +def eos_workgraph(structure=None, scales=None, scf_inputs=None): + wg = WorkGraph("eos") + scale_structure1 = wg.add_task( + scale_structure, name="scale_structure1", structure=structure, scales=scales + ) + all_scf1 = wg.add_task(all_scf, name="all_scf1", scf_inputs=scf_inputs) + eos1 = wg.add_task(eos, name="eos1") + wg.add_link(scale_structure1.outputs["structures"], all_scf1.inputs["structures"]) + wg.add_link(all_scf1.outputs["result"], eos1.inputs["datas"]) + return wg + + +# + +# ------------------------------------------------------- +wg = WorkGraph("relax_eos") +relax_task = wg.add_task(PwCalculation, name="relax1") +eos_wg_task = wg.add_task( + eos_workgraph, name="eos1", structure=relax_task.outputs["output_structure"] +) +wg.to_html() + + +# %% +# Useful tool: Web GUI +# ===================== +# Open a terminal, and run: +# +# .. code-block:: console +# +# workgraph web start +# +# Then visit the page `http://127.0.0.1:8000/workgraph`, where you can view all the workgraphs. +# +# What's Next +# =========== +# +# +-----------------------------------------------------+-----------------------------------------------------------------------------+ +# | `Concepts <../../concept/index.html>`_ | A brief introduction of WorkGraph’s main concepts. | +# +-----------------------------------------------------+-----------------------------------------------------------------------------+ +# | `HowTo <../../howto/index.html>`_ | Advanced topics, e.g., flow control using `if`, `while`, and `context`. | +# +-----------------------------------------------------+-----------------------------------------------------------------------------+ +# diff --git a/docs/index.rst b/docs/index.rst index 8ea7495e..14016ac6 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -14,6 +14,7 @@ The material is divided in 5 units: sections/running_processes/index sections/managing_data/index sections/writing_workflows/index + sections/writing_workflows_with_workgraph/index sections/creating_plugins/index These are also accessible via the sidebar on the left. diff --git a/docs/sections/writing_workflows_with_workgraph/index.rst b/docs/sections/writing_workflows_with_workgraph/index.rst new file mode 100644 index 00000000..ca37d3b8 --- /dev/null +++ b/docs/sections/writing_workflows_with_workgraph/index.rst @@ -0,0 +1,96 @@ +Writing workflows with WorkGraph +================================ + +.. rst-class:: header-text + + Recently a new way to create workflows has been developed, the *workgraph*. + The workgraph should simplify the creation of the and provide a more user + friendly GUI that provides information about your workflow before + execution. In this section you will learn how to write different examples + with the workgraph. + +.. panels:: + :header: panel-header-text + :body: bg-light + :footer: bg-light border-0 + + ------ + :column: col-lg-12 + + .. link-button:: autogen/zero_to_hero + :type: ref + :text: Zero to hero + :classes: btn-light text-left stretched-link font-weight-bold + ^^^^^^^^^^^^ + + A short module on how to write the basic type of workflows in AiiDA: work functions. + The module also revises the usage of calculation functions to add simple Python functions to the provenance. + + +++++++++++++ + .. list-table:: + :widths: 50 50 + :class: footer-table + :header-rows: 0 + + * - |time| 30 min + - |aiida| :aiida-green:`Basic` + +.. panels:: + :header: panel-header-text + :body: bg-light + :footer: bg-light border-0 + + ------ + :column: col-lg-12 + + .. link-button:: autogen/qe + :type: ref + :text: Computational materials science + :classes: btn-light text-left stretched-link font-weight-bold + ^^^^^^^^^^^^ + + A step-by-step introduction to the basics of writing work chains in AiiDA. + After completing this module, you will be ready to start writing your own scientific workflows! + + +++++++++++++ + .. list-table:: + :widths: 50 50 + :class: footer-table + :header-rows: 0 + + * - |time| 60 min + - |aiida| :aiida-green:`Intermediate` + + +.. panels:: + :header: panel-header-text + :body: bg-light + :footer: bg-light border-0 + + ------ + :column: col-lg-12 + + .. link-button:: autogen/eos + :type: ref + :text: A Real-world example - Equation of state + :classes: btn-light text-left stretched-link font-weight-bold + ^^^^^^^^^^^^ + + A step-by-step introduction to the basics of writing work chains in AiiDA. + After completing this module, you will be ready to start writing your own scientific workflows! + + +++++++++++++ + .. list-table:: + :widths: 50 50 + :class: footer-table + :header-rows: 0 + + * - |time| 60 min + - |aiida| :aiida-green:`Intermediate` + +.. toctree:: + :hidden: + + autogen/zero_to_hero + autogen/qe + autogen/eos diff --git a/environment.yml b/environment.yml new file mode 100644 index 00000000..d7f7ca5a --- /dev/null +++ b/environment.yml @@ -0,0 +1,11 @@ +name: base +channels: + - conda-forge + - defaults +dependencies: + - aiida-core + - aiida-core.services + - qe + - pip + - pip: + - -r requirements.txt diff --git a/requirements.txt b/requirements.txt index be019ba1..bd6cd99e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,9 @@ sphinx-book-theme~=1.1.3 sphinx-copybutton~=0.5.2 sphinx-panels~=0.4.1 sphinx-tabs~=3.4.5 +sphinx-gallery~=0.17.1 myst-nb~=1.1.1 +# to run notebooks +aiida-quantumespresso +aiida-pseudo +aiida-workgraph[widget]