84 lines
2.2 KiB
Python
84 lines
2.2 KiB
Python
from functools import partial
|
|
from typing import Callable, Literal
|
|
|
|
from matplotlib import animation, pyplot as plt
|
|
import numpy as np
|
|
|
|
from .carry import Carry
|
|
|
|
writer = animation.writers["ffmpeg"](fps=15, metadata={"artist": "Me"})
|
|
|
|
|
|
def times(base: np.ndarray, carry: Carry, val):
|
|
base *= val
|
|
carry.apply(base)
|
|
return base
|
|
|
|
|
|
def add(base: np.ndarray, carry: Carry, val=1, center=(0, 0)):
|
|
base[center] += val
|
|
carry.apply(base)
|
|
return base
|
|
|
|
|
|
def animate(
|
|
fig, frames: list[int] | None, **kwargs
|
|
) -> Callable[..., animation.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"] = "add",
|
|
op_val=1,
|
|
center: tuple[int, int] | None = None,
|
|
# Animation parameters
|
|
frames: list[int] | None = None,
|
|
interval=200,
|
|
) -> animation.FuncAnimation:
|
|
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))
|
|
expansion[center] = 1
|
|
title = [1]
|
|
|
|
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"{title[0]}")
|
|
image.set_clim(0, carry.overflow - 1)
|
|
fig.tight_layout()
|
|
|
|
@animate(fig, frames, interval=interval, init_func=init)
|
|
def ret(_):
|
|
plt.title(f"{title[0]}")
|
|
image.set_data(expansion)
|
|
fig.tight_layout()
|
|
|
|
next_func(expansion)
|
|
if operation == "add":
|
|
title[0] += op_val
|
|
else:
|
|
title[0] *= op_val
|
|
|
|
return ret
|