Graph Visualisation for TypeDB Data

# Convert all constraints to edges
from typing import Tuple
def to_query_edge(constraint: Constraint) -> Tuple[ConstraintVertex, str, ConstraintVertex]:
    if constraint.is_isa():
        isa = constraint.as_isa()
        return (isa.instance(), "isa", isa.type())
    elif constraint.is_has():
        has = constraint.as_has()
        return (has.owner(), "has", has.attribute())
    else:
        raise NotImplementedError("Not implemented in tutorial.")

query_edges = [to_query_edge(constraint) for constraint in constraints]
query_edges

import networkx
from matplotlib import pyplot

query_graph = networkx.MultiDiGraph()
for (u, label, v) in query_edges:
    query_graph.add_edge(u,v)
node_labels = {node: str(node) for node in query_graph.nodes()}
networkx.draw(query_graph, labels=node_labels)
pyplot.show()

from typing import Dict, List, Tuple
def node_style(pipeline: Pipeline, node: ConstraintVertex) -> Dict[str, any]:
    color = "g" if node.is_variable() else "b"
    label = ("$" + pipeline.get_variable_name(node.as_variable())) if node.is_variable() else str(node)
    shape = "s" if node.is_label() else "o"
    return {
        "color": color,
        "label": label,
        "shape": shape,
    }

def draw(edges: List[Tuple[ConstraintVertex, str, ConstraintVertex]], node_styles: Dict[ConstraintVertex, Dict[str, any]]):
    graph = networkx.MultiDiGraph()
    graph.add_edges_from((u,v, label) for (u, label, v) in edges)
    pos = networkx.forceatlas2_layout(graph) if hasattr(networkx, 'forceatlas2_layout') else networkx.planar_layout(graph)

    nodes_by_shape = {node_styles[n]["shape"]: [] for n in graph.nodes}
    for node in node_styles:
        nodes_by_shape[node_styles[node]["shape"]].append(node)

    for (shape, node_subset) in nodes_by_shape.items():
        node_colors = [node_styles[n]["color"] for n in node_subset]
        node_labels = {n: node_styles[n]["label"] for n in node_subset}
        networkx.draw_networkx_nodes(graph, pos, nodelist=nodes_by_shape[shape], node_color=node_colors, node_shape=shape)
        networkx.draw_networkx_labels(graph, pos, labels = node_labels)
    networkx.draw_networkx_edges(graph, pos)

    edge_labels = { (u,v): label for (u, label, v) in edges}
    networkx.draw_networkx_edge_labels(graph, pos, edge_labels=edge_labels)
    pyplot.show()

# Prepare node styles
nodes = set(u for (u,_,_) in query_edges).union(set(v for (_,_,v) in query_edges))
node_styles = { n: node_style(pipeline, n) for n in nodes}
# Draw
draw(query_edges, node_styles)

from typedb.driver import ConceptRow, Concept
def substitute(pipeline: Pipeline, vertex: ConstraintVertex, row: ConceptRow) -> Concept:
    if vertex.is_label():
        return vertex.as_label()
    elif vertex.is_variable():
        var_name = pipeline.get_variable_name(vertex.as_variable())
        return row.get(var_name) if var_name else None
    else:
        raise NotImplementedError("Not implemented in tutorial. See resolve_constraint_vertex")

answers_as_data_edges = [
    [(substitute(pipeline, u, row), label, substitute(pipeline, v, row)) for (u,label,v) in query_edges]
    for row in answers
]
answers_as_data_edges

# We need to update our node_style
def data_node_style(node: Concept) -> Dict[str, any]:
    if node.is_type():
        label = str(node)
        color = "c"
    elif node.is_attribute():
        label = f"{node.get_type().get_label()}:{node.get_value()}"
        color = "g"
    else:
        label = f"{node.get_type().get_label()}#{node.get_iid()[-4:]}"
        color = "b"

    shape = "s" if (node.is_attribute() or node.is_attribute_type()) else "o"
    return {
        "color": color,
        "label": label,
        "shape": shape,
    }

# Flatten them and remove any duplicate edges:
data_edges = set(e for answer_edges in answers_as_data_edges for e in answer_edges)

