✨ graphs: evaluate inline comparisons without a task

Author: GeigerJ2

src/aiida/engine/processes/graphs/__init__.py

@@ -12,6 +12,7 @@
 from aiida.engine.processes.graphs.build_source import GraphHandle, graph, lower_from_source, task
 from aiida.engine.processes.graphs.ir import (
     BranchControl,
+    Comparison,
     Dependency,
     Endpoint,
     GraphIR,
@@ -24,6 +25,7 @@
 
 __all__ = (
     'BranchControl',
+    'Comparison',
     'Dependency',
     'Endpoint',
     'GraphHandle',

src/aiida/engine/processes/graphs/build_source.py

@@ -30,6 +30,7 @@
 from aiida.engine.processes.graphs.ir import (
     CONDITION_PORT,
     BranchControl,
+    Comparison,
     Dependency,
     Endpoint,
     GraphIR,
@@ -503,10 +504,14 @@ def _lower_branch(state: _LoweringState, statement: ast.If) -> None:
         _reject(state, statement, 'condition conflicts with a branch input; rename one of them')
     instance = f'branch_{len(state.tasks) + 1}'
     given: dict[str, CapturedValue] = {}
-    _wire(state, instance, CONDITION_PORT, _lower_value(state, test), given)
+    comparison = _lower_comparison(state, instance, test, given)
+    if comparison is None:
+        _wire(state, instance, CONDITION_PORT, _lower_value(state, test), given)
     for name in sorted({path.split('.')[0] for side in (first, second) for path in side.child.inputs}):
         _wire(state, instance, name, state.names[name], given)
-    state.tasks.append(BranchControl(name=instance, inputs=given, body=first.body, otherwise=second.body))
+    state.tasks.append(
+        BranchControl(name=instance, inputs=given, body=first.body, otherwise=second.body, comparison=comparison)
+    )
     selected = first.child.names[first.name]
     state.names[first.name] = _Reference(instance, selected.port, selected.shape)
 
@@ -521,9 +526,11 @@ def _lower_loop(state: _LoweringState, statement: ast.While) -> None:
         ast.copy_location(update.targets[0], statement.test)
         body.append(update)
     test: ast.expr = bound if bound is not None else statement.test
-    if statement.orelse or not isinstance(test, ast.Name) or test.id not in state.names:
-        _reject(state, statement, 'while requires a bound condition name and no else')
-    condition = test.id
+    if statement.orelse or (
+        not isinstance(test, ast.Compare) and (not isinstance(test, ast.Name) or test.id not in state.names)
+    ):
+        _reject(state, statement, 'while requires a bound condition or comparison and no else')
+    condition = test.id if isinstance(test, ast.Name) else CONDITION_PORT
     child = _collect_region_inputs(state, body)
     # Loop state is read from the previous iteration, including values rebound in the body.
     assigned: set[str] = {
@@ -536,7 +543,7 @@ def _lower_loop(state: _LoweringState, statement: ast.While) -> None:
             _reject(state, statement, f'loop state {name!r} must be initialized before the loop')
         child.inputs.setdefault(name, [])
         child.names[name] = _Reference(None, name, state.names[name].shape)
-    if condition not in assigned:
+    if not isinstance(test, ast.Compare) and condition not in assigned:
         _reject(state, statement, 'while body must update its condition; rebind the condition name in the body')
     for node in body:
         _lower_assignment(child, node, rebind=True)
@@ -545,7 +552,12 @@ def _lower_loop(state: _LoweringState, statement: ast.While) -> None:
     instance = f'loop_{len(state.tasks) + 1}'
     given: dict[str, CapturedValue] = {}
     _route_region(state, child, instance, given)
-    state.tasks.append(LoopControl(name=instance, inputs=given, body=body_spec, condition_port=condition))
+    comparison = _lower_comparison(state, instance, test, given) if isinstance(test, ast.Compare) else None
+    if comparison is None and isinstance(test, ast.Compare):
+        _reject(state, test, _unsupported_reason(test))
+    state.tasks.append(
+        LoopControl(name=instance, inputs=given, body=body_spec, condition_port=condition, comparison=comparison)
+    )
     for name in assigned:
         state.names[name] = _Reference(instance, name, child.names[name].shape)
 
@@ -611,6 +623,38 @@ def _bind_condition(state: _LoweringState, test: ast.expr) -> ast.Name | None:
     return ast.Name(id=name, ctx=ast.Load())
 
 
+_COMPARISON_OPERATORS: dict[type[ast.cmpop], str] = {
+    ast.Eq: '==',
+    ast.NotEq: '!=',
+    ast.Lt: '<',
+    ast.LtE: '<=',
+    ast.Gt: '>',
+    ast.GtE: '>=',
+}
+
+
+def _lower_comparison(
+    state: _LoweringState, instance: str, test: ast.expr, given: dict[str, CapturedValue]
+) -> Comparison | None:
+    """Lower one comparison over names and literals, or decline anything richer."""
+    if not isinstance(test, ast.Compare) or len(test.ops) != 1 or type(test.ops[0]) not in _COMPARISON_OPERATORS:
+        return None
+    operands: list[str | int | float | bool | None] = []
+    for side in (test.left, *test.comparators):
+        if isinstance(side, ast.Constant) and (side.value is None or isinstance(side.value, (int, float, str, bool))):
+            operands.append(side.value)
+        elif isinstance(side, ast.Name):
+            lowered = _lower_value(state, side)
+            if not isinstance(lowered, _Reference):
+                return None
+            _wire(state, instance, side.id, lowered, given)
+            operands.append(side.id)
+        else:
+            return None
+    left, right = operands
+    return Comparison(operator=_COMPARISON_OPERATORS[type(test.ops[0])], left=left, right=right)
+
+
 def _snapshot_namespace_boundary(state: _LoweringState) -> dict[str, t.Any] | None:
     """Snapshot region boundaries only when routing a declared namespace field."""
     if not any('.' in path for path in state.inputs):

src/aiida/engine/processes/graphs/ir.py

@@ -26,6 +26,7 @@
 
 __all__ = (
     'BranchControl',
+    'Comparison',
     'Dependency',
     'Endpoint',
     'GraphIR',
@@ -35,7 +36,7 @@
     'TaskIR',
 )
 
-GRAPH_IR_VERSION = '0.1'
+GRAPH_IR_VERSION = '0.2'
 TASK_IR_VERSION = '0.1'
 CONDITION_PORT = 'condition'
 TaskKind = t.Literal['process', 'graph', 'branch', 'loop']
@@ -269,13 +270,15 @@ def from_dict(cls, data: Mapping[str, t.Any]) -> GraphTask:
                 body=body,
                 otherwise=GraphIR.from_dict(data['otherwise']) if data.get('otherwise') is not None else None,
                 condition_port=data.get('condition_port', CONDITION_PORT),
+                comparison=Comparison.from_dict(data['comparison']) if data.get('comparison') is not None else None,
                 **common,
             )
         if kind == 'loop':
             assert body is not None
             return LoopControl(
                 body=body,
                 condition_port=data.get('condition_port', CONDITION_PORT),
+                comparison=Comparison.from_dict(data['comparison']) if data.get('comparison') is not None else None,
                 max_iterations=data['max_iterations'],
                 **common,
             )
@@ -309,12 +312,44 @@ class SubgraphTask(BodyTask):
     KIND: t.ClassVar[TaskKind] = 'graph'
 
 
+@dataclass(frozen=True, kw_only=True)
+class Comparison:
+    """A comparison over control inputs or literals, evaluated without a task."""
+
+    operator: str
+    left: str | int | float | bool | None
+    right: str | int | float | bool | None
+
+    def __post_init__(self) -> None:
+        if self.operator not in COMPARISON_OPERATORS:
+            msg = f'unsupported comparison operator {self.operator!r}; use one of {sorted(COMPARISON_OPERATORS)}.'
+            raise GraphTypeCheckError(msg)
+
+    @property
+    def ports(self) -> tuple[str, ...]:
+        """Return the control input ports this comparison reads."""
+        return tuple(side for side in (self.left, self.right) if isinstance(side, str))
+
+    def to_dict(self) -> dict[str, t.Any]:
+        """Return the comparison without its evaluated values."""
+        return {'operator': self.operator, 'left': self.left, 'right': self.right}
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, t.Any]) -> Comparison:
+        """Rebuild a comparison from its serialized form."""
+        return cls(operator=data['operator'], left=data.get('left'), right=data.get('right'))
+
+
+COMPARISON_OPERATORS = ('==', '!=', '<', '<=', '>', '>=')
+
+
 @dataclass(frozen=True, kw_only=True)
 class BranchControl(BodyTask):
     """A condition dispatching exactly one of two nested templates."""
 
     KIND: t.ClassVar[TaskKind] = 'branch'
     condition_port: str = CONDITION_PORT
+    comparison: Comparison | None = None
     otherwise: GraphIR | None = None
 
     @property
@@ -331,6 +366,7 @@ def to_dict(self) -> dict[str, t.Any]:
         return {
             **super().to_dict(),
             'condition_port': self.condition_port,
+            'comparison': None if self.comparison is None else self.comparison.to_dict(),
             'otherwise': None if self.otherwise is None else self.otherwise.to_dict(),
         }
 
@@ -341,10 +377,16 @@ class LoopControl(BodyTask):
 
     KIND: t.ClassVar[TaskKind] = 'loop'
     condition_port: str = CONDITION_PORT
+    comparison: Comparison | None = None
     max_iterations: int = 1000
 
     def to_dict(self) -> dict[str, t.Any]:
-        return {**super().to_dict(), 'condition_port': self.condition_port, 'max_iterations': self.max_iterations}
+        return {
+            **super().to_dict(),
+            'condition_port': self.condition_port,
+            'comparison': None if self.comparison is None else self.comparison.to_dict(),
+            'max_iterations': self.max_iterations,
+        }
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -369,7 +411,11 @@ def to_dict(self) -> dict[str, t.Any]:
 
 def _accepts_condition(task: GraphTask, port: str | None) -> bool:
     """Return whether a control task wires its condition outside declared inputs."""
-    return isinstance(task, (BranchControl, LoopControl)) and port == task.condition_port
+    if not isinstance(task, (BranchControl, LoopControl)):
+        return False
+    if port == task.condition_port:
+        return True
+    return task.comparison is not None and port in task.comparison.ports
 
 
 def _stored_types(port: t.Any) -> tuple[type, ...]:
@@ -515,7 +561,9 @@ def serialize_inputs(self, inputs: Mapping[str, t.Any]) -> dict[str, t.Any]:
 
     @staticmethod
     def _input_ports_at(task: GraphTask, path: str) -> list[t.Any]:
-        if isinstance(task, (BranchControl, LoopControl)) and path == task.condition_port:
+        if isinstance(task, (BranchControl, LoopControl)) and (
+            path == task.condition_port or (task.comparison is not None and path in task.comparison.ports)
+        ):
             return []
         return [port for namespace in task.input_ports() if (port := port_at(namespace, path)) is not None]
 
@@ -630,7 +678,11 @@ def _validate_region(self, task: GraphTask) -> None:
                 msg = 'both branches must return the same outputs.'
                 raise GraphTypeCheckError(msg)
             self._check_compatible(task.body.output_spec(), task.otherwise.output_spec(), 'branch outputs')
-        if isinstance(task, LoopControl) and (task.condition_port not in task.body.outputs or task.max_iterations < 1):
+        if (
+            isinstance(task, LoopControl)
+            and task.comparison is None
+            and (task.condition_port not in task.body.outputs or task.max_iterations < 1)
+        ):
             msg = 'a loop must return its condition and allow at least one iteration.'
             raise GraphTypeCheckError(msg)
 

src/aiida/engine/processes/graphs/run.py

@@ -18,6 +18,7 @@
 from aiida.engine.processes.graphs.inputs import at, place
 from aiida.engine.processes.graphs.ir import (
     BranchControl,
+    Comparison,
     GraphIR,
     GraphTask,
     LoopControl,
@@ -37,6 +38,28 @@ def holds(condition: t.Any) -> bool:
     return bool(condition.value if isinstance(condition, BaseType) else condition)
 
 
+def evaluate(comparison: Comparison, values: Mapping[str, t.Any]) -> bool:
+    """Decide a comparison from control inputs or iteration outputs."""
+    left: t.Any = values[comparison.left] if isinstance(comparison.left, str) else comparison.left
+    right: t.Any = values[comparison.right] if isinstance(comparison.right, str) else comparison.right
+    left = left.value if isinstance(left, BaseType) else left
+    right = right.value if isinstance(right, BaseType) else right
+    if comparison.operator == '==':
+        return bool(left == right)
+    if comparison.operator == '!=':
+        return bool(left != right)
+    if comparison.operator == '<':
+        return bool(left < right)
+    if comparison.operator == '<=':
+        return bool(left <= right)
+    if comparison.operator == '>':
+        return bool(left > right)
+    if comparison.operator == '>=':
+        return bool(left >= right)
+    msg = f'unsupported comparison operator {comparison.operator!r}.'
+    raise RuntimeError(msg)
+
+
 def returned(node: Node) -> dict[str, t.Any]:
     """Read the leaf outputs of one graph iteration."""
     return {entry.link_label: entry.node for entry in node.base.links.get_outgoing(link_type=LinkType.RETURN).all()}
@@ -209,7 +232,11 @@ def _begin(self, name: str, step: Step) -> None:
 
     def _begin_branch(self, task: BranchControl, inputs: dict[str, t.Any], step: Step) -> None:
         """Dispatch the taken side of a branch, skipping when it has no body."""
-        taken = task.body if holds(inputs.pop(task.condition_port, None)) else task.otherwise
+        if task.comparison is not None:
+            decision = evaluate(task.comparison, inputs)
+        else:
+            decision = holds(inputs.pop(task.condition_port, None))
+        taken = task.body if decision else task.otherwise
         if taken is None:
             self.skipped.add(task.name)
             return
@@ -219,7 +246,11 @@ def _begin_branch(self, task: BranchControl, inputs: dict[str, t.Any], step: Ste
 
     def _begin_loop(self, task: LoopControl, inputs: dict[str, t.Any], step: Step) -> None:
         """Start iterating a loop, or pass through never-entered initial state."""
-        if not holds(inputs.get(task.condition_port, True)):
+        if task.comparison is not None:
+            if not evaluate(task.comparison, inputs):
+                self._pass_initial_state(task, inputs)
+                return
+        elif not holds(inputs.get(task.condition_port, True)):
             self._pass_initial_state(task, inputs)
             return
         self.instances[task.name] = []
@@ -236,7 +267,12 @@ def _continue_loop(self, task: LoopControl, step: Step) -> None:
         if not runs or any(instance not in self.done for instance in runs):
             return
         last = load_node(self.done[runs[-1]])
-        if not last.is_finished_ok or not holds(returned(last).get(task.condition_port)):
+        if not last.is_finished_ok:
+            return
+        if task.comparison is not None:
+            if not evaluate(task.comparison, {**self._inputs_for(task), **returned(last)}):
+                return
+        elif not holds(returned(last).get(task.condition_port)):
             return
         if len(runs) >= task.max_iterations:
             msg = f'loop `{task.name}` exceeded {task.max_iterations} iterations with its condition still true.'

tests/engine/processes/graphs/language/constructs/conditional/rejected.py

@@ -13,7 +13,6 @@
 
 __all__ = (
     'boolean_condition',
-    'compare_condition',
     'inline_conditional',
     'mismatched_branches',
     'walrus_condition',
@@ -47,23 +46,6 @@ def boolean_condition(value: int, flag: bool) -> int:
 #     return flag and value > 0
 
 
-@graph
-def compare_condition(value: int) -> int:
-    """Compare inline; conditions are calls or names."""
-    if value > 0:
-        chosen = add(x=value, y=1)
-    else:
-        chosen = add(x=value, y=2)
-    return chosen
-
-
-# Compute the predicate in a task instead:
-
-# @task
-# def positive(value: int) -> bool:
-#     return value > 0
-
-
 @graph
 def inline_conditional(value: int, flag: bool) -> int:
     """Select inline; branch explicitly instead."""

tests/engine/processes/graphs/language/constructs/conditional/supported.py

@@ -12,6 +12,7 @@
 from tests.engine.processes.graphs.language.tasks import add, is_positive
 
 __all__ = (
+    'compare_condition',
     'computed_choice',
     'external_choice',
     'inline_condition',
@@ -39,6 +40,16 @@ def computed_choice(value: int) -> int:
     return selected_branch
 
 
+@graph
+def compare_condition(value: int) -> int:
+    """Compare against a bound; the scheduler evaluates the comparison."""
+    if value > 0:
+        chosen = add(x=value, y=1)
+    else:
+        chosen = add(x=value, y=2)
+    return chosen
+
+
 @graph
 def inline_condition(value: int) -> int:
     """Branch on an inline predicate call, bound like the explicit form."""