Export to pillow graph

CHANGELOG.md

@@ -7,7 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 ### Added:
 - Docs: Tips for setting custom coordinates for plots.
+- Tree Exporter: `tree_to_pillow_graph` method to export tree to pillow image in graph format.
 ### Changed:
+- Plot: Allow `reverse` argument to allow top-bottom y coordinates in Reingold Tilford algorithm.
 - Docs: Add more elaboration for exporting to image for tree and dag.
 - Misc: Split tree/construct and tree/export into multiple files.
 

assets/tree.dot

@@ -0,0 +1,30 @@
+strict digraph G {
+	graph [bb="0,0,162,180",
+		rankdir=TB
+	];
+	node [label="\N"];
+	a0	[height=0.5,
+		label=a,
+		pos="99,162",
+		width=0.75];
+	b0	[height=0.5,
+		label=b,
+		pos="63,90",
+		width=0.75];
+	a0 -> b0	[pos="e,71.304,107.15 90.65,144.76 86.425,136.55 81.192,126.37 76.419,117.09"];
+	c0	[height=0.5,
+		label=c,
+		pos="135,90",
+		width=0.75];
+	a0 -> c0	[pos="e,126.7,107.15 107.35,144.76 111.58,136.55 116.81,126.37 121.58,117.09"];
+	d0	[height=0.5,
+		label=d,
+		pos="27,18",
+		width=0.75];
+	b0 -> d0	[pos="e,35.304,35.147 54.65,72.765 50.425,64.548 45.192,54.373 40.419,45.093"];
+	e0	[height=0.5,
+		label=e,
+		pos="99,18",
+		width=0.75];
+	b0 -> e0	[pos="e,90.696,35.147 71.35,72.765 75.575,64.548 80.808,54.373 85.581,45.093"];
+}

bigtree/__init__.py

@@ -37,6 +37,7 @@
     tree_to_nested_dict,
     tree_to_newick,
     tree_to_pillow,
+    tree_to_pillow_graph,
     tree_to_polars,
     yield_tree,
 )

bigtree/tree/export/__init__.py

@@ -1,6 +1,11 @@
 from .dataframes import tree_to_dataframe, tree_to_polars  # noqa
 from .dictionaries import tree_to_dict, tree_to_nested_dict  # noqa
-from .images import tree_to_dot, tree_to_mermaid, tree_to_pillow  # noqa
+from .images import (  # noqa
+    tree_to_dot,
+    tree_to_mermaid,
+    tree_to_pillow,
+    tree_to_pillow_graph,
+)
 from .stdout import (  # noqa
     hprint_tree,
     hyield_tree,
@@ -20,6 +25,7 @@
     "tree_to_nested_dict",
     "tree_to_dot",
     "tree_to_pillow",
+    "tree_to_pillow_graph",
     "tree_to_mermaid",
     "tree_to_newick",
 ]

bigtree/tree/export/images.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import collections
+import re
 from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
 
 from bigtree.node import node
@@ -24,6 +25,7 @@
 
 __all__ = [
     "tree_to_dot",
+    "tree_to_pillow_graph",
     "tree_to_pillow",
     "tree_to_mermaid",
 ]
