Animate the "flower" pattern of primes, using the Hungarian algorithm to ensure smooth interpolation

prime_dance.py

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from collections import defaultdict
import colorsys
import munkres  # provides Hungarian algorithm
import scipy.spatial.distance as dist
import functools
import os


def primesbelow(N):
    # http://stackoverflow.com/questions/2068372/fastest-way-to-list-all-primes-below-n-in-python/3035188#3035188
    # """ Input N>=6, Returns a list of primes, 2 <= p < N """
    correction = N % 6 > 1
    N = {0: N, 1: N - 1, 2: N + 4, 3: N + 3, 4: N + 2, 5: N + 1}[N % 6]
    sieve = [True] * (N // 3)
    sieve[0] = False
    for i in range(int(N ** 0.5) // 3 + 1):
        if sieve[i]:
            k = (3 * i + 1) | 1
            sieve[k * k // 3 :: 2 * k] = [False] * (
                (N // 6 - (k * k) // 6 - 1) // k + 1
            )
            sieve[(k * k + 4 * k - 2 * k * (i % 2)) // 3 :: 2 * k] = [False] * (
                (N // 6 - (k * k + 4 * k - 2 * k * (i % 2)) // 6 - 1) // k + 1
            )
    return [2, 3] + [(3 * i + 1) | 1 for i in range(1, N // 3 - correction) if sieve[i]]


max_prime = 100_000
smallprimes = primesbelow(max_prime)

# factorise an integer
# only works for n < max_prime
@functools.lru_cache(maxsize=None)
def primefactors(n):
    factors = []

    for fac in smallprimes:
        while n % fac == 0:
            factors.append(fac)
            n //= fac
        if fac > n:
            break
    return factors


# create a circular ring of points around a specific
# origin, optionally with a phase shift
def ring(x, y, radius, n, phase=0):
    i = np.linspace(0, np.pi * 2, n, endpoint=False)
    xs = x + radius * np.cos(i + phase)
    ys = y + radius * -np.sin(i + phase)
    return xs, ys


def color_factors(set_factors):
    h, s, v = 0.0, 0.5, 0.5
    for j, (i, factor) in enumerate(set_factors):
        x = i / factor
        if j == 1:
            h = x
        if j == 2:
            s = 0.15 + x * 0.85
        if j == 3:
            v = 0.2 + x * 0.6
    r, g, b = colorsys.hsv_to_rgb(h, s, v)
    return r, g, b, 1.0


def _flower_color(x, y, radius, factors, set_factors=[]):

    if len(factors) == 0:
        # just one point; fix the color
        r, g, b, a = color_factors(set_factors)
        return [(x, y, r, g, b, a)]
    else:
        # we have subrings remaining
        all_points = []
        factor = factors[0]

        # compute the phase offset
        if len(set_factors) == 0:
            phase = 0
        else:
            i, f = set_factors[0]
            phase = 2 * np.pi * (i / f + 0.5 / f)

        xs, ys = ring(x, y, radius, factor, phase)

        for i, (rx, ry) in enumerate(zip(xs, ys)):
            # punt factors from the "factors" list into "set_factors", passing on the specific
            # factor we have chosen for this subring, to produce pairs (i, n)
            all_points += _flower_color(
                rx, ry, radius / factor, factors[1:], [(i, factor)] + set_factors
            )
        return all_points


# return the [x,y,r,g,b,a] point layout corresponding to the
# factorisation of n
@functools.lru_cache(maxsize=None)
def flower_color(n):
    factors = primefactors(n)
    return _flower_color(0, 0, 1.0, factors)


# compute the node positions for two integers, ordered such that
# the distance between points with the same index is minimised
# to allow for
# memoize, so we only compute the node layout for points once
@functools.lru_cache(maxsize=None)
def align(n1, n2):
    # new black, transparent point at origin
    # x, y, r, g, b, a
    new_point = [0, 0, 0, 0, 0, 0]
    # compute distance matrix betwen flower points
    first_pts = np.array(flower_color(n1) + [new_point])
    second_pts = np.array(flower_color(n2))
    distance = dist.cdist(first_pts[:, 0:2], second_pts[:, 0:2])
    # compute optimal matching
    m = munkres.Munkres()
    indices = np.array(m.compute(distance))
    # get best indices in second array
    selected = indices[:, 1]
    second_order = second_pts[selected]
    return first_pts, second_order


# given a fractional number n, compute the interpolated
# layout of points between floor(n) and floor(n)+1
# Uses a modified cosine interpolation to get smooth
# interpolation with slowdowns on each integer
def interpolate(n_frac, whole_number_pause=5):
    n1 = np.floor(n_frac)
    n2 = n1 + 1
    frac = n_frac - n1

    # cosine interpolation function
    cosfrac = np.cos(frac * np.pi)
    frac = (1 - cosfrac) / 2
    frac = np.tanh((frac - 0.5) * whole_number_pause) / 2 + 0.5  # slow on whole numbers

    # get the positions, in aligned order (for minimal movement of points)
    first_pts, second_order = align(n1, n2)

    # linear interpolate aligned points
    return (1 - frac) * first_pts + frac * second_order


# animate the movement of points, from 2 up to n
def render_animation(
    n, fps=24, path=".", dpi=200, crop=1.8, base_name="prime_dance", preview=False
):
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    if preview:
        plt.ion()
        plt.show()
    for i in range((n - 2) * fps):
        ix = i / fps + 2.0
        ax.cla()
        coloured_points = interpolate(ix)
        # sizing function for points, to ensure enough room
        # as number of points increases
        sz = 5000.0 / (ix ** (3.0 / 2.0))

        # render the points
        ax.scatter(
            coloured_points[:, 0], coloured_points[:, 1], c=coloured_points[:, 2:], s=sz
        )

        number_label = "{index:.0f}".format(index=ix + 0.25)
        # and a label at the origin
        ax.text(
            0,
            0,
            number_label,
            horizontalalignment="center",
            verticalalignment="center",
            size=18,
            color=[0.45, 0.45, 0.45],
            zorder=-1,
            fontdict={"family": "GeoSansLight"},
        )

        # clean up axes and write the frame
        ax.set_aspect(1.0)
        ax.axis("off")
        ax.set_xlim(-crop, crop)
        ax.set_ylim(-crop, crop)

        if preview:
            plt.draw()
            plt.pause(0.0001)
        else:
            output_path = os.path.join(path, f"{base_name}_{i:05d}.png")
            print(f"Writing {ix:.2f} as {output_path}")
            fig.savefig(output_path, dpi=dpi)


if __name__ == "__main__":
    render_animation(10, preview=True)
    # render_animation(10, path="c:\\temp", base_name="prime_dance")