Extend annotation function capabilites

pyreason/pyreason.py

@@ -1413,16 +1413,71 @@ def add_fact_from_csv(csv_path: str, raise_errors = True) -> None:
 
 
 def add_annotation_function(function: Callable) -> None:
-    """Function to add annotation functions to PyReason. The added functions can be used in rules
+    """Function to add annotation functions to PyReason. The added functions can be used in rules.
 
-    :param function: Function to be added. This has to be under a numba `njit` decorator. function has signature: two parameters as input -- annotations, weights
+    The function must be ``@numba.njit``-decorated and must accept exactly one of the two
+    supported signatures:
+
+    - 2 args (legacy)::
+
+          def fn(annotations, weights) -> Tuple[float, float]
+
+      ``annotations`` is a per-clause list of bounds for the atoms that satisfied that
+      clause. The relational join across body variables is performed inside the engine
+      and projected away before this hook fires, so 2-arg functions cannot recover which
+      grounding produced which bound.
+
+    - 6 args (extended)::
+
+          def fn(annotations, weights,
+                 qualified_nodes, qualified_edges,
+                 clause_labels, clause_variables) -> Tuple[float, float]
+
+      Same as the legacy signature, plus four extra args that expose the per-clause
+      structure the engine already builds:
+
+      ====================  =====================================================
+      ``annotations[i]``    bounds of atoms that satisfied clause ``i``
+      ``qualified_nodes[i]``    nodes that satisfied clause ``i`` (empty for edge clauses)
+      ``qualified_edges[i]``    edges that satisfied clause ``i`` (empty for node clauses)
+      ``clause_labels[i]``  predicate label of clause ``i``
+      ``clause_variables[i]``  variable names of clause ``i``
+                                 (length 1 for node clauses, length 2 for edge clauses)
+      ====================  =====================================================
+
+      Comparison clauses are not surfaced, so the list lengths may be less than
+      ``len(rule.get_clauses())``. Match clauses by predicate name and variable role,
+      not by position in the rule body — the parser's ``reorder_clauses`` optimization
+      may rewrite the order.
+
+    Arity validation runs at registration time, so anything other than 2 or 6 raises
+    ``TypeError`` here rather than producing a confusing failure inside the reasoning
+    loop.
+
+    :param function: Function to be added. Must be ``@numba.njit``-decorated and must
+        match one of the two supported signatures above.
     :type function: Callable
     :return: None