@@ -67,8 +69,10 @@ def tree_to_dot(
 
         Export to image, dot file, etc.
 
-        >>> graph.write_png("assets/docstr/tree.png")
-        >>> graph.write_dot("assets/docstr/tree.dot")
+        >>> graph.write_png("assets/tree.png")
+        >>> graph.write_dot("assets/tree.dot")
+
+        ![Export to Dot](https://github.com/kayjan/bigtree/raw/master/assets/tree.png)
 
         Export to string
 
@@ -207,6 +211,193 @@ def _recursive_append(parent_name: Optional[str], child_node: T) -> None:
     return _graph
 
 
+@exceptions.optional_dependencies_image("Pillow")
+def tree_to_pillow_graph(
+    tree: T,
+    node_content: str = "{node_name}",
+    *,
+    margin: Optional[Dict[str, int]] = None,
+    height_buffer: Union[int, float] = 20,
+    width_buffer: Union[int, float] = 10,
+    font_family: str = "",
+    font_size: int = 12,
+    font_colour: Union[Tuple[int, int, int], str] = "black",
+    text_align: str = "center",
+    bg_colour: Union[Tuple[int, int, int], str] = "white",
+    rect_margin: Optional[Dict[str, int]] = None,
+    rect_fill: Union[Tuple[int, int, int], str] = "white",
+    rect_outline: Union[Tuple[int, int, int], str] = "black",
+    rect_width: Union[float, int] = 1,
+) -> Image.Image:
+    r"""Export tree to PIL.Image.Image object. Object can be
+    converted to other formats, such as jpg, or png. Image will look
+    like a tree/graph-like structure.
+
+    Customisations:
+
+        - To change the margin of tree within diagram, vary `margin`
+        - To change the margin of the text within node, vary `rect_margin`
+        - For more separation between nodes, change `height_buffer` and `width_buffer`
+
+    Examples:
+        >>> from bigtree import Node, tree_to_pillow_graph
+        >>> root = Node("a", age=90)
+        >>> b = Node("b", age=65, parent=root)
+        >>> c = Node("c", age=60, parent=root)
+        >>> d = Node("d", age=40, parent=b)
+        >>> e = Node("e", age=35, parent=b)
+        >>> pillow_image = tree_to_pillow_graph(root, node_content="{node_name}\nAge: {age}")
+
+        Export to image (PNG, JPG) file, etc.
+
+        >>> pillow_image.save("assets/tree_pillow_graph.png")
+        >>> pillow_image.save("assets/tree_pillow_graph.jpg")
+
+        ![Export to Pillow Graph](https://github.com/kayjan/bigtree/raw/master/assets/tree_pillow_graph.png)
+
+    Args:
+        tree (Node/List[Node]): tree or list of trees to be exported
+        node_content (str): display text in node
+        margin (Dict[str, int]): margin of diagram
+        height_buffer (Union[int, float]): height buffer between node layers, in pixels
+        width_buffer (Union[int, float]) : width buffer between sibling nodes, in pixels
+        font_family (str): file path of font family, requires .ttf file, defaults to DejaVuSans
+        font_size (int): font size, defaults to 12
+        font_colour (Union[Tuple[int, int, int], str]): font colour, accepts tuple of RGB values or string, defaults to black
+        text_align (str): text align for multi-line text
+        bg_colour (Union[Tuple[int, int, int], str]): background of image, accepts tuple of RGB values or string, defaults to white
+        rect_margin (Dict[str, int]): (for rectangle) margin of text to rectangle, in pixels
+        rect_fill (Union[Tuple[int, int, int], str]): (for rectangle) colour to use for fill
+        rect_outline (Union[Tuple[int, int, int], str]): (for rectangle) colour to use for outline
+        rect_width: Union[float, int]: (for rectangle) line width, in pixels
+
+    Returns:
+        (PIL.Image.Image)
+    """
+    default_margin = {"t": 10, "b": 10, "l": 10, "r": 10}
+    default_rect_margin = {"t": 5, "b": 5, "l": 5, "r": 5}
+    if not margin:
+        margin = default_margin
+    else:
+        margin = {**default_margin, **margin}
+    if not rect_margin:
+        rect_margin = default_rect_margin
+    else:
+        rect_margin = {**default_rect_margin, **rect_margin}
+
+    # Initialize font
+    if not font_family:
+        from urllib.request import urlopen
+
+        dejavusans_url = "https://github.com/kayjan/bigtree/raw/master/assets/DejaVuSans.ttf?raw=true"
+        font_family = urlopen(dejavusans_url)
+    try:
+        font = ImageFont.truetype(font_family, font_size)
+    except OSError:
+        raise ValueError(
+            f"Font file {font_family} is not found, set `font_family` parameter to point to a valid .ttf file."
+        )
+
+    # Calculate image dimension from text, otherwise override with argument
+    _max_text_width = 0
+    _max_text_height = 0
+    _image = Image.new("RGB", (0, 0))
+    _draw = ImageDraw.Draw(_image)
+
+    def get_node_text(_node: T, _node_content: str) -> str:
+        pattern = re.compile(r"\{(.*)\}")
+        matches = re.findall(pattern, _node_content)
+        for match in matches:
+            _node_content = _node_content.replace(
+                f"{{{match}}}",
+                str(_node.get_attr(match)) if _node.get_attr(match) else "",
+            )
+        return _node_content
+
+    for _, _, _node in yield_tree(tree):
+        l, t, r, b = _draw.multiline_textbbox(
+            (0, 0), get_node_text(_node, node_content), font=font
+        )
+        _max_text_width = max(
+            _max_text_width, l + r + rect_margin.get("l", 0) + rect_margin.get("r", 0)
+        )
+        _max_text_height = max(
+            _max_text_height, t + b + rect_margin.get("t", 0) + rect_margin.get("b", 0)
+        )
+
+    # Get x, y, coordinates and height, width of diagram
+    from bigtree.utils.plot import reingold_tilford
+
+    tree = tree.copy()
+    reingold_tilford(
+        tree,
+        subtree_separation=_max_text_width + width_buffer,
+        sibling_separation=_max_text_width + width_buffer,
+        level_separation=_max_text_height + height_buffer,
+        x_offset=0.5 * _max_text_width + margin.get("l", 0),
+        y_offset=0.5 * _max_text_height + margin.get("t", 0),
+        reverse=True,  # top-left corner is (0, 0)
+    )
+
+    _width, _height = 0, 0
+    _width_margin = 0.5 * _max_text_width + margin.get("r", 0)
+    _height_margin = 0.5 * _max_text_height + margin.get("b")
+    for _, _, _node in yield_tree(tree):
+        _width = max(_width, _node.get_attr("x") + _width_margin)
+        _height = max(_height, _node.get_attr("y") + _height_margin)
+    _width = _width.__ceil__()
+    _height = _height.__ceil__()
+
+    # Initialize and draw image
+    image = Image.new("RGB", (_width, _height), bg_colour)
+    image_draw = ImageDraw.Draw(image)
+
+    for _, _, _node in yield_tree(tree):
+        _x, _y = _node.get_attr("x"), _node.get_attr("y")
+        x1, x2 = _x - 0.5 * _max_text_width, _x + 0.5 * _max_text_width
+        y1, y2 = _y - 0.5 * _max_text_height, _y + 0.5 * _max_text_height
+        # Draw box
+        image_draw.rectangle(
+            [x1, y1, x2, y2], fill=rect_fill, outline=rect_outline, width=rect_width
+        )
+        # Draw text
+        image_draw.text(
+            (x1 + rect_margin.get("l", 0), y1 + rect_margin.get("t", 0)),
+            get_node_text(_node, node_content),
+            font=font,
+            fill=font_colour,
+            align=text_align,
+        )
+        # Draw line to parent
+        if _node.parent:
+            _child_x, _child_y = (
+                _node.get_attr("x"),
+                _node.get_attr("y") - 0.5 * _max_text_height,
+            )
+            _parent_x, _parent_y = (
+                _node.parent.get_attr("x"),
+                _node.parent.get_attr("y") + 0.5 * _max_text_height,
+            )
+            middle_y = (_child_y + _parent_y) / 2
+            image_draw.line(
+                (_parent_x, _parent_y, _parent_x, middle_y),
+                fill=rect_outline,
+                width=rect_width,
+            )
+            image_draw.line(
+                (_parent_x, middle_y, _child_x, middle_y),
+                fill=rect_outline,
+                width=rect_width,
+            )
+            image_draw.line(
+                (_child_x, _child_y, _child_x, middle_y),
+                fill=rect_outline,
+                width=rect_width,
+            )
+
+    return image
+
+
 @exceptions.optional_dependencies_image("Pillow")
 def tree_to_pillow(
     tree: T,
@@ -220,7 +411,7 @@ def tree_to_pillow(
     **kwargs: Any,
 ) -> Image.Image:
     """Export tree to PIL.Image.Image object. Object can be
-    converted to other format, such as jpg, or png. Image will be
+    converted to other formats, such as jpg, or png. Image will be
     similar format as `print_tree`, accepts additional keyword arguments
     as input to `yield_tree`.
 

bigtree/utils/plot.py

@@ -26,6 +26,7 @@ def reingold_tilford(
     level_separation: float = 1.0,
     x_offset: float = 0.0,
     y_offset: float = 0.0,
+    reverse: bool = False,
 ) -> None:
     """
     Algorithm for drawing tree structure, retrieves `(x, y)` coordinates for a tree structure.
@@ -77,9 +78,12 @@ def reingold_tilford(
         level_separation (float): fixed distance between adjacent levels of the tree
         x_offset (float): graph offset of x-coordinates
         y_offset (float): graph offset of y-coordinates
+        reverse (bool): graph begins bottom to top by default, set to True for top to bottom y coordinates
     """
     _first_pass(tree_node, sibling_separation, subtree_separation)
-    x_adjustment = _second_pass(tree_node, level_separation, x_offset, y_offset)
+    x_adjustment = _second_pass(
+        tree_node, level_separation, x_offset, y_offset, reverse
+    )
     _third_pass(tree_node, x_adjustment)
 
 
@@ -333,6 +337,7 @@ def _second_pass(
     level_separation: float,
     x_offset: float,
     y_offset: float,
+    reverse: bool,
     cum_mod: Optional[float] = 0.0,
     max_depth: Optional[int] = None,
     x_adjustment: Optional[float] = 0.0,
@@ -356,6 +361,7 @@ def _second_pass(
         level_separation (float): fixed distance between adjacent levels of the tree (constant across iteration)
         x_offset (float): graph offset of x-coordinates (constant across iteration)
         y_offset (float): graph offset of y-coordinates (constant across iteration)
+        reverse (bool): graph begins bottom to top by default, set to True for top to bottom y coordinates
         cum_mod (Optional[float]): cumulative `mod + shift` for tree/subtree from the ancestors
         max_depth (Optional[int]): maximum depth of tree (constant across iteration)
         x_adjustment (Optional[float]): amount of x-adjustment for third pass, in case any x-coordinates goes below 0
@@ -369,7 +375,10 @@ def _second_pass(
     final_x: float = (
         tree_node.get_attr("x") + tree_node.get_attr("shift") + cum_mod + x_offset
     )
-    final_y: float = (max_depth - tree_node.depth) * level_separation + y_offset
+    if reverse:
+        final_y: float = (tree_node.depth - 1) * level_separation + y_offset
+    else:
+        final_y = (max_depth - tree_node.depth) * level_separation + y_offset
     tree_node.set_attrs({"x": final_x, "y": final_y})
 
     # Pre-order iteration (NLR)
@@ -381,6 +390,7 @@ def _second_pass(
                     level_separation,
                     x_offset,
                     y_offset,
+                    reverse,
                     cum_mod + tree_node.get_attr("mod") + tree_node.get_attr("shift"),
                     max_depth,
                     x_adjustment,