Croft-Moser manim, work in progress

main.py

import numpy as np
from numpy import ndarray
from manim import *
from shapely.geometry import Point as ShapelyPoint
from shapely.geometry import Polygon as ShapelyPolygon
from shapely.ops import unary_union  # To compute the union of multiple geometries


def moser_spindle():
    delta = np.exp(np.pi * 1j / 3) # 60 degrees
    # solution to dist(alpha * (delta + 1), delta + 1) = 1
    rot = 2 * np.arcsin(1 / np.sqrt(12))
    alpha = np.exp(1j * rot)
    assert np.isclose(np.abs(alpha * (delta + 1) - (delta + 1)), 1)
    return np.array([
        0,
        1,
        delta,
        1 + delta,
        alpha,
        alpha * delta,
        alpha * (1 + delta)
    ])


def transform_complex(vs, rotation, translation):
    return vs * np.exp(1j * rotation) + translation


def get_adjacency_matrix(vertices: ndarray) -> ndarray:
    assert np.issubdtype(vertices.dtype, np.complexfloating)
    assert len(vertices.shape) == 1
    return np.isclose(np.abs(vertices[:, None] - vertices[None, :]), 1)


def get_independent_subsets_from_matrix(adjacency_matrix: ndarray) -> ndarray:
    n = adjacency_matrix.shape[0]
    assert adjacency_matrix.shape == (n, n)
    breadth = np.array([[0], [1]], dtype=bool)
    for i in range(1, n):
        row = adjacency_matrix[i, :i]

        can_add_vertex = ~np.any(
            breadth[:, row],
            axis=1
        )        
        new_breadth_len = len(breadth) + np.sum(can_add_vertex)

        new_breadth = np.empty(
            (new_breadth_len, i + 1),
            dtype=bool
        )
        new_breadth[:len(breadth), :i] = breadth
        new_breadth[:len(breadth), i] = False
        new_breadth[len(breadth):, :i] = breadth[can_add_vertex]
        new_breadth[len(breadth):, i] = True

        breadth = new_breadth

    return breadth


# almost the same as get_independent_subsets_from_matrix, but sorts the atoms,
# first by number of nonzeros, then lexicographically.
def get_atoms(udg):
    atoms = get_independent_subsets_from_matrix(get_adjacency_matrix(udg)).astype(int)

    true_counts = atoms.sum(axis=1)
    # Step 2: Convert each row to an integer for lexicographic comparison
    lexicographic_keys = atoms.dot(1 << np.arange(atoms.shape[1])[::-1])
    # Step 3: Use np.lexsort with true_counts as the primary key and lexicographic_keys as the secondary key
    sorted_indices = np.lexsort((-lexicographic_keys, true_counts))
    atoms = atoms[sorted_indices]
    return atoms



def sample_distribution(udg, atoms, union_of_tortoises):
    sample = np.zeros(len(atoms), dtype=int)
    malformed = 0
    for i in range(2000):
        probe = udg * np.exp(1j * np.random.uniform(0, 2*PI)) + np.random.uniform(-4, 1) + np.random.uniform(-2, 2) * 1j
        atom = []
        for pos in probe:
            bit = union_of_tortoises.contains(ShapelyPoint(pos.real, pos.imag))
            atom.append(bit)
        atom = np.array(atom, dtype=bool)
        indices = np.where((atoms == atom).all(axis=1))[0]
        if len(indices) != 1:
            assert len(indices) < 2, "there cannot be two matches"
            # assert len(indices) > 0, f"there must be a match {atom}"
            malformed += 1
        else:
            indx = indices[0]
            sample[indx] += 1
    for cnt, atom in zip(sample, atoms):
        print(cnt, atom.astype(int))
    print("malformed", malformed)
    return sample


# override these to my favorite color scheme
BLUE = '#6699cc'
RED = '#ff3333'


