zenzicubi.co/posts/polycount/sand-1/carry2d/anim.py

from functools import partial
import time
from threading import Thread
from typing import Callable, Literal

from matplotlib import animation, pyplot as plt
import numpy as np
from sympy.plotting import plot_implicit

from .carry import Carry
from .expansion import add, times, poly_from_array


def animate(
    fig, frames: list[int] | None, **kwargs
) -> Callable[..., animation.FuncAnimation]:
    """Decorator-style wrapper for FuncAnimation."""
    # Why isn't this how the function is exposed by default?
    return lambda func: animation.FuncAnimation(fig, func, frames, **kwargs)


def animate_carry_count(
    carry: Carry,
    dims: int = 100,
    operation: Literal["add", "multiply", "log"] = "add",
    op_val=1,
    center: tuple[int, int] | None = None,
    # Animation parameters
    frames: list[int] | None = None,
    interval=200,
) -> animation.FuncAnimation:
    """
    Create a video of integer expansions using a `carry`.
    Expansions use a `dims` x `dims` numpy array.

    The video starts from the expansion of one, and subsequent frames come from
    either adding `op_val` to the constant term or multiplying each term by `op_val`,
    both followed by carrying.
    `operation` decides which of these to use.

    The constant term of the expansion is specified with `center`.
    print("Starting animation")
    By default, the constant term is at position (0, 0), in the upper left corner.
    If the carry expands leftward or upward, it is instead in the center of the array.
    """
    if center is None:
        if carry.over_pos != (0, 0):
            center = (dims // 2, dims // 2)
        else:
            center = (0, 0)

    if operation == "add" or operation == "log":
        next_func = partial(add, carry=carry, center=center)
    elif operation == "multiply":
        if op_val == 1:
            raise ValueError(f"too small value {op_val} for repeated multiplication")
        next_func = partial(times, carry=carry)
    else:
        raise ValueError(f"Cannot use {repr(operation)} for animation")

    expansion = np.zeros((dims, dims))
    expansion[center] = 1
    state = [1, op_val, op_val]

    fig = plt.gcf()

    image = plt.imshow(
        expansion,
        extent=[-center[0], dims - center[0], dims - center[1], -center[1]],  # type: ignore
    )
    def init():
        plt.title(f"{state[0]}")
        image.set_clim(0, carry.overflow - 1)
        fig.tight_layout()

    @animate(fig, frames, interval=interval, init_func=init)
    def ret(_):
        title, op_val, old_op_val = state
        plt.title(f"{title}")
        image.set_data(expansion)
        fig.tight_layout()

        next_func(expansion, val=op_val)
        if operation == "add":
            state[0] += op_val
        elif operation == "multiply":
            state[0] *= op_val
        elif operation == "log":
            state[0] += op_val
            if state[0] >= op_val * old_op_val * old_op_val:
                state[1] *= old_op_val

    return ret

# writer = animation.writers["ffmpeg"](fps=15, metadata={"artist": "Me"})

def bindfig(fig):
    """
    Create a function which is compatible with the signature of `pyplot.figure`,
    but alters the already-existing figure `fig` instead.
    """
    def ret(**kwargs):
        for i, j in kwargs.items():
            if i == "figsize":
                if j is not None:
                    fig.set_figwidth(j[0])
                    fig.set_figheight(j[1])
                continue
            fig.__dict__["set_" + i](j)
        return fig
    return ret


def animate_carry_curves(
    carry: Carry,
    dims: int = 100,
    operation: Literal["add", "multiply"] = "add",
    op_val=1,
    center: tuple[int, int] | None = None,
    # Animation parameters
    frames: list[int] | None = None,
    # interval=200,
) -> animation.FuncAnimation:
    """
    Create a video of implicit plots of "incremented curves" based on a `carry`.
    These are curves for which we re-interpret an expansion as a polynomial equal
    equal to the integer it represents.
    """
    if center is None:
        if carry.over_pos != (0, 0):
            center = (dims // 2, dims // 2)
        else:
            center = (0, 0)

    if operation == "add":
        next_func = partial(add, carry=carry, val=op_val, center=center)
    elif operation == "multiply":
        if op_val == 1:
            raise ValueError(f"too small value {op_val} for repeated multiplication")
        next_func = partial(times, carry=carry, val=op_val)
    else:
        raise ValueError(f"Cannot use {repr(operation)} for animation")

    expansion = np.zeros((dims, dims), dtype=np.int32)
    expansion[center] = 0
    count = [0]

    # I hate doing it, but there's no other way to control the figure uses
    fig = plt.gcf()
    plt.figure = bindfig(fig)
    plt.tight_layout()

    # More sympy.plot_implicit nonsense
    t = Thread(target=lambda: time.sleep(1) or plt.close())
    t.run()
    plt.title(f"{count[0]}")
    plot_implicit(
        poly_from_array(carry._arr),
        backend="matplotlib",
    )

    @animate(fig, frames)
    def ret(_):
        next_poly = poly_from_array(expansion) - count[0]

        if next_poly != 0:
            fig.clf()
            plot_implicit(
                next_poly,
                backend="matplotlib",
            )
            plt.title(f"{count[0]}")

        if operation == "add":
            count[0] += op_val
        else:
            count[0] *= op_val
        next_func(expansion)

    return ret