[Docs] Update tutorial

docs/gallery/howto/autogen/annotate_inputs_outputs.py

@@ -129,6 +129,7 @@ def AddMultiplyInputs(x: int, y: int):
 # %%
 # We can see that even though we passed the inputs as a single dictionary, they were serialized as two separate ``Int`` nodes, ``x`` and ``y``, before being passed to the task.
 #
+# .. _dynamic_namespaces:
 # Dynamic namespaces
 # ~~~~~~~~~~~~~~~~~~
 #

docs/gallery/howto/autogen/parallel.py

@@ -1,78 +1,100 @@
 """
-Run tasks in parallel
-=====================
+Run tasks in parallel (Scatter-Gather)
+=====================================
 """
 
 # %%
 # Introduction
 # ------------
 #
-# Once you have developed a correct and functioning workflow, the next step is often to scale it up for large datasets.
-# This typically involves applying the same workflow to many independent data points.
-# In this how-to, we show how to run the workflow in parallel for each data point to improve performance and scalability.
+# A common pattern in scientific workflows is "scatter-gather," where a collection of inputs is processed in parallel (the "scatter" phase), and the results are then collected for a final processing step (the "gather" phase). This is a powerful way to parallelize work *within* a single, larger workflow.
+#
+# For example, applying a squaring operation `x²` to each number in a list `[x₁, x₂, ..., xₙ]`. Each squaring operation can be performed independently of the others.
+#
+# This how-to demonstrates how to implement the scatter-gather pattern to leverage this feature.
 
 import typing as t
-
 from aiida import load_profile
-
 from aiida_workgraph import namespace, task, dynamic
 
-
 load_profile()
 
+
 # %%
-# Perfectly parallelizable problem
-# --------------------------------
-#
-# A perfectly parallelizable problem can be broken down into smaller, independent subproblems that require no shared resources.
-# For example, consider an addition operation ``x + y`` applied element-wise to two lists: ``[x₁, ..., xₙ]`` and ``[y₁, ..., yₙ]``.
-# Each individual addition can be performed independently of the others.
-# ``WorkGraph`` automatically parallelizes task execution when there are no data dependencies between tasks (for more details on this concept, refer to `WorkGraph Engine <../../concept/autogen/engine>`_).
+# Scatter
+# ---------------------
 #
-# We will take advantage of this concept and show three different ways of how one can parallelize the add operation over the list with ``WorkGraph``.
+# The "scatter" phase involves creating and running multiple independent tasks. We can achieve this by simply iterating over our inputs within a `WorkGraph` and creating a task for each item.
 #
-# .. note::
-#
-#    In practice, a simple element-wise addition like this would typically be parallelized at a lower level, such as using NumPy vectorization or multithreading.
-#    We use it here for illustrative purposes.
-#    The concepts demonstrated in this guide can be applied to any workflow that is perfectly parallelizable.
+# Let's define a simple task to square a number.
 
-# %%
-# Conventional for-loop
-# ~~~~~~~~~~~~~~~~~~~~~
 
+@task
+def square(x: int) -> int:
+    """Square an integer."""
+    return x * x
+
+# %%
+# Define a helper task that generates a dictionary of numbers.
+# This will serve as input to the scatter phase.
 
 @task
-def add(x: int, y: int) -> int:
-    return x + y
+def generate_numbers(
+    n: int,
+) -> t.Annotated[dict[str, int], namespace(data=dynamic(int))]:
+    """Generate a dictionary of numbers from 1 to n."""
+    return {"data": {f"number_{i+1}": i + 1 for i in range(n)}}
+
+
+# %%
+# At first glance, one might try to write the scatter logic directly as follows:
+#
+# .. code:: python
+#
+#     data = generate_numbers(n=n).data
+#     squares = {}
+#     # Since these tasks have no dependencies on each other, they run in parallel.
+#     for key, value in data.items():
+#         squares[key] = square(x=value).result
+#
+# However, **this will not work as expected**.
+#
+# The reason is that ``data`` is not immediately available when constructing the graph.
+# Instead, ``generate_numbers(n=n).data`` is a *future output*, a placeholder that will only be resolved at runtime.
+#
+# To correctly handle this, we must wrap the loop inside another task graph. This ensures that the graph engine knows how to schedule and parallelize the tasks.
 
 
 @task.graph
