226 lines
6.3 KiB
Python

import argparse
from itertools import accumulate, chain, repeat
from pathlib import Path
from typing import Any, Generator, TypeVar
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from matplotlib.animation import FuncAnimation
import numpy as np
from numpy.typing import NDArray
T = TypeVar("T")
def to_base(b: int, x: int) -> list[int]:
"""Convert x to an integer base b."""
if x == 0:
return [0]
ret = []
while x > 0:
ret.append(x % b)
x = x // b
return ret
def from_base(b: float, xs: list[int]) -> float:
"""
Interpret xs as a string in base `b`.
For genericness, the base and return are both floats.
"""
ret = 0
plval = 1
for i in xs:
ret += i * plval
plval *= b
return ret
class Expansions:
"""List of base-b expansions of a number"""
def __init__(self, base: float, name: str, xs: list[list[int]]):
self.xs = xs
self.name = name
self.base = base
@classmethod
def from_file(cls, base: float, expansions_file: Path, name: str | None = None):
ret = []
with open(expansions_file) as f:
for line in f.readlines():
ret.append(eval(line))
# Extract name of base from filename
if name is None:
name_from_file = expansions_file.name.removesuffix(expansions_file.suffix)
second_underscore = name_from_file.find("_", name_from_file.find("_") + 1)
name = (
name_from_file[:second_underscore]
if second_underscore >= 0
else name_from_file
)
return cls(base, name, ret)
@classmethod
def from_integer(cls, base: int, count: int, name: str | None = None):
return cls(
float(base),
str(base) if name is None else name,
[to_base(base, i) for i in range(count)],
)
def slices(xs):
for i in range(1, len(xs) + 1):
yield xs[:i]
def extend_last(xs: Generator[T, Any, Any]) -> Generator[T, Any, Any]:
i = next(xs)
yield i
for i in xs:
yield i
while True:
yield i
# More efficient embedding
def to_angle(xs: list[int], len_: int, b: float) -> complex:
return np.exp(2j * np.pi * from_base(b, xs) / b ** (len_ + 1))
def _adic_angles(b: float, expansion: list[int], maxpow: int = 10) -> list[complex]:
return [
to_angle(x, i, b) for i, x in zip(range(maxpow), extend_last(slices(expansion)))
]
def geosum(c, phis: list[complex]) -> complex:
return sum(
i * j
for i, j in zip(phis, accumulate(chain([1], repeat(c)), lambda x, y: x * y))
)
def interactive(expansions: Expansions):
phiss = [
[_adic_angles(expansions.base, expansion, 200) for expansion in expansions.xs]
]
embedding: NDArray[np.complex64] = np.array(
[geosum(0.9, phis) for phis in phiss[0]]
)
# Create the figure and the line that we will manipulate
fig, ax = plt.subplots()
line = ax.plot(embedding.real, embedding.imag, linestyle="None", marker=".")[0]
# adjust the main plot to make room for the sliders
fig.subplots_adjust(bottom=0.25)
# Make a horizontal slider to control the frequency.
axfreq = fig.add_axes([0.1, 0.1, 0.75, 0.03]) # type: ignore
freq_slider = Slider(
ax=axfreq,
label="$c$",
valmin=0.001,
valmax=0.999,
valinit=0.9,
)
# The function to be called anytime a slider's value changes
def update(val):
embedding: NDArray[np.complex64] = np.array(
[geosum(val, phis) for phis in phiss[0]]
)
line.set_data(embedding.real, embedding.imag) # type: ignore
ax.set_xlim(min(embedding.real), max(embedding.real))
ax.set_ylim(min(embedding.imag), max(embedding.imag))
fig.canvas.draw_idle()
# register the update function with each slider
freq_slider.on_changed(update)
plt.show()
def anim(expansions: Expansions):
phiss = [
[_adic_angles(expansions.base, expansion, 200) for expansion in expansions.xs]
]
embedding: NDArray[np.complex64] = np.array(
[geosum(0.9, phis) for phis in phiss[0]]
)
# Create the figure and the line that we will manipulate
fig, ax = plt.subplots()
line = ax.plot(embedding.real, embedding.imag, linestyle="None", marker=".")[0]
vals = np.linspace(0.9, 0.1, 100)
i = [0]
# The function to be called anytime a slider's value changes
def update(_):
embedding: NDArray[np.complex64] = np.array(
[geosum(vals[i[0]], phis) for phis in phiss[0]]
)
line.set_data(embedding.real, embedding.imag) # type: ignore
ax.set_xlim(min(embedding.real), max(embedding.real))
ax.set_ylim(min(embedding.imag), max(embedding.imag))
ax.set_title(f"{vals[i[0]]}")
i[0] = i[0] + 1
FuncAnimation(fig, update, frames=99).save( # type: ignore
f"{expansions.name}-adics_count_{len(expansions.xs)}_{len(expansions.xs[-1])}_digits.mp4"
)
cendree_divmod = Expansions.from_file(
1 + 3**0.5, Path("./cendree_DivMod_count_1024_256_digits.txt"), "divMod_cendree"
)
cendree_quotrem = Expansions.from_file(
1 + 3**0.5, Path("./cendree_QuotRem_count_1024_256_digits.txt"), "quotRem_cendree"
)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--interactive",
"-i",
action="store_true",
help="Interactively set the geometric constant of the Fourier series.",
)
parser.add_argument("--base", "-p", type=int, help="Use p-adic integers")
parser.add_argument(
"--count",
"-c",
default=1024,
type=int,
help="Number of numbers to use. Only applies when using -p.",
)
parser.add_argument("--cendree", "-k", type=Path, help="Use cendree-adic integers")
args = parser.parse_args()
expansions: Expansions
if args.cendree:
if not args.cendree.exists():
raise ValueError(f"Received cendree-adic expansion file: {args.cendree}")
expansions = Expansions.from_file(1 + 3**0.5, args.cendree)
else:
expansions = Expansions.from_integer(args.base, args.count)
if args.interactive:
interactive(expansions)
else:
anim(expansions)
if __name__ == "__main__":
main()