Skip to content

graflo.plot

Plotting utilities for graph visualization.

This module provides tools for visualizing graph schemas and structures. It includes functionality for creating visual representations of graph databases, their vertices, edges, and relationships.

Key Components
  • SchemaPlotter: Creates visual representations of graph schemas
Example

plotter = SchemaPlotter(schema) plotter.plot("schema.png")

SchemaPlotter

Main class for schema visualization.

This class provides methods to visualize different aspects of a graph database schema, including vertex collections, resources, and their relationships.

Attributes:

Name Type Description
fig_path

Path to save visualizations

config

Schema configuration

schema

Schema instance

name

Schema name

prefix

Prefix for output files

Source code in graflo/plot/plotter.py
class SchemaPlotter:
    """Main class for schema visualization.

    This class provides methods to visualize different aspects of a graph database
    schema, including vertex collections, resources, and their relationships.

    Attributes:
        fig_path: Path to save visualizations
        config: Schema configuration
        schema: Schema instance
        name: Schema name
        prefix: Prefix for output files
    """

    def __init__(self, config_filename, fig_path):
        """Initialize the schema plotter.

        Args:
            config_filename: Path to schema configuration file
            fig_path: Path to save visualizations
        """
        self.fig_path = fig_path

        self.config = FileHandle.load(fpath=config_filename)

        self.schema = Schema.from_dict(self.config)

        self.name = self.schema.general.name
        self.prefix = self.name

    def _discover_edges_from_resources(self):
        """Discover edges from resources by walking through ActorWrappers.

        This method finds all EdgeActors in resources and extracts their edges,
        which may include edges with dynamic relations (relation_field, relation_from_key)
        that aren't fully represented in edge_config.

        Returns:
            dict: Dictionary mapping (source, target, purpose) to Edge objects
        """
        discovered_edges = {}

        for resource in self.schema.resources:
            # Collect all actors from the resource's ActorWrapper
            actors = resource.root.collect_actors()

            for actor in actors:
                if isinstance(actor, EdgeActor):
                    edge = actor.edge
                    edge_id = edge.edge_id
                    # Store the edge, preferring already discovered edges from edge_config
                    # but allowing resource edges to supplement
                    if edge_id not in discovered_edges:
                        discovered_edges[edge_id] = edge

        return discovered_edges

    def plot_vc2fields(self):
        """Plot vertex collections and their fields.

        Creates a visualization showing the relationship between vertex collections
        and their fields, including index fields. The visualization is saved as
        a PDF file.
        """
        g = nx.DiGraph()
        nodes = []
        edges = []
        vconf = self.schema.vertex_config
        vertex_prefix_dict = lto_dict([v for v in self.schema.vertex_config.vertex_set])

        kwargs = {"vfield": True, "vertex_sh": vertex_prefix_dict}
        for k in vconf.vertex_set:
            index_fields = vconf.index(k)
            fields = vconf.fields_names(k)
            kwargs["vertex"] = k
            nodes_collection = [
                (
                    get_auxnode_id(AuxNodeType.VERTEX, **kwargs),
                    {
                        "type": AuxNodeType.VERTEX,
                        "label": get_auxnode_id(
                            AuxNodeType.VERTEX, label=True, **kwargs
                        ),
                    },
                )
            ]
            nodes_fields = [
                (
                    get_auxnode_id(AuxNodeType.FIELD, field=item, **kwargs),
                    {
                        "type": (
                            AuxNodeType.FIELD_DEFINITION
                            if item in index_fields
                            else AuxNodeType.FIELD
                        ),
                        "label": get_auxnode_id(
                            AuxNodeType.FIELD, field=item, label=True, **kwargs
                        ),
                    },
                )
                for item in fields
            ]
            nodes += nodes_collection
            nodes += nodes_fields
            edges += [(x[0], y[0]) for x, y in product(nodes_collection, nodes_fields)]

        g.add_nodes_from(nodes)
        g.add_edges_from(edges)

        for n in g.nodes():
            props = g.nodes()[n]
            upd_dict = props.copy()
            if "type" in upd_dict:
                upd_dict["shape"] = map_type2shape[props["type"]]
                upd_dict["color"] = map_type2color[props["type"]]
            if "label" in upd_dict:
                upd_dict["forcelabel"] = True
            upd_dict["style"] = "filled"

            for k, v in upd_dict.items():
                g.nodes[n][k] = v

        for e in g.edges(data=True):
            s, t, _ = e
            upd_dict = {"style": "solid", "arrowhead": "vee"}
            for k, v in upd_dict.items():
                g.edges[s, t][k] = v

        ag = nx.nx_agraph.to_agraph(g)

        for k in vconf.vertex_set:
            level_index = [
                get_auxnode_id(
                    AuxNodeType.FIELD,
                    vertex=k,
                    field=item,
                    vfield=True,
                    vertex_sh=vertex_prefix_dict,
                )
                for item in vconf.index(k)
            ]
            index_subgraph = ag.add_subgraph(level_index, name=f"cluster_{k}:def")
            index_subgraph.node_attr["style"] = "filled"
            index_subgraph.node_attr["label"] = "definition"

        ag = ag.unflatten("-l 5 -f -c 3")
        ag.draw(
            os.path.join(self.fig_path, f"{self.prefix}_vc2fields.pdf"),
            "pdf",
            prog="dot",
        )

    def plot_resources(self):
        """Plot resource relationships.

        Creates visualizations for each resource in the schema, showing their
        internal structure and relationships. Each resource is saved as a
        separate PDF file.
        """
        resource_prefix_dict = lto_dict(
            [resource.name for resource in self.schema.resources]
        )
        vertex_prefix_dict = lto_dict([v for v in self.schema.vertex_config.vertex_set])
        kwargs = {"vertex_sh": vertex_prefix_dict, "resource_sh": resource_prefix_dict}

        for resource in self.schema.resources:
            kwargs["resource"] = resource.name
            assemble_tree(
                resource.root,
                os.path.join(
                    self.fig_path,
                    f"{self.schema.general.name}.resource-{resource.resource_name}.pdf",
                ),
            )

    def plot_source2vc(self):
        """Plot source to vertex collection mappings.

        Creates a visualization showing the relationship between source resources
        and vertex collections. The visualization is saved as a PDF file.
        """
        nodes = []
        g = nx.MultiDiGraph()
        edges = []
        resource_prefix_dict = lto_dict(
            [resource.name for resource in self.schema.resources]
        )
        vertex_prefix_dict = lto_dict([v for v in self.schema.vertex_config.vertex_set])
        kwargs = {"vertex_sh": vertex_prefix_dict, "resource_sh": resource_prefix_dict}

        for resource in self.schema.resources:
            kwargs["resource"] = resource.name

            g = assemble_tree(resource.root)

            vertices = []
            nodes_resource = [
                (
                    get_auxnode_id(AuxNodeType.RESOURCE, **kwargs),
                    {
                        "type": AuxNodeType.RESOURCE,
                        "label": get_auxnode_id(
                            AuxNodeType.RESOURCE, label=True, **kwargs
                        ),
                    },
                )
            ]
            nodes_vertex = [
                (
                    get_auxnode_id(AuxNodeType.VERTEX, vertex=v, **kwargs),
                    {
                        "type": AuxNodeType.VERTEX,
                        "label": get_auxnode_id(
                            AuxNodeType.VERTEX, vertex=v, label=True, **kwargs
                        ),
                    },
                )
                for v in vertices
            ]
            nodes += nodes_resource
            nodes += nodes_vertex
            edges += [
                (nt[0], nc[0]) for nt, nc in product(nodes_resource, nodes_vertex)
            ]

        g.add_nodes_from(nodes)

        g.add_edges_from(edges)

        for n in g.nodes():
            props = g.nodes()[n]
            upd_dict = {
                "shape": map_type2shape[props["type"]],
                "color": map_type2color[props["type"]],
                "style": "filled",
            }
            if "label" in props:
                upd_dict["forcelabel"] = True
            if "name" in props:
                upd_dict["label"] = props["name"]
            for resource, v in upd_dict.items():
                g.nodes[n][resource] = v

        ag = nx.nx_agraph.to_agraph(g)
        ag.draw(
            os.path.join(self.fig_path, f"{self.prefix}_source2vc.pdf"),
            "pdf",
            prog="dot",
        )

    def plot_vc2vc(self, prune_leaves=False):
        """Plot vertex collection relationships.

        Creates a visualization showing the relationships between vertex collections.
        Optionally prunes leaf nodes from the visualization.

        This method discovers edges from both edge_config and resources to ensure
        all relationships are visualized, including those with dynamic relations.

        Args:
            prune_leaves: Whether to remove leaf nodes from the visualization

        Example:
            >>> plotter.plot_vc2vc(prune_leaves=True)
        """
        g = nx.MultiDiGraph()
        nodes = []
        edges = []

        # Discover edges from resources (may include edges not in edge_config)
        discovered_edges = self._discover_edges_from_resources()

        # Collect all edges: from edge_config and discovered from resources
        all_edges = {}
        for edge_id, e in self.schema.edge_config.edges_items():
            all_edges[edge_id] = e
        # Add discovered edges (they may already be in edge_config, but that's fine)
        for edge_id, e in discovered_edges.items():
            if edge_id not in all_edges:
                all_edges[edge_id] = e

        # Create graph edges with relation labels
        for (source, target, purpose), e in all_edges.items():
            # Determine label based on relation configuration
            label = None
            if e.relation is not None:
                # Static relation
                label = e.relation
            elif e.relation_field is not None:
                # Dynamic relation from field - show indicator
                label = f"[{e.relation_field}]"
            elif e.relation_from_key:
                # Dynamic relation from key - show indicator
                label = "[key]"

            if label is not None:
                ee = (
                    get_auxnode_id(AuxNodeType.VERTEX, vertex=source),
                    get_auxnode_id(AuxNodeType.VERTEX, vertex=target),
                    {"label": label},
                )
            else:
                ee = (
                    get_auxnode_id(AuxNodeType.VERTEX, vertex=source),
                    get_auxnode_id(AuxNodeType.VERTEX, vertex=target),
                )
            edges += [ee]

        # Create nodes for all vertices involved in edges
        for (source, target, purpose), e in all_edges.items():
            for v in (source, target):
                nodes += [
                    (
                        get_auxnode_id(AuxNodeType.VERTEX, vertex=v),
                        {
                            "type": AuxNodeType.VERTEX,
                            "label": get_auxnode_id(
                                AuxNodeType.VERTEX, vertex=v, label=True
                            ),
                        },
                    )
                ]

        for nid, weight in nodes:
            g.add_node(nid, **weight)

        g.add_nodes_from(nodes)
        g.add_edges_from(edges)

        if prune_leaves:
            out_deg = g.out_degree()
            in_deg = g.in_degree()

            nodes_to_remove = set([k for k, v in out_deg if v == 0]) & set(
                [k for k, v in in_deg if v < 2]
            )
            g.remove_nodes_from(nodes_to_remove)

        for n in g.nodes():
            props = g.nodes()[n]
            upd_dict = {
                "shape": map_type2shape[props["type"]],
                "color": map_type2color[props["type"]],
                "style": "filled",
            }
            for k, v in upd_dict.items():
                g.nodes[n][k] = v

        for e in g.edges:
            s, t, ix = e
            target_props = g.nodes[s]
            edge_data = g.edges[s, t, ix]
            upd_dict = {
                "style": edge_status[target_props["type"]],
                "arrowhead": "vee",
            }
            # Preserve existing label if present (for relation display)
            if "label" in edge_data:
                upd_dict["label"] = edge_data["label"]
            for k, v in upd_dict.items():
                g.edges[s, t, ix][k] = v

        ag = nx.nx_agraph.to_agraph(g)
        ag.draw(
            os.path.join(self.fig_path, f"{self.prefix}_vc2vc.pdf"),
            "pdf",
            prog="dot",
        )

