Change shape and color of the leaves based on their values

[Miha Gazvoda]

Hi,
I'd like my leaves to be different shape than other nodes. I'd also like them to be either different color based on their value (green if 1, red if 0).
My tree looks like:

from bigtree import list_to_tree, tree_to_dot

root = list_to_tree([
    "Everyone/A/X/xxx/0",
    "Everyone/A/X/yyy/1",
    "Everyone/A/X/zzz/0",
    "Everyone/A/Y/xxx/1",
    "Everyone/A/Y/yyy/0",
    "Everyone/A/Y/zzz/1",
    "Everyone/A/Z/xxx/0",
    "Everyone/A/Z/yyy/1",
    "Everyone/A/Z/zzz/0",
    "Everyone/B/X/xxx/1",
    "Everyone/B/X/yyy/0",
    "Everyone/B/X/zzz/0",
    "Everyone/B/Y/xxx/1",
    "Everyone/B/Y/yyy/0",
    "Everyone/B/Y/zzz/0",
    "Everyone/B/Z/xxx/0",
    "Everyone/B/Z/yyy/1",
    "Everyone/B/Z/zzz/1",
])
tree_to_dot(root).write_png("before.png")

It looks like this:

Again, I'd like my leaves to be different shape (square, let's say) and their color should be based on their value.

Is it possible to achieve that? What's the best way to do it? Thanks a lot!

[Kay Jan W.]

Hi, thanks for your question.

There is documentation on how to set custom node and/or edge design here, but it does require some knowledge on pydot terminology on how to style the colour, shape, etc. Below is a working example of the design you described:

from bigtree import Node, clone_tree, tree_to_dot

# Define style
colour_red = dict(style="filled", fillcolor="red")
colour_green = dict(style="filled", fillcolor="green")
shape_square = dict(shape="square")

# Define customized node to have `node_attr` property
class CustomNode(Node):
     def __init__(self, name, **kwargs):
         super().__init__(name, **kwargs)

     @property
     def node_attr(self):
         node_design = {}
         if self.is_leaf:
             node_design.update(shape_square)
         if self.node_name == "0":
             node_design.update(colour_red)
         elif self.node_name == "1":
             node_design.update(colour_green)
         return node_design

# Clone tree to `CustomNode` type
root = clone_tree(root, CustomNode)
tree_to_dot(root, node_attr="node_attr").write_png("tree.png")

Resulting tree:

Hope this helps!