relationalai 0.11.4__py3-none-any.whl → 0.12.0__py3-none-any.whl

relationalai/semantics/reasoners/graph/core.py

@@ -20,6 +20,7 @@ from relationalai.semantics import (
     count, sum, avg,
 )
 from relationalai.docutils import include_in_docs
+from relationalai.semantics.internal import annotations
 from relationalai.semantics.std.math import abs, isnan, isinf, maximum, natural_log, sqrt
 
 Numeric = Union[int, float, Decimal]
@@ -1250,6 +1251,8 @@ class Graph():
     def _num_nodes(self):
         """Lazily define and cache the self._num_nodes relationship."""
         _num_nodes_rel = self._model.Relationship("The graph has {num_nodes:Integer} nodes")
+        _num_nodes_rel.annotate(annotations.track("graphs", "num_nodes"))
+
         define(_num_nodes_rel(count(self.Node) | 0))
         return _num_nodes_rel
 
@@ -1316,6 +1319,7 @@ class Graph():
     def _num_edges(self):
         """Lazily define and cache the self._num_edges relationship."""
         _num_edges_rel = self._model.Relationship("The graph has {num_edges:Integer} edges")
+        _num_edges_rel.annotate(annotations.track("graphs", "num_edges"))
 
         src, dst = self.Node.ref(), self.Node.ref()
         if self.directed:
@@ -1468,7 +1472,9 @@ class Graph():
     @cached_property
     def _neighbor(self):
         """Lazily define and cache the self._neighbor relationship."""
-        return self._create_neighbor_relationship(nodes_subset=None)
+        _neighbor_rel = self._create_neighbor_relationship(nodes_subset=None)
+        _neighbor_rel.annotate(annotations.track("graphs", "neighbor"))
+        return _neighbor_rel
 
     def _neighbor_of(self, nodes_subset: Relationship):
         """
@@ -1476,7 +1482,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_neighbor_relationship(nodes_subset=nodes_subset)
+        _neighbor_rel = self._create_neighbor_relationship(nodes_subset=nodes_subset)
+        _neighbor_rel.annotate(annotations.track("graphs", "neighbor_of"))
+        return _neighbor_rel
 
     def _create_neighbor_relationship(self, *, nodes_subset: Optional[Relationship]):
         _neighbor_rel = self._model.Relationship(f"{{src:{self._NodeConceptStr}}} has neighbor {{dst:{self._NodeConceptStr}}}")
@@ -1626,7 +1634,9 @@ class Graph():
     @cached_property
     def _inneighbor(self):
         """Lazily define and cache the self._inneighbor relationship."""
-        return self._create_inneighbor_relationship(nodes_subset=None)
+        _inneighbor_rel = self._create_inneighbor_relationship(nodes_subset=None)
+        _inneighbor_rel.annotate(annotations.track("graphs", "inneighbor"))
+        return _inneighbor_rel
 
     def _inneighbor_of(self, nodes_subset: Relationship):
         """
@@ -1634,7 +1644,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_inneighbor_relationship(nodes_subset=nodes_subset)
+        _inneighbor_rel = self._create_inneighbor_relationship(nodes_subset=nodes_subset)
+        _inneighbor_rel.annotate(annotations.track("graphs", "inneighbor_of"))
+        return _inneighbor_rel
 
     def _create_inneighbor_relationship(self, *, nodes_subset: Optional[Relationship]):
         _inneighbor_rel = self._model.Relationship(f"{{dst:{self._NodeConceptStr}}} has inneighbor {{src:{self._NodeConceptStr}}}")
@@ -1777,7 +1789,9 @@ class Graph():
     @cached_property
     def _outneighbor(self):
         """Lazily define and cache the self._outneighbor relationship."""