class TortoiseWithTriangle(Scene):
    def __init__(self):
        # this is a global state of the animation, a bitmask
        # describing which of the UDG vertices is over Croft at the moment.
        self.global_atom = None
        super().__init__()

    def construct(self):
        self.camera.background_color = WHITE

        udg = moser_spindle()

        atoms = get_atoms(udg)

        self.global_atom = atoms[0]

        tortoises, union_of_tortoises = self.create_tortoises()

        point_circles, graph_edges = self.create_unit_distance_graph(udg)
        # caricature of the big one: half absolute edge width, double relative dot size.
        for p in point_circles:
            p.scale(2)
        for e in graph_edges:
            e.set_stroke_width(1)
        mini_graph_prototype = VGroup(graph_edges, point_circles).scale(0.3)

        # we redo them for the big unique UDG
        point_circles, graph_edges = self.create_unit_distance_graph(udg)

        sample = sample_distribution(udg, atoms, union_of_tortoises)

        assert len(atoms) == 18, "we only deal with the Moser spindle, really."
        grid_x = 0
        grid_y = 0
        graph_grid = VGroup()
        box_grid = VGroup()
        bar_grid = VGroup()
        for atom in atoms:
            mini_graph = mini_graph_prototype.copy().move_to((grid_x + 1.85, (5 - grid_y) * 0.7 - 1.8, 0))
            self.set_atom(mini_graph[1], atom) # mini_graph[1] is the vgroup of circles
            graph_grid.add(mini_graph)
            box = SurroundingRectangle(mini_graph, color=WHITE, buff=0.05)
            box_grid.add(box)

            bar = Rectangle(height=0.00001, width=0.1, color=BLACK, fill_opacity=0.8)
            bar.next_to(box, RIGHT, buff=0.05)
            bar.shift((0, -0.1, 0))
            bar_grid.add(bar)

            grid_x += 1
            if grid_x >= 3:
                grid_x = 0
                grid_y += 1

        def make_box_updater(local_atom):
            def box_updater(b):
                fire = np.allclose(local_atom, self.global_atom)
                b.set_color(RED if fire else WHITE)
            return box_updater

        def make_bar_updater(local_atom):
            def bar_updater(rect):
                fire = np.allclose(local_atom, self.global_atom)
                if fire:
                    new_height = rect.get_height() + 0.001
                    bottom_y = rect.get_bottom()[1]  # Get the current y-coordinate of the bottom edge

                    # Stretch the rectangle to the new height
                    rect.stretch_to_fit_height(new_height)

                    # Move the rectangle down so the bottom edge stays in place
                    rect.move_to([rect.get_center()[0], bottom_y + new_height / 2, 0])

            return bar_updater

        for i, atom in enumerate(atoms):
            box_grid[i].add_updater(make_box_updater(atom))
            bar_grid[i].add_updater(make_bar_updater(atom))

        # Set camera frame dimensions relative to Manim's default aspect ratio
        target_frame_height = 5  # Adjust this to zoom in or out
        self.camera.frame_height = target_frame_height
        self.camera.frame_width = target_frame_height * config.frame_width / config.frame_height

        self.add(tortoises)
        # we only add the grid now, because we want it in the front
        self.add(graph_grid)
        self.add(box_grid)
        self.add(bar_grid)


        graph = VGroup(graph_edges, point_circles)

        self.add(graph)

        def get_atom():
            atom = []
            for p in point_circles:
                pos = p.get_center()
                bit = union_of_tortoises.contains(ShapelyPoint(pos[0], pos[1]))
                atom.append(bit)
            return np.array(atom, dtype=bool)

        def vertex_color_updater(p):
            pos = p.get_center()
            if union_of_tortoises.contains(ShapelyPoint(pos[0], pos[1])):
                return p.set_fill(BLACK)
            else:
                return p.set_fill(WHITE)

        for p in point_circles:
            p.add_updater(vertex_color_updater)

        def spiral_fn(t01):
            t_range = (0.5, 2.0)
            t = t01 * (t_range[1] - t_range[0]) + t_range[0]
            return (0.5 * t * np.cos(0.7 * t ** 2 * PI) - 1, 0.5 * t * np.sin(0.3 * t ** 3 * PI), 0)

        spiral = ParametricFunction(spiral_fn, t_range=(0, 1), color=PINK)
        # we only show it until we find a nice pattern
        # self.add(spiral)

        # the graph with an invisible arrow stuck to it, thank you uwezi!
        obj = VGroup(
            graph,
            Line(ORIGIN,1e-3*RIGHT,stroke_width=0,stroke_opacity=0)
        )
        posTracker = ValueTracker(1e-3)
        def smooth_updater(mobj):
            pos = spiral_fn(posTracker.get_value())
            lastpos = spiral_fn(posTracker.get_value() - 1e-3)

            secangle = Line(lastpos,pos).get_angle()
            # mobj.rotate(secangle-mobj[1].get_angle(), about_point=mobj.get_center())
            mobj.rotate(0.01)
            mobj.move_to(pos)
            self.global_atom = get_atom()

        obj.add_updater(smooth_updater, call_updater=True)
        self.add(obj)
        self.play(
            posTracker.animate.set_value(1),
            rate_func=rate_functions.linear,
            run_time=15
        )
        obj.remove_updater(smooth_updater)
        def jumpy_updater(mobj):
            mobj.move_to((np.random.uniform(-4, 1), np.random.uniform(-2, 2), 0))
            mobj.rotate(np.random.uniform(0, 2*PI))
            self.global_atom = get_atom()
        obj.add_updater(jumpy_updater, call_updater=True)
        self.play(
            posTracker.animate.set_value(1),
            rate_func=rate_functions.linear,
            run_time=10
        )

    @staticmethod
    def set_atom(point_circles, atom):
        assert len(point_circles) == len(atom)
        for p, bit in zip(point_circles, atom):
            if bit:
                p.set_fill(BLACK)
            else:
                p.set_fill(WHITE)

    def create_tortoises(self):
        # Define parameters for the tortoise pattern
        x = 0.96553
        R6 = np.exp(1j * np.pi / 3)
        ROOTS = R6 ** np.arange(6)

        # Create a hexagon and scale it
        hexagon_points = [(q.real, q.imag, 0) for q in ROOTS]
        hexagon = Polygon(*hexagon_points, stroke_color=BLACK, stroke_width=1, fill_color=BLUE, fill_opacity=1.0)
        scale_factor = x / np.sqrt(3)
        hexagon.scale(scale_factor)
        
        # Define the circle parameters
        radius = 0.5
        circle = Circle(radius=radius, stroke_color=BLACK, stroke_width=1, fill_color=BLUE, fill_opacity=1.0)
        
        # Intersect circle and hexagon by clipping to create tortoise shape
        tortoise = Intersection(circle, hexagon, stroke_color=BLACK, stroke_width=1, fill_color=BLUE, fill_opacity=1.0)
        
        # Convert tortoise to Shapely geometry for point containment check
        tortoise_shapely_points = [
            (point[0], point[1]) for point in tortoise.get_all_points()[:, :2]
        ]
        tortoise_polygon = ShapelyPolygon(tortoise_shapely_points)

        tortoises = VGroup()

        # Create multiple tortoises and add them to the scene
        for i in range(-4, 3): # range(-2, 3):
            for j in range(-4, 2): # range(-1, 2):
                direction = (ROOTS[0] * i + ROOTS[2] * j) * (x + 1)
                if direction.real > 1:
                    continue
                tortoise_i = tortoise.copy().rotate(np.pi / 6).shift(direction.real * RIGHT + direction.imag * UP)
                tortoises.add(tortoise_i)

        # Compute the union of all tortoises for containment checks
        union_of_tortoises = unary_union([
            ShapelyPolygon([(point[0], point[1]) for point in tortoise_i.get_all_points()[:, :2]])
            for tortoise_i in tortoises
        ])
        return tortoises, union_of_tortoises

    def create_unit_distance_graph(self, complex_points):
        # Create a VGroup to hold points and edges
        graph_edges = VGroup()
        point_circles = VGroup()

        # Render points as circles
        for point in complex_points:
            circle = Circle(radius=0.05, color=BLACK, stroke_width=1, fill_opacity=1).move_to(point.real * RIGHT + point.imag * UP)
            point_circles.add(circle)
            # self.add(circle)  # Add point to scene

        # Connect points with edges if they are approximately unit distance apart
        for i in range(len(complex_points)):
            for j in range(i + 1, len(complex_points)):
                if np.isclose(np.abs(complex_points[i] - complex_points[j]), 1.0):
                    line = Line(complex_points[i].real * RIGHT + complex_points[i].imag * UP,
                                 complex_points[j].real * RIGHT + complex_points[j].imag * UP,
                                 stroke_color=BLACK, stroke_width=2)
                    graph_edges.add(line)
                    # self.add(line)  # Add line to scene

        # Optional: Add points and edges to the scene (already added during the loop)
        # self.add(point_circles)  # Not needed since points are added during the loop
        # self.add(graph_edges)  # Not needed since edges are added during the loop
        return point_circles, graph_edges