-def ParallelAdd(
-    data: t.Annotated[dict[str, dict[str, int]], dynamic(namespace(x=int, y=int))],
-) -> t.Annotated[dict, namespace(sums=dynamic(int))]:
-    sums = {}
-    for i, list_i in enumerate(data.values()):
-        x, y = list_i.values()
-        sums[f"sum_{i}"] = add(x=x, y=y).result
-    return {"sums": sums}
+def ParallelSquare(
+    data: t.Annotated[dict[str, int], dynamic(int)],
+) -> t.Annotated[dict, namespace(squares=dynamic(int))]:
+    """Applies the square task to each number in the input data in parallel."""
+    squares = {}
+    for key, value in data.items():
+        squares[key] = square(x=value).result
+    return {"squares": squares}
 
 
+# %%
+# .. tip::
+#
+#     - **Parallelization**: In `WorkGraph`, tasks without data dependencies between them are automatically scheduled to run in parallel.
+#     - **Dynamic inputs/outputs**: Our use of dynamic type annotations, such as `dynamic(int)`, allows AiiDA to create a distinct node for each input and output in the collection, which is essential for data provenance tracking. For more details, please refer to the section on :ref:`Dynamic namespaces <dynamic_namespaces>`.
+#
+
 # %%
 # Let's run it with some sample data.
 
-data = {f"list_{i}": {"x": i, "y": i} for i in range(1, 5)}
+data = {f"number_{i}": i for i in range(1, 5)}
 
-wg = ParallelAdd.build(data)
+wg = ParallelSquare.build(data)
 wg.run()
 
 print("\nResults:")
-lists = list(data.values())
-for i, sum_ in enumerate(wg.outputs.sums):
-    print(f"{lists[i]['x']} + {lists[i]['y']} = {sum_.value}")
-
-# %%
-# Each addition was executed independently, yielding an AiiDA node for each result.
+for i, result_node in enumerate(wg.outputs.squares):
+    original_value = list(data.values())[i]
+    print(f"{original_value}² = {result_node.value}")
 
 # %%
 # Workflow view