-        return self._create_outneighbor_relationship(nodes_subset=None)
+        _outneighbor_rel = self._create_outneighbor_relationship(nodes_subset=None)
+        _outneighbor_rel.annotate(annotations.track("graphs", "outneighbor"))
+        return _outneighbor_rel
 
     def _outneighbor_of(self, nodes_subset: Relationship):
         """
@@ -1785,7 +1799,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_outneighbor_relationship(nodes_subset=nodes_subset)
+        _outneighbor_rel = self._create_outneighbor_relationship(nodes_subset=nodes_subset)
+        _outneighbor_rel.annotate(annotations.track("graphs", "outneighbor_of"))
+        return _outneighbor_rel
 
     def _create_outneighbor_relationship(self, *, nodes_subset: Optional[Relationship]):
         _outneighbor_rel = self._model.Relationship(f"{{src:{self._NodeConceptStr}}} has outneighbor {{dst:{self._NodeConceptStr}}}")
@@ -1914,6 +1930,8 @@ class Graph():
     def _common_neighbor(self):
         """Lazily define and cache the self._common_neighbor relationship."""
         _common_neighbor_rel = self._model.Relationship(f"{{node_a:{self._NodeConceptStr}}} and {{node_b:{self._NodeConceptStr}}} have common neighbor {{node_c:{self._NodeConceptStr}}}")
+        _common_neighbor_rel.annotate(annotations.track("graphs", "common_neighbor"))
+
         node_a, node_b, node_c = self.Node.ref(), self.Node.ref(), self.Node.ref()
         where(self._neighbor(node_a, node_c), self._neighbor(node_b, node_c)).define(_common_neighbor_rel(node_a, node_b, node_c))
         return _common_neighbor_rel
@@ -2072,7 +2090,9 @@ class Graph():
     @cached_property
     def _degree(self):
         """Lazily define and cache the self._degree relationship."""
-        return self._create_degree_relationship(nodes_subset=None)
+        _degree_rel = self._create_degree_relationship(nodes_subset=None)
+        _degree_rel.annotate(annotations.track("graphs", "degree"))
+        return _degree_rel
 
     def _degree_of(self, nodes_subset: Relationship):
         """
@@ -2080,7 +2100,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_degree_relationship(nodes_subset=nodes_subset)
+        _degree_rel = self._create_degree_relationship(nodes_subset=nodes_subset)
+        _degree_rel.annotate(annotations.track("graphs", "degree_of"))
+        return _degree_rel
 
     def _create_degree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _degree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has degree {{count:Integer}}")
@@ -2263,7 +2285,9 @@ class Graph():
     @cached_property
     def _indegree(self):
         """Lazily define and cache the self._indegree relationship."""
-        return self._create_indegree_relationship(nodes_subset=None)
+        _indegree_rel = self._create_indegree_relationship(nodes_subset=None)
+        _indegree_rel.annotate(annotations.track("graphs", "indegree"))
+        return _indegree_rel
 
     def _indegree_of(self, nodes_subset: Relationship):
         """
@@ -2271,7 +2295,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_indegree_relationship(nodes_subset=nodes_subset)
+        _indegree_rel = self._create_indegree_relationship(nodes_subset=nodes_subset)
+        _indegree_rel.annotate(annotations.track("graphs", "indegree_of"))
+        return _indegree_rel
 
     def _create_indegree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _indegree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has indegree {{count:Integer}}")
@@ -2443,7 +2469,9 @@ class Graph():
     @cached_property
     def _outdegree(self):
         """Lazily define and cache the self._outdegree relationship."""
-        return self._create_outdegree_relationship(nodes_subset=None)
+        _outdegree_rel = self._create_outdegree_relationship(nodes_subset=None)
+        _outdegree_rel.annotate(annotations.track("graphs", "outdegree"))
+        return _outdegree_rel
 
     def _outdegree_of(self, nodes_subset: Relationship):
         """
@@ -2451,7 +2479,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_outdegree_relationship(nodes_subset=nodes_subset)
+        _outdegree_rel = self._create_outdegree_relationship(nodes_subset=nodes_subset)
+        _outdegree_rel.annotate(annotations.track("graphs", "outdegree_of"))
+        return _outdegree_rel
 
     def _create_outdegree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _outdegree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has outdegree {{count:Integer}}")