__init__(config_filename, fig_path)

Initialize the schema plotter.

Parameters:

Name Type Description Default
config_filename

Path to schema configuration file

required
fig_path

Path to save visualizations

required
Source code in graflo/plot/plotter.py
def __init__(self, config_filename, fig_path):
    """Initialize the schema plotter.

    Args:
        config_filename: Path to schema configuration file
        fig_path: Path to save visualizations
    """
    self.fig_path = fig_path

    self.config = FileHandle.load(fpath=config_filename)

    self.schema = Schema.from_dict(self.config)

    self.name = self.schema.general.name
    self.prefix = self.name

plot_resources()

Plot resource relationships.

Creates visualizations for each resource in the schema, showing their internal structure and relationships. Each resource is saved as a separate PDF file.

Source code in graflo/plot/plotter.py
def plot_resources(self):
    """Plot resource relationships.

    Creates visualizations for each resource in the schema, showing their
    internal structure and relationships. Each resource is saved as a
    separate PDF file.
    """
    resource_prefix_dict = lto_dict(
        [resource.name for resource in self.schema.resources]
    )
    vertex_prefix_dict = lto_dict([v for v in self.schema.vertex_config.vertex_set])
    kwargs = {"vertex_sh": vertex_prefix_dict, "resource_sh": resource_prefix_dict}

    for resource in self.schema.resources:
        kwargs["resource"] = resource.name
        assemble_tree(
            resource.root,
            os.path.join(
                self.fig_path,
                f"{self.schema.general.name}.resource-{resource.resource_name}.pdf",
            ),
        )

