collatz_coral.py

#!/usr/bin/env python3

import random
import svgwrite
import math


def lerp(t, a, b):
    """
    Perform a linear interpolation with parameter t and interpolants a and b
    """
    return (1-t)*a + t*b


def make_coral(outfile, threshold, options={}):
    def random_color():
        """Generate a random color with components in [0.3,0.9]"""
        r, g, b = [lerp(random.random(), 0.3, 0.9) for _ in range(3)]
        return (r, g, b)

    def write_color(c):
        """Construct a color object."""
        r, g, b = [math.floor(100*x) for x in c]
        return svgwrite.rgb(r, g, b, '%')

    default_options = {
            'dimx': 1920,
            'dimy': 1080,
            'even angle': 2*math.pi/360,
            'odd angle': 2*math.pi/360,
            'line width': 1e-3,
            'length': 2e-1,
            'curve length': 2e-1,
            'initial angle': 45*2*math.pi/360,
            'initial position': (0, 0),
            'use bezier': True,
        }
    default_options.update(options)
    options = default_options

    dimx = options['dimx']
    dimy = options['dimy']

    # This is a set of previously encountered numbers.
    # We do not want to revisit them.
    encountered = set([0])
    # We will extend these numbers in each iteration. They are the front
    # edge of the branches of the tree.
    front = set([1])

    # Start with the initial angles.
    angles = {}
    angles[0] = options['initial angle']
    angles[1] = angles[0] + options['even angle']

    # ... and initial position.
    positions = {}
    positions[0] = options['initial position']

    # Choose a random color for the initial branch.
    colors = {}
    colors[0] = random_color()
    colors[1] = colors[0]

    # Construct the SVG drawing.
    dwg = svgwrite.Drawing(profile='full',
                           size=('{}px'.format(dimx), '{}px'.format(dimy)))
    dwg.viewbox(width=1, height=dimy/dimx)

    # Make a white background.
    rect = dwg.rect(insert=(0, 0), size=('100%', '100%'), fill='white')
    dwg.add(rect)

    # This is the group containing all of the branches.
    g = dwg.g()
    dwg.add(g)
    # Set the line width.
    g.stroke(width=options['line width'])
    # We do not need the interiors filled.
    g.fill(color='none')

    # This is a map of the polylines. Branches are segments of polylines.
    polylines = {}

    # Construct the first line. It extents from item k to item j.
    k = 0
    j = 1
    # This is the originating position on the canvas.
    x1, y1 = positions[k]

    # This is the destination position. It is the originating position
    # extended by one length unit with the given angle.
    x2 = math.cos(angles[j])*options['length']
    y2 = math.sin(angles[j])*options['length']
    x2 += x1
    y2 += y1
    positions[j] = (x2, y2)

    if options['use bezier']:
        # Construct a Bezier curve.
        polylines[0] = dwg.path()
        polylines[1] = polylines[0]

        # This point is on the tangent line of the originating point.
        cx = math.cos(angles[k])*options['curve length']
        cy = math.sin(angles[k])*options['curve length']
        cx += x1
        cy += y1

        # This point is on the tangent line of the destination point.
        out_angle = angles[j] + options['even angle']
        dx = math.cos(out_angle)*options['curve length']
        dy = math.sin(out_angle)*options['curve length']
        dx = x2 - dx
        dy = y2 - dy

        # Insert the Bezier curve parameters.
        polylines[j].push('M {} {}'.format(x1, y1))
        polylines[j].push('C {} {}, {} {}, {} {}'.format(
            cx, cy, dx, dy, x2, y2))
    else:
        # Construct a polyline and insert the two points.
        polylines[0] = dwg.polyline()
        polylines[1] = polylines[0]
        polylines[0].points.append((x1, y1))
        polylines[0].points.append((x2, y2))

    # Set the color of the polyline and add the polyline to the group.
    polylines[0].stroke(color=write_color(colors[0]))
    g.add(polylines[0])

    for iteration in range(threshold):
        newfront = set()

        def add_node(j):
            """Attempt to insert a node j"""

            # The next front will contain this item.
            newfront.add(j)

            # Choose the angle depending on the branch type.
            # Add it to the originating angle
            if j % 2 == 0:
                angles[j] = angles[k] + options['even angle']
            else:
                angles[j] = angles[k] + options['odd angle']

            # Get the originating point coordinates.
            x1, y1 = positions[k]

            # Extend the originating point along the line with the given angle.
            x2 = math.cos(angles[j])*options['length']
            y2 = math.sin(angles[j])*options['length']
            x2 += x1
            y2 += y1
            positions[j] = (x2, y2)

            if options['use bezier']:
                # This point is on the tangent line of the originating point.
                cx = math.cos((angles[k]+angles[j])/2)*options['curve length']
                cy = math.sin((angles[k]+angles[j])/2)*options['curve length']
                cx += x1
                cy += y1

                # This point is on the tangent line of the destination point.
                out_angle = angles[j] + options['even angle']/2
                dx = math.cos(out_angle)*options['curve length']
                dy = math.sin(out_angle)*options['curve length']
                dx = x2 - dx
                dy = y2 - dy

            if j % 2 == 0:
                # We do not bifurcate for even branches.
                colors[j] = colors[k]
                polylines[j] = polylines[k]

                if options['use bezier']:
                    polylines[j].push('C {} {}, {} {}, {} {}'.format(
                        cx, cy, dx, dy, x2, y2))
                else:
                    polylines[j].points.append((x2, y2))
            else:
                # Construct a new branch. Select a random color first.
                colors[j] = random_color()

                if options['use bezier']:
                    # Use the previous tangent for the new Bezier curve.
                    polylines[j] = dwg.path()
                    polylines[j].push('M {} {}'.format(x1, y1))
                    polylines[j].push('C {} {}, {} {}, {} {}'.format(
                        cx, cy, dx, dy, x2, y2))
                else:
                    # Construct a new polyline with the branch point
                    # as the starting point.
                    polylines[j] = dwg.polyline()
                    polylines[j].points.append((x1, y1))
                    polylines[j].points.append((x2, y2))
                # Set the color of the branch.
                polylines[j].stroke(color=write_color(colors[j]))

                # Insert the branch into the group.
                g.add(polylines[j])

        for k in front:
            # Make an even branch. This is always possible.
            add_node(k*2)

            # Make an odd branch if this node can be reached by a 3*x+1 move.
            if (k-1) % 3 == 0 and ((k-1)//3) % 2 == 1:
                add_node((k-1)//3)

        # Insert the processed points into the set of encountered points.
        # We will not deal with them in the future.
        encountered = encountered | front

        # Remove encountered points from the next front.
        front = newfront - encountered

    print('point count: {}'.format(len(positions)))

    # Get the extents of the drawing and set the canvas
    # boundaries to coincide.
    max_x, max_y = next(iter(positions.values()))
    min_x, min_y = (max_x, max_y)
    for (x, y) in positions.values():
        max_x = max(x, max_x)
        max_y = max(y, max_y)
        min_x = min(x, min_x)
        min_y = min(y, min_y)

    # Perform the coordinate transforms.
    # We first map the drawing to [0,1]^2, then flip it,
    # then translate it by 1.
    g.translate(1, 1)
    g.scale(-1, -1)
    g.scale(1/(max_x-min_x), 1/(max_y-min_y))
    g.translate(-min_x, -min_y)

    # Write to file3
    dwg.write(outfile)


def main():
    with open('out.svg', 'w') as f:
        dim = 1024
        options = {
            'even angle': 7 * 2*math.pi/360,
            'odd angle': -17 * 2*math.pi/360,
            'line width': 5 / dim,
            'length': 100 / dim,
            'curve length': 60 / dim,
            'dimx': dim,
            'dimy': dim,
            'use bezier': False
        }
        make_coral(outfile=f, threshold=45, options=options)

if __name__ == '__main__':
    main()