@@ -2588,7 +2618,9 @@ class Graph():
     @cached_property
     def _weighted_degree(self):
         """Lazily define and cache the self._weighted_degree relationship."""
-        return self._create_weighted_degree_relationship(nodes_subset=None)
+        _weighted_degree_rel = self._create_weighted_degree_relationship(nodes_subset=None)
+        _weighted_degree_rel.annotate(annotations.track("graphs", "weighted_degree"))
+        return _weighted_degree_rel
 
     def _weighted_degree_of(self, nodes_subset: Relationship):
         """
@@ -2596,7 +2628,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_weighted_degree_relationship(nodes_subset=nodes_subset)
+        _weighted_degree_rel = self._create_weighted_degree_relationship(nodes_subset=nodes_subset)
+        _weighted_degree_rel.annotate(annotations.track("graphs", "weighted_degree_of"))
+        return _weighted_degree_rel
 
     def _create_weighted_degree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _weighted_degree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has weighted degree {{weight:Float}}")
@@ -2744,7 +2778,9 @@ class Graph():
     @cached_property
     def _weighted_indegree(self):
         """Lazily define and cache the self._weighted_indegree relationship."""
-        return self._create_weighted_indegree_relationship(nodes_subset=None)
+        _weighted_indegree_rel = self._create_weighted_indegree_relationship(nodes_subset=None)
+        _weighted_indegree_rel.annotate(annotations.track("graphs", "weighted_indegree"))
+        return _weighted_indegree_rel
 
     def _weighted_indegree_of(self, nodes_subset: Relationship):
         """
@@ -2752,7 +2788,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_weighted_indegree_relationship(nodes_subset=nodes_subset)
+        _weighted_indegree_rel = self._create_weighted_indegree_relationship(nodes_subset=nodes_subset)
+        _weighted_indegree_rel.annotate(annotations.track("graphs", "weighted_indegree_of"))
+        return _weighted_indegree_rel
 
     def _create_weighted_indegree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _weighted_indegree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has weighted indegree {{weight:Float}}")
@@ -2892,7 +2930,9 @@ class Graph():
     @cached_property
     def _weighted_outdegree(self):
         """Lazily define and cache the self._weighted_outdegree relationship."""
-        return self._create_weighted_outdegree_relationship(nodes_subset=None)
+        _weighted_outdegree_rel = self._create_weighted_outdegree_relationship(nodes_subset=None)
+        _weighted_outdegree_rel.annotate(annotations.track("graphs", "weighted_outdegree"))
+        return _weighted_outdegree_rel
 
     def _weighted_outdegree_of(self, nodes_subset: Relationship):
         """
@@ -2900,7 +2940,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_weighted_outdegree_relationship(nodes_subset=nodes_subset)
+        _weighted_outdegree_rel = self._create_weighted_outdegree_relationship(nodes_subset=nodes_subset)
+        _weighted_outdegree_rel.annotate(annotations.track("graphs", "weighted_outdegree_of"))
+        return _weighted_outdegree_rel
 
     def _create_weighted_outdegree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _weighted_outdegree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has weighted outdegree {{weight:Float}}")
@@ -3067,7 +3109,9 @@ class Graph():
     @cached_property
     def _degree_centrality(self):
         """Lazily define and cache the self._degree_centrality relationship."""
-        return self._create_degree_centrality_relationship(nodes_subset=None)
+        _degree_centrality_rel = self._create_degree_centrality_relationship(nodes_subset=None)
+        _degree_centrality_rel.annotate(annotations.track("graphs", "degree_centrality"))
+        return _degree_centrality_rel
 
     def _degree_centrality_of(self, nodes_subset: Relationship):
         """
@@ -3075,7 +3119,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_degree_centrality_relationship(nodes_subset=nodes_subset)
+        _degree_centrality_rel = self._create_degree_centrality_relationship(nodes_subset=nodes_subset)
+        _degree_centrality_rel.annotate(annotations.track("graphs", "degree_centrality_of"))
+        return _degree_centrality_rel
 
     def _create_degree_centrality_relationship(self, *, nodes_subset: Optional[Relationship]):
         """Create a degree centrality relationship, optionally constrained to a subset of nodes."""
@@ -3572,6 +3618,7 @@ class Graph():
     def _triangle(self):
         """Lazily define and cache the self._triangle relationship."""
         _triangle_rel = self._model.Relationship(f"{{node_a:{self._NodeConceptStr}}} and {{node_b:{self._NodeConceptStr}}} and {{node_c:{self._NodeConceptStr}}} form a triangle")
+        _triangle_rel.annotate(annotations.track("graphs", "triangle"))
 
         a, b, c = self.Node.ref(), self.Node.ref(), self.Node.ref()
 
@@ -3714,6 +3761,7 @@ class Graph():
     def _unique_triangle(self):
         """Lazily define and cache the self._unique_triangle relationship."""
         _unique_triangle_rel = self._model.Relationship(f"{{node_a:{self._NodeConceptStr}}} and {{node_b:{self._NodeConceptStr}}} and {{node_c:{self._NodeConceptStr}}} form unique triangle")
+        _unique_triangle_rel.annotate(annotations.track("graphs", "unique_triangle"))
 
         node_a, node_b, node_c = self.Node.ref(), self.Node.ref(), self.Node.ref()
 
@@ -3849,6 +3897,7 @@ class Graph():
     def _num_triangles(self):
         """Lazily define and cache the self._num_triangles relationship."""
         _num_triangles_rel = self._model.Relationship("The graph has {num_triangles:Integer} triangles")
+        _num_triangles_rel.annotate(annotations.track("graphs", "num_triangles"))
 
         _num_triangles = Integer.ref()
         node_a, node_b, node_c = self.Node.ref(), self.Node.ref(), self.Node.ref()
@@ -3973,7 +4022,9 @@ class Graph():
     @cached_property
     def _triangle_count(self):
         """Lazily define and cache the self._triangle_count relationship."""
-        return self._create_triangle_count_relationship(nodes_subset=None)
+        _triangle_count_rel = self._create_triangle_count_relationship(nodes_subset=None)
+        _triangle_count_rel.annotate(annotations.track("graphs", "triangle_count"))
+        return _triangle_count_rel
 
     def _triangle_count_of(self, nodes_subset: Relationship):
         """
@@ -3981,7 +4032,9 @@ class Graph():
         in `nodes_subset`. Note this relationship is not cached; it is
         specific to the callsite.
         """
-        return self._create_triangle_count_relationship(nodes_subset=nodes_subset)
+        _triangle_count_rel = self._create_triangle_count_relationship(nodes_subset=nodes_subset)
+        _triangle_count_rel.annotate(annotations.track("graphs", "triangle_count_of"))
+        return _triangle_count_rel
 
     def _create_triangle_count_relationship(self, *, nodes_subset: Optional[Relationship]):
         """Create a triangle count relationship, optionally constrained to a subset of nodes."""
@@ -4113,7 +4166,7 @@ class Graph():
 
 
     @include_in_docs
