Collatz tree generator. Output as text or png (png requires graphviz). Produces regular binary collatz tree or trees without even numbers.

collatz-tree.py

# Collatz tree generator. Output as text or png (png requires graphviz).
# Produces regular binary collatz tree or trees without even numbers.
#
# Note: The tree starts at 8 to avoid the 4-2-1 loop.
# Note': See the end of the file for two example text trees.
#
# Setup - on linux or wsl:
# ```shell
# sudo apt install python3-pip ipython3 graphviz
# sudo -H pip3 install anytree
# ```
# ipython3 isn't mandatory.
#
# Launch:
# ```
# $ ipython3 -i collatz.py
# >>> pngRellabeledTree()
# >>> # The above commands generates a file named
# >>> # "collatz-sophisticated-re-200.png" in the currend working directory.
# >>>
# >>> pngRellabeledTree(400, "Sophisticated")
# >>> # > "collatz-sophisticated-re-400.png"
# >>>
# >>> pngTree(150, "Binary")
# >>> # > "collatz-binary-150.png"
# >>>
# >>> printTree(100, "Binary")
# >>> # Example output at the end of the file (`print(example[0])`)
# >>>
# >>> printTree(200, "Sophisticated")
# >>> # Example output at the end of the file (`print(example[1])`)
# ```


from abc import abstractmethod

import anytree as a
from anytree import Node, RenderTree, AbstractStyle
from anytree.exporter import DotExporter

from sys import argv

mode = "Binary Sophisticated".split()[1]
# * Binary is the regular Collatz tree. Nodes have either 1 or 2 children.
# * Sophisticated is a Collatz tree without even numbers (except 8). Nodes have
#   either 0 or infinite children.

try:
    z = int(argv[1])
except Exception:
    z = 200
    def main():
        pass
else:
    def main():
        printTree()

# Frontend
def printTree(z=None, mode=None):
    """
    Produce text tree with labeled elements. See `pngRellabledTree` for the
    explanation about labels.
    """
    __printTree(z, mode)

def pngTree(z=None, mode=None):
    """
    Produce a png image of the tree, using graphviz.
    """
    __pngTree(z, mode)

def pngRellabeledTree(z=None, mode=None):
    """
    Produce a png image of the tree, including labels '_/\'.
    The labels indicate for each element, whether it is congruent
    to 0 (_), 1 (/) or 2 (\) modulo 3.
    """
    __pngRellabeledTree(z, mode)


# Backend
def __info(v, patterns="_/\\"):
    return " " + patterns[v % 3]

def __label(node, info=__info):
    return "{}{}".format(node.name, info(node.name))

semiStyle = AbstractStyle("│ ", "├─", "└─")
def __render(tree):
    r = RenderTree(tree, style=semiStyle)
    return "\n".join("{}{}".format(pre, __label(node)) for pre, _, node in r)


def __getClass(mode):
    return globals()[f"CollatzTree{mode}"]

def __printTree(_z=None, _mode=None):
    if _z is None:
        _z = z
    if _mode is None:
        _mode = mode
    tree = Node(8)
    __getClass(_mode)(_z).addChildrenOf(tree)
    print(__render(tree))

def __pngTree(_z=None, _mode=None):
    tree = Node(8)
    __getClass(_mode)(_z).addChildrenOf(tree)
    DotExporter(tree).to_picture(f"collatz-{_mode.lower()}-{_z}.png")

def __relabelTree(tree):
    reInfo = lambda v: __info(v, patterns=["_", "/", "\\\\"])
    for _, _, node in tuple(RenderTree(tree)):
        node.name = __label(node, info=reInfo)

def __pngRellabeledTree(_z=None, _mode=None):
    if _z is None:
        _z = z
    if _mode is None:
        _mode = mode
    tree = Node(8)
    __getClass(_mode)(_z).addChildrenOf(tree)
    __relabelTree(tree)
    DotExporter(tree).to_picture(f"collatz-{_mode.lower()}-re-{_z}.png")


class AbstractCollatzTree:
    def __init__(self, z):
        self.z = z
        self.y = 3 * z + 1

    @abstractmethod
    def childrenOf(self, node):
        pass

    def addChildrenOf(self, node):
        for child in self.childrenOf(node):
            if child >= self.z:
                break
            try:
                self.addChildrenOf(Node(child, parent=node))
            except RecursionError:
                import pdb
                pdb.post_mortem()


