make kadic.py have parity with showKappaAdic.ojs

posts/polycount/4/appendix/kadic.py

@@ -1,10 +1,12 @@
 import argparse
+import csv
 from itertools import accumulate, chain, repeat
 from pathlib import Path
+import random
 from typing import Any, Generator, TypeVar
 
 import matplotlib.pyplot as plt
-from matplotlib.widgets import Slider
+from matplotlib.widgets import Slider, RadioButtons
 from matplotlib.animation import FuncAnimation
 import numpy as np
 from numpy.typing import NDArray
@@ -49,8 +51,11 @@ class Expansions:
     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():
+            for row in csv.reader(f.readlines()):
-                ret.append(eval(line))
+                row_even = [int(i) for i in row]
+                row_odd = row_even.copy()
+                row_odd[0] = 1
+                ret.extend([row_even, row_odd])
 
         # Extract name of base from filename
         if name is None:
@@ -72,6 +77,22 @@ class Expansions:
             [to_base(base, i) for i in range(count)],
         )
 
+    @classmethod
+    def from_random(cls, base: int, count: int, digits_count: int):
+        return cls(
+            float(base),
+            f"{base}_(random)",
+            [[random.randint(0, 1) for _ in range(digits_count)] for _ in range(count)]
+        )
+
+
+cendree_divmod = Expansions.from_file(
+    1 + 3**0.5, Path("./cendree_DivMod_count_1024_256_digits.csv"), "divMod_cendree"
+)
+cendree_quotrem = Expansions.from_file(
+    1 + 3**0.5, Path("./cendree_QuotRem_count_1024_256_digits.csv"), "quotRem_cendree"
+)
+
 
 def slices(xs):
     for i in range(1, len(xs) + 1):
@@ -105,45 +126,143 @@ def geosum(c, phis: list[complex]) -> complex:
     )
 
 