-    def local_clustering_coefficient(self):
+    def local_clustering_coefficient(self, *, of: Optional[Relationship] = None):
         """Returns a binary relationship containing the local clustering coefficient of each node.
 
         The local clustering coefficient quantifies how close a node's neighbors
@@ -4122,6 +4175,14 @@ class Graph():
         directly connecting them, and 1.0 indicates all neighbors have edges
         directly connecting them.
 
+        Parameters
+        ----------
+        of : Relationship, optional
+            A unary relationship containing a subset of the graph's nodes. When
+            provided, constrains the domain of the local clustering coefficient
+            computation: only coefficients of nodes in this relationship are
+            computed and returned.
+
         Returns
         -------
         Relationship
@@ -4148,17 +4209,6 @@ class Graph():
         | Directed   | No        | Undirected only.   |
         | Weighted   | Yes       | Weights are ignored. |
 
-        Notes
-        -----
-        The formal definition of the local clustering coefficient (`C`) for a
-        node (`v`) can be given as::
-
-            C(v) = (2 * num_edges) / (degree(v) * (degree(v) - 1))
-
-        Here, `num_edges` represents the number of edges between the
-        neighbors of node `v`, and `degree(v)` represents the degree of the
-        node, i.e., the number of edges connected to the node.
-
         Examples
         --------
         >>> from relationalai.semantics import Model, define, select, Float
@@ -4194,6 +4244,41 @@ class Graph():
         3   4  0.333333
         4   5  0.000000
 
+        >>> # 4. Use 'of' parameter to constrain the set of nodes to compute local clustering coefficients of
+        >>> # Define a subset containing only nodes 1 and 3
+        >>> subset = model.Relationship(f"{{node:{Node}}} is in subset")
+        >>> node = Node.ref()
+        >>> where(union(node.id == 1, node.id == 3)).define(subset(node))
+        >>>
+        >>> # Get local clustering coefficients only of nodes in the subset
+        >>> constrained_lcc = graph.local_clustering_coefficient(of=subset)
+        >>> select(node.id, coeff).where(constrained_lcc(node, coeff)).inspect()
+        ▰▰▰▰ Setup complete
+           id     coeff
+        0   1  1.000000
+        1   3  0.666667
+
+        Notes
+        -----
+        The local clustering coefficient for node `v` is::
+
+            (2 * num_neighbor_edges(v)) / (ext_degree(v) * (ext_degree(v) - 1))
+
+        where `num_neighbor_edges(v)` is the number of edges between
+        the neighbors of node `v`, and `ext_degree(v)` is the degree of the
+        node excluding self-loops. If `ext_degree(v)` is less than 2,
+        the local clustering coefficient is 0.0.
+
+        The ``local_clustering_coefficient()`` method, called with no parameters, computes
+        and caches the full local clustering coefficient relationship, providing efficient
+        reuse across multiple calls to ``local_clustering_coefficient()``. In contrast,
+        ``local_clustering_coefficient(of=subset)`` computes a constrained relationship
+        specific to the passed-in ``subset`` and that call site. When a significant fraction
+        of the local clustering coefficient relation is needed across a program,
+        ``local_clustering_coefficient()`` is typically more efficient; this is the typical
+        case. Use ``local_clustering_coefficient(of=subset)`` only when small subsets of the
+        local clustering coefficient relationship are needed collectively across the program.
+
 
         See Also
         --------
@@ -4206,29 +4291,51 @@ class Graph():
             raise NotImplementedError(
                 "`local_clustering_coefficient` is not applicable to directed graphs"
             )
+
+        if of is not None:
+            self._validate_node_subset_parameter(of)
+            return self._local_clustering_coefficient_of(of)
         return self._local_clustering_coefficient
 
     @cached_property
     def _local_clustering_coefficient(self):
+        """Lazily define and cache the self._local_clustering_coefficient relationship."""
+        _local_clustering_coefficient_rel = self._create_local_clustering_coefficient_relationship(nodes_subset=None)
+        _local_clustering_coefficient_rel.annotate(annotations.track("graphs", "local_clustering_coefficient"))
+        return _local_clustering_coefficient_rel
+
+    def _local_clustering_coefficient_of(self, nodes_subset: Relationship):
         """