@@ -87,47 +109,51 @@ def ParallelAdd(
 wg.generate_provenance_graph()
 
 # %%
-# .. note::
+# Gather
+# --------------------
 #
-#    Due to our explicit use of type annotation, AiiDA yields a node per input/output.
-#    For more on leveraging type annotations, refer to the :doc:`/howto/autogen/annotate_inputs_outputs` how-to section.
-
-# %%
-# Gather results
-# ~~~~~~~~~~~~~~
+# The "gather" phase involves collecting the results from the parallel tasks and performing a final operation.
 #
-# We now extend the workflow by adding a task that sums the intermediate results.
-# This step is commonly known as a gather, aggregate, or reduce operation.
-# It is often used to automatically analyze or summarize the output of parallel computations.
-#
-# We will extend it the whole workflow only by the ``aggregate_sum`` task
+# We will now extend the workflow by adding a task that sums the results from the `square` tasks.
 
 
 @task
-def aggregate_sums(data: t.Annotated[dict, dynamic(int)]) -> int:
+def gather_and_sum(data: t.Annotated[dict, dynamic(int)]) -> int:
+    """Sums the values of a dictionary of integers."""
     return sum(data.values())
 
 
+# %%
+# We create a new `WorkGraph` that orchestrates the full scatter-gather pattern.
+# It first calls our `generate_numbers` and `ParallelSquare` graphs (scatter),
+# and then feeds the collected outputs into the `gather_and_sum` task (gather).
+
+
 @task.graph
-def ParallelAddAggregate(
-    data: t.Annotated[dict[str, dict[str, int]], dynamic(namespace(x=int, y=int))],
-) -> int:
-    sums = ParallelAdd(data=data).sums
-    return aggregate_sums(data=sums).result
+def ScatterGatherSquare(n: int) -> int:
+    """A full scatter-gather workflow to generate numbers, square them in parallel, and sum the results."""
+    # Generate inputs
+    data = generate_numbers(n=n).data
+    # Scatter phase
+    squares = ParallelSquare(data=data).squares
+    # Gather phase
+    return gather_and_sum(data=squares).result
 
 
-wg = ParallelAddAggregate.build(data)
+wg = ScatterGatherSquare.build(4)
 wg.run()
 
-print("\nResult:", wg.outputs.result.value)
+print("\nAggregated Result:", wg.outputs.result.value)
 
-# (1+1) + (2+2) + (3+3) + (4+4) = 20
-assert wg.outputs.result.value == 20
+assert wg.outputs.result.value == 30
 
 # %%
 # Conclusion
 # ----------
 #
-# In this how-to, we demonstrated how to run tasks in parallel using ``WorkGraph``.
-# We illustrated this with a simple example of element-wise addition, showcasing how to structure a perfectly parallelizable workflow.
-# We also showed how to gather the results of parallel computations using an aggregation task.
+# In this how-to, we demonstrated how to implement the powerful scatter-gather pattern using `WorkGraph`.
+#
+# - The **scatter** phase is achieved by creating multiple independent tasks within a graph, which the engine automatically runs in parallel.
+# - The **gather** phase collects the results from the parallel tasks for a final processing step.
+# - We also highlighted an important concept: **future outputs**. When a value is the output of another task, it cannot be used in a Python loop directly at graph-construction time. Instead, the loop must be wrapped in another task graph.
+#

docs/gallery/howto/autogen/remote_job.py

@@ -5,6 +5,7 @@
 """
 
 # %%
+# .. _remote_calculations:
 # Introduction
 # ------------
 #

docs/gallery/tutorial/autogen/materials_science_ase.py

@@ -8,7 +8,7 @@
 We'll explore two key examples that highlight the flexibility of AiiDA-WorkGraph:
 
 1.  **Atomization energy**: A simple, linear workflow to calculate the atomization energy of a diatomic molecule.
-2.  **Equation of state (EOS)**: A more advanced workflow for bulk structure that showcases how to handle loops and dynamic inputs/outputs, a common pattern in dynamic workflows.
+2.  **Equation of state (EOS)**: A more advanced workflow for bulk structure that showcases how to handle ``if`` condition, parallel execution, and dynamic inputs/outputs, a common pattern in dynamic workflows.
 
 """
 
@@ -45,6 +45,8 @@
 #
 #    We're using the ASE EMT (Effective Medium Theory) calculator because it's exceptionally fast and perfect for demonstrations.
 #    You can easily swap it with any other ASE-compatible calculator, like Quantum ESPRESSO, VASP, or GPAW, for your research.
+#    For realistic simulations, especially with DFT codes, you would typically run these calculations on a remote computer.
+#    For more details on running calculations remotely, please refer to the section on :ref:`Run calculations remotely <remote_calculations>`.
 #
 from aiida_workgraph import task, spec
 from ase import Atoms
@@ -122,7 +124,7 @@ def atomization_energy_workflow(molecule_obj: Atoms, atom_obj: Atoms) -> float:
 # Now for a more complex and practical example: calculating the Equation of State (EOS) for a bulk material.
 # The process involves several steps:
 #
-# 1.  **Relax** the initial atomic structure to its lowest-energy state.
+# 1.  **Relax**(optional) the initial atomic structure to its lowest-energy state.
 # 2.  **Strain** the relaxed structure by applying a series of scaling factors.
 # 3.  **Calculate** the total energy and volume for each strained structure.
 # 4.  **Fit** the resulting energy-volume data to an EOS model to find properties like the equilibrium volume and bulk modulus.
@@ -131,6 +133,7 @@ def atomization_energy_workflow(molecule_obj: Atoms, atom_obj: Atoms) -> float:
 
 from ase.calculators.emt import EMT
 from ase.optimize import BFGS
+from typing import Annotated
 
 
 @task
@@ -142,10 +145,18 @@ def relax_structure(atoms: Atoms) -> Atoms:
     return atoms
 
 
+# %%
+# .. note::
+#
+#   If you want to run the ``relax_structure`` task on a remote computer, you can use the ``@task.pythonjob`` decorator.
+#   Please refer to the section on :ref:`Run calculations remotely <remote_calculations>`.
+#
+
+
 @task
 def create_strained_structures(
     atoms: Atoms, scales: list
-) -> spec.namespace(scaled_structures=spec.dynamic(Atoms)):
+) -> Annotated[dict, spec.namespace(scaled_structures=spec.dynamic(Atoms))]:
     """Generate a series of strained structures from a list of scaling factors."""
     scaled_structures = {}
     for i, scale in enumerate(scales):