class CollatzTreeBinary(AbstractCollatzTree):
    """
    Binary Collatz Tree.

    Contains both even and odd numbers.
    """
    def childrenOf(self, node):
        assert node not in [1, 2, 4]
        if node.name % 6 == 4:
            yield (node.name - 1) // 3
        yield 2 * node.name
        # Note that the two yields must be in increasing order


class CollatzTreeSophisticated(AbstractCollatzTree):
    """
    Sophisticated Collatz Tree.

    Contains only odd numbers. Each node theorically has an infinit number of
    decendents.
    """
    def childrenOf(self, node):
        epsilon = node.name % 3
        if epsilon == 0: return
        child = (2 ** (3 - epsilon) * node.name - 1) // 3
        while True:
            yield child
            child = 4 * child + 1


if __name__ == "__main__":
    main()

example = {}
example[0] = example["""printTree(100, "Binary")"""] = r"""
8 \
└─16 /
  ├─5 \
  │ └─10 /
  │   ├─3 _
  │   │ └─6 _
  │   │   └─12 _
  │   │     └─24 _
  │   │       └─48 _
  │   │         └─96 _
  │   └─20 \
  │     └─40 /
  │       ├─13 /
  │       │ └─26 \
  │       │   └─52 /
  │       │     └─17 \
  │       │       └─34 /
  │       │         ├─11 \
  │       │         │ └─22 /
  │       │         │   ├─7 /
  │       │         │   │ └─14 \
  │       │         │   │   └─28 /
  │       │         │   │     ├─9 _
  │       │         │   │     │ └─18 _
  │       │         │   │     │   └─36 _
  │       │         │   │     │     └─72 _
  │       │         │   │     └─56 \
  │       │         │   └─44 \
  │       │         │     └─88 /
  │       │         │       └─29 \
  │       │         │         └─58 /
  │       │         │           └─19 /
  │       │         │             └─38 \
  │       │         │               └─76 /
  │       │         │                 └─25 /
  │       │         │                   └─50 \
  │       │         └─68 \
  │       └─80 \
  └─32 \
    └─64 /
      └─21 _
        └─42 _
          └─84 _
"""

example[1] = example["""printTree(200, "Sophisticated")"""] = r"""
8 \
├─5 \
│ ├─3 _
│ ├─13 /
│ │ ├─17 \
│ │ │ ├─11 \
│ │ │ │ ├─7 /
│ │ │ │ │ ├─9 _
│ │ │ │ │ ├─37 /
│ │ │ │ │ │ ├─49 /
│ │ │ │ │ │ │ └─65 \
│ │ │ │ │ │ │   ├─43 /
│ │ │ │ │ │ │   │ └─57 _
│ │ │ │ │ │ │   └─173 \
│ │ │ │ │ │ │     └─115 /
│ │ │ │ │ │ │       └─153 _
│ │ │ │ │ │ └─197 \
│ │ │ │ │ │   └─131 \
│ │ │ │ │ │     └─87 _
│ │ │ │ │ └─149 \
│ │ │ │ │   └─99 _
│ │ │ │ ├─29 \
│ │ │ │ │ ├─19 /
│ │ │ │ │ │ ├─25 /
│ │ │ │ │ │ │ ├─33 _
│ │ │ │ │ │ │ └─133 /
│ │ │ │ │ │ │   └─177 _
│ │ │ │ │ │ └─101 \
│ │ │ │ │ │   └─67 /
│ │ │ │ │ │     └─89 \
│ │ │ │ │ │       └─59 \
│ │ │ │ │ │         ├─39 _
│ │ │ │ │ │         └─157 /
│ │ │ │ │ └─77 \
│ │ │ │ │   └─51 _
│ │ │ │ └─117 _
│ │ │ ├─45 _
│ │ │ └─181 /
│ │ └─69 _
│ └─53 \
│   ├─35 \
│   │ ├─23 \
│   │ │ ├─15 _
│   │ │ └─61 /
│   │ │   └─81 _
│   │ └─93 _
│   └─141 _
├─21 _
└─85 /
  └─113 \
    └─75 _
"""