-        Lazily define and cache the self._local_clustering_coefficient relationship,
-        which only applies to undirected graphs.
+        Create a local clustering coefficient relationship constrained to the subset of nodes
+        in `nodes_subset`. Note this relationship is not cached; it is
+        specific to the callsite.
         """
-        _local_clustering_coefficient_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has local clustering coefficient {{coefficient:Float}}")
+        _local_clustering_coefficient_rel = self._create_local_clustering_coefficient_relationship(nodes_subset=nodes_subset)
+        _local_clustering_coefficient_rel.annotate(annotations.track("graphs", "local_clustering_coefficient_of"))
+        return _local_clustering_coefficient_rel
 
-        if self.directed:
-            raise NotImplementedError(
-                "`local_clustering_coefficient is not defined for directed graphs."
-            )
+    def _create_local_clustering_coefficient_relationship(self, *, nodes_subset: Optional[Relationship]):
+        """Create a local clustering coefficient relationship, optionally constrained to a subset of nodes."""
+        _local_clustering_coefficient_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has local clustering coefficient {{coefficient:Float}}")
 
         node = self.Node.ref()
+
+        if nodes_subset is None:
+            degree_no_self_rel = self._degree_no_self
+            triangle_count_rel = self._triangle_count
+            node_constraint = node  # No constraint on nodes.
+        else:
+            degree_no_self_rel = self._degree_no_self_of(nodes_subset)
+            triangle_count_rel = self._triangle_count_of(nodes_subset)
+            node_constraint = nodes_subset(node)  # Nodes constrained to given subset.
+
         degree_no_self = Integer.ref()
         triangle_count = Integer.ref()
         where(
-            node,
+            node_constraint,
             _lcc := where(
-                self._degree_no_self(node, degree_no_self),
-                self._triangle_count(node, triangle_count),
+                degree_no_self_rel(node, degree_no_self),
+                triangle_count_rel(node, triangle_count),
                 degree_no_self > 1
             ).select(
                 2.0 * triangle_count / (degree_no_self * (degree_no_self - 1.0))
@@ -4243,11 +4350,32 @@ class Graph():
         Lazily define and cache the self._degree_no_self relationship,
         a non-public helper for local_clustering_coefficient.
         """
+        return self._create_degree_no_self_relationship(nodes_subset=None)
+
+    def _degree_no_self_of(self, nodes_subset: Relationship):
+        """
+        Create a self-loop-exclusive degree relationship constrained to
+        the subset of nodes in `nodes_subset`. Note this relationship
+        is not cached; it is specific to the callsite.
+        """
+        return self._create_degree_no_self_relationship(nodes_subset=nodes_subset)
+
+    def _create_degree_no_self_relationship(self, *, nodes_subset: Optional[Relationship]):
+        """
+        Create a self-loop-exclusive degree relationship,
+        optionally constrained to a subset of nodes.
+        """
         _degree_no_self_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has degree excluding self loops {{num:Integer}}")
 
         node, neighbor = self.Node.ref(), self.Node.ref()
+
+        if nodes_subset is None:
+            node_constraint = node  # No constraint on nodes.
+        else:
+            node_constraint = nodes_subset(node)  # Nodes constrained to given subset.
+
         where(
-            self.Node(node),
+            node_constraint,
             _dns := count(neighbor).per(node).where(self._no_loop_edge(node, neighbor)) | 0,
         ).define(_degree_no_self_rel(node, _dns))
 
@@ -4331,6 +4459,7 @@ class Graph():
         which only applies to undirected graphs.
         """
         _average_clustering_coefficient_rel = self._model.Relationship("The graph has average clustering coefficient {{coefficient:Float}}")
+        _average_clustering_coefficient_rel.annotate(annotations.track("graphs", "average_clustering_coefficient"))
 
         if self.directed:
             raise NotImplementedError(
@@ -4471,6 +4600,7 @@ class Graph():
     def _reachable_from(self):
         """Lazily define and cache the self._reachable_from relationship."""
         _reachable_from_rel = self._model.Relationship(f"{{node_a:{self._NodeConceptStr}}} reaches {{node_b:{self._NodeConceptStr}}}")
+        _reachable_from_rel.annotate(annotations.track("graphs", "reachable_from"))
 
         node_a, node_b, node_c = self.Node.ref(), self.Node.ref(), self.Node.ref()
         define(_reachable_from_rel(node_a, node_a))
@@ -4613,9 +4743,12 @@ class Graph():
     def _distance(self):
         """Lazily define and cache the self._distance relationship."""
         if not self.weighted:
-            return self._distance_non_weighted
+            _distance_rel = self._distance_non_weighted
         else:
-            return self._distance_weighted
+            _distance_rel = self._distance_weighted
+
+        _distance_rel.annotate(annotations.track("graphs", "distance"))
+        return _distance_rel
 
     @cached_property
     def _distance_weighted(self):
@@ -4741,6 +4874,7 @@ class Graph():
     def _weakly_connected_component(self):
         """Lazily define and cache the self._weakly_connected_component relationship."""
         _weakly_connected_component_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} is in the connected component {{id:{self._NodeConceptStr}}}")
+        _weakly_connected_component_rel.annotate(annotations.track("graphs", "weakly_connected_component"))
 
         node, node_v, component = self.Node.ref(), self.Node.ref(), self.Node.ref()
         node, component = union(
@@ -4864,6 +4998,8 @@ class Graph():
         """
         _diameter_range_min_rel = self._model.Relationship("The graph has a min diameter range of {value:Integer}")
         _diameter_range_max_rel = self._model.Relationship("The graph has a max diameter range of {value:Integer}")