-def interactive(expansions: Expansions):
+class PadicPlot:
-    phiss = [
+    EXPANSIONS_LENGTH = 1024
-        [_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
+    def __init__(self, expansions: Expansions, axes=None, plot=False, interactive=False, **kwargs):
-    fig, ax = plt.subplots()
+        self._expansions: Expansions = expansions
+        self._angles: NDArray[np.complex64]
+        self._embedding: NDArray[np.complex64]
 
-    line = ax.plot(embedding.real, embedding.imag, linestyle="None", marker=".")[0]
+        self._list_base: int = 2 if isinstance(expansions.base, float) else expansions.base
+        self._embedding_base: float = 2 if isinstance(expansions.base, float) else expansions.base
+        self._geometric: float = 0.9
+        self._update(set_data=False)
 
-    # adjust the main plot to make room for the sliders
+        self._axes = {}
-    fig.subplots_adjust(bottom=0.25)
+        self._inputs = {}
+        self._plot = None
 
-    # Make a horizontal slider to control the frequency.
+        axes = axes if axes is not None else plt.gca()
-    axfreq = fig.add_axes([0.1, 0.1, 0.75, 0.03])  # type: ignore
+        if interactive:
-    freq_slider = Slider(
+            self.plot(axes=axes, **kwargs)
-        ax=axfreq,
+            self._interactive(axes=axes)
-        label="$c$",
+        elif plot:
-        valmin=0.001,
+            self.plot(axes=axes, **kwargs)
-        valmax=0.999,
-        valinit=0.9,
-    )
 
-    # The function to be called anytime a slider's value changes
+    def plot(self, axes, linestyle="None", marker=".", **kwargs):
-    def update(val):
+        self._axes["plot"] = axes
-        embedding: NDArray[np.complex64] = np.array(
+        self._plot = axes.plot(
-            [geosum(val, phis) for phis in phiss[0]]
+            self._embedding.real,
+            self._embedding.imag,
+            linestyle=linestyle,
+            marker=marker,
+            **kwargs,
+        )[0]
+
+        plt.tight_layout()
+
+    def _interactive(self, axes):
+        # adjust the main plot to make room for the sliders
+        fig = axes.get_figure()
+        fig.subplots_adjust(bottom=0.25)
+
+        self._axes["adic"] = fig.add_axes([0.65, 0.025, 0.3, 0.175])  # type: ignore
+        self._inputs["adic"] = RadioButtons(
+            labels=["b-adics", "$\\kappa$-adic (balanced)", "$\\kappa$-adic (binary)", "Random Binary"],
+            ax=self._axes["adic"],
         )
-        line.set_data(embedding.real, embedding.imag)  # type: ignore
+        self._inputs["adic"].on_clicked(self._update_adic)
-        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
+        self._axes["list_base"] = fig.add_axes([0.1, 0.15, 0.45, 0.03])  # type: ignore
-    freq_slider.on_changed(update)
+        self._inputs["list_base"] = Slider(
+            ax=self._axes["list_base"],
+            label="$b$",
+            valmin=2,
+            valmax=10,
+            valstep=1,
+            valinit=2,
+        )
+        self._inputs["list_base"].on_changed(self._update_list_base)
 
+        self._axes["embedding_base"] = fig.add_axes([0.1, 0.1, 0.45, 0.03])  # type: ignore
+        self._inputs["embedding_base"] = Slider(
+            ax=self._axes["embedding_base"],
+            label="$p$",
+            valmin=1.1,
+            valmax=10,
+            valstep=0.1,
+            valinit=2,
+        )
+        self._inputs["embedding_base"].on_changed(self._update_embedding_base)
+
+        self._axes["geometric"] = fig.add_axes([0.1, 0.05, 0.45, 0.03])  # type: ignore
+        self._inputs["geometric"] = Slider(
+            ax=self._axes["geometric"],
+            label="$c$",
+            valmin=0.001,
+            valmax=0.999,
+            valinit=0.9,
+        )
+        self._inputs["geometric"].on_changed(self._update_geometric)
+
+    def _update_adic(self, val):
+        if val == "b-adics":
+            expansions = Expansions.from_integer(int(self._list_base), self.EXPANSIONS_LENGTH)
+        elif val == "$\\kappa$-adic (balanced)":
+            expansions = cendree_quotrem
+        elif val == "$\\kappa$-adic (binary)":
+            expansions = cendree_divmod
+        elif val == "Random Binary":
+            expansions = Expansions.from_random(2, self.EXPANSIONS_LENGTH, self.EXPANSIONS_LENGTH // 10)
+        else:
+            raise ValueError("Impossible")
+
+        self._expansions = expansions
+        self._update()
+
+    def _update_list_base(self, val):
+        self._list_base = val
+        if val == "b-adics":
+            self._expansions = Expansions.from_integer(int(val), self.EXPANSIONS_LENGTH)
+            self._update()
+
+    def _update_embedding_base(self, val):
+        self._embedding_base = val
+        self._update()
+
+    def _update_geometric(self, val):
+        self._geometric = val
+        self._update(False)
+
+    def _update(self, regen_angles: bool = True, set_data: bool = True):
+        if regen_angles:
+            self._angles = np.array([
+                _adic_angles(self._embedding_base, expansion, 15) for expansion in self._expansions.xs
+            ])
+        # TODO: vectorize
+        self._embedding: NDArray[np.complex64] = np.array(
+            [geosum(self._geometric, phis) for phis in self._angles]
+        )
+
+        if set_data:
+            self._set_data()
+
+    def _set_data(self):
+        self._plot.set_data(self._embedding.real, self._embedding.imag)  # type: ignore
+        self._axes["plot"].set_xlim(
+            min(self._embedding.real),
+            max(self._embedding.real)
+        )
+        self._axes["plot"].set_ylim(
+            min(self._embedding.imag),
+            max(self._embedding.imag)
+        )
+        self._axes["plot"].get_figure().canvas.draw_idle()
+
+
+def interactive(expansions: Expansions):
+    PadicPlot(expansions=expansions, interactive=True)
     plt.show()
 
 
@@ -179,13 +298,6 @@ def anim(expansions: Expansions):
     )
 
 
-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()
@@ -195,7 +307,13 @@ def main():
         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(
+        "--base",
+        "-p",
+        default=2,
+        type=int,
+        help="Use p-adic integers"
+    )
     parser.add_argument(
         "--count",
         "-c",