plot_source2vc()

Plot source to vertex collection mappings.

Creates a visualization showing the relationship between source resources and vertex collections. The visualization is saved as a PDF file.

Source code in graflo/plot/plotter.py
def plot_source2vc(self):
    """Plot source to vertex collection mappings.

    Creates a visualization showing the relationship between source resources
    and vertex collections. The visualization is saved as a PDF file.
    """
    nodes = []
    g = nx.MultiDiGraph()
    edges = []
    resource_prefix_dict = lto_dict(
        [resource.name for resource in self.schema.resources]
    )
    vertex_prefix_dict = lto_dict([v for v in self.schema.vertex_config.vertex_set])
    kwargs = {"vertex_sh": vertex_prefix_dict, "resource_sh": resource_prefix_dict}

    for resource in self.schema.resources:
        kwargs["resource"] = resource.name

        g = assemble_tree(resource.root)

        vertices = []
        nodes_resource = [
            (
                get_auxnode_id(AuxNodeType.RESOURCE, **kwargs),
                {
                    "type": AuxNodeType.RESOURCE,
                    "label": get_auxnode_id(
                        AuxNodeType.RESOURCE, label=True, **kwargs
                    ),
                },
            )
        ]
        nodes_vertex = [
            (
                get_auxnode_id(AuxNodeType.VERTEX, vertex=v, **kwargs),
                {
                    "type": AuxNodeType.VERTEX,
                    "label": get_auxnode_id(
                        AuxNodeType.VERTEX, vertex=v, label=True, **kwargs
                    ),
                },
            )
            for v in vertices
        ]
        nodes += nodes_resource
        nodes += nodes_vertex
        edges += [
            (nt[0], nc[0]) for nt, nc in product(nodes_resource, nodes_vertex)
        ]

    g.add_nodes_from(nodes)

    g.add_edges_from(edges)

    for n in g.nodes():
        props = g.nodes()[n]
        upd_dict = {
            "shape": map_type2shape[props["type"]],
            "color": map_type2color[props["type"]],
            "style": "filled",
        }
        if "label" in props:
            upd_dict["forcelabel"] = True
        if "name" in props:
            upd_dict["label"] = props["name"]
        for resource, v in upd_dict.items():
            g.nodes[n][resource] = v

    ag = nx.nx_agraph.to_agraph(g)
    ag.draw(
        os.path.join(self.fig_path, f"{self.prefix}_source2vc.pdf"),
        "pdf",
        prog="dot",
    )

plot_vc2fields()

Plot vertex collections and their fields.

Creates a visualization showing the relationship between vertex collections and their fields, including index fields. The visualization is saved as a PDF file.

Source code in graflo/plot/plotter.py
def plot_vc2fields(self):
    """Plot vertex collections and their fields.

    Creates a visualization showing the relationship between vertex collections
    and their fields, including index fields. The visualization is saved as
    a PDF file.
    """
    g = nx.DiGraph()
    nodes = []
    edges = []
    vconf = self.schema.vertex_config
    vertex_prefix_dict = lto_dict([v for v in self.schema.vertex_config.vertex_set])

    kwargs = {"vfield": True, "vertex_sh": vertex_prefix_dict}
    for k in vconf.vertex_set:
        index_fields = vconf.index(k)
        fields = vconf.fields_names(k)
        kwargs["vertex"] = k
        nodes_collection = [
            (
                get_auxnode_id(AuxNodeType.VERTEX, **kwargs),
                {
                    "type": AuxNodeType.VERTEX,
                    "label": get_auxnode_id(
                        AuxNodeType.VERTEX, label=True, **kwargs
                    ),
                },
            )
        ]
        nodes_fields = [
            (
                get_auxnode_id(AuxNodeType.FIELD, field=item, **kwargs),
                {
                    "type": (
                        AuxNodeType.FIELD_DEFINITION
                        if item in index_fields
                        else AuxNodeType.FIELD
                    ),
                    "label": get_auxnode_id(
                        AuxNodeType.FIELD, field=item, label=True, **kwargs
                    ),
                },
            )
            for item in fields
        ]
        nodes += nodes_collection
        nodes += nodes_fields
        edges += [(x[0], y[0]) for x, y in product(nodes_collection, nodes_fields)]

    g.add_nodes_from(nodes)
    g.add_edges_from(edges)

    for n in g.nodes():
        props = g.nodes()[n]
        upd_dict = props.copy()
        if "type" in upd_dict:
            upd_dict["shape"] = map_type2shape[props["type"]]
            upd_dict["color"] = map_type2color[props["type"]]
        if "label" in upd_dict:
            upd_dict["forcelabel"] = True
        upd_dict["style"] = "filled"

        for k, v in upd_dict.items():
            g.nodes[n][k] = v

    for e in g.edges(data=True):
        s, t, _ = e
        upd_dict = {"style": "solid", "arrowhead": "vee"}
        for k, v in upd_dict.items():
            g.edges[s, t][k] = v

    ag = nx.nx_agraph.to_agraph(g)

    for k in vconf.vertex_set:
        level_index = [
            get_auxnode_id(
                AuxNodeType.FIELD,
                vertex=k,
                field=item,
                vfield=True,
                vertex_sh=vertex_prefix_dict,
            )
            for item in vconf.index(k)
        ]
        index_subgraph = ag.add_subgraph(level_index, name=f"cluster_{k}:def")
        index_subgraph.node_attr["style"] = "filled"
        index_subgraph.node_attr["label"] = "definition"

    ag = ag.unflatten("-l 5 -f -c 3")
    ag.draw(
        os.path.join(self.fig_path, f"{self.prefix}_vc2fields.pdf"),
        "pdf",
        prog="dot",
    )