+    :raises TypeError: if ``function`` does not have exactly 2 or 6 positional
+        arguments.
     """
     # Make sure that the functions are jitted so that they can be passed around in other jitted functions
     # TODO: Remove if necessary
     # assert hasattr(function, 'nopython_signatures'), 'The function to be added has to be under a `numba.njit` decorator'
 
+    # Arity gate: only 2-arg and 6-arg signatures are supported by `annotate`.
+    # Validating here keeps the error close to the user's call site and avoids
+    # `raise` inside numba.objmode (which would fail with_lifting).
+    py_func = getattr(function, 'py_func', function)
+    nargs = py_func.__code__.co_argcount
+    if nargs != 2 and nargs != 6:
+        raise TypeError(
+            f"Annotation function {py_func.__name__!r} must accept exactly 2 positional "
+            f"args (annotations, weights) or exactly 6 positional args (annotations, "
+            f"weights, qualified_nodes, qualified_edges, clause_labels, clause_variables); "
+            f"got {nargs}."
+        )
+
     __annotation_functions.append(function)
 
 

setup.py

@@ -8,7 +8,7 @@
 
 setup(
     name='pyreason',
-    version='3.5.1',
+    version='3.6.0',
     author='Dyuman Aditya',
     author_email='[email protected]',
     description='An explainable inference software supporting annotated, real valued, graph based and temporal logic',

tests/functional/test_advanced_features.py

@@ -43,6 +43,72 @@ def identity_func(annotations):
     return result
 
 
+@numba.njit
+def ann_fn_paired(annotations, weights, qualified_nodes, qualified_edges, clause_labels, clause_variables):
+    # 6-arg annotation function: pair hasLabel(CB1,X) and hasLabel(CB2,Y) atoms
+    # via conn(X,Y) groundings. Identify clauses by predicate + variable role
+    # rather than body position (reorder_clauses may rewrite clause order).
+    head_var = "CB2"
+    conn_idx = -1
+    has_label_cb1_idx = -1
+    has_label_cb2_idx = -1
+    for i in range(len(clause_labels)):
+        name = clause_labels[i].value
+        if name == "conn":
+            conn_idx = i
+        elif name == "hasLabel":
+            if clause_variables[i][0] == head_var:
+                has_label_cb2_idx = i
+            else:
+                has_label_cb1_idx = i
+
+    if conn_idx < 0 or has_label_cb1_idx < 0 or has_label_cb2_idx < 0:
+        return 0.0, 1.0
+
+    best_lower = 0.0
+    best_upper = 0.0
+    found_any = False
+    conn_pairs = qualified_edges[conn_idx]
+    for ci in range(len(conn_pairs)):
+        x_val = conn_pairs[ci][0]
+        y_val = conn_pairs[ci][1]
+
+        x_lower = -1.0
+        x_upper = -1.0
+        for k in range(len(qualified_edges[has_label_cb1_idx])):
+            if qualified_edges[has_label_cb1_idx][k][1] == x_val:
+                x_lower = annotations[has_label_cb1_idx][k].lower
+                x_upper = annotations[has_label_cb1_idx][k].upper
+                break
+        if x_lower < 0.0:
+            continue
+
+        y_lower = -1.0
+        y_upper = -1.0
+        for k in range(len(qualified_edges[has_label_cb2_idx])):
+            if qualified_edges[has_label_cb2_idx][k][1] == y_val:
+                y_lower = annotations[has_label_cb2_idx][k].lower
+                y_upper = annotations[has_label_cb2_idx][k].upper
+                break
+        if y_lower < 0.0:
+            continue
+
+        pair_lower = min(x_lower, y_lower)
+        pair_upper = min(x_upper, y_upper)
+        if not found_any or pair_lower > best_lower:
+            best_lower = pair_lower
+            best_upper = pair_upper
+            found_any = True
+
+    if not found_any:
+        return 0.0, 1.0
+    lower = min(best_lower, 1.0)
+    upper = min(best_upper, 1.0)
+    if lower > upper:
+        return 0.0, 1.0
+    return lower, upper
+
+
 @pytest.mark.parametrize("mode", ["regular", "fp", "parallel"])
 def test_probability_func_consistency(mode):
     """Ensure annotation function behaves the same with and without JIT."""
@@ -86,15 +152,48 @@ def test_head_functions(mode):
 @pytest.mark.slow
 @pytest.mark.parametrize("mode", ["regular", "fp", "parallel"])
 def test_annotation_function(mode):
-    """Test annotation function usage in reasoning."""
+    """Annotation function usage in reasoning.
+
+    Exercises both supported signatures in a single reason() call:
+      - ``probability_func`` (2-arg legacy: annotations, weights)
+      - ``ann_fn_paired``    (6-arg extended: + qualified_nodes,
+                              qualified_edges, clause_labels,
+                              clause_variables)
+
+    Mixing the two signatures on different rules confirms per-rule dispatch
+    in ``annotate`` and the per-rule metadata gate in ``_ground_rule``
+    (``extended_ann_fn_flags[i]``) route each rule to the correct path.
+    ``atom_trace`` is left at its default (off) so the 6-arg path is
+    exercised through the perf gate alone, not via the atom_trace branch.
+    """
     setup_mode(mode)
 
     pr.settings.allow_ground_rules = True
 
+    # 2-arg path: simple disjoint probability
     pr.add_fact(pr.Fact('P(A) : [0.01, 1]'))
     pr.add_fact(pr.Fact('P(B) : [0.2, 1]'))
+
+    # 6-arg path: hasLabel + conn join. Correct answer for hackerAt(b) is
+    # [0.5, 1] from the conn-paired (l1, l5) grounding, not [0.3, 1] from
+    # the positional weakest-link aggregation a 2-arg fn would yield.
+    pr.add_fact(pr.Fact("hasLabel(a, l1):[0.5,1]"))
+    pr.add_fact(pr.Fact("hasLabel(a, l2):[0.6,1]"))
+    pr.add_fact(pr.Fact("hasLabel(a, l_unused):[0.8,1]"))
+    pr.add_fact(pr.Fact("hasLabel(b, l4):[0.3,1]"))
+    pr.add_fact(pr.Fact("hasLabel(b, l5):[0.8,1]"))
+    pr.add_fact(pr.Fact("conn(l1, l4)"))
+    pr.add_fact(pr.Fact("conn(l2, l4)"))
+    pr.add_fact(pr.Fact("conn(l1, l5)"))
+    pr.add_fact(pr.Fact("conn(l1, l6)"))
+
     pr.add_annotation_function(probability_func)
+    pr.add_annotation_function(ann_fn_paired)
+
     pr.add_rule(pr.Rule('union_probability(A, B):probability_func <- P(A):[0, 1], P(B):[0, 1]', infer_edges=True))
+    pr.add_rule(pr.Rule(
+        "hackerAt(CB2):ann_fn_paired <- hasLabel(CB1, X):[0.001,1], hasLabel(CB2,Y):[0.001,1], conn(X,Y)"
+    ))
 
     interpretation = pr.reason(timesteps=1)
 
@@ -105,6 +204,8 @@ def test_annotation_function(mode):
         print()
 
     assert interpretation.query(pr.Query('union_probability(A, B) : [0.21, 1]')), 'Union probability should be 0.21'
+    assert interpretation.query(pr.Query('hackerAt(b) : [0.5, 1]')), \
+        'hackerAt(b) should be [0.5, 1] (best conn-paired grounding: hasLabel(a,l1) + hasLabel(b,l5))'
 
 
 @pytest.mark.slow

tests/unit/disable_jit/interpretations/test_ground_rule_helpers.py

@@ -1039,7 +1039,7 @@ def __init__(self, d): self.world = d
 
     # One applicable rule instance for node head
     assert len(apps_node) == 1 and apps_edge == []
-    head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_add = apps_node[0]
+    head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_add, _, _ = apps_node[0]
 
     assert head_grounding == "H"
     # _add_node called to materialize head node (ground rule)
@@ -1093,7 +1093,7 @@ def test_ground_rule_edge_infer_adds_nodes_and_unlabeled_edge(monkeypatch):
     # One applicable edge instance
     assert apps_node == []
     assert len(apps_edge) == 1
-    (e, annotations, qn, qe, edges_to_add) = apps_edge[0]
+    (e, annotations, qn, qe, edges_to_add, _, _) = apps_edge[0]
     assert e == ("S","T")
 
     # infer_edges → edges_to_add lists get S and T
@@ -1160,7 +1160,7 @@ def __init__(self, d): self.world = d
     # One applicable instance using existing edge
     assert apps_node == []
     assert len(apps_edge) == 1
-    (e, annotations, qn, qe, edges_to_add) = apps_edge[0]
+    (e, annotations, qn, qe, edges_to_add, _, _) = apps_edge[0]
 
     # The head grounding should be the concrete edge
     assert e == ("A","B")
@@ -1255,7 +1255,7 @@ def __init__(self, d): self.world = d
 
     # Node-head: expect exactly one applicable rule instance
     assert len(apps_node) == 1 and apps_edge == []
-    head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_add = apps_node[0]
+    head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_add, _, _ = apps_node[0]
     assert head_grounding == "H"
 
     # We have 3 edge clauses → 3 entries added for both trace and annotations
@@ -1376,7 +1376,7 @@ def add_label(edges, lbl):
     # Exactly one head pair (H1,H2) → one applicable edge
     assert apps_node == []
     assert len(apps_edge) == 1
-    (e, annotations, qn, qe, edges_to_add) = apps_edge[0]
+    (e, annotations, qn, qe, edges_to_add, _, _) = apps_edge[0]
     assert e == (hv1, hv2)
 
     # 7 clauses → 7 entries
@@ -1463,7 +1463,7 @@ def __init__(self, d):
 
     assert apps_edge == []
     assert len(apps_node) == 1
-    head_grounding, annotations, qn, qe, edges_to_add = apps_node[0]
+    head_grounding, annotations, qn, qe, edges_to_add, _, _ = apps_node[0]
     assert head_grounding == "A"
     assert qn[0] == ["A"]
     assert annotations[0] == ["ANN_A"]
@@ -1555,7 +1555,7 @@ def test_ground_rule_edge_head_vars_use_existing_nodes_when_allowed(monkeypatch)
 
     assert apps_node == []
     assert len(apps_edge) == 1
-    e, annotations, qn, qe, edges_to_add = apps_edge[0]
+    e, annotations, qn, qe, edges_to_add, _, _ = apps_edge[0]
     assert e == ("A", "B")
     assert annotations == []
     assert qn == []
@@ -1969,7 +1969,7 @@ def __init__(self, d):
 
     assert apps_node == []
     assert len(apps_edge) == 1
-    (e, annotations, qn, qe, edges_to_add) = apps_edge[0]
+    (e, annotations, qn, qe, edges_to_add, _, _) = apps_edge[0]
     assert e == ("a1", "b1")
     assert qn[0] == ["a1"]
     assert qn[1] == ["b1"]