From 755e03a33abcfff7422b5b9e3a52eed56b2e787e Mon Sep 17 00:00:00 2001
From: superstar54 <xingwang1991@gmail.com>
Date: Sun, 25 Aug 2024 20:21:19 +0200
Subject: [PATCH] Add workgraph.py

---
 docs/gallery/concept/autogen/workgraph.py     | 78 ++++++++++++++++++
 docs/gallery/tutorial/autogen/zero_to_hero.py | 11 ++-
 docs/source/concept/workgraph.rst             | 82 -------------------
 3 files changed, 83 insertions(+), 88 deletions(-)
 create mode 100644 docs/gallery/concept/autogen/workgraph.py
 delete mode 100644 docs/source/concept/workgraph.rst

diff --git a/docs/gallery/concept/autogen/workgraph.py b/docs/gallery/concept/autogen/workgraph.py
new file mode 100644
index 00000000..09632982
--- /dev/null
+++ b/docs/gallery/concept/autogen/workgraph.py
@@ -0,0 +1,78 @@
+"""
+WorkGraph Example
+=================
+
+This example demonstrates how to create, configure, and use a
+:class:`~aiida_workgraph.workgraph.WorkGraph` object, which is a collection of tasks and links.
+
+Create and launch workgraph
+----------------------------
+"""
+
+# %%
+# First, create an empty workgraph:
+#
+
+from aiida_workgraph import WorkGraph, task
+
+wg = WorkGraph(name="my_first_workgraph")
+
+# %%
+# Define and use tasks
+# --------------------
+#
+
+# Define a task using a `calcfunction`:
+@task.calcfunction()
+def add(x, y):
+    return x + y
+
+
+# Add tasks to the workgraph
+add1 = wg.add_task(add, name="add1")
+add2 = wg.add_task(add, name="add2")
+
+# %%
+# Add a link between tasks:
+
+wg.add_link(add1.outputs["result"], add2.inputs["x"])
+wg.to_html()
+
+# %%
+# Submit the workgraph:
+
+wg.submit(inputs={"add1": {"x": 1, "y": 2}, "add2": {"y": 3}}, wait=True)
+
+# %%
+# Load workgraph from the AiiDA process
+# -------------------------------------
+#
+# WorkGraph saves its data as an extra attribute in its process, allowing reconstruction of the WorkGraph from the process.
+
+from aiida_workgraph import WorkGraph
+
+wg_loaded = WorkGraph.load(wg.pk)
+
+# %%
+# Execute order
+# -------------
+# The tasks will be executed under the following conditions:
+#
+# - No input task
+# - All input tasks are finished.
+# Group outputs
+# -------------
+# You can output the results of the tasks as the output of the WorkGraph.
+
+wg = WorkGraph("test_workgraph_group_outputs")
+wg.add_task(add, "add1", x=2, y=3)
+wg.group_outputs = [{"name": "sum", "from": "add1.result"}]
+wg.submit(wait=True)
+assert wg.process.outputs.sum.value == 5
+
+# %%
+# List of all Methods
+# ===================
+#
+# .. autoclass:: aiida_workgraph.workgraph.WorkGraph
+#    :members:
diff --git a/docs/gallery/tutorial/autogen/zero_to_hero.py b/docs/gallery/tutorial/autogen/zero_to_hero.py
index 30f6743d..0bc9d5a8 100644
--- a/docs/gallery/tutorial/autogen/zero_to_hero.py
+++ b/docs/gallery/tutorial/autogen/zero_to_hero.py
@@ -39,16 +39,15 @@
 # In AiiDA, we can define two `calcfunction` to do the `add` and `mutiply`:
 #
 
-from aiida.engine import calcfunction
+from aiida_workgraph import task
 
-#
-@calcfunction
+
+@task.calcfunction()
 def add(x, y):
     return x + y
 
 
-#
-@calcfunction
+@task.calcfunction()
 def multiply(x, y):
     return x * y
 
@@ -58,7 +57,7 @@ def multiply(x, y):
 # --------------------
 # Three steps:
 #
-# - create a empty `WorkGraph`
+# - create a empty WorkGraph
 # - add tasks: `add` and `multiply`.
 # - link the output of the `add` task to the `x` input of the `multiply` task.
 #
diff --git a/docs/source/concept/workgraph.rst b/docs/source/concept/workgraph.rst
deleted file mode 100644
index 15a05cc7..00000000
--- a/docs/source/concept/workgraph.rst
+++ /dev/null
@@ -1,82 +0,0 @@
-.. _workgraph:
-
-.. module:: workgraph
-
-
-===========================================
-WorkGraph
-===========================================
-The :class:`~aiida_workgraph.workgraph.WorkGraph` object is a collection of tasks and links.
-
-Create and launch workgraph
-============================
-- Create a empty workgraph:
-
-.. code-block:: python
-
-    from aiida_workgraph import WorkGraph, task
-    wg = WorkGraph(name="my_first_workgraph")
-
-
-Create and use `task`.
-
-.. code:: python
-
-    # define add calcfunction task
-    @task.calcfunction()
-    def add(x, y):
-       return x + y
-
-    add1 = wg.add_task(add, name="add1")
-    add2 = wg.add_task(add, name="add2")
-
-
-- Add link between tasks:
-
-.. code-block:: python
-
-    wg.add_link(add1.outputs["result"], add2.inputs["x"])
-
-- Submit the workgraph:
-
-.. code-block:: python
-
-    wg.submit()
-
-Load workgraph from the AiiDA process
-=====================================
-WorkGraph save its data as a extra attribute into its process, so that one can rebuild the WorkGraph from the process.
-
-
-.. code-block:: python
-
-    from aiida_workgraph import WorkGraph
-    # pk is the process id of a WorkGraph
-    WorkGraph.load(pk)
-
-Execute order
-===============
-The tasks will be executed when:
-
-- No input task
-- All input tasks finish.
-
-
-Group outputs
-=====================================
-One can output the results of the tasks as the output of the WorkGraph.
-
-.. code-block:: python
-
-    wg = WorkGraph("test_workgraph_group_outputs")
-    wg.add_task(add, "add1", x=2, y=3)
-    wg.group_outputs = [{"name": "sum", "from": "add1.result"}]
-    wg.submit(wait=True)
-    assert wg.process.outputs.sum.value == 5
-
-
-List of all Methods
-===================
-
-.. autoclass:: aiida_workgraph.workgraph.WorkGraph
-   :members: