diff --git a/pyreason/pyreason.py b/pyreason/pyreason.py index 7de69aca..52fa3cf6 100755 --- a/pyreason/pyreason.py +++ b/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) diff --git a/pyreason/scripts/interpretation/interpretation.py b/pyreason/scripts/interpretation/interpretation.py index 680bb5fe..c2e1a187 100755 --- a/pyreason/scripts/interpretation/interpretation.py +++ b/pyreason/scripts/interpretation/interpretation.py @@ -28,13 +28,16 @@ facts_to_be_applied_edge_type = numba.types.Tuple((numba.types.uint16, edge_type, label.label_type, interval.interval_type, numba.types.boolean, numba.types.boolean)) # Type for returning list of applicable rules for a certain rule -# node/edge, annotations, qualified nodes, qualified edges, edges to be added +# node/edge, annotations, qualified nodes, qualified edges, edges to be added, +# clause labels (per-clause predicate label), clause variables (per-clause variable names) node_applicable_rule_type = numba.types.Tuple(( node_type, numba.types.ListType(numba.types.ListType(interval.interval_type)), numba.types.ListType(numba.types.ListType(node_type)), numba.types.ListType(numba.types.ListType(edge_type)), - numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)) + numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)), + numba.types.ListType(label.label_type), + numba.types.ListType(numba.types.ListType(numba.types.string)) )) edge_applicable_rule_type = numba.types.Tuple(( @@ -42,7 +45,9 @@ numba.types.ListType(numba.types.ListType(interval.interval_type)), numba.types.ListType(numba.types.ListType(node_type)), numba.types.ListType(numba.types.ListType(edge_type)), - numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)) + numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)), + numba.types.ListType(label.label_type), + numba.types.ListType(numba.types.ListType(numba.types.string)) )) rules_to_be_applied_node_type = numba.types.Tuple((numba.types.uint16, node_type, label.label_type, interval.interval_type, numba.types.boolean)) @@ -216,7 +221,15 @@ def _start_fp(self, rules, max_facts_time, verbose, again, restart): if restart: self.time = 0 self.prev_reasoning_data[0] = 0 - fp_cnt, t = self.reason(self.interpretations_node, self.interpretations_edge, self.predicate_map_node, self.predicate_map_edge, self.tmax, self.prev_reasoning_data, rules, self.nodes, self.edges, self.neighbors, self.reverse_neighbors, self.rules_to_be_applied_node, self.rules_to_be_applied_edge, self.edges_to_be_added_node_rule, self.edges_to_be_added_edge_rule, self.rules_to_be_applied_node_trace, self.rules_to_be_applied_edge_trace, self.facts_to_be_applied_node, self.facts_to_be_applied_edge, self.facts_to_be_applied_node_trace, self.facts_to_be_applied_edge_trace, self.ipl, self.rule_trace_node, self.rule_trace_edge, self.rule_trace_node_atoms, self.rule_trace_edge_atoms, self.reverse_graph, self.atom_trace, self.save_graph_attributes_to_rule_trace, self.persistent, self.inconsistency_check, self.store_interpretation_changes, self.update_mode, self.allow_ground_rules, max_facts_time, self.annotation_functions, self.head_functions, self._convergence_mode, self._convergence_delta, self.num_ga, verbose, again, self.closed_world_predicates) + # Per-rule flag: True iff the rule's annotation function is registered + # with the extended 6-arg signature. Gates the per-grounding metadata + # build in _ground_rule so legacy 2-arg ann_fns pay zero extra cost. + ann_fn_arity = {f.__name__: getattr(f, 'py_func', f).__code__.co_argcount for f in self.annotation_functions} + extended_ann_fn_flags = numba.typed.List.empty_list(numba.types.boolean) + for r in rules: + fn_name = r.get_annotation_function() + extended_ann_fn_flags.append(fn_name != '' and ann_fn_arity.get(fn_name, 0) == 6) + fp_cnt, t = self.reason(self.interpretations_node, self.interpretations_edge, self.predicate_map_node, self.predicate_map_edge, self.tmax, self.prev_reasoning_data, rules, self.nodes, self.edges, self.neighbors, self.reverse_neighbors, self.rules_to_be_applied_node, self.rules_to_be_applied_edge, self.edges_to_be_added_node_rule, self.edges_to_be_added_edge_rule, self.rules_to_be_applied_node_trace, self.rules_to_be_applied_edge_trace, self.facts_to_be_applied_node, self.facts_to_be_applied_edge, self.facts_to_be_applied_node_trace, self.facts_to_be_applied_edge_trace, self.ipl, self.rule_trace_node, self.rule_trace_edge, self.rule_trace_node_atoms, self.rule_trace_edge_atoms, self.reverse_graph, self.atom_trace, self.save_graph_attributes_to_rule_trace, self.persistent, self.inconsistency_check, self.store_interpretation_changes, self.update_mode, self.allow_ground_rules, max_facts_time, self.annotation_functions, extended_ann_fn_flags, self.head_functions, self._convergence_mode, self._convergence_delta, self.num_ga, verbose, again, self.closed_world_predicates) self.time = t - 1 # If we need to reason again, store the next timestep to start from self.prev_reasoning_data[0] = t @@ -226,7 +239,7 @@ def _start_fp(self, rules, max_facts_time, verbose, again, restart): @staticmethod @numba.njit(cache=True, parallel=False) - def reason(interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, tmax, prev_reasoning_data, rules, nodes, edges, neighbors, reverse_neighbors, rules_to_be_applied_node, rules_to_be_applied_edge, edges_to_be_added_node_rule, edges_to_be_added_edge_rule, rules_to_be_applied_node_trace, rules_to_be_applied_edge_trace, facts_to_be_applied_node, facts_to_be_applied_edge, facts_to_be_applied_node_trace, facts_to_be_applied_edge_trace, ipl, rule_trace_node, rule_trace_edge, rule_trace_node_atoms, rule_trace_edge_atoms, reverse_graph, atom_trace, save_graph_attributes_to_rule_trace, persistent, inconsistency_check, store_interpretation_changes, update_mode, allow_ground_rules, max_facts_time, annotation_functions, head_functions, convergence_mode, convergence_delta, num_ga, verbose, again, closed_world_predicates): + def reason(interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, tmax, prev_reasoning_data, rules, nodes, edges, neighbors, reverse_neighbors, rules_to_be_applied_node, rules_to_be_applied_edge, edges_to_be_added_node_rule, edges_to_be_added_edge_rule, rules_to_be_applied_node_trace, rules_to_be_applied_edge_trace, facts_to_be_applied_node, facts_to_be_applied_edge, facts_to_be_applied_node_trace, facts_to_be_applied_edge_trace, ipl, rule_trace_node, rule_trace_edge, rule_trace_node_atoms, rule_trace_edge_atoms, reverse_graph, atom_trace, save_graph_attributes_to_rule_trace, persistent, inconsistency_check, store_interpretation_changes, update_mode, allow_ground_rules, max_facts_time, annotation_functions, extended_ann_fn_flags, head_functions, convergence_mode, convergence_delta, num_ga, verbose, again, closed_world_predicates): t = prev_reasoning_data[0] fp_cnt = prev_reasoning_data[1] max_rules_time = 0 @@ -561,14 +574,14 @@ def reason(interpretations_node, interpretations_edge, predicate_map_node, predi # Only go through if the rule can be applied within the given timesteps, or we're running until convergence delta_t = rule.get_delta() if t + delta_t <= tmax or tmax == -1 or again: - applicable_node_rules, applicable_edge_rules = _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, allow_ground_rules, num_ga, t, head_functions, closed_world_predicates) + applicable_node_rules, applicable_edge_rules = _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, extended_ann_fn_flags[i], allow_ground_rules, num_ga, t, head_functions, closed_world_predicates) # Loop through applicable rules and add them to the rules to be applied for later or next fp operation for applicable_rule in applicable_node_rules: - n, annotations, qualified_nodes, qualified_edges, _ = applicable_rule + n, annotations, qualified_nodes, qualified_edges, _, clause_labels, clause_variables = applicable_rule # If there is an edge to add or the predicate doesn't exist or the interpretation is not static if rule.get_target() not in interpretations_node[n].world or not interpretations_node[n].world[rule.get_target()].is_static(): - bnd = annotate(annotation_functions, rule, annotations, rule.get_weights()) + bnd = annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, rule.get_weights()) # If the rule head was negated AND an ann_fn produced the bound, invert it: # ~[l,u] = [1-u, 1-l]. For non-ann_fn negation the parser already folded # the inversion into target_bound, so we must NOT re-invert here. @@ -589,10 +602,10 @@ def reason(interpretations_node, interpretations_edge, predicate_map_node, predi update_threadsafe[i] = False for applicable_rule in applicable_edge_rules: - e, annotations, qualified_nodes, qualified_edges, edges_to_add = applicable_rule + e, annotations, qualified_nodes, qualified_edges, edges_to_add, clause_labels, clause_variables = applicable_rule # If there is an edge to add or the predicate doesn't exist or the interpretation is not static if len(edges_to_add[0]) > 0 or rule.get_target() not in interpretations_edge[e].world or not interpretations_edge[e].world[rule.get_target()].is_static(): - bnd = annotate(annotation_functions, rule, annotations, rule.get_weights()) + bnd = annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, rule.get_weights()) # If the rule head was negated AND an ann_fn produced the bound, invert it: # ~[l,u] = [1-u, 1-l]. For non-ann_fn negation the parser already folded # the inversion into target_bound, so we must NOT re-invert here. @@ -794,7 +807,7 @@ def query(self, query, return_bool=True) -> Union[bool, Tuple[float, float]]: @numba.njit(cache=True) -def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, allow_ground_rules, num_ga, t, head_functions, closed_world_predicates): +def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, extended_ann_fn, allow_ground_rules, num_ga, t, head_functions, closed_world_predicates): # Extract rule params rule_type = rule.get_type() head_variables = rule.get_head_variables() @@ -959,10 +972,16 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map add_head_var_node_to_graph = True groundings[head_var_1] = numba.typed.List([head_var_1]) + # TODO(perf): when extended_ann_fn is True, `numba.typed.List(clause_variables)` + # is allocated K*N times (K head groundings * N clauses) below, but the data + # is rule-static. Hoist a precomputed `clause_variables_precomputed` list + # here and append references in the inner loop instead of fresh copies. for head_grounding in groundings[head_var_1]: qualified_nodes = numba.typed.List.empty_list(numba.typed.List.empty_list(node_type)) qualified_edges = numba.typed.List.empty_list(numba.typed.List.empty_list(edge_type)) annotations = numba.typed.List.empty_list(numba.typed.List.empty_list(interval.interval_type)) + clause_labels_out = numba.typed.List.empty_list(label.label_type) + clause_variables_out = numba.typed.List.empty_list(numba.typed.List.empty_list(numba.types.string)) edges_to_be_added = (numba.typed.List.empty_list(node_type), numba.typed.List.empty_list(node_type), rule_edges[-1]) # Check for satisfaction one more time in case the refining process has changed the groundings @@ -980,12 +999,15 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map clause_var_1 = clause_variables[0] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: if clause_var_1 == head_var_1: qualified_nodes.append(numba.typed.List([head_grounding])) else: qualified_nodes.append(numba.typed.List(groundings[clause_var_1])) qualified_edges.append(numba.typed.List.empty_list(edge_type)) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -999,7 +1021,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map elif clause_type == 'edge': clause_var_1, clause_var_2 = clause_variables[0], clause_variables[1] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: # Cases: Both equal, one equal, none equal qualified_nodes.append(numba.typed.List.empty_list(node_type)) if clause_var_1 == head_var_1: @@ -1010,6 +1032,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map qualified_edges.append(es) else: qualified_edges.append(numba.typed.List(groundings_edges[(clause_var_1, clause_var_2)])) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -1034,7 +1059,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map _add_node(head_var_1, neighbors, reverse_neighbors, nodes, interpretations_node) # For each grounding add a rule to be applied - applicable_rules_node.append((head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_be_added)) + applicable_rules_node.append((head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_be_added, clause_labels_out, clause_variables_out)) elif rule_type == 'edge': head_var_1 = head_variables[0] @@ -1090,11 +1115,17 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map valid_edge_groundings.append((g1, g2)) # Loop through the head variable groundings + # TODO(perf): when extended_ann_fn is True, `numba.typed.List(clause_variables)` + # is allocated K*N times (K edge groundings * N clauses) below, but the data + # is rule-static. Hoist a precomputed `clause_variables_precomputed` list + # here and append references in the inner loop instead of fresh copies. for valid_e in valid_edge_groundings: head_var_1_grounding, head_var_2_grounding = valid_e[0], valid_e[1] qualified_nodes = numba.typed.List.empty_list(numba.typed.List.empty_list(node_type)) qualified_edges = numba.typed.List.empty_list(numba.typed.List.empty_list(edge_type)) annotations = numba.typed.List.empty_list(numba.typed.List.empty_list(interval.interval_type)) + clause_labels_out = numba.typed.List.empty_list(label.label_type) + clause_variables_out = numba.typed.List.empty_list(numba.typed.List.empty_list(numba.types.string)) edges_to_be_added = (numba.typed.List.empty_list(node_type), numba.typed.List.empty_list(node_type), rule_edges[-1]) # Containers to keep track of groundings to make sure that the edge pair is valid @@ -1144,7 +1175,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map if clause_type == 'node': clause_var_1 = clause_variables[0] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: if clause_var_1 == head_var_1: qualified_nodes.append(numba.typed.List([head_var_1_grounding])) elif clause_var_1 == head_var_2: @@ -1152,6 +1183,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map else: qualified_nodes.append(numba.typed.List(temp_groundings[clause_var_1])) qualified_edges.append(numba.typed.List.empty_list(edge_type)) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -1167,7 +1201,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map elif clause_type == 'edge': clause_var_1, clause_var_2 = clause_variables[0], clause_variables[1] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: # Cases: # 1. Both equal (cv1 = hv1 and cv2 = hv2 or cv1 = hv2 and cv2 = hv1) # 2. One equal (cv1 = hv1 or cv2 = hv1 or cv1 = hv2 or cv2 = hv2) @@ -1193,6 +1227,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map qualified_edges.append(es) else: qualified_edges.append(numba.typed.List(temp_groundings_edges[(clause_var_1, clause_var_2)])) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': @@ -1238,7 +1275,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map # Only if all the clauses have valid groundings # if satisfaction: e = (head_var_1_grounding, head_var_2_grounding) - applicable_rules_edge.append((e, annotations, qualified_nodes, qualified_edges, edges_to_be_added)) + applicable_rules_edge.append((e, annotations, qualified_nodes, qualified_edges, edges_to_be_added, clause_labels_out, clause_variables_out)) # Return the applicable rules return applicable_rules_node, applicable_rules_edge @@ -1829,7 +1866,51 @@ def is_satisfied_edge_comparison(interpretations, comp, na): @numba.njit(cache=True) -def annotate(annotation_functions, rule, annotations, weights): +def annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, weights): + """Resolve and invoke the rule's annotation function. + + If the rule has no annotation function attached, returns the rule's static + bound directly. Otherwise looks up the user-registered function by name + and dispatches based on arity (arity is validated at registration; see + ``add_annotation_function``). + + Two supported signatures: + + - 2-arg legacy:: + + def fn(annotations, weights) -> Tuple[float, float] + + - 6-arg extended:: + + def fn(annotations, weights, + qualified_nodes, qualified_edges, + clause_labels, clause_variables) -> Tuple[float, float] + + The extended args expose the per-clause join structure the engine + already builds, so user code can identify clauses by predicate name + + variable role (robust to ``reorder_clauses``) and walk the actual + per-grounding pairings rather than the flattened per-atom bounds. + + Per-clause alignment (extended signature) + ----------------------------------------- + The five list-shaped args are indexed by the same per-clause position ``i``:: + + annotations[i] # List[Interval] — bounds of qualifying atoms + qualified_nodes[i] # List[Node] — nodes that qualified + qualified_edges[i] # List[Edge] — edges that qualified + clause_labels[i] # Label — predicate of this clause + clause_variables[i] # List[str] — variable names of this clause + + Per-clause arity: + + - Node clause: ``qualified_edges[i]`` is empty; ``len(clause_variables[i]) == 1``. + - Edge clause: ``qualified_nodes[i]`` is empty; ``len(clause_variables[i]) == 2``. + + Comparison clauses are skipped during metadata collection, so + ``len(clause_labels)`` may be less than ``len(rule.get_clauses())``. + Do not assume index parity with the original rule body — match clauses + by predicate name and variable role instead. + """ func_name = rule.get_annotation_function() if func_name == '': return rule.get_bnd().lower, rule.get_bnd().upper @@ -1837,7 +1918,16 @@ def annotate(annotation_functions, rule, annotations, weights): with numba.objmode(annotation='Tuple((float64, float64))'): for func in annotation_functions: if func.__name__ == func_name: - annotation = func(annotations, weights) + # Arity is gated at registration time by + # `add_annotation_function`, so only 2 and 6 reach here. + # (Raising inside numba.objmode is unsupported, so we + # rely on the registration-time check for validation.) + py_func = getattr(func, 'py_func', func) + nargs = py_func.__code__.co_argcount + if nargs == 6: + annotation = func(annotations, weights, qualified_nodes, qualified_edges, clause_labels, clause_variables) + else: + annotation = func(annotations, weights) return annotation diff --git a/pyreason/scripts/interpretation/interpretation_fp.py b/pyreason/scripts/interpretation/interpretation_fp.py index 3944732e..2cf174d0 100755 --- a/pyreason/scripts/interpretation/interpretation_fp.py +++ b/pyreason/scripts/interpretation/interpretation_fp.py @@ -32,13 +32,16 @@ facts_to_be_applied_edge_type = numba.types.Tuple((numba.types.uint16, edge_type, label.label_type, interval.interval_type, numba.types.boolean, numba.types.boolean)) # Type for returning list of applicable rules for a certain rule -# node/edge, annotations, qualified nodes, qualified edges, edges to be added +# node/edge, annotations, qualified nodes, qualified edges, edges to be added, +# clause labels (per-clause predicate label), clause variables (per-clause variable names) node_applicable_rule_type = numba.types.Tuple(( node_type, numba.types.ListType(numba.types.ListType(interval.interval_type)), numba.types.ListType(numba.types.ListType(node_type)), numba.types.ListType(numba.types.ListType(edge_type)), - numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)) + numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)), + numba.types.ListType(label.label_type), + numba.types.ListType(numba.types.ListType(numba.types.string)) )) edge_applicable_rule_type = numba.types.Tuple(( @@ -46,7 +49,9 @@ numba.types.ListType(numba.types.ListType(interval.interval_type)), numba.types.ListType(numba.types.ListType(node_type)), numba.types.ListType(numba.types.ListType(edge_type)), - numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)) + numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)), + numba.types.ListType(label.label_type), + numba.types.ListType(numba.types.ListType(numba.types.string)) )) rules_to_be_applied_node_type = numba.types.Tuple((numba.types.uint16, node_type, label.label_type, interval.interval_type, numba.types.boolean)) @@ -225,7 +230,15 @@ def _start_fp(self, rules, max_facts_time, verbose, again, restart): if restart: self.time = 0 self.prev_reasoning_data[0] = 0 - fp_cnt, t = self.reason(self.interpretations_node, self.interpretations_edge, self.predicate_map_node, self.predicate_map_edge, self.tmax, self.prev_reasoning_data, rules, self.nodes, self.edges, self.neighbors, self.reverse_neighbors, self.rules_to_be_applied_node, self.rules_to_be_applied_edge, self.edges_to_be_added_node_rule, self.edges_to_be_added_edge_rule, self.rules_to_be_applied_node_trace, self.rules_to_be_applied_edge_trace, self.facts_to_be_applied_node, self.facts_to_be_applied_edge, self.facts_to_be_applied_node_trace, self.facts_to_be_applied_edge_trace, self.ipl, self.rule_trace_node, self.rule_trace_edge, self.rule_trace_node_atoms, self.rule_trace_edge_atoms, self.reverse_graph, self.atom_trace, self.save_graph_attributes_to_rule_trace, self.persistent, self.inconsistency_check, self.store_interpretation_changes, self.update_mode, self.allow_ground_rules, max_facts_time, self.annotation_functions, self.head_functions, self._convergence_mode, self._convergence_delta, verbose, again, self.closed_world_predicates) + # Per-rule flag: True iff the rule's annotation function is registered + # with the extended 6-arg signature. Gates the per-grounding metadata + # build in _ground_rule so legacy 2-arg ann_fns pay zero extra cost. + ann_fn_arity = {f.__name__: getattr(f, 'py_func', f).__code__.co_argcount for f in self.annotation_functions} + extended_ann_fn_flags = numba.typed.List.empty_list(numba.types.boolean) + for r in rules: + fn_name = r.get_annotation_function() + extended_ann_fn_flags.append(fn_name != '' and ann_fn_arity.get(fn_name, 0) == 6) + fp_cnt, t = self.reason(self.interpretations_node, self.interpretations_edge, self.predicate_map_node, self.predicate_map_edge, self.tmax, self.prev_reasoning_data, rules, self.nodes, self.edges, self.neighbors, self.reverse_neighbors, self.rules_to_be_applied_node, self.rules_to_be_applied_edge, self.edges_to_be_added_node_rule, self.edges_to_be_added_edge_rule, self.rules_to_be_applied_node_trace, self.rules_to_be_applied_edge_trace, self.facts_to_be_applied_node, self.facts_to_be_applied_edge, self.facts_to_be_applied_node_trace, self.facts_to_be_applied_edge_trace, self.ipl, self.rule_trace_node, self.rule_trace_edge, self.rule_trace_node_atoms, self.rule_trace_edge_atoms, self.reverse_graph, self.atom_trace, self.save_graph_attributes_to_rule_trace, self.persistent, self.inconsistency_check, self.store_interpretation_changes, self.update_mode, self.allow_ground_rules, max_facts_time, self.annotation_functions, extended_ann_fn_flags, self.head_functions, self._convergence_mode, self._convergence_delta, verbose, again, self.closed_world_predicates) self.time = t - 1 # If we need to reason again, store the next timestep to start from self.prev_reasoning_data[0] = t @@ -235,7 +248,7 @@ def _start_fp(self, rules, max_facts_time, verbose, again, restart): @staticmethod @numba.njit(cache=True, parallel=False) - def reason(interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, tmax, prev_reasoning_data, rules, nodes, edges, neighbors, reverse_neighbors, rules_to_be_applied_node, rules_to_be_applied_edge, edges_to_be_added_node_rule, edges_to_be_added_edge_rule, rules_to_be_applied_node_trace, rules_to_be_applied_edge_trace, facts_to_be_applied_node, facts_to_be_applied_edge, facts_to_be_applied_node_trace, facts_to_be_applied_edge_trace, ipl, rule_trace_node, rule_trace_edge, rule_trace_node_atoms, rule_trace_edge_atoms, reverse_graph, atom_trace, save_graph_attributes_to_rule_trace, persistent, inconsistency_check, store_interpretation_changes, update_mode, allow_ground_rules, max_facts_time, annotation_functions, head_functions, convergence_mode, convergence_delta, verbose, again, closed_world_predicates): + def reason(interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, tmax, prev_reasoning_data, rules, nodes, edges, neighbors, reverse_neighbors, rules_to_be_applied_node, rules_to_be_applied_edge, edges_to_be_added_node_rule, edges_to_be_added_edge_rule, rules_to_be_applied_node_trace, rules_to_be_applied_edge_trace, facts_to_be_applied_node, facts_to_be_applied_edge, facts_to_be_applied_node_trace, facts_to_be_applied_edge_trace, ipl, rule_trace_node, rule_trace_edge, rule_trace_node_atoms, rule_trace_edge_atoms, reverse_graph, atom_trace, save_graph_attributes_to_rule_trace, persistent, inconsistency_check, store_interpretation_changes, update_mode, allow_ground_rules, max_facts_time, annotation_functions, extended_ann_fn_flags, head_functions, convergence_mode, convergence_delta, verbose, again, closed_world_predicates): t = prev_reasoning_data[0] max_t = t # Keeps track of the max time in each fp operation max_t_changes = t @@ -521,11 +534,11 @@ def reason(interpretations_node, interpretations_edge, predicate_map_node, predi # Only go through if the rule can be applied within the given timesteps, or we're running until convergence delta_t = rule.get_delta() if t + delta_t <= tmax or tmax == -1 or again: - applicable_node_rules, applicable_edge_rules = _ground_rule(rule, interpretations_node[t], interpretations_edge[t], predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, allow_ground_rules, t, head_functions, closed_world_predicates) + applicable_node_rules, applicable_edge_rules = _ground_rule(rule, interpretations_node[t], interpretations_edge[t], predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, extended_ann_fn_flags[i], allow_ground_rules, t, head_functions, closed_world_predicates) # Loop through applicable rules and add them to the rules to be applied for later or next fp operation for applicable_rule in applicable_node_rules: - n, annotations, qualified_nodes, qualified_edges, _ = applicable_rule + n, annotations, qualified_nodes, qualified_edges, _, clause_labels, clause_variables = applicable_rule # Check if this edge rule should be applied should_apply_rule = False @@ -543,7 +556,7 @@ def reason(interpretations_node, interpretations_edge, predicate_map_node, predi should_apply_rule = True if should_apply_rule: - bnd = annotate(annotation_functions, rule, annotations, rule.get_weights()) + bnd = annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, rule.get_weights()) # If the rule head was negated AND an ann_fn produced the bound, invert it: # ~[l,u] = [1-u, 1-l]. For non-ann_fn negation the parser already folded # the inversion into target_bound, so we must NOT re-invert here. @@ -563,7 +576,7 @@ def reason(interpretations_node, interpretations_edge, predicate_map_node, predi update = False for applicable_rule in applicable_edge_rules: - e, annotations, qualified_nodes, qualified_edges, edges_to_add = applicable_rule + e, annotations, qualified_nodes, qualified_edges, edges_to_add, clause_labels, clause_variables = applicable_rule # Check if this edge rule should be applied should_apply_rule = False @@ -585,7 +598,7 @@ def reason(interpretations_node, interpretations_edge, predicate_map_node, predi should_apply_rule = True if should_apply_rule: - bnd = annotate(annotation_functions, rule, annotations, rule.get_weights()) + bnd = annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, rule.get_weights()) # If the rule head was negated AND an ann_fn produced the bound, invert it: # ~[l,u] = [1-u, 1-l]. For non-ann_fn negation the parser already folded # the inversion into target_bound, so we must NOT re-invert here. @@ -916,7 +929,7 @@ def query(self, query, t=-1, return_bool=True) -> Union[bool, Tuple[float, float @numba.njit(cache=True) -def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, allow_ground_rules, t, head_functions, closed_world_predicates): +def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, extended_ann_fn, allow_ground_rules, t, head_functions, closed_world_predicates): # Extract rule params rule_type = rule.get_type() head_variables = rule.get_head_variables() @@ -1081,10 +1094,16 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map add_head_var_node_to_graph = True groundings[head_var_1] = numba.typed.List([head_var_1]) + # TODO(perf): when extended_ann_fn is True, `numba.typed.List(clause_variables)` + # is allocated K*N times (K head groundings * N clauses) below, but the data + # is rule-static. Hoist a precomputed `clause_variables_precomputed` list + # here and append references in the inner loop instead of fresh copies. for head_grounding in groundings[head_var_1]: qualified_nodes = numba.typed.List.empty_list(numba.typed.List.empty_list(node_type)) qualified_edges = numba.typed.List.empty_list(numba.typed.List.empty_list(edge_type)) annotations = numba.typed.List.empty_list(numba.typed.List.empty_list(interval.interval_type)) + clause_labels_out = numba.typed.List.empty_list(label.label_type) + clause_variables_out = numba.typed.List.empty_list(numba.typed.List.empty_list(numba.types.string)) edges_to_be_added = (numba.typed.List.empty_list(node_type), numba.typed.List.empty_list(node_type), rule_edges[-1]) # Check for satisfaction one more time in case the refining process has changed the groundings @@ -1101,12 +1120,15 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map clause_var_1 = clause_variables[0] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: if clause_var_1 == head_var_1: qualified_nodes.append(numba.typed.List([head_grounding])) else: qualified_nodes.append(numba.typed.List(groundings[clause_var_1])) qualified_edges.append(numba.typed.List.empty_list(edge_type)) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -1120,7 +1142,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map elif clause_type == 'edge': clause_var_1, clause_var_2 = clause_variables[0], clause_variables[1] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: # Cases: Both equal, one equal, none equal qualified_nodes.append(numba.typed.List.empty_list(node_type)) if clause_var_1 == head_var_1: @@ -1131,6 +1153,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map qualified_edges.append(es) else: qualified_edges.append(numba.typed.List(groundings_edges[(clause_var_1, clause_var_2)])) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -1155,7 +1180,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map _add_node(head_var_1, neighbors, reverse_neighbors, nodes, interpretations_node) # For each grounding add a rule to be applied - applicable_rules_node.append((head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_be_added)) + applicable_rules_node.append((head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_be_added, clause_labels_out, clause_variables_out)) elif rule_type == 'edge': head_var_1 = head_variables[0] @@ -1211,11 +1236,17 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map valid_edge_groundings.append((g1, g2)) # Loop through the head variable groundings + # TODO(perf): when extended_ann_fn is True, `numba.typed.List(clause_variables)` + # is allocated K*N times (K edge groundings * N clauses) below, but the data + # is rule-static. Hoist a precomputed `clause_variables_precomputed` list + # here and append references in the inner loop instead of fresh copies. for valid_e in valid_edge_groundings: head_var_1_grounding, head_var_2_grounding = valid_e[0], valid_e[1] qualified_nodes = numba.typed.List.empty_list(numba.typed.List.empty_list(node_type)) qualified_edges = numba.typed.List.empty_list(numba.typed.List.empty_list(edge_type)) annotations = numba.typed.List.empty_list(numba.typed.List.empty_list(interval.interval_type)) + clause_labels_out = numba.typed.List.empty_list(label.label_type) + clause_variables_out = numba.typed.List.empty_list(numba.typed.List.empty_list(numba.types.string)) edges_to_be_added = (numba.typed.List.empty_list(node_type), numba.typed.List.empty_list(node_type), rule_edges[-1]) # Containers to keep track of groundings to make sure that the edge pair is valid @@ -1265,7 +1296,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map if clause_type == 'node': clause_var_1 = clause_variables[0] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: if clause_var_1 == head_var_1: qualified_nodes.append(numba.typed.List([head_var_1_grounding])) elif clause_var_1 == head_var_2: @@ -1273,6 +1304,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map else: qualified_nodes.append(numba.typed.List(temp_groundings[clause_var_1])) qualified_edges.append(numba.typed.List.empty_list(edge_type)) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -1288,7 +1322,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map elif clause_type == 'edge': clause_var_1, clause_var_2 = clause_variables[0], clause_variables[1] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: # Cases: # 1. Both equal (cv1 = hv1 and cv2 = hv2 or cv1 = hv2 and cv2 = hv1) # 2. One equal (cv1 = hv1 or cv2 = hv1 or cv1 = hv2 or cv2 = hv2) @@ -1314,6 +1348,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map qualified_edges.append(es) else: qualified_edges.append(numba.typed.List(temp_groundings_edges[(clause_var_1, clause_var_2)])) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': @@ -1359,7 +1396,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map # Only if all the clauses have valid groundings # if satisfaction: e = (head_var_1_grounding, head_var_2_grounding) - applicable_rules_edge.append((e, annotations, qualified_nodes, qualified_edges, edges_to_be_added)) + applicable_rules_edge.append((e, annotations, qualified_nodes, qualified_edges, edges_to_be_added, clause_labels_out, clause_variables_out)) # Return the applicable rules return applicable_rules_node, applicable_rules_edge @@ -1939,7 +1976,51 @@ def is_satisfied_edge_comparison(interpretations, comp, na): @numba.njit(cache=True) -def annotate(annotation_functions, rule, annotations, weights): +def annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, weights): + """Resolve and invoke the rule's annotation function. + + If the rule has no annotation function attached, returns the rule's static + bound directly. Otherwise looks up the user-registered function by name + and dispatches based on arity (arity is validated at registration; see + ``add_annotation_function``). + + Two supported signatures: + + - 2-arg legacy:: + + def fn(annotations, weights) -> Tuple[float, float] + + - 6-arg extended:: + + def fn(annotations, weights, + qualified_nodes, qualified_edges, + clause_labels, clause_variables) -> Tuple[float, float] + + The extended args expose the per-clause join structure the engine + already builds, so user code can identify clauses by predicate name + + variable role (robust to ``reorder_clauses``) and walk the actual + per-grounding pairings rather than the flattened per-atom bounds. + + Per-clause alignment (extended signature) + ----------------------------------------- + The five list-shaped args are indexed by the same per-clause position ``i``:: + + annotations[i] # List[Interval] — bounds of qualifying atoms + qualified_nodes[i] # List[Node] — nodes that qualified + qualified_edges[i] # List[Edge] — edges that qualified + clause_labels[i] # Label — predicate of this clause + clause_variables[i] # List[str] — variable names of this clause + + Per-clause arity: + + - Node clause: ``qualified_edges[i]`` is empty; ``len(clause_variables[i]) == 1``. + - Edge clause: ``qualified_nodes[i]`` is empty; ``len(clause_variables[i]) == 2``. + + Comparison clauses are skipped during metadata collection, so + ``len(clause_labels)`` may be less than ``len(rule.get_clauses())``. + Do not assume index parity with the original rule body — match clauses + by predicate name and variable role instead. + """ func_name = rule.get_annotation_function() if func_name == '': return rule.get_bnd().lower, rule.get_bnd().upper @@ -1947,7 +2028,16 @@ def annotate(annotation_functions, rule, annotations, weights): with numba.objmode(annotation='Tuple((float64, float64))'): for func in annotation_functions: if func.__name__ == func_name: - annotation = func(annotations, weights) + # Arity is gated at registration time by + # `add_annotation_function`, so only 2 and 6 reach here. + # (Raising inside numba.objmode is unsupported, so we + # rely on the registration-time check for validation.) + py_func = getattr(func, 'py_func', func) + nargs = py_func.__code__.co_argcount + if nargs == 6: + annotation = func(annotations, weights, qualified_nodes, qualified_edges, clause_labels, clause_variables) + else: + annotation = func(annotations, weights) return annotation diff --git a/pyreason/scripts/interpretation/interpretation_parallel.py b/pyreason/scripts/interpretation/interpretation_parallel.py index 16d48776..251713b8 100644 --- a/pyreason/scripts/interpretation/interpretation_parallel.py +++ b/pyreason/scripts/interpretation/interpretation_parallel.py @@ -28,13 +28,16 @@ facts_to_be_applied_edge_type = numba.types.Tuple((numba.types.uint16, edge_type, label.label_type, interval.interval_type, numba.types.boolean, numba.types.boolean)) # Type for returning list of applicable rules for a certain rule -# node/edge, annotations, qualified nodes, qualified edges, edges to be added +# node/edge, annotations, qualified nodes, qualified edges, edges to be added, +# clause labels (per-clause predicate label), clause variables (per-clause variable names) node_applicable_rule_type = numba.types.Tuple(( node_type, numba.types.ListType(numba.types.ListType(interval.interval_type)), numba.types.ListType(numba.types.ListType(node_type)), numba.types.ListType(numba.types.ListType(edge_type)), - numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)) + numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)), + numba.types.ListType(label.label_type), + numba.types.ListType(numba.types.ListType(numba.types.string)) )) edge_applicable_rule_type = numba.types.Tuple(( @@ -42,7 +45,9 @@ numba.types.ListType(numba.types.ListType(interval.interval_type)), numba.types.ListType(numba.types.ListType(node_type)), numba.types.ListType(numba.types.ListType(edge_type)), - numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)) + numba.types.Tuple((numba.types.ListType(node_type), numba.types.ListType(node_type), label.label_type)), + numba.types.ListType(label.label_type), + numba.types.ListType(numba.types.ListType(numba.types.string)) )) rules_to_be_applied_node_type = numba.types.Tuple((numba.types.uint16, node_type, label.label_type, interval.interval_type, numba.types.boolean)) @@ -216,7 +221,15 @@ def _start_fp(self, rules, max_facts_time, verbose, again, restart): if restart: self.time = 0 self.prev_reasoning_data[0] = 0 - fp_cnt, t = self.reason(self.interpretations_node, self.interpretations_edge, self.predicate_map_node, self.predicate_map_edge, self.tmax, self.prev_reasoning_data, rules, self.nodes, self.edges, self.neighbors, self.reverse_neighbors, self.rules_to_be_applied_node, self.rules_to_be_applied_edge, self.edges_to_be_added_node_rule, self.edges_to_be_added_edge_rule, self.rules_to_be_applied_node_trace, self.rules_to_be_applied_edge_trace, self.facts_to_be_applied_node, self.facts_to_be_applied_edge, self.facts_to_be_applied_node_trace, self.facts_to_be_applied_edge_trace, self.ipl, self.rule_trace_node, self.rule_trace_edge, self.rule_trace_node_atoms, self.rule_trace_edge_atoms, self.reverse_graph, self.atom_trace, self.save_graph_attributes_to_rule_trace, self.persistent, self.inconsistency_check, self.store_interpretation_changes, self.update_mode, self.allow_ground_rules, max_facts_time, self.annotation_functions, self.head_functions, self._convergence_mode, self._convergence_delta, self.num_ga, verbose, again, self.closed_world_predicates) + # Per-rule flag: True iff the rule's annotation function is registered + # with the extended 6-arg signature. Gates the per-grounding metadata + # build in _ground_rule so legacy 2-arg ann_fns pay zero extra cost. + ann_fn_arity = {f.__name__: getattr(f, 'py_func', f).__code__.co_argcount for f in self.annotation_functions} + extended_ann_fn_flags = numba.typed.List.empty_list(numba.types.boolean) + for r in rules: + fn_name = r.get_annotation_function() + extended_ann_fn_flags.append(fn_name != '' and ann_fn_arity.get(fn_name, 0) == 6) + fp_cnt, t = self.reason(self.interpretations_node, self.interpretations_edge, self.predicate_map_node, self.predicate_map_edge, self.tmax, self.prev_reasoning_data, rules, self.nodes, self.edges, self.neighbors, self.reverse_neighbors, self.rules_to_be_applied_node, self.rules_to_be_applied_edge, self.edges_to_be_added_node_rule, self.edges_to_be_added_edge_rule, self.rules_to_be_applied_node_trace, self.rules_to_be_applied_edge_trace, self.facts_to_be_applied_node, self.facts_to_be_applied_edge, self.facts_to_be_applied_node_trace, self.facts_to_be_applied_edge_trace, self.ipl, self.rule_trace_node, self.rule_trace_edge, self.rule_trace_node_atoms, self.rule_trace_edge_atoms, self.reverse_graph, self.atom_trace, self.save_graph_attributes_to_rule_trace, self.persistent, self.inconsistency_check, self.store_interpretation_changes, self.update_mode, self.allow_ground_rules, max_facts_time, self.annotation_functions, extended_ann_fn_flags, self.head_functions, self._convergence_mode, self._convergence_delta, self.num_ga, verbose, again, self.closed_world_predicates) self.time = t - 1 # If we need to reason again, store the next timestep to start from self.prev_reasoning_data[0] = t @@ -226,7 +239,7 @@ def _start_fp(self, rules, max_facts_time, verbose, again, restart): @staticmethod @numba.njit(cache=True, parallel=True) - def reason(interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, tmax, prev_reasoning_data, rules, nodes, edges, neighbors, reverse_neighbors, rules_to_be_applied_node, rules_to_be_applied_edge, edges_to_be_added_node_rule, edges_to_be_added_edge_rule, rules_to_be_applied_node_trace, rules_to_be_applied_edge_trace, facts_to_be_applied_node, facts_to_be_applied_edge, facts_to_be_applied_node_trace, facts_to_be_applied_edge_trace, ipl, rule_trace_node, rule_trace_edge, rule_trace_node_atoms, rule_trace_edge_atoms, reverse_graph, atom_trace, save_graph_attributes_to_rule_trace, persistent, inconsistency_check, store_interpretation_changes, update_mode, allow_ground_rules, max_facts_time, annotation_functions, head_functions, convergence_mode, convergence_delta, num_ga, verbose, again, closed_world_predicates): + def reason(interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, tmax, prev_reasoning_data, rules, nodes, edges, neighbors, reverse_neighbors, rules_to_be_applied_node, rules_to_be_applied_edge, edges_to_be_added_node_rule, edges_to_be_added_edge_rule, rules_to_be_applied_node_trace, rules_to_be_applied_edge_trace, facts_to_be_applied_node, facts_to_be_applied_edge, facts_to_be_applied_node_trace, facts_to_be_applied_edge_trace, ipl, rule_trace_node, rule_trace_edge, rule_trace_node_atoms, rule_trace_edge_atoms, reverse_graph, atom_trace, save_graph_attributes_to_rule_trace, persistent, inconsistency_check, store_interpretation_changes, update_mode, allow_ground_rules, max_facts_time, annotation_functions, extended_ann_fn_flags, head_functions, convergence_mode, convergence_delta, num_ga, verbose, again, closed_world_predicates): t = prev_reasoning_data[0] fp_cnt = prev_reasoning_data[1] max_rules_time = 0 @@ -561,14 +574,14 @@ def reason(interpretations_node, interpretations_edge, predicate_map_node, predi # Only go through if the rule can be applied within the given timesteps, or we're running until convergence delta_t = rule.get_delta() if t + delta_t <= tmax or tmax == -1 or again: - applicable_node_rules, applicable_edge_rules = _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, allow_ground_rules, num_ga, t, head_functions, closed_world_predicates) + applicable_node_rules, applicable_edge_rules = _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, extended_ann_fn_flags[i], allow_ground_rules, num_ga, t, head_functions, closed_world_predicates) # Loop through applicable rules and add them to the rules to be applied for later or next fp operation for applicable_rule in applicable_node_rules: - n, annotations, qualified_nodes, qualified_edges, _ = applicable_rule + n, annotations, qualified_nodes, qualified_edges, _, clause_labels, clause_variables = applicable_rule # If there is an edge to add or the predicate doesn't exist or the interpretation is not static if rule.get_target() not in interpretations_node[n].world or not interpretations_node[n].world[rule.get_target()].is_static(): - bnd = annotate(annotation_functions, rule, annotations, rule.get_weights()) + bnd = annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, rule.get_weights()) # If the rule head was negated AND an ann_fn produced the bound, invert it: # ~[l,u] = [1-u, 1-l]. For non-ann_fn negation the parser already folded # the inversion into target_bound, so we must NOT re-invert here. @@ -589,10 +602,10 @@ def reason(interpretations_node, interpretations_edge, predicate_map_node, predi update_threadsafe[i] = False for applicable_rule in applicable_edge_rules: - e, annotations, qualified_nodes, qualified_edges, edges_to_add = applicable_rule + e, annotations, qualified_nodes, qualified_edges, edges_to_add, clause_labels, clause_variables = applicable_rule # If there is an edge to add or the predicate doesn't exist or the interpretation is not static if len(edges_to_add[0]) > 0 or rule.get_target() not in interpretations_edge[e].world or not interpretations_edge[e].world[rule.get_target()].is_static(): - bnd = annotate(annotation_functions, rule, annotations, rule.get_weights()) + bnd = annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, rule.get_weights()) # If the rule head was negated AND an ann_fn produced the bound, invert it: # ~[l,u] = [1-u, 1-l]. For non-ann_fn negation the parser already folded # the inversion into target_bound, so we must NOT re-invert here. @@ -794,7 +807,7 @@ def query(self, query, return_bool=True) -> Union[bool, Tuple[float, float]]: @numba.njit(cache=True) -def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, allow_ground_rules, num_ga, t, head_functions, closed_world_predicates): +def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map_node, predicate_map_edge, nodes, edges, neighbors, reverse_neighbors, atom_trace, extended_ann_fn, allow_ground_rules, num_ga, t, head_functions, closed_world_predicates): # Extract rule params rule_type = rule.get_type() head_variables = rule.get_head_variables() @@ -959,10 +972,16 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map add_head_var_node_to_graph = True groundings[head_var_1] = numba.typed.List([head_var_1]) + # TODO(perf): when extended_ann_fn is True, `numba.typed.List(clause_variables)` + # is allocated K*N times (K head groundings * N clauses) below, but the data + # is rule-static. Hoist a precomputed `clause_variables_precomputed` list + # here and append references in the inner loop instead of fresh copies. for head_grounding in groundings[head_var_1]: qualified_nodes = numba.typed.List.empty_list(numba.typed.List.empty_list(node_type)) qualified_edges = numba.typed.List.empty_list(numba.typed.List.empty_list(edge_type)) annotations = numba.typed.List.empty_list(numba.typed.List.empty_list(interval.interval_type)) + clause_labels_out = numba.typed.List.empty_list(label.label_type) + clause_variables_out = numba.typed.List.empty_list(numba.typed.List.empty_list(numba.types.string)) edges_to_be_added = (numba.typed.List.empty_list(node_type), numba.typed.List.empty_list(node_type), rule_edges[-1]) # Check for satisfaction one more time in case the refining process has changed the groundings @@ -980,12 +999,15 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map clause_var_1 = clause_variables[0] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: if clause_var_1 == head_var_1: qualified_nodes.append(numba.typed.List([head_grounding])) else: qualified_nodes.append(numba.typed.List(groundings[clause_var_1])) qualified_edges.append(numba.typed.List.empty_list(edge_type)) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -999,7 +1021,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map elif clause_type == 'edge': clause_var_1, clause_var_2 = clause_variables[0], clause_variables[1] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: # Cases: Both equal, one equal, none equal qualified_nodes.append(numba.typed.List.empty_list(node_type)) if clause_var_1 == head_var_1: @@ -1010,6 +1032,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map qualified_edges.append(es) else: qualified_edges.append(numba.typed.List(groundings_edges[(clause_var_1, clause_var_2)])) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -1034,7 +1059,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map _add_node(head_var_1, neighbors, reverse_neighbors, nodes, interpretations_node) # For each grounding add a rule to be applied - applicable_rules_node.append((head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_be_added)) + applicable_rules_node.append((head_grounding, annotations, qualified_nodes, qualified_edges, edges_to_be_added, clause_labels_out, clause_variables_out)) elif rule_type == 'edge': head_var_1 = head_variables[0] @@ -1090,11 +1115,17 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map valid_edge_groundings.append((g1, g2)) # Loop through the head variable groundings + # TODO(perf): when extended_ann_fn is True, `numba.typed.List(clause_variables)` + # is allocated K*N times (K edge groundings * N clauses) below, but the data + # is rule-static. Hoist a precomputed `clause_variables_precomputed` list + # here and append references in the inner loop instead of fresh copies. for valid_e in valid_edge_groundings: head_var_1_grounding, head_var_2_grounding = valid_e[0], valid_e[1] qualified_nodes = numba.typed.List.empty_list(numba.typed.List.empty_list(node_type)) qualified_edges = numba.typed.List.empty_list(numba.typed.List.empty_list(edge_type)) annotations = numba.typed.List.empty_list(numba.typed.List.empty_list(interval.interval_type)) + clause_labels_out = numba.typed.List.empty_list(label.label_type) + clause_variables_out = numba.typed.List.empty_list(numba.typed.List.empty_list(numba.types.string)) edges_to_be_added = (numba.typed.List.empty_list(node_type), numba.typed.List.empty_list(node_type), rule_edges[-1]) # Containers to keep track of groundings to make sure that the edge pair is valid @@ -1144,7 +1175,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map if clause_type == 'node': clause_var_1 = clause_variables[0] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: if clause_var_1 == head_var_1: qualified_nodes.append(numba.typed.List([head_var_1_grounding])) elif clause_var_1 == head_var_2: @@ -1152,6 +1183,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map else: qualified_nodes.append(numba.typed.List(temp_groundings[clause_var_1])) qualified_edges.append(numba.typed.List.empty_list(edge_type)) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': a = numba.typed.List.empty_list(interval.interval_type) @@ -1167,7 +1201,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map elif clause_type == 'edge': clause_var_1, clause_var_2 = clause_variables[0], clause_variables[1] # 1. - if atom_trace: + if atom_trace or extended_ann_fn: # Cases: # 1. Both equal (cv1 = hv1 and cv2 = hv2 or cv1 = hv2 and cv2 = hv1) # 2. One equal (cv1 = hv1 or cv2 = hv1 or cv1 = hv2 or cv2 = hv2) @@ -1193,6 +1227,9 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map qualified_edges.append(es) else: qualified_edges.append(numba.typed.List(temp_groundings_edges[(clause_var_1, clause_var_2)])) + if extended_ann_fn: + clause_labels_out.append(clause_label) + clause_variables_out.append(numba.typed.List(clause_variables)) # 2. if ann_fn != '': @@ -1238,7 +1275,7 @@ def _ground_rule(rule, interpretations_node, interpretations_edge, predicate_map # Only if all the clauses have valid groundings # if satisfaction: e = (head_var_1_grounding, head_var_2_grounding) - applicable_rules_edge.append((e, annotations, qualified_nodes, qualified_edges, edges_to_be_added)) + applicable_rules_edge.append((e, annotations, qualified_nodes, qualified_edges, edges_to_be_added, clause_labels_out, clause_variables_out)) # Return the applicable rules return applicable_rules_node, applicable_rules_edge @@ -1829,7 +1866,51 @@ def is_satisfied_edge_comparison(interpretations, comp, na): @numba.njit(cache=True) -def annotate(annotation_functions, rule, annotations, weights): +def annotate(annotation_functions, rule, annotations, qualified_nodes, qualified_edges, clause_labels, clause_variables, weights): + """Resolve and invoke the rule's annotation function. + + If the rule has no annotation function attached, returns the rule's static + bound directly. Otherwise looks up the user-registered function by name + and dispatches based on arity (arity is validated at registration; see + ``add_annotation_function``). + + Two supported signatures: + + - 2-arg legacy:: + + def fn(annotations, weights) -> Tuple[float, float] + + - 6-arg extended:: + + def fn(annotations, weights, + qualified_nodes, qualified_edges, + clause_labels, clause_variables) -> Tuple[float, float] + + The extended args expose the per-clause join structure the engine + already builds, so user code can identify clauses by predicate name + + variable role (robust to ``reorder_clauses``) and walk the actual + per-grounding pairings rather than the flattened per-atom bounds. + + Per-clause alignment (extended signature) + ----------------------------------------- + The five list-shaped args are indexed by the same per-clause position ``i``:: + + annotations[i] # List[Interval] — bounds of qualifying atoms + qualified_nodes[i] # List[Node] — nodes that qualified + qualified_edges[i] # List[Edge] — edges that qualified + clause_labels[i] # Label — predicate of this clause + clause_variables[i] # List[str] — variable names of this clause + + Per-clause arity: + + - Node clause: ``qualified_edges[i]`` is empty; ``len(clause_variables[i]) == 1``. + - Edge clause: ``qualified_nodes[i]`` is empty; ``len(clause_variables[i]) == 2``. + + Comparison clauses are skipped during metadata collection, so + ``len(clause_labels)`` may be less than ``len(rule.get_clauses())``. + Do not assume index parity with the original rule body — match clauses + by predicate name and variable role instead. + """ func_name = rule.get_annotation_function() if func_name == '': return rule.get_bnd().lower, rule.get_bnd().upper @@ -1837,7 +1918,16 @@ def annotate(annotation_functions, rule, annotations, weights): with numba.objmode(annotation='Tuple((float64, float64))'): for func in annotation_functions: if func.__name__ == func_name: - annotation = func(annotations, weights) + # Arity is gated at registration time by + # `add_annotation_function`, so only 2 and 6 reach here. + # (Raising inside numba.objmode is unsupported, so we + # rely on the registration-time check for validation.) + py_func = getattr(func, 'py_func', func) + nargs = py_func.__code__.co_argcount + if nargs == 6: + annotation = func(annotations, weights, qualified_nodes, qualified_edges, clause_labels, clause_variables) + else: + annotation = func(annotations, weights) return annotation diff --git a/setup.py b/setup.py index 0269d25f..54f08ef2 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ setup( name='pyreason', - version='3.5.1', + version='3.6.0', author='Dyuman Aditya', author_email='dyuman.aditya@gmail.com', description='An explainable inference software supporting annotated, real valued, graph based and temporal logic', diff --git a/tests/functional/test_advanced_features.py b/tests/functional/test_advanced_features.py index 559fe1da..c67736ba 100644 --- a/tests/functional/test_advanced_features.py +++ b/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 diff --git a/tests/unit/disable_jit/interpretations/test_ground_rule_helpers.py b/tests/unit/disable_jit/interpretations/test_ground_rule_helpers.py index a24013fa..cb538966 100644 --- a/tests/unit/disable_jit/interpretations/test_ground_rule_helpers.py +++ b/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"] diff --git a/tests/unit/disable_jit/interpretations/test_interpretation_common.py b/tests/unit/disable_jit/interpretations/test_interpretation_common.py index 2c0723ae..ca1f9461 100644 --- a/tests/unit/disable_jit/interpretations/test_interpretation_common.py +++ b/tests/unit/disable_jit/interpretations/test_interpretation_common.py @@ -73,11 +73,13 @@ def add_edge(*args): def ground_rule(*args, **kwargs): kwargs.setdefault('head_functions', ()) kwargs.setdefault('closed_world_predicates', []) + kwargs.setdefault('extended_ann_fn', False) return _ground_rule_fn(*args, num_ga=[0], **kwargs) else: def ground_rule(*args, **kwargs): kwargs.setdefault('head_functions', ()) kwargs.setdefault('closed_world_predicates', []) + kwargs.setdefault('extended_ann_fn', False) return _ground_rule_fn(*args, **kwargs) ns.ground_rule = ground_rule ns.update_rule_trace = _py(interpretation._update_rule_trace) @@ -152,51 +154,58 @@ def init_interpretations_edge(edges, specific_labels): ) _reason_fn = _py(interpretation.Interpretation.reason) - if "num_ga" in inspect.signature(_reason_fn).parameters: - def reason( - interpretations_node, - interpretations_edge, - predicate_map_node, - predicate_map_edge, - tmax, - prev_reasoning_data, - rules, - nodes, - edges, - neighbors, - reverse_neighbors, - rules_to_be_applied_node, - rules_to_be_applied_edge, - edges_to_be_added_node_rule, - edges_to_be_added_edge_rule, - rules_to_be_applied_node_trace, - rules_to_be_applied_edge_trace, - facts_to_be_applied_node, - facts_to_be_applied_edge, - facts_to_be_applied_node_trace, - facts_to_be_applied_edge_trace, - ipl, - rule_trace_node, - rule_trace_edge, - rule_trace_node_atoms, - rule_trace_edge_atoms, - reverse_graph, - atom_trace, - save_graph_attributes_to_rule_trace, - persistent, - inconsistency_check, - store_interpretation_changes, - update_mode, - allow_ground_rules, - max_facts_time, - annotation_functions, - head_functions, - convergence_mode, - convergence_delta, - verbose, - again, - closed_world_predicates, - ): + _reason_params = inspect.signature(_reason_fn).parameters + _has_num_ga = "num_ga" in _reason_params + + # Both branches wrap so we can inject `extended_ann_fn_flags` (one bool per + # rule, default False since the disable_jit tests don't use 6-arg ann_fns). + # The num_ga branch additionally injects `[0]` for num_ga. + def reason( + interpretations_node, + interpretations_edge, + predicate_map_node, + predicate_map_edge, + tmax, + prev_reasoning_data, + rules, + nodes, + edges, + neighbors, + reverse_neighbors, + rules_to_be_applied_node, + rules_to_be_applied_edge, + edges_to_be_added_node_rule, + edges_to_be_added_edge_rule, + rules_to_be_applied_node_trace, + rules_to_be_applied_edge_trace, + facts_to_be_applied_node, + facts_to_be_applied_edge, + facts_to_be_applied_node_trace, + facts_to_be_applied_edge_trace, + ipl, + rule_trace_node, + rule_trace_edge, + rule_trace_node_atoms, + rule_trace_edge_atoms, + reverse_graph, + atom_trace, + save_graph_attributes_to_rule_trace, + persistent, + inconsistency_check, + store_interpretation_changes, + update_mode, + allow_ground_rules, + max_facts_time, + annotation_functions, + head_functions, + convergence_mode, + convergence_delta, + verbose, + again, + closed_world_predicates, + ): + extended_ann_fn_flags = [False] * len(rules) + if _has_num_ga: return _reason_fn( interpretations_node[0], interpretations_edge[0], @@ -234,6 +243,7 @@ def reason( allow_ground_rules, max_facts_time, annotation_functions, + extended_ann_fn_flags, head_functions, convergence_mode, convergence_delta, @@ -242,8 +252,52 @@ def reason( again, closed_world_predicates, ) - else: - reason = _reason_fn + else: + return _reason_fn( + interpretations_node, + interpretations_edge, + predicate_map_node, + predicate_map_edge, + tmax, + prev_reasoning_data, + rules, + nodes, + edges, + neighbors, + reverse_neighbors, + rules_to_be_applied_node, + rules_to_be_applied_edge, + edges_to_be_added_node_rule, + edges_to_be_added_edge_rule, + rules_to_be_applied_node_trace, + rules_to_be_applied_edge_trace, + facts_to_be_applied_node, + facts_to_be_applied_edge, + facts_to_be_applied_node_trace, + facts_to_be_applied_edge_trace, + ipl, + rule_trace_node, + rule_trace_edge, + rule_trace_node_atoms, + rule_trace_edge_atoms, + reverse_graph, + atom_trace, + save_graph_attributes_to_rule_trace, + persistent, + inconsistency_check, + store_interpretation_changes, + update_mode, + allow_ground_rules, + max_facts_time, + annotation_functions, + extended_ann_fn_flags, + head_functions, + convergence_mode, + convergence_delta, + verbose, + again, + closed_world_predicates, + ) ns.reason = reason class FakeLabel: @@ -604,7 +658,7 @@ def get_bnd(self): def test_annotate_returns_bounds_when_no_function(): bnd = _Interval(0.2, 0.3) rule = AnnRule("", bnd) - lo, up = annotate([], rule, [], []) + lo, up = annotate([], rule, [], [], [], [], [], []) assert (lo, up) == (0.2, 0.3) @@ -615,7 +669,7 @@ def test_annotate_calls_named_function(): def foo(ann, wts): return (len(ann), len(wts)) - out = annotate([foo], rule, [1, 2], [3]) + out = annotate([foo], rule, [1, 2], [], [], [], [], [3]) assert out == (2, 1) diff --git a/tests/unit/disable_jit/interpretations/test_old_interp_file.py b/tests/unit/disable_jit/interpretations/test_old_interp_file.py index ea5cd4af..29cf5024 100644 --- a/tests/unit/disable_jit/interpretations/test_old_interp_file.py +++ b/tests/unit/disable_jit/interpretations/test_old_interp_file.py @@ -241,7 +241,7 @@ def get_weights(self): rule = Rule() reason_env["rules"] = [rule] - applicable_rule = (reason_env["node"], [], [], [], None) + applicable_rule = (reason_env["node"], [], [], [], None, [], []) mock_ground = Mock(side_effect=[([applicable_rule], []), ([], [])]) monkeypatch.setattr(interpretation, "_ground_rule", mock_ground) monkeypatch.setattr(interpretation, "annotate", Mock(return_value=(0.0, 1.0))) @@ -291,7 +291,7 @@ def get_weights(self): return () reason_env["rules"] = [Rule()] - applicable_rule = (reason_env["node"], [], [], [], None) + applicable_rule = (reason_env["node"], [], [], [], None, [], []) monkeypatch.setattr(interpretation, "_ground_rule", Mock(return_value=([applicable_rule], []))) mock_annotate = Mock(return_value=(0.0, 1.0)) monkeypatch.setattr(interpretation, "annotate", mock_annotate) @@ -379,7 +379,7 @@ def get_weights(self): reason_env["rules"] = [rule1, rule2] edges_to_add = ([node], [other], edge_lbl) - applicable_edge_rule = (edge, [], [], [], edges_to_add) + applicable_edge_rule = (edge, [], [], [], edges_to_add, [], []) mock_ground = Mock( side_effect=[ ([], [applicable_edge_rule]), diff --git a/tests/unit/disable_jit/interpretations/test_reason_core.py b/tests/unit/disable_jit/interpretations/test_reason_core.py index 9cc80c59..9e2b980c 100644 --- a/tests/unit/disable_jit/interpretations/test_reason_core.py +++ b/tests/unit/disable_jit/interpretations/test_reason_core.py @@ -944,7 +944,7 @@ def test_reason_node_rule_delta_zero_traces_and_applies(monkeypatch, reason_env) def ground_rule_stub(*args, **kwargs): ground_calls["n"] += 1 if ground_calls["n"] == 1: - return ([(new_node, [], ["qn"], ["qe"], None)], []) + return ([(new_node, [], ["qn"], ["qe"], None, [], [])], []) return ([], []) monkeypatch.setattr(interpretation, "_ground_rule", ground_rule_stub) @@ -1029,7 +1029,7 @@ def test_reason_node_rule_skips_when_static(monkeypatch, reason_env): rule.get_annotation_function.return_value = "" rule.get_name.return_value = "r" - mock_ground = Mock(return_value=([(node, [], [], [], None)], [])) + mock_ground = Mock(return_value=([(node, [], [], [], None, [], [])], [])) monkeypatch.setattr(interpretation, "_ground_rule", mock_ground) rules_list = [] @@ -1064,7 +1064,7 @@ def test_reason_node_rule_resolves_inconsistency(monkeypatch, reason_env): rule.get_name.return_value = "r" monkeypatch.setattr( - interpretation, "_ground_rule", lambda *a, **k: ([(reason_env["node"], [], [], [], None)], []) + interpretation, "_ground_rule", lambda *a, **k: ([(reason_env["node"], [], [], [], None, [], [])], []) ) monkeypatch.setattr(interpretation, "check_consistent_node", lambda *a, **k: False) @@ -1163,7 +1163,7 @@ def test_reason_edge_rule_records_trace(monkeypatch, reason_env): rule.get_bnd.return_value = Mock(lower=0, upper=1) edge_placeholder = type("EL", (), {"value": ""})() - applicable = [(edge, [], [node], [edge], ([], [], edge_placeholder))] + applicable = [(edge, [], [node], [edge], ([], [], edge_placeholder), [], [])] monkeypatch.setattr(interpretation, "_ground_rule", lambda *a, **k: ([], applicable)) def add_edge_to_interp(comp, interp_edge): @@ -1217,7 +1217,7 @@ def test_reason_edge_rule_delta_zero_applies(monkeypatch, reason_env): rule.get_bnd.return_value = Mock(lower=0, upper=1) edge_placeholder = type("EL", (), {"value": ""})() - applicable = [(edge, [], [], [], ([], [], edge_placeholder))] + applicable = [(edge, [], [], [], ([], [], edge_placeholder), [], [])] monkeypatch.setattr(interpretation, "_ground_rule", lambda *a, **k: ([], applicable)) def add_edge_to_interp(comp, interp_edge):