✨ aiida: PWD via @task and @graph

binder/environment.yml

@@ -11,7 +11,9 @@ dependencies:
 - pyiron_base =0.15.19
 - pyiron_workflow =0.18.0
 - pygraphviz =1.14
-- aiida-workgraph =0.9.0
 - conda_subprocess =0.0.12
 - networkx =3.6.1
 - cwltool =3.2.20260413085819
+- pip
+- pip:
+  - git+https://github.com/aiidateam/aiida-core.git@main

example_workflows/arithmetic/aiida.ipynb

@@ -23,115 +23,132 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "source": "# Aiida\n\nThis notebook defines the arithmetic workflow with [`aiida-workgraph`](https://github.com/aiidateam/aiida-workgraph) and then loads the resulting `workflow.json` into `jobflow`, `pyiron_base` and `pyiron_workflow`, to demonstrate that the same workflow definition can be executed by several different engines.",
+   "source": [
+    "# Aiida\n",
+    "\n",
+    "This notebook defines the arithmetic workflow with the aiida-core graph authoring layer (`@task` and `@graph`) and then loads the resulting `workflow.json` into `jobflow`, `pyiron_base` and `pyiron_workflow`, to demonstrate that the same workflow definition can be executed by several different engines."
+   ],
    "metadata": {}
   },
   {
    "cell_type": "markdown",
-   "source": "## Define workflow with aiida\n\n`aiida-workgraph` represents a workflow as a `WorkGraph` of `Task`s. Tasks are executed and stored by the AiiDA engine, so we first need to connect to an AiiDA profile with `load_profile()` before any task can run.",
+   "source": [
+    "## Define workflow with aiida\n",
+    "\n",
+    "Tasks are plain Python functions annotated with `@task`; the `@graph` function wires their calls without executing them. Building the graph needs no profile; only execution does."
+   ],
    "metadata": {}
   },
   {
    "cell_type": "code",
    "source": [
     "from python_workflow_definition.aiida import write_workflow_json\n",
     "\n",
-    "from aiida_workgraph import WorkGraph, task\n",
-    "from aiida import orm, load_profile\n",
-    "load_profile()\n",
+    "from aiida.engine import graph, task\n",
     "\n",
-    "workflow_json_filename =  \"aiida_simple.json\""
+    "workflow_json_filename = \"aiida_simple.json\""
    ],
    "metadata": {
     "trusted": true
    },
    "outputs": [],
-   "execution_count": 1
+   "execution_count": null
   },
   {
    "cell_type": "code",
    "source": [
+    "import typing as t\n",
+    "from typing import Any\n",
+    "\n",
     "from workflow import (\n",
     "    get_sum as _get_sum,\n",
     "    get_prod_and_div as _get_prod_and_div,\n",
     "    get_square as _get_square,\n",
-    ")"
+    ")\n",
+    "\n",
+    "\n",
+    "class ProdDiv(t.TypedDict):\n",
+    "    prod: Any\n",
+    "    div: Any\n",
+    "\n",
+    "\n",
+    "def get_prod_and_div(x, y) -> ProdDiv:\n",
+    "    return _get_prod_and_div(x, y)\n",
+    "\n",
+    "\n",
+    "get_prod_and_div.__module__ = \"workflow\"\n",
+    "get_prod_and_div = task(get_prod_and_div)\n",
+    "\n",
+    "\n",
+    "def get_sum(x, y) -> Any:\n",
+    "    return _get_sum(x, y)\n",
+    "\n",
+    "\n",
+    "get_sum.__module__ = \"workflow\"\n",
+    "get_sum = task(get_sum)\n",
+    "\n",
+    "\n",
+    "def get_square(x) -> Any:\n",
+    "    return _get_square(x)\n",
+    "\n",
+    "\n",
+    "get_square.__module__ = \"workflow\"\n",
+    "get_square = task(get_square)"
    ],
    "metadata": {
     "trusted": true
    },
    "outputs": [],
-   "execution_count": 2
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "source": "The arithmetic functions are imported under a leading underscore because they are still plain Python functions at this point; `wg.add_task` below turns each one into an AiiDA `Task` that can be linked to other tasks' inputs and outputs. `get_prod_and_div` returns a dictionary, so it is wrapped with `task(outputs=['prod', 'div'])` to expose `prod` and `div` as separate output sockets.",
+   "source": [
+    "The arithmetic functions are imported under a leading underscore because they are still plain Python functions at this point; the wrappers below turn each one into an AiiDA task with declared outputs. `get_prod_and_div` returns a dictionary, so a `TypedDict` declares `prod` and `div` as its output ports. Setting `__module__` keeps the exported `module.function` values pointing at `workflow.py`."
+   ],
    "metadata": {}
   },
   {
    "cell_type": "code",
-   "source": "wg = WorkGraph(\"arithmetic\")",
-   "metadata": {
-    "trusted": true
-   },
-   "outputs": [],
-   "execution_count": 3
-  },
-  {
-   "cell_type": "code",
-   "source": "get_prod_and_div_task = wg.add_task(\n    task(outputs=['prod', 'div'])(_get_prod_and_div),\n    x=orm.Float(1),\n    y=orm.Float(2),\n)",
-   "metadata": {
-    "trusted": true
-   },
-   "outputs": [],
-   "execution_count": 4
-  },
-  {
-   "cell_type": "code",
-   "source": "get_sum_task = wg.add_task(\n    _get_sum,\n    x=get_prod_and_div_task.outputs.prod,\n    y=get_prod_and_div_task.outputs.div,\n)",
-   "metadata": {
-    "trusted": true
-   },
-   "outputs": [],
-   "execution_count": 5
-  },
-  {
-   "cell_type": "code",
-   "source": "get_square_task = wg.add_task(\n    _get_square,\n    x=get_sum_task.outputs.result,\n)",
+   "source": [
+    "@graph\n",
+    "def arithmetic(x=1, y=2):\n",
+    "    prod_div = get_prod_and_div(x=x, y=y)\n",
+    "    total = get_sum(x=prod_div[\"prod\"], y=prod_div[\"div\"])\n",
+    "    return {\"result\": get_square(x=total)}"
+   ],
    "metadata": {
     "trusted": true
    },
    "outputs": [],
-   "execution_count": 6
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "source": "With all tasks connected, `write_workflow_json` walks the `WorkGraph`'s tasks and links and serializes them into the PWD `workflow.json` format: each task becomes a `function` node and each link becomes an edge between an input and an output port.",
+   "source": [
+    "With all tasks connected, `write_workflow_json` builds the `GraphIR` and serializes it into the PWD `workflow.json` format: each task becomes a `function` node, each dependency becomes an edge, and the graph defaults become `input` values."
+   ],
    "metadata": {}
   },
   {
    "cell_type": "code",
-   "source": "write_workflow_json(wg=wg, file_name=workflow_json_filename)",
+   "source": [
+    "write_workflow_json(arithmetic, workflow_json_filename)"
+   ],
    "metadata": {
     "trusted": true
    },
    "outputs": [],
-   "execution_count": 7
+   "execution_count": null
   },
   {
    "cell_type": "code",
    "source": "!cat {workflow_json_filename}",
    "metadata": {
     "trusted": true
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": "{\n  \"version\": \"0.1.0\",\n  \"nodes\": [\n    {\n      \"id\": 0,\n      \"type\": \"function\",\n      \"value\": \"workflow.get_prod_and_div\"\n    },\n    {\n      \"id\": 1,\n      \"type\": \"function\",\n      \"value\": \"workflow.get_sum\"\n    },\n    {\n      \"id\": 2,\n      \"type\": \"function\",\n      \"value\": \"workflow.get_square\"\n    },\n    {\n      \"id\": 3,\n      \"type\": \"input\",\n      \"name\": \"x\",\n      \"value\": 1.0\n    },\n    {\n      \"id\": 4,\n      \"type\": \"input\",\n      \"name\": \"y\",\n      \"value\": 2.0\n    },\n    {\n      \"id\": 5,\n      \"type\": \"output\",\n      \"name\": \"result\"\n    }\n  ],\n  \"edges\": [\n    {\n      \"target\": 1,\n      \"targetPort\": \"x\",\n      \"source\": 0,\n      \"sourcePort\": \"prod\"\n    },\n    {\n      \"target\": 1,\n      \"targetPort\": \"y\",\n      \"source\": 0,\n      \"sourcePort\": \"div\"\n    },\n    {\n      \"target\": 2,\n      \"targetPort\": \"x\",\n      \"source\": 1,\n      \"sourcePort\": null\n    },\n    {\n      \"target\": 0,\n      \"targetPort\": \"x\",\n      \"source\": 3,\n      \"sourcePort\": null\n    },\n    {\n      \"target\": 0,\n      \"targetPort\": \"y\",\n      \"source\": 4,\n      \"sourcePort\": null\n    },\n    {\n      \"target\": 5,\n      \"targetPort\": null,\n      \"source\": 2,\n      \"sourcePort\": null\n    }\n  ]\n}"
-    }
-   ],
-   "execution_count": 8
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "markdown",