@@ -176,8 +187,8 @@ def calculate_energy_and_volume(atoms: Atoms) -> dict:
 
 @task.graph
 def calc_all_structures(
-    **scaled_structures,
-) -> spec.namespace(results=spec.dynamic(dict)):
+    scaled_structures: Annotated[dict, spec.dynamic(Atoms)]
+) -> Annotated[dict, spec.namespace(results=spec.dynamic(dict))]:
     """Sub-workflow to calculate energy and volume for all strained structures in parallel."""
     results = {}
     for key, atoms in scaled_structures.items():
@@ -190,7 +201,7 @@ def calc_all_structures(
 
 
 @task
-def fit_eos_model(**data) -> dict:
+def fit_eos_model(data: Annotated[dict, spec.dynamic(dict)]) -> dict:
     """Fit Energy-Volume data to a Birch-Murnaghan Equation of State."""
     from ase.eos import EquationOfState
     from ase.units import kJ
@@ -208,33 +219,15 @@ def fit_eos_model(**data) -> dict:
 
 
 @task.graph()
-def eos_workflow(atoms: Atoms, scales: list) -> dict:
+def eos_workflow(atoms: Atoms, scales: list, run_relax: bool = True) -> dict:
     """The complete EOS workflow graph."""
-    relaxed_atoms = relax_structure(atoms=atoms).result
-    strained = create_strained_structures(atoms=relaxed_atoms, scales=scales)
+    if run_relax:
+        atoms = relax_structure(atoms=atoms).result
+    strained = create_strained_structures(atoms=atoms, scales=scales)
     emt_outputs = calc_all_structures(scaled_structures=strained.scaled_structures)
     return fit_eos_model(data=emt_outputs.results).result
 
 
-# %%
-#
-# .. important::
-#
-#    When linking task outputs to a keyword argument of the task, you **must** explicitly name the input argument, e.g.:
-#
-#    .. code-block:: python
-#
-#        calc_all_structures(scaled_structures=strained.scaled_structures)
-#
-#    rather than using argument unpacking like:
-#
-#    .. code-block:: python
-#
-#        calc_all_structures(**strained.scaled_structures)
-#
-#    For further details, please refer to the section on :ref:`Dynamic Namespaces <dynamic_namespaces>`.
-#
-
 # %%
 # Build and Run the EOS Workflow
 # ------------------------------
@@ -277,7 +270,7 @@ def eos_workflow(atoms: Atoms, scales: list) -> dict:
 #
 # -   Transform any Python function into a robust, provenance-tracked task with the ``@task`` decorator.
 # -   Compose tasks into complete workflows using ``@task.graph``, from simple linear chains to complex graphs.
-# -   Manage advanced patterns like loops and parallel execution using **dynamic namespaces**.
+# -   Manage advanced patterns like ``if`` condition, parallel execution, and **dynamic namespaces**.
 # -   Build, execute, and visualize both the workflow plan and its final, rich **provenance graph**.
 #
 # These powerful concepts are the foundation for building sophisticated, automated, and fully reproducible simulation pipelines for your own research projects. Happy computing! 🚀