PyFunction

Default and custom outputs

The @pyfunction decorator turns a Python function into an AiiDA PyFunction. By default, the entire return value is stored in a single output node named result.

from aiida import load_profile
from aiida.engine import run_get_node

from aiida_pythonjob import pyfunction, spec
from typing import Any

load_profile()


@pyfunction()
def add_and_subtract(x, y):
    return {"sum": x + y, "diff": x - y}


result, node = run_get_node(add_and_subtract, x=1, y=2)
print("Default result: ", result)

Default result:  uuid: 95d1b49e-3ade-475c-80b4-46612862849f (pk: 69)

You can specify the outputs parameter to unpack a returned dictionary into separate output nodes.

@pyfunction(outputs=spec.namespace(sum=Any, diff=Any))
def add_and_subtract(x, y):
    return {"sum": x + y, "diff": x - y}


result, node = run_get_node(add_and_subtract, x=1, y=2)
print("Unpacked results: ")
print("sum: ", result["sum"])
print("diff: ", result["diff"])

Unpacked results:
sum:  uuid: 717cb451-a372-46ed-a7f2-a068cdaaaa3c (pk: 73) value: 3
diff:  uuid: 068a8393-4145-44d6-974d-3c13e0b1e64c (pk: 74) value: -1

Advanced features

pyfunction supports many advanced features for data handling and workflow control. These functionalities are shared with PythonJob.

See also

For a detailed guide on the following topics, please refer to the Common Concepts:

• Dynamic and Nested Namespaces: For handling complex or variable outputs.

• Custom Exit Codes: For robust error handling and workflow control.

• Data Serialization and Deserialization: For working with custom data types like ase.Atoms or other AiiDA data nodes.

Here is a more complex example demonstrating a dynamic namespace output with custom data types.

from ase import Atoms  # noqa: E402
from ase.build import bulk  # noqa: E402


@pyfunction(outputs=spec.dynamic(Atoms))
def generate_structures(element: str, factors: list) -> dict:
    """Scale a bulk structure by the given factors."""
    scaled_structures = {}
    initial_structure = bulk(element)
    for i, factor in enumerate(factors):
        atoms = initial_structure.copy()
        atoms.set_cell(atoms.cell * factor, scale_atoms=True)
        scaled_structures[f"s_{i}"] = atoms
    return scaled_structures


result, node = run_get_node(generate_structures, element="Al", factors=[0.95, 1.0, 1.05])

print("Generated scaled structures:")
for key, value in result.items():
    print(key, value)

Generated scaled structures:
s_0 uuid: 7a8f5ff3-2b3d-4b8f-84b2-808132b22303 (pk: 78)
s_1 uuid: a3fa8353-4237-4e82-b615-f0fc5a987490 (pk: 79)
s_2 uuid: 0d0ca77a-f053-4b01-a14f-5f3c00ba5a8a (pk: 80)

Async functions

pyfunction also supports Python’s async functions. This is a powerful feature for tasks that are I/O-bound (e.g., waiting for network requests, file operations) or for running multiple tasks concurrently without blocking the AiiDA daemon.

When you submit an async function, the call returns immediately with a process node, allowing your script to continue running while the function executes in the background.

from aiida.engine import submit
import datetime
from aiida_pythonjob import prepare_pyfunction_inputs, PyFunction


@pyfunction()
async def add_async(x, y, time: float):
    """A simple function that adds two numbers."""
    import asyncio

    # Simulate asynchronous I/O or computation
    await asyncio.sleep(time)
    return x + y


inputs = prepare_pyfunction_inputs(
    add_async,
    function_inputs={"x": 2, "y": 3, "time": 2.0},
)

node = submit(PyFunction, **inputs)

Killing an async process

Since async functions run as regular AiiDA processes, they can be controlled and killed programmatically. This is useful for managing long-running or stuck tasks. You can kill a running async function using the AiiDA command line interface.

$ verdi process kill <pk>

Monitor external events

Async functions are particularly useful for monitoring external events or conditions without blocking the AiiDA daemon. Here is an example that waits until a specified time.

async def monitor_time(time: datetime.datetime, interval: float = 0.5, timeout: float = 60.0):
    """Monitor the current time until it reaches the specified target time."""
    import asyncio

    start_time = datetime.datetime.now()
    while datetime.datetime.now() < time:
        print("Waiting...")
        await asyncio.sleep(interval)
        if (datetime.datetime.now() - start_time).total_seconds() > timeout:
            raise TimeoutError("Monitoring timed out.")


inputs = prepare_pyfunction_inputs(
    monitor_time,
    function_inputs={"time": datetime.datetime.now() + datetime.timedelta(seconds=5), "interval": 1.0},
)

node = submit(PyFunction, **inputs)

For user’s convenience, we provide a dedicated MonitorFunction class that inherits from PyFunction. User only need to write normal function, which returns True when the monitoring condition is met.

from aiida_pythonjob import MonitorPyFunction


def monitor_time(time: datetime.datetime):
    # return True when the current time is greater than the target time
    return datetime.datetime.now() > time


inputs = prepare_pyfunction_inputs(
    monitor_time,
    function_inputs={"time": datetime.datetime.now() + datetime.timedelta(seconds=5)},
    interval=1.0,
    timeout=20.0,
)

node = submit(MonitorPyFunction, **inputs)