+        _diameter_range_min_rel.annotate(annotations.track("graphs", "diameter_range_min"))
+        _diameter_range_max_rel.annotate(annotations.track("graphs", "diameter_range_max"))
 
         component_node_pairs = self._model.Relationship(f"component id {{cid:{self._NodeConceptStr}}} has node id {{nid:{self._NodeConceptStr}}}")
         nodeid, cid, degreevalue = self.Node.ref(), self.Node.ref(), Integer.ref()
@@ -4924,16 +5060,22 @@ class Graph():
 
     @include_in_docs
     def is_connected(self):
-        """Returns a query fragment that is satisfied if the graph is connected.
+        """Returns a unary relationship containing whether the graph is connected.
 
         A graph is considered connected if every node is reachable from every
         other node in the underlying undirected graph.
 
         Returns
         -------
-        Fragment
-            A query fragment that can be used as a condition in other
-            queries to assert that the graph is connected.
+        Relationship
+            A unary relationship containing a boolean indicator of whether the graph
+            is connected.
+
+        Relationship Schema
+        -------------------
+        ``is_connected(connected)``
+
+        * **connected** (*Boolean*): Whether the graph is connected.
 
         Supported Graph Types
         ---------------------
@@ -4951,8 +5093,6 @@ class Graph():
         --------
         **Connected Graph Example**
 
-        The following query will produce a result because the graph is connected.
-
         >>> from relationalai.semantics import Model, define, select
         >>> from relationalai.semantics.reasoners.graph import Graph
         >>>
@@ -4970,17 +5110,14 @@ class Graph():
         ...     Edge.new(src=n4, dst=n3),
         ... )
         >>>
-        >>> # 3. Use the fragment as a condition in a query
-        >>> select("Graph is connected").where(graph.is_connected()).inspect()
+        >>> # 3. Select and inspect the relation
+        >>> select(graph.is_connected()).inspect()
         ▰▰▰▰ Setup complete
-                            v
-        0  Graph is connected
+           is_connected
+        0          True
 
         **Disconnected Graph Example**
 
-        The following query will produce no results because the graph is not
-        connected.
-
         >>> from relationalai.semantics import Model, define, select
         >>> from relationalai.semantics.reasoners.graph import Graph
         >>>
@@ -4998,22 +5135,31 @@ class Graph():
         ...     Edge.new(src=n4, dst=n5),  # This edge creates a separate component
         ... )
         >>>
-        >>> # 3. The conditional query produces no output
-        >>> select("Graph is connected").where(graph.is_connected()).inspect()
+        >>> # 3. Select and inspect the relation
+        >>> select(graph.is_connected()).inspect()
         ▰▰▰▰ Setup complete