# Prepare node styles
nodes = set(u for (u,_,_) in data_edges).union(set(v for (_,_,v) in data_edges))
node_styles = { n: data_node_style(n) for n in nodes}
# Draw
draw(data_edges, node_styles)

BRANCHED_QUERY = """
match
$x isa person;
{ $x has name $y; } or { $x has age $y; };
"""

with driver.transaction(DB_NAME, TransactionType.READ) as tx:
    answers = list(tx.query(BRANCHED_QUERY, QUERY_OPTIONS).resolve())
assert 4 == len(answers), "TypeDB answer count mismatch"
answers

pipeline = answers[0].query_structure()
stages = list(pipeline.stages())
trunk_id = stages[0].as_match().block()
trunk = list(pipeline.conjunction(trunk_id).constraints())
(trunk_id, trunk)

branches = { branch_id: list(pipeline.conjunction(branch_id).constraints()) for branch_id in trunk[1].as_or().branches() }
branches

[list(answer.involved_conjunctions()) for answer in answers]

def flatten(list_of_lists: List[List[any]]) -> List[any]:
    return [x for l in list_of_lists for x in l]

answers_as_data_edges = []
for row in answers:
    involved_constraints = flatten([list(pipeline.conjunction(conj_id).constraints()) for conj_id in row.involved_conjunctions()])
    as_query_edges = [to_query_edge(c) for c in involved_constraints if not c.is_or()] # Exclude or constraint
    as_data_edges = [(substitute(pipeline, u, row), label, substitute(pipeline, v, row)) for (u,label,v) in as_query_edges]
    answers_as_data_edges.append(as_data_edges)

answers_as_data_edges

# Prepare node styles
flattened_data_edges = flatten(answers_as_data_edges)
nodes = set(u for (u,_,_) in data_edges).union(set(v for (_,_,v) in flattened_data_edges))
node_styles = { n: data_node_style(n) for n in nodes }
# Draw
draw(flattened_data_edges, node_styles)

# From library
class Isa(DataConstraint):
    def instance(self) -> DataVertex
    def type(self) -> DataVertex
    def exactness(self) -> ConstraintExactness

# From driver
class Isa(Constraint, ABC):
    def instance(self) -> ConstraintVertex
    def type(self) -> ConstraintVertex
    def exactness(self) -> ConstraintExactness # isa or isa!

from typedb_graph_utils import data_constraint, TypeDBAnswerConverter, MatplotlibVisualizer
from typedb.common.enums import ConstraintExactness
from networkx import MultiDiGraph
class MyTutorialBuilder(TypeDBAnswerConverter[MultiDiGraph]):

    def __init__(self, pipeline: Pipeline):
        super().__init__(pipeline)
        self.graph = MultiDiGraph()

    def finish(self) -> MultiDiGraph:
        return self.graph

    def add_isa(self, isa: data_constraint.Isa):
        edge_type = "isa!" if isa.exactness() == ConstraintExactness.Exact else "isa"
         # Use the edge attributes to store metadata. The visualiser uses it.
        self.graph.add_edge(isa.instance(), isa.type(), label=edge_type)

    def add_has(self, has: data_constraint.Has):
        if has.owner() is None or has.attribute() is None:
            return
        edge_type = "has!" if has.exactness() == ConstraintExactness.Exact else "has"
        self.graph.add_edge(has.owner(), has.attribute(), label = edge_type)

    # We leave the remaining unimplemented for brevity. Check the NetworkXConverter
    def add_comparison(c): pass
    def add_expression(c): pass
    def  add_function_call(c): pass
    def  add_iid(c): pass
    def  add_is(c): pass
    def  add_kind(c): pass
    def  add_label(c): pass
    def  add_links(c): pass
    def  add_owns(c): pass
    def  add_plays(c): pass
    def  add_relates(c): pass
    def  add_sub(c): pass
    def  add_value(c): pass

with driver.transaction(DB_NAME, TransactionType.READ) as tx:
    answers = list(tx.query(BRANCHED_QUERY, QUERY_OPTIONS).resolve())

builder = MyTutorialBuilder(answers[0].query_structure())
for (i, answer) in enumerate(answers):
    builder.add_answer(i, answer)

graph = builder.finish()
MatplotlibVisualizer.draw(graph)