plot_vc2vc(prune_leaves=False)

Plot vertex collection relationships.

Creates a visualization showing the relationships between vertex collections. Optionally prunes leaf nodes from the visualization.

This method discovers edges from both edge_config and resources to ensure all relationships are visualized, including those with dynamic relations.

Parameters:

Name Type Description Default
prune_leaves

Whether to remove leaf nodes from the visualization

False
Example

plotter.plot_vc2vc(prune_leaves=True)

Source code in graflo/plot/plotter.py
def plot_vc2vc(self, prune_leaves=False):
    """Plot vertex collection relationships.

    Creates a visualization showing the relationships between vertex collections.
    Optionally prunes leaf nodes from the visualization.

    This method discovers edges from both edge_config and resources to ensure
    all relationships are visualized, including those with dynamic relations.

    Args:
        prune_leaves: Whether to remove leaf nodes from the visualization

    Example:
        >>> plotter.plot_vc2vc(prune_leaves=True)
    """
    g = nx.MultiDiGraph()
    nodes = []
    edges = []

    # Discover edges from resources (may include edges not in edge_config)
    discovered_edges = self._discover_edges_from_resources()

    # Collect all edges: from edge_config and discovered from resources
    all_edges = {}
    for edge_id, e in self.schema.edge_config.edges_items():
        all_edges[edge_id] = e
    # Add discovered edges (they may already be in edge_config, but that's fine)
    for edge_id, e in discovered_edges.items():
        if edge_id not in all_edges:
            all_edges[edge_id] = e

    # Create graph edges with relation labels
    for (source, target, purpose), e in all_edges.items():
        # Determine label based on relation configuration
        label = None
        if e.relation is not None:
            # Static relation
            label = e.relation
        elif e.relation_field is not None:
            # Dynamic relation from field - show indicator
            label = f"[{e.relation_field}]"
        elif e.relation_from_key:
            # Dynamic relation from key - show indicator
            label = "[key]"

        if label is not None:
            ee = (
                get_auxnode_id(AuxNodeType.VERTEX, vertex=source),
                get_auxnode_id(AuxNodeType.VERTEX, vertex=target),
                {"label": label},
            )
        else:
            ee = (
                get_auxnode_id(AuxNodeType.VERTEX, vertex=source),
                get_auxnode_id(AuxNodeType.VERTEX, vertex=target),
            )
        edges += [ee]

    # Create nodes for all vertices involved in edges
    for (source, target, purpose), e in all_edges.items():
        for v in (source, target):
            nodes += [
                (
                    get_auxnode_id(AuxNodeType.VERTEX, vertex=v),
                    {
                        "type": AuxNodeType.VERTEX,
                        "label": get_auxnode_id(
                            AuxNodeType.VERTEX, vertex=v, label=True
                        ),
                    },
                )
            ]

    for nid, weight in nodes:
        g.add_node(nid, **weight)

    g.add_nodes_from(nodes)
    g.add_edges_from(edges)

    if prune_leaves:
        out_deg = g.out_degree()
        in_deg = g.in_degree()

        nodes_to_remove = set([k for k, v in out_deg if v == 0]) & set(
            [k for k, v in in_deg if v < 2]
        )
        g.remove_nodes_from(nodes_to_remove)

    for n in g.nodes():
        props = g.nodes()[n]
        upd_dict = {
            "shape": map_type2shape[props["type"]],
            "color": map_type2color[props["type"]],
            "style": "filled",
        }
        for k, v in upd_dict.items():
            g.nodes[n][k] = v

    for e in g.edges:
        s, t, ix = e
        target_props = g.nodes[s]
        edge_data = g.edges[s, t, ix]
        upd_dict = {
            "style": edge_status[target_props["type"]],
            "arrowhead": "vee",
        }
        # Preserve existing label if present (for relation display)
        if "label" in edge_data:
            upd_dict["label"] = edge_data["label"]
        for k, v in upd_dict.items():
            g.edges[s, t, ix][k] = v

    ag = nx.nx_agraph.to_agraph(g)
    ag.draw(
        os.path.join(self.fig_path, f"{self.prefix}_vc2vc.pdf"),
        "pdf",
        prog="dot",
    )