-        Empty DataFrame
-        Columns: []
-        Index: []
+           is_connected
+        0         False
 
         """
-        # TODO (dba) This method is inconsistent with the other,
-        # public methods. It does not return a `Relationship`. Revisit
-        # this. See GH thread:
-        # https://github.com/RelationalAI/relationalai-python/pull/2077#discussion_r2190538074
-        return where(
+        return self._is_connected
+
+    @cached_property
+    def _is_connected(self):
+        """Lazily define and cache the self._is_connected relationship."""
+        _is_connected_rel = self._model.Relationship("'The graph is connected' is {is_connected:Boolean}")
+        _is_connected_rel.annotate(annotations.track("graphs", "is_connected"))
+
+        where(
             self._num_nodes(0) |
             count(self._reachable_from_min_node(self.Node.ref())) == self._num_nodes(Integer.ref())
-        )
+        ).define(_is_connected_rel(True))
+
+        where(
+            not_(_is_connected_rel(True))
+        ).define(_is_connected_rel(False))
+
+        return _is_connected_rel
 
 
     @include_in_docs
@@ -5179,6 +5325,7 @@ class Graph():
     def _jaccard_similarity(self):
         """Lazily define and cache the self._jaccard_similarity relationship."""
         _jaccard_similarity_rel = self._model.Relationship(f"{{node_u:{self._NodeConceptStr}}} has a similarity to {{node_v:{self._NodeConceptStr}}} of {{similarity:Float}}")
+        _jaccard_similarity_rel.annotate(annotations.track("graphs", "jaccard_similarity"))
 
         if not self.weighted:
             node_u, node_v = self.Node.ref(), self.Node.ref()
@@ -5430,6 +5577,7 @@ class Graph():
     def _cosine_similarity(self):
         """Lazily define and cache the self._cosine_similarity relationship."""
         _cosine_similarity_rel = self._model.Relationship(f"{{node_u:{self._NodeConceptStr}}} has a cosine similarity to {{node_v:{self._NodeConceptStr}}} of {{score:Float}}")
+        _cosine_similarity_rel.annotate(annotations.track("graphs", "cosine_similarity"))
 
         if not self.weighted:
             node_u, node_v = self.Node.ref(), self.Node.ref()
@@ -5550,6 +5698,7 @@ class Graph():
     def _adamic_adar(self):
         """Lazily define and cache the self._adamic_adar relationship."""
         _adamic_adar_rel = self._model.Relationship(f"{{node_u:{self._NodeConceptStr}}} and {{node_v:{self._NodeConceptStr}}} have adamic adar score {{score:Float}}")
+        _adamic_adar_rel.annotate(annotations.track("graphs", "adamic_adar"))
 
         node_u, node_v, common_neighbor = self.Node.ref(), self.Node.ref(), self.Node.ref()
         neighbor_count = Integer.ref()
@@ -5648,6 +5797,7 @@ class Graph():
     def _preferential_attachment(self):
         """Lazily define and cache the self._preferential_attachment relationship."""
         _preferential_attachment_rel = self._model.Relationship(f"{{node_u:{self._NodeConceptStr}}} and {{node_v:{self._NodeConceptStr}}} have preferential attachment score {{score:Integer}}")
+        _preferential_attachment_rel.annotate(annotations.track("graphs", "preferential_attachment"))
 
         node_u, node_v = self.Node.ref(), self.Node.ref()
         count_u, count_v = Integer.ref(), Integer.ref()

relationalai/semantics/rel/builtins.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 from relationalai.semantics.metamodel import types, factory as f
 from relationalai.semantics.metamodel.util import OrderedSet
+from relationalai.semantics.metamodel import builtins
 
 # Rel Annotations as IR Relations (to be used in IR Annotations)
 
@@ -28,7 +29,10 @@ inner_loop_non_stratified_annotation = f.annotation(inner_loop_non_stratified, [
 
 # collect all supported builtin rel annotations
 builtin_annotations = OrderedSet.from_iterable([
-    arrow, no_diagnostics, no_inline, function, inner_loop, inner_loop_non_stratified
+    arrow, no_diagnostics, no_inline, function, inner_loop, inner_loop_non_stratified,
+    # track annotations on relations do not currently propagate into Rel
+    # TODO: from Thiago, ensure annotation goes from the Logical into the proper declaration
+    builtins.track,
 ])
 
 builtin_annotation_names = OrderedSet.from_iterable([a.name for a in builtin_annotations])