Deploy app with precomputed results for p=15,23,29,31

.gitattributes

@@ -33,3 +33,44 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_015/p015_lineplot_in.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_015/p015_lineplot_out.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_015/p015_phase_align_approx1.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_015/p015_phase_align_approx2.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_015/p015_single_freq_quad.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_015/p015_single_freq_relu.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_full_training_para_origin.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_grokk_decoded_weights_dynamic.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_grokk_memorization_common_to_rare.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_lineplot_in.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_lineplot_out.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_output_logits.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_overview_loss_ipr.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_phase_align_approx1.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_phase_align_approx2.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_single_freq_quad.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_023/p023_single_freq_relu.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_full_training_para_origin.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_grokk_decoded_weights_dynamic.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_grokk_memorization_accuracy.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_grokk_memorization_common_to_rare.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_lineplot_in.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_lineplot_out.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_output_logits.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_overview_loss_ipr.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_phase_align_approx1.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_phase_align_approx2.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_single_freq_quad.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_029/p029_single_freq_relu.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_full_training_para_origin.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_grokk_decoded_weights_dynamic.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_grokk_memorization_accuracy.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_grokk_memorization_common_to_rare.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_lineplot_in.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_lineplot_out.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_output_logits.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_overview_loss_ipr.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_phase_align_approx1.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_phase_align_approx2.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_single_freq_quad.png filter=lfs diff=lfs merge=lfs -text
+precomputed_results/p_031/p031_single_freq_relu.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,18 @@
+src/wandb/
+notebooks/simulate_dynamics.ipynb
+
+# Claude AI files
+.claude/
+
+# Python cache files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Model checkpoints (too large for git; regenerate with precompute/run_pipeline.sh)
+trained_models/
+saved_models/
+
+# OS files
+.DS_Store
+tmp/

README.md

@@ -1,14 +1,202 @@
+# On the Mechanism and Dynamics of Modular Addition
+
+### Fourier Features, Lottery Ticket, and Grokking
+
+**Jianliang He, Leda Wang, Siyu Chen, Zhuoran Yang**
+*Department of Statistics and Data Science, Yale University*
+
+[[arXiv (coming soon)](#)] [[Blog (coming soon)](#)] [[Interactive Demo](https://huggingface.co/spaces/y-agent/modular-addition-feature-learning)]
+
 ---
-title: Modular Addition Feature Learning
-emoji: 😻
-colorFrom: green
-colorTo: green
-sdk: gradio
-sdk_version: 6.6.0
-app_file: app.py
-pinned: false
-license: mit
-short_description: Interactive Demo of Paper on Modular Addition
----
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## Overview
+
+This repository provides the code for studying how a two-layer neural network learns modular arithmetic $f(x,y) = (x+y) \bmod p$. We analyze three phenomena:
+
+1. **Fourier Feature Learning** — Each neuron independently discovers a cosine wave at a single frequency, collectively implementing a discrete Fourier transform that the network was never taught.
+2. **Lottery Ticket Dynamics** — Random initialization determines which frequency each neuron will specialize in: the frequency with the best initial phase alignment wins a winner-take-all competition.
+3. **Grokking** — Under partial data with weight decay, the network first memorizes, then suddenly generalizes through a three-stage process: memorization → sparsification → cleanup.
+
+## Interactive Demo
+
+An interactive Gradio app visualizes all results with math explanations and interactive Plotly charts:
+
+- **9 analysis tabs** covering mechanism, dynamics, grokking, and analytical simulations
+- **Interactive features**: neuron frequency inspector, logit explorer, grokking epoch slider
+- **On-demand training**: generate results for any odd $p \geq 3$ directly from the app
+- **Pre-computed examples** included for $p = 15, 23, 29, 31$
+
+### Launch Locally
+
+```bash
+pip install -r requirements.txt
+python hf_app/app.py
+# Opens at http://localhost:7860
+```
+
+### Deploy to Hugging Face Spaces
+
+1. Create a new Space at [huggingface.co/new-space](https://huggingface.co/new-space) (SDK: Gradio)
+2. Push the repo:
+   ```bash
+   git remote add hf https://huggingface.co/spaces/YOUR_USERNAME/YOUR_SPACE_NAME
+   git push hf main
+   ```
+3. The app reads from `precomputed_results/` — the included examples (p=15, 23, 29, 31) work out of the box
+4. Users can generate results for additional $p$ values on-demand via the "Generate" button. New results are auto-committed back to the Space repo so they persist.
+
+> **Tip:** For GPU-accelerated on-demand training, select a GPU runtime in your Space settings.
+
+## Pre-computation Pipeline
+
+The `precompute/` directory trains 5 model configurations per modulus and generates all plots + interactive JSON data. See [`precompute/README.md`](precompute/README.md) for full documentation.
+
+### Quick Start
+
+```bash
+# Full pipeline for a single modulus (train → plots → analytical → verify)
+bash precompute/run_pipeline.sh 23
+
+# With custom d_mlp
+bash precompute/run_pipeline.sh 23 --d_mlp 128
+
+# Delete checkpoints after generating plots (saves disk space)
+CLEANUP=1 bash precompute/run_pipeline.sh 23
+
+# Batch: all odd p in [3, 99]
+bash precompute/run_all.sh
+
+# Or up to p=199
+MAX_P=199 bash precompute/run_all.sh
+```
+
+### Manual Steps
+
+```bash
+# Step 1: Train all 5 configurations
+python precompute/train_all.py --p 23 --output ./trained_models --resume
+
+# Step 2: Generate model-based plots (21 PNGs + 7 JSONs)
+python precompute/generate_plots.py --p 23 --input ./trained_models --output ./precomputed_results
+
+# Step 3: Generate analytical simulation plots (2 PNGs, no model needed)
+python precompute/generate_analytical.py --p 23 --output ./precomputed_results
+```
+
+### Output
+
+Each modulus produces ~33 files in `precomputed_results/p_XXX/`:
+
+| Category | Files | Description |
+|----------|-------|-------------|
+| Overview (Tab 1) | 2 PNGs + 1 JSON | Loss, IPR, phase scatter |
+| Fourier Weights (Tab 2) | 3 PNGs + 1 JSON | DFT heatmaps, cosine fits, neuron spectra |
+| Phase Analysis (Tab 3) | 3 PNGs | Phase distribution, alignment, magnitudes |
+| Output Logits (Tab 4) | 1 PNG + 1 JSON | Logit heatmap, interactive explorer |
+| Lottery Mechanism (Tab 5) | 3 PNGs | Magnitude race, phase convergence, contour |
+| Grokking (Tab 6) | 5 PNGs + 3 JSONs | Loss/acc curves, memorization, weight evolution |
+| Gradient Dynamics (Tab 7) | 4 PNGs | Phase alignment + DFT for Quad and ReLU |
+| Decoupled Simulation (Tab 8) | 2 PNGs | Analytical ODE integration |
+| Metadata | 2 JSONs | Config + training log |
+
+> **Note:** Grokking results (Tab 6) require $p \geq 19$. Smaller values of $p$ have too few data points for a meaningful train/test split.
+
+## The 5 Training Configurations
+
+| Config | Activation | Optimizer | LR | Weight Decay | Data | Epochs | Used In |
+|--------|-----------|-----------|-----|-------------|------|--------|---------|
+| `standard` | ReLU | AdamW | 5e-5 | 0 | 100% | 5,000 | Tabs 1–4 |
+| `grokking` | ReLU | AdamW | 1e-4 | 2.0 | 75% | 50,000 | Tabs 1, 6 |
+| `quad_random` | Quad | AdamW | 5e-5 | 0 | 100% | 5,000 | Tab 5 |
+| `quad_single_freq` | Quad | SGD | 0.1 | 0 | 100% | 5,000 | Tab 7 |
+| `relu_single_freq` | ReLU | SGD | 0.01 | 0 | 100% | 5,000 | Tab 7 |
+
+## Running a Single Experiment
+
+For custom experiments outside the pre-computation pipeline:
+
+```bash
+cd src
+
+# Train with default config (p=97, d_mlp=1024, ReLU, 5000 epochs)
+python module_nn.py
+
+# Train with specific parameters
+python module_nn.py --p 23 --d_mlp 512 --num_epochs 5000 --lr 5e-5
+
+# Dry run: see config without training
+python module_nn.py --dry_run --p 23 --d_mlp 512
+```
+
+## Notebooks
+
+Interactive analysis notebooks in `notebooks/`:
+
+| Notebook | Description |
+|----------|-------------|
+| `empirical_insight_standard.ipynb` | Fourier weight analysis, phase distributions, output logits |
+| `empirical_insight_grokk.ipynb` | Grokking stages, weight dynamics, IPR evolution |
+| `lottery_mechanism.ipynb` | Neuron specialization, frequency magnitude/phase tracking |
+| `interprete_gd_dynamics.ipynb` | Phase alignment under single-frequency initialization |
+| `decouple_dynamics_simulation.ipynb` | Analytical gradient flow simulation |
+
+## Setup
+
+### Requirements
+
+- Python 3.8+
+- PyTorch 2.0+
+- CUDA-capable GPU (recommended for $p > 50$; CPU works for small $p$)
+
+### Installation
+
+```bash
+git clone https://github.com/Y-Agent/modular-addition-feature-learning.git
+cd modular-addition-feature-learning
+pip install -r requirements.txt
+```
+
+## Project Structure
+
+```
+modular-addition-feature-learning/
+├── src/                          # Core source code
+│   ├── module_nn.py             # Training script with CLI
+│   ├── nnTrainer.py             # Training loop and optimization
+│   ├── model_base.py            # Neural network architecture (EmbedMLP)
+│   ├── mechanism_base.py        # Fourier analysis and decomposition
+│   ├── utils.py                 # Configuration and helpers
+│   └── configs.yaml             # Default hyperparameters
+├── precompute/                   # Batch training and plot generation
+│   ├── run_pipeline.sh          # Full pipeline for one modulus
+│   ├── run_all.sh               # Batch pipeline for all odd p
+│   ├── train_all.py             # Train 5 configurations
+│   ├── generate_plots.py        # Generate model-based plots + JSONs
+│   ├── generate_analytical.py   # Analytical ODE simulation plots
+│   └── prime_config.py          # Configurations and sizing formula
+├── hf_app/                       # Gradio web application
+│   └── app.py                   # Interactive visualization app
+├── precomputed_results/          # Pre-computed plots and data
+│   ├── p_015/                   # Results for p=15
+│   ├── p_023/                   # Results for p=23
+│   ├── p_029/                   # Results for p=29
+│   └── p_031/                   # Results for p=31
+├── notebooks/                    # Analysis and visualization notebooks
+├── requirements.txt              # Python dependencies
+└── README.md
+```
+
+## Citation
+
+```bibtex
+@article{he2025modular,
+  title={On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking},
+  author={He, Jianliang and Wang, Leda and Chen, Siyu and Yang, Zhuoran},
+  journal={arXiv preprint arXiv:XXXX.XXXXX},
+  year={2025}
+}
+```
+
+## License
+
+[MIT License](LICENSE)

hf_app/app.py

@@ -0,0 +1,1375 @@
+#!/usr/bin/env python3
+"""
+Gradio app for Modular Addition Feature Learning visualization.
+Serves pre-computed results for odd moduli p in [3, 199].
+
+All results are pre-computed as PNG images and JSON data files.
+No GPU needed at serving time.
+
+Tab structure:
+  Core Interpretability:
+    1. Training Overview    -- loss + IPR sparsity
+    2. Fourier Weights      -- decoded W_in/W_out heatmaps + line plots + neuron inspector
+    3. Phase Analysis       -- phase distribution, 2phi vs psi, magnitudes
+    4. Output Logits        -- predicted logit heatmap + interactive logit explorer
+    5. Lottery Mechanism    -- neuron specialization, magnitude/phase, contour
+  Grokking:
+    6. Grokking             -- loss/acc, phase alignment, IPR, memorization, epoch slider
+  Theory:
+    7. Gradient Dynamics    -- phase alignment for Quad & ReLU single-freq init
+    8. Decoupled Simulation -- analytical gradient flow (no model needed)
+  Diagnostics:
+    9. Training Log         -- per-run hyperparameters and epoch-by-epoch metrics
+"""
+import gradio as gr
+import json
+import logging
+import os
+import shutil
+import subprocess
+import sys
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+# Force pandas to be fully imported before plotly lazily imports it
+# (avoids "partially initialized module 'pandas'" in threaded callbacks)
+import pandas  # noqa: F401
+import plotly.graph_objects as go
+
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+RESULTS_DIR = os.path.join(PROJECT_ROOT, "precomputed_results")
+TRAINED_MODELS_DIR = os.path.join(PROJECT_ROOT, "trained_models")
+
+# Max p for on-demand training (d_mlp grows as O(p^2), memory limit)
+MAX_P_ON_DEMAND = 97
+
+COLORS = ['#0D2758', '#60656F', '#DEA54B', '#A32015', '#347186']
+STAGE_COLORS = ['rgba(212,175,55,0.15)', 'rgba(139,115,85,0.15)', 'rgba(192,192,192,0.15)']
+
+# KaTeX delimiters for Gradio Markdown
+LATEX_DELIMITERS = [
+    {"left": "$$", "right": "$$", "display": True},
+    {"left": "$", "right": "$", "display": False},
+]
+
+# Custom CSS for Palatino font and styling
+CUSTOM_CSS = r"""
+@import url('https://fonts.googleapis.com/css2?family=Libre+Baskerville:ital,wght@0,400;0,700;1,400&display=swap');
+
+* {
+    font-family: "Palatino Linotype", "Book Antiqua", Palatino, "Libre Baskerville", Georgia, serif !important;
+}
+code, pre, .code, .monospace {
+    font-family: "Menlo", "Consolas", "Monaco", monospace !important;
+}
+.katex, .katex * {
+    font-family: KaTeX_Main, "Times New Roman", serif !important;
+}
+h1 {
+    font-family: "Palatino Linotype", "Book Antiqua", Palatino, "Libre Baskerville", Georgia, serif !important;
+    text-align: center !important;
+    margin-bottom: 0.1em !important;
+}
+h3 {
+    font-family: "Palatino Linotype", "Book Antiqua", Palatino, "Libre Baskerville", Georgia, serif !important;
+    text-align: center !important;
+    color: var(--neutral-500) !important;
+    font-weight: normal !important;
+    margin-top: 0 !important;
+}
+h2, h4 {
+    font-family: "Palatino Linotype", "Book Antiqua", Palatino, "Libre Baskerville", Georgia, serif !important;
+}
+blockquote {
+    border-left: 3px solid var(--color-accent) !important;
+    background-color: var(--block-background-fill) !important;
+    padding: 0.5em 1em !important;
+    margin: 0.5em 0 !important;
+}
+"""
+
+# ---------------------------------------------------------------------------
+# Math explanation text for each tab (following the paper precisely)
+# ---------------------------------------------------------------------------
+
+MATH_TAB1 = r"""
+### Overview
+
+We study how a two-layer neural network learns to compute modular addition $f(x,y) = (x+y) \bmod p$. The network has $M$ hidden neurons. Each input integer $x$ is represented as a one-hot vector, and the network produces a score for each of the $p$ possible answers. During training, the network learns two weight vectors per neuron: an **input weight** $\theta_m$ and an **output weight** $\xi_m$, both vectors of length $p$.
+
+#### Two Training Setups
+
+1. **Full-data (Tabs 1--5, 7).** Train on all $p^2$ input pairs with no held-out data and no regularization. This produces clean features ideal for studying what the network learns and how.
+
+2. **Grokking (Tab 6).** Train on only 75% of input pairs with weight decay $\lambda = 2.0$ (a penalty that shrinks weights over time). These two ingredients -- incomplete data + weight decay -- cause the network to first memorize, then suddenly generalize, a phenomenon called **grokking**.
+
+#### What the Network Learns
+
+Each neuron's weight vectors turn into **cosine waves** at a single frequency -- the network independently rediscovers the Discrete Fourier Transform. The neurons collectively cover all frequencies with balanced strengths, enabling them to "vote" together and identify the correct answer $(x+y) \bmod p$.
+
+#### How It Learns (Dynamics)
+
+Frequencies **compete** within each neuron during training. The frequency whose input and output phases happen to start best-aligned grows fastest -- a **lottery ticket mechanism** where the random initialization determines the outcome before training begins.
+
+#### Grokking (Three Stages)
+
+When trained on partial data with weight decay: **(I) Memorization** -- the network fits the training data using noisy, multi-frequency features. **(II) Generalization** -- weight decay prunes away the noise, leaving clean single-frequency features; test accuracy jumps. **(III) Cleanup** -- weight decay slowly polishes the solution.
+
+#### Progress Measures on These Plots
+
+- **Loss**: Cross-entropy loss (lower = better predictions). We show both training loss and test loss.
+
+- **IPR (Inverse Participation Ratio)**: Measures how concentrated a neuron's energy is across frequencies. We decompose each neuron's weights into Fourier components, measure the strength $A_k$ at each frequency $k$, and compute:
+
+$$\text{IPR} = \frac{\sum_k A_k^4}{\left(\sum_k A_k^2\right)^2}.$$
+
+When a neuron uses only **one frequency**, IPR $= 1$ (fully specialized). When energy is spread across **many frequencies**, IPR is close to $0$. Watching IPR rise toward 1 during training shows the network specializing.
+
+- **Phase scatter**: Each neuron has an input phase $\phi_m$ and output phase $\psi_m$. The theory predicts the output phase equals twice the input phase ($\psi_m = 2\phi_m$). The scatter plot checks this: all points should fall on the diagonal.
+"""
+
+MATH_TAB2 = r"""
+### Every Neuron is a Cosine Wave
+> **Setup:** ReLU activation, full data, no weight decay.
+
+After training, each neuron's weight vectors become clean **cosine waves** at a single frequency. Concretely, the input weight of neuron $m$ looks like:
+
+$$\underbrace{\theta_m[j]}_{\text{input weight at position } j} = \underbrace{\alpha_m}_{\text{input magnitude}} \cdot \cos\!\left(\underbrace{\frac{2\pi k}{p}}_{\text{frequency}} \cdot j + \underbrace{\phi_m}_{\text{input phase}}\right),$$
+
+and the output weight has the same form with its own magnitude $\beta_m$ (output magnitude) and phase $\psi_m$ (output phase). Each neuron picks **one frequency** $k$ out of the $(p{-}1)/2$ possible frequencies. No one told the network about Fourier analysis -- it rediscovered this representation on its own through training.
+
+**Heatmap**: Each row is a neuron, each column is a Fourier component (cosine and sine at each frequency). If a row has only one bright cell, that neuron is using a single frequency -- and that's exactly what we see.
+
+**Line Plots**: The dots are the actual learned weights; the dashed curves are best-fit cosines. The near-perfect fits confirm each neuron is well-described by a single cosine at a single frequency.
+
+**Neuron Inspector**: Select a neuron from the dropdown to see how its energy is distributed across all frequencies (for both input and output weights).
+"""
+
+MATH_TAB3 = r"""
+### Phase Alignment and Collective Diversification
+> **Setup:** ReLU activation, full data, no weight decay.
+
+#### The Input and Output Phases Lock Together
+
+Each neuron has an input phase $\phi_m$ and an output phase $\psi_m$ (the "shift" of each cosine wave). These are not independent -- training drives them into a precise relationship:
+
+$$\underbrace{\psi_m}_{\text{output phase}} = 2 \times \underbrace{\phi_m}_{\text{input phase}}.$$
+
+**Why "doubled"?** The activation function squares (or, for ReLU, roughly squares) the sum of two cosines. Squaring a cosine at phase $\phi$ naturally produces terms at phase $2\phi$. The output layer learns to match this by setting its own phase to $2\phi$, so the two layers work together coherently.
+
+The **scatter plot** checks this: we plot $2\phi_m$ (horizontal) vs. $\psi_m$ (vertical) for every neuron. If the relationship holds, all points land on the diagonal. This relationship is not built into the architecture -- it **emerges from training** (see Tab 7 for why).
+
+#### Neurons Organize Themselves into a Balanced Ensemble
+
+The neurons don't just specialize to single frequencies -- they also organize *collectively*:
+
+1. **Frequency balance:** Every frequency gets roughly the same number of neurons.
+2. **Phase spread:** Within each frequency group, the phases are spread uniformly around the circle. This is what enables **noise cancellation** -- the random noise from individual neurons averages out when their phases are evenly spaced.
+3. **Magnitude balance:** All neurons contribute roughly equally to the output (no single neuron dominates).
+
+The **polar plot** shows phases at multiples ($1\times, 2\times, 3\times, 4\times$) on concentric rings -- uniform spread confirms the cancellation condition. The **violin plots** show the distribution of input magnitudes ($\alpha$) and output magnitudes ($\beta$) -- tight concentration confirms magnitude balance.
+"""
+
+MATH_TAB4 = r"""
+### The Mechanism: Majority Voting in Fourier Space
+> **Setup:** ReLU activation, full data, no weight decay.
+
+#### How Neurons Vote for the Correct Answer
+
+Each neuron produces a score for every possible output $j \in \{0, 1, \ldots, p{-}1\}$. Thanks to the phase alignment ($\psi = 2\phi$, see Tab 3), each neuron's score has a **signal** component that peaks at the correct answer $j = (x+y) \bmod p$, plus **noise** that depends on that neuron's particular phase.
+
+When we sum over many neurons within a frequency group, the signal adds up (every neuron agrees on the right answer) while the noise cancels out (different neurons have different phases, and the phase spread from Tab 3 ensures the noise averages to zero). This is **majority voting** -- each neuron casts a noisy vote, but the consensus is correct.
+
+#### The "Flawed Indicator"
+
+After summing over all neurons and all frequency groups, the network's output simplifies to:
+
+$$\text{score for answer } j \;\propto\; \underbrace{\frac{p}{2} \cdot \mathbf{1}[j = (x{+}y) \bmod p]}_{\text{correct answer (strongest)}} \;+\; \underbrace{\frac{p}{4} \cdot \bigl(\mathbf{1}[j = 2x \bmod p] + \mathbf{1}[j = 2y \bmod p]\bigr)}_{\text{two "ghost" peaks (half strength)}}.$$
+
+The correct answer gets score $p/2$, but two **spurious ghost peaks** appear at $2x \bmod p$ and $2y \bmod p$ with score $p/4$. The correct answer always wins because $p/2 > p/4$, so the network always predicts correctly despite the ghosts.
+
+**Heatmap**: The network's output scores for all inputs with $x = 0$. The bright diagonal is the correct answer. The faint lines are the ghost peaks.
+
+**Logit Explorer**: Pick an input pair $(x, y)$ to see the full score distribution. The correct answer (highlighted) should be the tallest bar.
+"""
+
+MATH_TAB5 = r"""
+### The Lottery Ticket: How Each Neuron Picks Its Frequency
+> **Setup:** Quadratic activation ($\sigma(x) = x^2$), full data, random initialization.
+
+#### The Competition
+
+At the start of training, every neuron has a tiny bit of energy at **every** frequency -- nothing is specialized yet. But the input and output phases at each frequency start at random values, so some frequencies happen to be better aligned (input phase and output phase closer to the $\psi = 2\phi$ relationship) than others.
+
+The key insight: **a frequency grows faster when its phases are better aligned.** The growth rate of a frequency's magnitude depends on how close it is to alignment:
+
+$$\text{growth rate} \;\propto\; \cos(\underbrace{2\phi - \psi}_{\text{phase misalignment }\mathcal{D}}).$$
+
+When the misalignment $\mathcal{D}$ is small (phases nearly aligned), $\cos(\mathcal{D}) \approx 1$ and the frequency grows quickly. When $\mathcal{D}$ is large, growth stalls.
+
+#### Winner Takes All
+
+This creates a **positive feedback loop**: the best-aligned frequency grows a little, which helps it align even better, which makes it grow even faster. The gap compounds exponentially until one frequency completely dominates -- **the winner takes all.**
+
+The winning frequency is simply the one that started closest to alignment:
+
+$$\text{winning frequency} = \text{the } k \text{ with smallest initial misalignment } |\mathcal{D}_m^k|.$$
+
+This is a **lottery ticket**: the outcome is determined by the random initialization before training even begins. Since each neuron draws independent random phases, different neurons pick different winning frequencies, naturally producing the balanced frequency coverage seen in Tab 3.
+
+**Phase plot:** Shows how the misalignment $\mathcal{D}$ evolves over training for each frequency within one neuron. The winner (red) converges to zero first; the others barely move.
+
+**Magnitude plot:** Shows how the output magnitude $\beta$ (strength of each frequency) evolves. All start equal. Once the winner aligns, it grows explosively while the others stay frozen.
+
+**Contour plot:** Final magnitude as a function of (initial magnitude, initial misalignment). Largest values appear at small misalignment -- confirming that alignment determines the winner.
+"""
+
+MATH_TAB6 = r"""
+### Grokking: From Memorization to Generalization
+> **Setup:** ReLU activation, 75% training fraction, weight decay $\lambda = 2.0$.
+
+Under the train-test split setup, the network quickly memorizes the training set but takes much longer to generalize. Our analysis reveals grokking is a **three-stage process**, each driven by a different balance of forces.
+
+**Stage I -- Memorization (loss gradient dominates).** The loss gradient dominates and the network rapidly memorizes training data. Training accuracy reaches 100% while test accuracy reaches only ~70%. The ~70% figure (not ~50%) arises because the architecture is symmetric in $x$ and $y$: since $\theta_m[x] + \theta_m[y]$ is invariant under swapping $(x,y) \leftrightarrow (y,x)$, memorizing $(x,y)$ automatically gives the correct answer for $(y,x)$. The lottery mechanism runs on incomplete data, producing a "noisy" multi-frequency representation. We also observe a **common-to-rare ordering**: the network first memorizes symmetric pairs (both $(i,j)$ and $(j,i)$ in training) while actively *suppressing* rare pairs, before eventually memorizing them too.
+
+**Stage II -- Fast Generalization (loss + weight decay).** Weight decay penalizes all magnitudes equally, but the dominant frequency has much larger magnitude and can "afford" the penalty. Non-feature frequencies are driven to zero -- a **sparsification** effect visible as the sharp IPR increase. This transforms the noisy memorization solution into clean single-frequency-per-neuron features. Test accuracy jumps steeply.
+
+**Stage III -- Slow Cleanup (weight decay dominates).** The loss gradient becomes negligible (both losses $\approx 0$). Weight decay alone slowly shrinks norms at rate $\partial_t \|w\| = -\lambda \|w\|$. The feature frequencies are already identified; this stage fine-tunes magnitudes. The network transitions from a lookup table to a generalizing algorithm implementing the indicator function from the mechanism (Tab 4).
+
+**Four progress measures**: (a) Loss -- train drops in Stage I, test drops in Stage II. (b) Accuracy -- train reaches 100% early, test jumps in Stage II. (c) Phase alignment -- $|\sin(\mathcal{D}_m^\star)|$ decreases throughout. (d) IPR + parameter norms -- IPR increases sharply in Stage II, norms shrink in Stage III.
+
+**Epoch Slider**: Use the slider below to see how the accuracy grid evolves across the three stages.
+"""
+
+MATH_TAB7 = r"""
+### Training Dynamics: Phase Alignment and Single-Frequency Preservation
+> **Setup:** Quadratic and ReLU activations, full data, single-frequency initialization, SGD.
+
+#### The Four-Variable ODE
+
+Under small initialization ($\kappa_{\mathrm{init}} \ll 1$), the dynamics decouple: each neuron evolves independently, and within each neuron, different Fourier modes evolve independently (because $\sum_{x \in \mathbb{Z}_p} \cos(\omega_k x) \cos(\omega_\tau x) = \frac{p}{2}\delta_{k,\tau}$). The full dynamics reduce to independent four-variable ODEs per (neuron, frequency):
+
+$$\partial_t \alpha \approx 2p \cdot \alpha \cdot \beta \cdot \cos(\mathcal{D}), \qquad \partial_t \beta \approx p \cdot \alpha^2 \cdot \cos(\mathcal{D}),$$
+$$\partial_t \phi \approx 2p \cdot \beta \cdot \sin(\mathcal{D}), \qquad \partial_t \psi \approx -p \cdot \frac{\alpha^2}{\beta} \cdot \sin(\mathcal{D}),$$
+
+where $\mathcal{D} = (2\phi - \psi) \bmod 2\pi$ is the **phase misalignment**. This system has a clear physical interpretation: **magnitudes grow when phases are aligned** ($\cos(\mathcal{D}) \approx 1$), and **phases rotate toward alignment** ($\sin(\mathcal{D}) \to 0$). The dynamics self-coordinate: phases align first (while magnitudes are small), then magnitudes explode.
+
+#### Phase Alignment Theorem
+
+$\mathcal{D}(t) \to 0$ from any initial condition except the measure-zero unstable point $\mathcal{D} = \pi$. The dynamics on the circle behave like an **overdamped pendulum**: $\mathcal{D} = 0$ is a stable attractor, $\mathcal{D} = \pi$ is an unstable repeller. This is not a coincidence or a property of initialization -- it is an **inevitable consequence of the training dynamics**. It explains Observation 2 ($\psi = 2\phi$).
+
+#### Single-Frequency Preservation Theorem
+
+Under the decoupled flow, if a neuron starts at a single frequency, it remains there for all time. The Fourier orthogonality on $\mathbb{Z}_p$ prevents energy from leaking between modes.
+
+**Quadratic** (left panels): Theory matches experiment almost exactly. The DFT heatmap shows the dominant frequency growing while all others stay at zero.
+
+**ReLU** (right panels): Same qualitative behavior with minor quantitative differences. Small energy "leaks" to harmonic multiples ($3k^\star, 5k^\star, \ldots$ for input; $2k^\star, 3k^\star, \ldots$ for output). The leakage decays as $O(r^{-2})$ where $r$ is the harmonic order (third harmonic has $1/9$ the strength, fifth has $1/25$), keeping the dominant frequency overwhelmingly dominant.
+"""
+
+MATH_TAB9 = r"""
+### Training Log
+
+This tab shows the training logs for each of the 5 configurations run for the selected modulo $p$. Select a run from the dropdown to view its hyperparameters and per-epoch training metrics.
+
+The 5 training runs are:
+- **standard**: ReLU, full data, no weight decay -- produces the clean Fourier features analyzed in Tabs 1--5
+- **grokking**: ReLU, 75% data, weight decay $\lambda = 2.0$ -- demonstrates the memorization $\to$ generalization transition (Tab 6)
+- **quad_random**: Quadratic activation, full data, random init -- the lottery ticket mechanism (Tab 5)
+- **quad_single_freq**: Quadratic activation, single-frequency init, SGD -- verifies single-frequency preservation (Tab 7)
+- **relu_single_freq**: ReLU, single-frequency init, SGD -- ReLU variant of the dynamics (Tab 7)
+"""
+
+MATH_TAB8 = r"""
+### Decoupled Gradient Flow Simulation
+> **Setup:** Analytical ODE integration (no neural network training).
+
+This tab shows a pure mathematical simulation of the multi-frequency gradient flow, **without training any neural network**. We numerically integrate the four-variable ODEs for all frequency modes simultaneously within a single neuron:
+
+$$\partial_t \alpha_k \approx 2p \cdot \alpha_k \cdot \beta_k \cdot \cos(\mathcal{D}_k), \qquad \partial_t \beta_k \approx p \cdot \alpha_k^2 \cdot \cos(\mathcal{D}_k),$$
+$$\partial_t \phi_k \approx 2p \cdot \beta_k \cdot \sin(\mathcal{D}_k), \qquad \partial_t \psi_k \approx -p \cdot \frac{\alpha_k^2}{\beta_k} \cdot \sin(\mathcal{D}_k),$$
+
+for each frequency $k = 1, \ldots, (p{-}1)/2$, with random initial conditions.
+
+The simulation confirms the theoretical predictions from Tab 7:
+
+- **Phase alignment:** Phase misalignments $\mathcal{D}_k = (2\phi_k - \psi_k) \bmod 2\pi$ converge to $0$ for most frequencies, or linger near $\pi$ (the unstable repeller) before eventually escaping.
+- **Magnitude competition:** Magnitudes grow explosively for the frequency where $\mathcal{D}_k \approx 0$ first, while others remain near their initial level.
+- **Lottery outcome:** The winning frequency (smallest initial $\mathcal{D}_k$) dominates all others, reproducing the full lottery ticket mechanism without any neural network -- just ODEs.
+
+Two cases are shown with different initial conditions to illustrate that the mechanism is robust: regardless of the random starting point, the frequency with the best initial phase alignment always wins.
+"""
+
+
+# ---------------------------------------------------------------------------
+# Data loading helpers
+# ---------------------------------------------------------------------------
+
+MIN_P = 3  # p=2 has 0 non-DC Fourier frequencies; analysis is degenerate
+
+
+def get_available_moduli():
+    """Discover which p values have pre-computed results (odd p >= 3)."""
+    moduli = []
+    if os.path.exists(RESULTS_DIR):
+        for d in sorted(os.listdir(RESULTS_DIR)):
+            if d.startswith("p_"):
+                try:
+                    p = int(d.split("_")[1])
+                    if p >= MIN_P:
+                        moduli.append(p)
+                except ValueError:
+                    pass
+    return moduli
+
+
+def _prime_dir(p):
+    return os.path.join(RESULTS_DIR, f"p_{p:03d}")
+
+
+def load_json_file(p, filename):
+    """Load a JSON file from the prime's directory."""
+    path = os.path.join(_prime_dir(p), f"p{p:03d}_{filename}")
+    if os.path.exists(path):
+        with open(path) as f:
+            return json.load(f)
+    return None
+
+
+def safe_img(p, filename):
+    """Return image path or None (Gradio handles None gracefully)."""
+    path = os.path.join(_prime_dir(p), f"p{p:03d}_{filename}")
+    return path if os.path.exists(path) else None
+
+
+# ---------------------------------------------------------------------------
+# Interactive Plotly chart builders
+# ---------------------------------------------------------------------------
+
+def _to_np(v):
+    """Convert a list/value to a numpy array (bypasses plotly's pandas check)."""
+    if v is None:
+        return None
+    return np.asarray(v)
+
+
+def make_loss_chart(data, title="Training Loss"):
+    """Build an interactive Plotly loss chart from JSON data."""
+    if data is None:
+        return None
+    fig = go.Figure()
+    n = len(data.get('train_losses', []))
+    epochs = np.arange(n)
+
+    fig.add_trace(go.Scatter(
+        x=epochs, y=_to_np(data['train_losses']),
+        name='Train Loss', line=dict(color=COLORS[0]),
+    ))
+    if 'test_losses' in data:
+        fig.add_trace(go.Scatter(
+            x=epochs, y=_to_np(data['test_losses']),
+            name='Test Loss', line=dict(color=COLORS[3]),
+        ))
+
+    s1 = data.get('stage1_end')
+    s2 = data.get('stage2_end')
+    if s1 is not None:
+        fig.add_vrect(x0=0, x1=s1, fillcolor=STAGE_COLORS[0],
+                      line_width=0, annotation_text="Memorization",
+                      annotation_position="top left")
+    if s1 is not None and s2 is not None:
+        fig.add_vrect(x0=s1, x1=s2, fillcolor=STAGE_COLORS[1],
+                      line_width=0, annotation_text="Transition",
+                      annotation_position="top left")
+    if s2 is not None:
+        fig.add_vrect(x0=s2, x1=n, fillcolor=STAGE_COLORS[2],
+                      line_width=0, annotation_text="Generalization",
+                      annotation_position="top left")
+
+    fig.update_layout(
+        title=title, xaxis_title='Epoch', yaxis_title='Loss',
+        template='plotly_white', height=400,
+        legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.99),
+    )
+    return fig
+
+
+def make_acc_chart(data, title="Training Accuracy"):
+    """Build an interactive Plotly accuracy chart."""
+    if data is None:
+        return None
+    fig = go.Figure()
+    epochs = _to_np(data.get('epochs', list(range(len(data.get('train_accs', []))))))
+
+    fig.add_trace(go.Scatter(
+        x=epochs, y=_to_np(data['train_accs']),
+        name='Train Acc', line=dict(color=COLORS[0]),
+    ))
+    if 'test_accs' in data:
+        fig.add_trace(go.Scatter(
+            x=epochs, y=_to_np(data['test_accs']),
+            name='Test Acc', line=dict(color=COLORS[3]),
+        ))
+
+    s1 = data.get('stage1_end')
+    s2 = data.get('stage2_end')
+    if s1 is not None:
+        fig.add_vrect(x0=0, x1=s1, fillcolor=STAGE_COLORS[0], line_width=0)
+    if s1 is not None and s2 is not None:
+        fig.add_vrect(x0=s1, x1=s2, fillcolor=STAGE_COLORS[1], line_width=0)
+    if s2 is not None:
+        n = int(epochs.max()) if len(epochs) > 0 else len(data.get('train_accs', []))
+        fig.add_vrect(x0=s2, x1=n, fillcolor=STAGE_COLORS[2], line_width=0)
+
+    fig.update_layout(
+        title=title, xaxis_title='Epoch', yaxis_title='Accuracy',
+        template='plotly_white', height=400,
+        legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.99),
+    )
+    return fig
+
+
+
+def make_neuron_spectrum_chart(data, neuron_key):
+    """Build a Plotly bar chart for a single neuron's Fourier spectrum."""
+    if data is None or neuron_key not in data.get('neurons', {}):
+        return None
+    neuron = data['neurons'][neuron_key]
+    names = data.get('fourier_basis_names', [])
+    mags_in = _to_np(neuron['fourier_magnitudes_in'])
+    mags_out = _to_np(neuron['fourier_magnitudes_out'])
+    dom_freq = neuron.get('dominant_freq', '?')
+
+    fig = go.Figure()
+    fig.add_trace(go.Bar(
+        x=names, y=mags_in, name='W_in magnitude',
+        marker_color=COLORS[0], opacity=0.8,
+    ))
+    fig.add_trace(go.Bar(
+        x=names, y=mags_out, name='W_out magnitude',
+        marker_color=COLORS[3], opacity=0.8,
+    ))
+    fig.update_layout(
+        title=f"Neuron {neuron_key} (dominant freq={dom_freq})",
+        xaxis_title='Fourier Component',
+        yaxis_title='Magnitude',
+        barmode='group',
+        template='plotly_white', height=350,
+    )
+    return fig
+
+
+def make_logit_bar_chart(data, pair_index):
+    """Build a Plotly bar chart of logits for a specific (a,b) pair."""
+    if data is None:
+        return None
+    pairs = data.get('pairs', [])
+    logits_all = data.get('logits', [])
+    correct = data.get('correct_answers', [])
+    classes = data.get('output_classes', [])
+
+    if pair_index >= len(pairs):
+        return None
+
+    a, b = pairs[pair_index]
+    logits = _to_np(logits_all[pair_index])
+    correct_ans = correct[pair_index]
+
+    bar_colors = [COLORS[3] if c == correct_ans else COLORS[0] for c in classes]
+
+    fig = go.Figure()
+    fig.add_trace(go.Bar(
+        x=[str(c) for c in classes], y=logits,
+        marker_color=bar_colors,
+        hovertemplate='Class %{x}: %{y:.3f}<extra></extra>',
+    ))
+    fig.update_layout(
+        title=f"Logits for ({a}, {b}) -- correct = {correct_ans}",
+        xaxis_title='Output Class',
+        yaxis_title='Logit Value',
+        template='plotly_white', height=350,
+    )
+    return fig
+
+
+def make_grokk_heatmap(data, epoch_index):
+    """Build a Plotly heatmap of accuracy grid at a grokking checkpoint."""
+    if data is None:
+        return None
+    epochs = data.get('epochs', [])
+    grids = data.get('grids', [])
+    if epoch_index >= len(grids):
+        return None
+
+    grid = _to_np(grids[epoch_index])
+    ep = epochs[epoch_index]
+
+    fig = go.Figure(data=go.Heatmap(
+        z=grid,
+        colorscale=[[0, 'white'], [1, COLORS[0]]],
+        zmin=0, zmax=1,
+        hovertemplate='a=%{y}, b=%{x}: %{z:.0f}<extra></extra>',
+    ))
+    fig.update_layout(
+        title=f"Accuracy Grid at Epoch {ep}",
+        xaxis_title='Second Input (b)',
+        yaxis_title='First Input (a)',
+        template='plotly_white', height=450,
+        yaxis=dict(autorange='reversed'),
+    )
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# Tab update functions
+# ---------------------------------------------------------------------------
+
+def update_tab1(p):
+    """Overview: standard + grokking loss/IPR, phase scatter."""
+    img_overview = safe_img(p, "overview_loss_ipr.png")
+    img_phase = safe_img(p, "overview_phase_scatter.png")
+    # Also build interactive charts from overview.json
+    data = load_json_file(p, "overview.json")
+    std_loss_chart = None
+    grokk_loss_chart = None
+    std_ipr_chart = None
+    grokk_ipr_chart = None
+
+    if data:
+        # Standard loss chart
+        std_ep = data.get('std_epochs', [])
+        std_tl = data.get('std_train_loss', [])
+        if std_tl:
+            fig = go.Figure()
+            fig.add_trace(go.Scatter(
+                x=_to_np(std_ep[:len(std_tl)]), y=_to_np(std_tl),
+                name='Train Loss', line=dict(color=COLORS[0]),
+            ))
+            fig.update_layout(
+                title='Standard: Training Loss (ReLU, full data)',
+                xaxis_title='Step', yaxis_title='Loss',
+                template='plotly_white', height=350,
+            )
+            std_loss_chart = fig
+
+        # Standard IPR chart
+        std_ipr = data.get('std_ipr', [])
+        if std_ipr:
+            fig = go.Figure()
+            fig.add_trace(go.Scatter(
+                x=_to_np(std_ep[:len(std_ipr)]), y=_to_np(std_ipr),
+                name='Avg IPR', line=dict(color=COLORS[3]),
+            ))
+            fig.update_layout(
+                title='Standard: IPR (Fourier Sparsity)',
+                xaxis_title='Step', yaxis_title='IPR',
+                yaxis=dict(range=[0, 1.05]),
+                template='plotly_white', height=350,
+            )
+            std_ipr_chart = fig
+
+        # Grokking loss chart
+        grokk_ep = data.get('grokk_epochs', [])
+        grokk_tl = data.get('grokk_train_loss', [])
+        grokk_tel = data.get('grokk_test_loss', [])
+        if grokk_tl or grokk_tel:
+            fig = go.Figure()
+            if grokk_tl:
+                fig.add_trace(go.Scatter(
+                    x=_to_np(grokk_ep[:len(grokk_tl)]), y=_to_np(grokk_tl),
+                    name='Train Loss', line=dict(color=COLORS[0]),
+                ))
+            if grokk_tel:
+                fig.add_trace(go.Scatter(
+                    x=_to_np(grokk_ep[:len(grokk_tel)]), y=_to_np(grokk_tel),
+                    name='Test Loss', line=dict(color=COLORS[3]),
+                ))
+            fig.update_layout(
+                title='Grokking: Loss (ReLU, 75% data, WD)',
+                xaxis_title='Step', yaxis_title='Loss',
+                template='plotly_white', height=350,
+            )
+            grokk_loss_chart = fig
+
+        # Grokking IPR chart
+        grokk_ipr = data.get('grokk_ipr', [])
+        if grokk_ipr:
+            fig = go.Figure()
+            fig.add_trace(go.Scatter(
+                x=_to_np(grokk_ep[:len(grokk_ipr)]), y=_to_np(grokk_ipr),
+                name='Avg IPR', line=dict(color=COLORS[3]),
+            ))
+            fig.update_layout(
+                title='Grokking: IPR (weight decay drives sparsification)',
+                xaxis_title='Step', yaxis_title='IPR',
+                yaxis=dict(range=[0, 1.05]),
+                template='plotly_white', height=350,
+            )
+            grokk_ipr_chart = fig
+
+    return (img_overview, std_loss_chart, grokk_loss_chart,
+            std_ipr_chart, grokk_ipr_chart, img_phase)
+
+
+def update_tab2(p):
+    """Fourier Weights: heatmap + line plots."""
+    return (
+        safe_img(p, "full_training_para_origin.png"),
+        safe_img(p, "lineplot_in.png"),
+        safe_img(p, "lineplot_out.png"),
+    )
+
+
+def update_tab3(p):
+    """Phase Analysis: distribution, relationship, magnitude."""
+    return (
+        safe_img(p, "phase_distribution.png"),
+        safe_img(p, "phase_relationship.png"),
+        safe_img(p, "magnitude_distribution.png"),
+    )
+
+
+def update_tab4(p):
+    """Output Logits."""
+    return safe_img(p, "output_logits.png")
+
+
+def update_tab5(p):
+    """Lottery Mechanism: magnitude, phase, contour."""
+    return (
+        safe_img(p, "lottery_mech_magnitude.png"),
+        safe_img(p, "lottery_mech_phase.png"),
+        safe_img(p, "lottery_beta_contour.png"),
+    )
+
+
+def update_tab6(p):
+    """Grokking: loss/acc charts + analysis images."""
+    loss_data = load_json_file(p, "grokk_loss.json")
+    acc_data = load_json_file(p, "grokk_acc.json")
+    loss_chart = make_loss_chart(loss_data, title="Grokking: Loss")
+    acc_chart = make_acc_chart(acc_data, title="Grokking: Accuracy")
+    return (
+        loss_chart,
+        acc_chart,
+        safe_img(p, "grokk_abs_phase_diff.png"),
+        safe_img(p, "grokk_avg_ipr.png"),
+        safe_img(p, "grokk_memorization_accuracy.png"),
+        safe_img(p, "grokk_memorization_common_to_rare.png"),
+        safe_img(p, "grokk_decoded_weights_dynamic.png"),
+    )
+
+
+def update_tab7(p):
+    """Gradient Dynamics: Quad and ReLU single-freq."""
+    return (
+        safe_img(p, "phase_align_quad.png"),
+        safe_img(p, "single_freq_quad.png"),
+        safe_img(p, "phase_align_relu.png"),
+        safe_img(p, "single_freq_relu.png"),
+    )
+
+
+def update_tab8(p):
+    """Decoupled Simulation: 2 analytical gradient flow plots."""
+    return (
+        safe_img(p, "phase_align_approx1.png"),
+        safe_img(p, "phase_align_approx2.png"),
+    )
+
+
+def update_tab9(p):
+    """Training Log: return available run names and initial log."""
+    data = load_json_file(p, "training_log.json")
+    if data is None:
+        return [], None, "", ""
+    run_names = list(data.keys())
+    # Show first run by default
+    first_run = run_names[0] if run_names else None
+    if first_run:
+        run_data = data[first_run]
+        config = run_data.get('config', {})
+        config_text = _format_config_md(first_run, config)
+        log_text = run_data.get('log_text', 'No log available.')
+    else:
+        config_text = ""
+        log_text = ""
+    return run_names, first_run, config_text, log_text
+
+
+def _format_config_md(run_name, config):
+    """Format a run config as a Markdown summary."""
+    lines = [f"**Run: {run_name}**\n"]
+    key_labels = {
+        'prime': 'Modulo (p)', 'd_mlp': 'd_mlp',
+        'act_type': 'Activation', 'init_type': 'Init Type',
+        'init_scale': 'Init Scale', 'optimizer': 'Optimizer',
+        'lr': 'Learning Rate', 'weight_decay': 'Weight Decay',
+        'frac_train': 'Frac Train', 'num_epochs': 'Num Epochs',
+        'seed': 'Seed',
+    }
+    for key, label in key_labels.items():
+        val = config.get(key, 'N/A')
+        lines.append(f"- **{label}**: `{val}`")
+    return "\n".join(lines)
+
+
+def update_info(p):
+    meta = load_json_file(p, "metadata.json")
+    if not meta:
+        return f"**p = {p}** | No metadata available"
+    d_mlp = meta.get('d_mlp', '?')
+    parts = [f"**p = {p}**", f"d_mlp = {d_mlp}"]
+    std_metrics = meta.get('final_metrics', {}).get('standard', {})
+    if 'train_acc' in std_metrics:
+        parts.append(f"Train Acc = {std_metrics['train_acc']:.4f}")
+    if 'test_acc' in std_metrics:
+        parts.append(f"Test Acc = {std_metrics['test_acc']:.4f}")
+    if 'train_loss' in std_metrics:
+        parts.append(f"Train Loss = {std_metrics['train_loss']:.6f}")
+    return " | ".join(parts)
+
+
+# ---------------------------------------------------------------------------
+# Interactive callback helpers
+# ---------------------------------------------------------------------------
+
+def _get_neuron_choices(p):
+    """Return list of neuron keys from neuron_spectra.json."""
+    data = load_json_file(p, "neuron_spectra.json")
+    if data is None:
+        return []
+    return list(data.get('neurons', {}).keys())
+
+
+def _get_pair_choices(p):
+    """Return list of (a,b) pair labels from logits_interactive.json."""
+    data = load_json_file(p, "logits_interactive.json")
+    if data is None:
+        return []
+    pairs = data.get('pairs', [])
+    return [f"({a}, {b})" for a, b in pairs]
+
+
+def _get_grokk_epochs(p):
+    """Return list of epoch values from grokk_epoch_data.json."""
+    data = load_json_file(p, "grokk_epoch_data.json")
+    if data is None:
+        return []
+    return data.get('epochs', [])
+
+
+# ---------------------------------------------------------------------------
+# Markdown helper -- ensures latex_delimiters are set
+# ---------------------------------------------------------------------------
+
+def _md(text, **kwargs):
+    """Create a gr.Markdown with KaTeX delimiters enabled."""
+    return gr.Markdown(text, latex_delimiters=LATEX_DELIMITERS, **kwargs)
+
+
+# ---------------------------------------------------------------------------
+# Main app
+# ---------------------------------------------------------------------------
+
+def on_p_change(p_str):
+    """Called when the prime dropdown changes. Returns all outputs."""
+    p = int(p_str)
+
+    info = update_info(p)
+
+    # Overview
+    (t1_img_overview, t1_std_loss, t1_grokk_loss,
+     t1_std_ipr, t1_grokk_ipr, t1_phase_scatter) = update_tab1(p)
+    # Core Interpretability
+    t2_heatmap, t2_line_in, t2_line_out = update_tab2(p)
+    t3_phase_dist, t3_phase_rel, t3_magnitude = update_tab3(p)
+    t4_logits = update_tab4(p)
+    t5_mag, t5_phase, t5_contour = update_tab5(p)
+    # Grokking
+    (t6_loss, t6_acc, t6_phase_diff, t6_ipr,
+     t6_memo, t6_memo_rare, t6_decoded) = update_tab6(p)
+    # Theory
+    t7_pa_quad, t7_sf_quad, t7_pa_relu, t7_sf_relu = update_tab7(p)
+    t8_approx1, t8_approx2 = update_tab8(p)
+
+    # Training Log
+    t9_run_names, t9_default_run, t9_config_text, t9_log = update_tab9(p)
+    t9_run_dd_update = gr.update(
+        choices=t9_run_names,
+        value=t9_default_run,
+    )
+
+    # Interactive widget updates
+    neuron_choices = _get_neuron_choices(p)
+    neuron_dd_update = gr.update(
+        choices=neuron_choices,
+        value=neuron_choices[0] if neuron_choices else None,
+    )
+    neuron_spectra_data = load_json_file(p, "neuron_spectra.json")
+    neuron_chart = make_neuron_spectrum_chart(
+        neuron_spectra_data, neuron_choices[0]
+    ) if neuron_choices else None
+
+    pair_choices = _get_pair_choices(p)
+    pair_dd_update = gr.update(
+        choices=pair_choices,
+        value=pair_choices[0] if pair_choices else None,
+    )
+    logit_data = load_json_file(p, "logits_interactive.json")
+    logit_chart = make_logit_bar_chart(logit_data, 0) if pair_choices else None
+
+    grokk_epochs = _get_grokk_epochs(p)
+    if grokk_epochs:
+        slider_update = gr.update(
+            minimum=0, maximum=len(grokk_epochs) - 1, value=0, step=1,
+            visible=True,
+        )
+    else:
+        slider_update = gr.update(minimum=0, maximum=0, value=0, visible=False)
+    grokk_slider_data = load_json_file(p, "grokk_epoch_data.json")
+    grokk_heatmap = make_grokk_heatmap(grokk_slider_data, 0) if grokk_epochs else None
+    epoch_label = f"Epoch: {grokk_epochs[0]}" if grokk_epochs else ""
+
+    return [
+        info,
+        # Tab 1: Overview
+        t1_img_overview, t1_std_loss, t1_grokk_loss,
+        t1_std_ipr, t1_grokk_ipr, t1_phase_scatter,
+        # Tab 2: Fourier Weights
+        t2_heatmap, t2_line_in, t2_line_out,
+        neuron_dd_update, neuron_chart,
+        # Tab 3: Phase Analysis
+        t3_phase_dist, t3_phase_rel, t3_magnitude,
+        # Tab 4: Output Logits
+        t4_logits,
+        pair_dd_update, logit_chart,
+        # Tab 5: Lottery Mechanism
+        t5_mag, t5_phase, t5_contour,
+        # Tab 6: Grokking
+        t6_loss, t6_acc, t6_phase_diff, t6_ipr,
+        t6_memo, t6_memo_rare, t6_decoded,
+        slider_update, grokk_heatmap, epoch_label,
+        # Tab 7: Gradient Dynamics
+        t7_pa_quad, t7_sf_quad, t7_pa_relu, t7_sf_relu,
+        # Tab 8: Decoupled Simulation
+        t8_approx1, t8_approx2,
+        # Tab 9: Training Log
+        t9_run_dd_update, t9_config_text, t9_log,
+    ]
+
+
+def _commit_results_to_repo(p):
+    """Try to commit new precomputed results back to the HF Space repo.
+
+    On HF Spaces, the repo is writable via the huggingface_hub API.
+    This allows results to accumulate as users generate them.
+    Returns (success, message).
+    """
+    try:
+        from huggingface_hub import HfApi
+    except ImportError:
+        return False, "huggingface_hub not installed"
+
+    space_id = os.environ.get("SPACE_ID")  # e.g. "username/space-name"
+    if not space_id:
+        return False, "Not running on HF Spaces (SPACE_ID not set)"
+
+    result_dir = os.path.join(RESULTS_DIR, f"p_{p:03d}")
+    if not os.path.isdir(result_dir):
+        return False, "No results directory found"
+
+    try:
+        api = HfApi()
+        api.upload_folder(
+            folder_path=result_dir,
+            path_in_repo=f"precomputed_results/p_{p:03d}",
+            repo_id=space_id,
+            repo_type="space",
+            commit_message=f"Add precomputed results for p={p}",
+        )
+        return True, f"Committed results for p={p} to {space_id}"
+    except Exception as e:
+        logger.warning(f"Failed to commit results for p={p}: {e}")
+        return False, str(e)
+
+
+def _run_step_streaming(cmd, env, label):
+    """Run a subprocess, yielding (line, error_flag) for each output line."""
+    proc = subprocess.Popen(
+        cmd, cwd=PROJECT_ROOT, env=env,
+        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+        text=True, bufsize=1,
+    )
+    for line in proc.stdout:
+        yield line.rstrip("\n"), False
+    proc.wait()
+    if proc.returncode != 0:
+        yield f"[ERROR] {label} failed (exit code {proc.returncode})", True
+
+
+def run_pipeline_for_p_streaming(p):
+    """Generator: run full pipeline for p, yielding log lines.
+
+    Yields (log_line: str, is_error: bool, is_done: bool).
+    Deletes model checkpoints after plot generation to save space.
+    """
+    if p < 3 or p % 2 == 0:
+        yield f"Error: p must be an odd number >= 3, got {p}", True, True
+        return
+    if p > MAX_P_ON_DEMAND:
+        yield f"Error: p={p} exceeds on-demand limit of {MAX_P_ON_DEMAND}", True, True
+        return
+
+    result_dir = os.path.join(RESULTS_DIR, f"p_{p:03d}")
+    if os.path.isdir(result_dir) and len(os.listdir(result_dir)) > 5:
+        yield f"Results for p={p} already exist ({len(os.listdir(result_dir))} files)", False, True
+        return
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = PROJECT_ROOT + ":" + env.get("PYTHONPATH", "")
+
+    steps = [
+        ("Step 1/3: Training 5 configurations", [
+            sys.executable, "precompute/train_all.py",
+            "--p", str(p), "--output", TRAINED_MODELS_DIR, "--resume",
+        ]),
+        ("Step 2/3: Generating model-based plots", [
+            sys.executable, "precompute/generate_plots.py",
+            "--p", str(p), "--input", TRAINED_MODELS_DIR,
+            "--output", RESULTS_DIR,
+        ]),
+        ("Step 3/3: Generating analytical plots", [
+            sys.executable, "precompute/generate_analytical.py",
+            "--p", str(p), "--output", RESULTS_DIR,
+        ]),
+    ]
+
+    for label, cmd in steps:
+        yield f"\n{'='*60}", False, False
+        yield f"  {label} (p={p})", False, False
+        yield f"{'='*60}", False, False
+        for line, is_err in _run_step_streaming(cmd, env, label):
+            if is_err:
+                yield line, True, True
+                return
+            yield line, False, False
+
+    # Cleanup checkpoints
+    model_dir = os.path.join(TRAINED_MODELS_DIR, f"p_{p:03d}")
+    if os.path.isdir(model_dir):
+        shutil.rmtree(model_dir)
+        yield f"Cleaned up checkpoints: {model_dir}", False, False
+
+    n_files = len(os.listdir(result_dir)) if os.path.isdir(result_dir) else 0
+
+    # Try to commit results back to the HF repo
+    ok_commit, commit_msg = _commit_results_to_repo(p)
+    if ok_commit:
+        yield f"Results saved to HF repo.", False, False
+
+    yield f"\nDone! Generated {n_files} files for p={p}.", False, True
+
+
+def create_app():
+    moduli = get_available_moduli()
+    p_choices = [str(p) for p in moduli]
+    default_p = p_choices[0] if p_choices else None
+
+    with gr.Blocks(
+        title="Modular Addition Feature Learning",
+    ) as app:
+        _md(
+            r"# On the Mechanism and Dynamics of Modular Addition" "\n"
+            r"### Fourier Features, Lottery Ticket, and Grokking" "\n\n"
+            r"**Jianliang He, Leda Wang, Siyu Chen, Zhuoran Yang**" "\n"
+            r"*Department of Statistics and Data Science, Yale University*" "\n\n"
+            r'<a href="#">[arXiv]</a> &nbsp; '
+            r'<a href="#">[Blog]</a> &nbsp; '
+            r'<a href="https://github.com/Y-Agent/modular-addition-feature-learning">[Code]</a>' "\n\n"
+            r"---" "\n\n"
+            r"This interactive explorer visualizes how a two-layer neural network "
+            r"learns modular arithmetic $f(x,y) = (x + y) \bmod p$ through "
+            r"Fourier feature learning, lottery ticket dynamics, and the grokking "
+            r"phenomenon. Select a modulo $p$ (any odd number $\geq 3$) below to view pre-computed results." "\n\n"
+            r"> **Note:** Grokking experiments (Tab 6) require $p \geq 19$ to have enough data for a meaningful train/test split. "
+            r"For $p < 19$, grokking plots will not be generated."
+        )
+
+        # Hidden state for current modulo
+        current_p = gr.State(value=int(default_p) if default_p else 3)
+
+        with gr.Row():
+            p_dropdown = gr.Dropdown(
+                choices=p_choices,
+                value=default_p,
+                label="Select Modulo (p)",
+                interactive=True,
+                scale=2,
+            )
+            info_md = _md(
+                update_info(int(default_p)) if default_p else ""
+            )
+
+        with gr.Accordion("Generate results for a new p", open=False):
+            _md(
+                f"Enter any odd number $p \\geq 3$ (max {MAX_P_ON_DEMAND} "
+                f"for on-demand training). This will train 5 model "
+                f"configurations and generate all plots. Training logs "
+                f"are streamed below in real time."
+            )
+            with gr.Row():
+                new_p_input = gr.Number(
+                    value=None, label="New p (odd, ≥ 3)",
+                    precision=0, scale=1,
+                )
+                generate_btn = gr.Button(
+                    "Generate", variant="primary", scale=1,
+                )
+            generate_status = _md("")
+            generate_log = gr.Code(
+                value="", language=None, label="Pipeline Log",
+                lines=15, interactive=False, visible=False,
+            )
+
+        # ----- Tabs -----
+        with gr.Tabs():
+
+            # === Core Interpretability ===
+
+            # Tab 1: Overview
+            with gr.Tab("1. Overview"):
+                _md(MATH_TAB1)
+                t1_img_overview = gr.Image(
+                    label="Loss & IPR Overview (Static)", type="filepath"
+                )
+                with gr.Row():
+                    t1_std_loss = gr.Plot(label="Standard: Loss")
+                    t1_grokk_loss = gr.Plot(label="Grokking: Loss")
+                with gr.Row():
+                    t1_std_ipr = gr.Plot(label="Standard: IPR")
+                    t1_grokk_ipr = gr.Plot(label="Grokking: IPR")
+                t1_phase_scatter = gr.Image(
+                    label="Phase Alignment: \u03c8 = 2\u03c6", type="filepath"
+                )
+
+            # Tab 2: Fourier Weights
+            with gr.Tab("2. Fourier Weights"):
+                _md(MATH_TAB2)
+                t2_heatmap = gr.Image(label="Decoded W_in / W_out Heatmap", type="filepath")
+                with gr.Row():
+                    t2_line_in = gr.Image(label="First-Layer Line Plots (with cosine fit)", type="filepath")
+                    t2_line_out = gr.Image(label="Second-Layer Line Plots (with cosine fit)", type="filepath")
+                _md("#### Neuron Frequency Inspector")
+                t2_neuron_dd = gr.Dropdown(
+                    choices=[], value=None,
+                    label="Select Neuron", interactive=True,
+                )
+                t2_neuron_chart = gr.Plot(label="Neuron Fourier Spectrum")
+
+            # Tab 3: Phase Analysis
+            with gr.Tab("3. Phase Analysis"):
+                _md(MATH_TAB3)
+                with gr.Row():
+                    t3_phase_dist = gr.Image(label="Phase Distribution", type="filepath")
+                    t3_phase_rel = gr.Image(
+                        label="Phase Relationship (2\u03c6 vs \u03c8)", type="filepath"
+                    )
+                t3_magnitude = gr.Image(label="Magnitude Distribution", type="filepath")
+
+            # Tab 4: Output Logits
+            with gr.Tab("4. Output Logits"):
+                _md(MATH_TAB4)
+                t4_logits = gr.Image(label="Output Logits Heatmap", type="filepath")
+                _md("#### Logit Explorer")
+                t4_pair_dd = gr.Dropdown(
+                    choices=[], value=None,
+                    label="Select Input Pair (a, b)", interactive=True,
+                )
+                t4_logit_chart = gr.Plot(label="Logit Distribution")
+
+            # Tab 5: Lottery Mechanism
+            with gr.Tab("5. Lottery Mechanism"):
+                _md(MATH_TAB5)
+                _md(r"""**Magnitude plot** below: Each curve tracks one frequency's output magnitude $\beta_k$ within a single neuron over training. All frequencies start with equal magnitude (from random initialization). The winning frequency (best initial phase alignment) grows explosively while others remain frozen.""")
+                t5_mag = gr.Image(label="Frequency Magnitude Evolution", type="filepath")
+                _md(r"""**Phase plot** below: Each curve shows the phase misalignment $\mathcal{D}_k = 2\phi_k - \psi_k$ for one frequency within the same neuron. The winning frequency (colored) converges to $\mathcal{D} = 0$ (perfect alignment) first; other frequencies barely change because their magnitudes remain small.""")
+                t5_phase = gr.Image(label="Phase Misalignment Convergence", type="filepath")
+                _md(r"""**Contour plot** below: Final output magnitude as a function of initial magnitude and initial phase misalignment, across all neurons. The largest final magnitudes (brightest regions) appear at small initial misalignment $|\mathcal{D}|$, confirming that initial phase alignment -- not initial magnitude -- determines which frequency wins.""")
+                t5_contour = gr.Image(label="Final Magnitude Contour", type="filepath")
+
+            # === Grokking ===
+
+            # Tab 6: Grokking
+            with gr.Tab("6. Grokking"):
+                _md(MATH_TAB6)
+
+                _md(r"""#### (a) Loss and (b) Accuracy
+
+**(a) Loss:** Training loss (blue) drops rapidly in Stage I as the network memorizes training data. Test loss (red) stays high until Stage II, when weight decay forces the network to find a generalizing solution, causing test loss to plummet. The three colored bands mark the three stages.
+
+**(b) Accuracy:** Training accuracy reaches 100% early (Stage I). Test accuracy stays at ~70% during memorization (not 50% -- the built-in symmetry $f(a,b) = f(b,a)$ gives "free" correct answers for the swapped pair). Test accuracy jumps sharply in Stage II when the network transitions from memorization to Fourier features.""")
+                with gr.Row():
+                    t6_loss = gr.Plot(label="Grokking Loss (Interactive)")
+                    t6_acc = gr.Plot(label="Grokking Accuracy (Interactive)")
+
+                _md(r"""#### (c) Phase Alignment and (d) IPR & Norms
+
+**(c) Phase alignment:** Average $|\sin(\mathcal{D}_m^\star)|$ over all neurons, where $\mathcal{D}_m^\star = 2\phi_m^\star - \psi_m^\star$ is the phase misalignment at each neuron's dominant frequency. This measures how far the network is from the ideal relationship $\psi = 2\phi$. It decreases throughout training as phases align, with the steepest drop during Stage II.
+
+**(d) IPR and parameter norms:** IPR (Fourier sparsity) increases sharply in Stage II -- this is the "aha" moment where multi-frequency noise collapses into clean single-frequency features. Parameter norms shrink steadily in Stage III as weight decay slowly polishes the solution.""")
+                with gr.Row():
+                    t6_phase_diff = gr.Image(
+                        label="Phase Difference |sin(D*)|", type="filepath"
+                    )
+                    t6_ipr = gr.Image(label="IPR & Parameter Norms", type="filepath")
+
+                _md(r"""#### (e) Memorization Accuracy Grid
+
+Each cell $(i,j)$ in the grid shows whether the network correctly predicts $(i+j) \bmod p$ at a given training epoch. **White = correct, dark = incorrect.** Training pairs are marked with dots.
+
+During Stage I, the network first memorizes **symmetric pairs** -- pairs where both $(i,j)$ and $(j,i)$ are in the training set (they appear on both sides of the diagonal). These are learned first because the architecture treats inputs symmetrically: $\theta_m[i] + \theta_m[j] = \theta_m[j] + \theta_m[i]$, so learning one automatically gives the other.
+
+**Asymmetric pairs** (where only one of $(i,j)$ or $(j,i)$ is in training) are harder to memorize and are learned later. Some test pairs may even be *actively suppressed* (the network gets them wrong on purpose) before eventually being memorized.""")
+                t6_memo = gr.Image(label="Memorization Accuracy", type="filepath")
+
+                _md(r"""#### (f) Common-to-Rare Ordering
+
+This plot reorders the accuracy grid to reveal the **memorization sequence**. Instead of plotting by input value, it sorts pairs by how "common" they are in the training set:
+
+- **Common pairs** (top-left): Both $(i,j)$ and $(j,i)$ in training set. These are memorized first.
+- **Rare pairs** (bottom-right): Only one ordering in training set. These are memorized last, and may be temporarily suppressed before being learned.
+
+The plot shows a clear **top-left to bottom-right** progression, confirming that the network memorizes common pairs before rare ones.""")
+                t6_memo_rare = gr.Image(label="Memorization: Common to Rare", type="filepath")
+
+                _md(r"""#### (g) Decoded Weights Across Stages
+
+DFT heatmaps of the network's weights at key epochs through the three stages. Each row is a neuron; each column is a Fourier frequency component.
+
+- **Stage I (memorization):** Weights are noisy with energy spread across many frequencies -- the network is using all available capacity to memorize.
+- **Stage II (generalization):** Weight decay kills the weak frequencies. Each neuron's energy concentrates into a single frequency -- clean Fourier features emerge.
+- **Stage III (cleanup):** Features are already clean; weight decay slowly shrinks overall magnitude without changing the structure.""")
+                t6_decoded = gr.Image(label="Decoded Weights Across Stages", type="filepath")
+
+                _md(r"""#### Accuracy Grid Across Training (Interactive)
+
+Use the slider to scrub through training epochs and watch the accuracy grid evolve. In Stage I, you'll see the symmetric pairs (along both diagonals) light up first, then asymmetric pairs fill in, and finally the entire grid becomes white in Stage II as the network generalizes.""")
+                t6_slider = gr.Slider(
+                    minimum=0, maximum=0, value=0, step=1,
+                    label="Epoch Snapshot Index", interactive=True,
+                    visible=False,
+                )
+                t6_heatmap = gr.Plot(label="Accuracy Heatmap")
+                t6_epoch_label = _md("")
+
+            # === Theory ===
+
+            # Tab 7: Gradient Dynamics
+            with gr.Tab("7. Gradient Dynamics"):
+                _md(MATH_TAB7)
+                _md(r"""#### Quadratic Activation ($\sigma(x) = x^2$)
+
+**Left -- Phase alignment:** Tracks the input phase $\phi_m^\star$, output phase $\psi_m^\star$, and doubled input phase $2\phi_m^\star$ of the dominant frequency in a single neuron over training. The theory predicts $\psi \to 2\phi$; here we see $\psi$ (red) and $2\phi$ (blue) converge and overlap, confirming phase alignment. The phases lock in early while magnitudes are still small.
+
+**Right -- DFT heatmaps:** Decoded weights in Fourier space at key training steps. At step 0, the neuron starts with energy at a single frequency (by construction -- single-frequency initialization). At later steps, the dominant frequency grows while all other frequencies stay at zero. This confirms the **single-frequency preservation theorem**: Fourier orthogonality prevents energy leakage between modes.""")
+                with gr.Row():
+                    t7_pa_quad = gr.Image(label="Phase Alignment (Quad)", type="filepath")
+                    t7_sf_quad = gr.Image(label="Decoded Weights (Quad)", type="filepath")
+                _md(r"""#### ReLU Activation ($\sigma(x) = \max(0, x)$)
+
+**Left -- Phase alignment:** Same as quadratic above, but with ReLU. The qualitative behavior is identical: $\psi$ converges to $2\phi$. Minor quantitative differences arise because ReLU is not exactly $x^2$.
+
+**Right -- DFT heatmaps:** Unlike quadratic, ReLU leaks small amounts of energy to **harmonic multiples** of the dominant frequency ($3k^\star, 5k^\star, \ldots$ for input weights; $2k^\star, 3k^\star, \ldots$ for output weights). This leakage decays as $O(r^{-2})$ where $r$ is the harmonic order, so the dominant frequency remains overwhelmingly dominant. The faint "stripes" at harmonic positions are this leakage.""")
+                with gr.Row():
+                    t7_pa_relu = gr.Image(label="Phase Alignment (ReLU)", type="filepath")
+                    t7_sf_relu = gr.Image(label="Decoded Weights (ReLU)", type="filepath")
+
+            # Tab 8: Decoupled Simulation
+            with gr.Tab("8. Decoupled Simulation"):
+                _md(MATH_TAB8)
+                _md(r"""Each 3-panel figure below shows one simulation run. The gray curves are non-winning frequencies; the colored curves are the winning frequency $k^\star$.
+
+**Top panel -- Phase alignment:** $\psi_{k^\star}$ (red) and $2\phi_{k^\star}$ (blue) converge toward each other, confirming that training drives phases into the $\psi = 2\phi$ relationship even in this pure ODE setting (no neural network).
+
+**Middle panel -- Phase difference $D_{k^\star}$:** The misalignment $\mathcal{D}_{k^\star} = 2\phi_{k^\star} - \psi_{k^\star}$ converges toward $0$ (or $\pi/2$ transiently in Case 1). The dashed horizontal line marks $\pi/2$. Non-winning frequencies (gray) remain scattered because their magnitudes are too small to drive phase alignment.
+
+**Bottom panel -- Magnitude evolution:** The winning frequency's magnitudes ($\alpha_{k^\star}$ and $\beta_{k^\star}$) grow explosively once phase alignment is achieved, while all other frequencies remain near their initial values. This is the lottery ticket effect in pure form.""")
+                with gr.Row():
+                    t8_approx1 = gr.Image(
+                        label="Gradient Flow (Case 1: with annotations)", type="filepath"
+                    )
+                    t8_approx2 = gr.Image(label="Gradient Flow (Case 2)", type="filepath")
+
+            # Tab 9: Training Log
+            with gr.Tab("9. Training Log"):
+                _md(MATH_TAB9)
+                t9_run_dd = gr.Dropdown(
+                    choices=[], value=None,
+                    label="Select Training Run", interactive=True,
+                )
+                t9_config_md = _md("")
+                t9_log_text = gr.Code(
+                    value="", language=None, label="Training Log",
+                    lines=30, interactive=False,
+                )
+
+        # All outputs for prime change
+        all_outputs = [
+            info_md,
+            # Tab 1: Overview
+            t1_img_overview, t1_std_loss, t1_grokk_loss,
+            t1_std_ipr, t1_grokk_ipr, t1_phase_scatter,
+            # Tab 2
+            t2_heatmap, t2_line_in, t2_line_out,
+            t2_neuron_dd, t2_neuron_chart,
+            # Tab 3
+            t3_phase_dist, t3_phase_rel, t3_magnitude,
+            # Tab 4
+            t4_logits,
+            t4_pair_dd, t4_logit_chart,
+            # Tab 5
+            t5_mag, t5_phase, t5_contour,
+            # Tab 6
+            t6_loss, t6_acc, t6_phase_diff, t6_ipr,
+            t6_memo, t6_memo_rare, t6_decoded,
+            t6_slider, t6_heatmap, t6_epoch_label,
+            # Tab 7
+            t7_pa_quad, t7_sf_quad, t7_pa_relu, t7_sf_relu,
+            # Tab 8
+            t8_approx1, t8_approx2,
+            # Tab 9
+            t9_run_dd, t9_config_md, t9_log_text,
+        ]
+
+        # --- Prime change handler ---
+        def p_change_and_store(p_str):
+            p = int(p_str)
+            results = on_p_change(p_str)
+            return [p] + results
+
+        p_dropdown.change(
+            fn=p_change_and_store,
+            inputs=[p_dropdown],
+            outputs=[current_p] + all_outputs,
+        )
+
+        # --- Neuron dropdown handler ---
+        def on_neuron_change(neuron_key, p):
+            data = load_json_file(p, "neuron_spectra.json")
+            return make_neuron_spectrum_chart(data, neuron_key)
+
+        t2_neuron_dd.change(
+            fn=on_neuron_change,
+            inputs=[t2_neuron_dd, current_p],
+            outputs=[t2_neuron_chart],
+        )
+
+        # --- Logit pair dropdown handler ---
+        def on_pair_change(pair_str, p):
+            data = load_json_file(p, "logits_interactive.json")
+            if data is None or not pair_str:
+                return None
+            pairs = data.get('pairs', [])
+            pair_labels = [f"({a}, {b})" for a, b in pairs]
+            if pair_str in pair_labels:
+                idx = pair_labels.index(pair_str)
+            else:
+                idx = 0
+            return make_logit_bar_chart(data, idx)
+
+        t4_pair_dd.change(
+            fn=on_pair_change,
+            inputs=[t4_pair_dd, current_p],
+            outputs=[t4_logit_chart],
+        )
+
+        # --- Grokking slider handler ---
+        def on_grokk_slider(slider_val, p):
+            data = load_json_file(p, "grokk_epoch_data.json")
+            if data is None:
+                return None, ""
+            idx = int(slider_val)
+            epochs = data.get('epochs', [])
+            label = f"**Epoch: {epochs[idx]}**" if idx < len(epochs) else ""
+            return make_grokk_heatmap(data, idx), label
+
+        t6_slider.change(
+            fn=on_grokk_slider,
+            inputs=[t6_slider, current_p],
+            outputs=[t6_heatmap, t6_epoch_label],
+        )
+
+        # --- Training log run dropdown handler ---
+        def on_log_run_change(run_name, p):
+            data = load_json_file(p, "training_log.json")
+            if data is None or run_name not in data:
+                return "", ""
+            run_data = data[run_name]
+            config = run_data.get('config', {})
+            config_text = _format_config_md(run_name, config)
+            log_text = run_data.get('log_text', 'No log available.')
+            return config_text, log_text
+
+        t9_run_dd.change(
+            fn=on_log_run_change,
+            inputs=[t9_run_dd, current_p],
+            outputs=[t9_config_md, t9_log_text],
+        )
+
+        # --- On-demand training handler (streaming) ---
+        def on_generate_click(new_p):
+            if new_p is None:
+                yield (
+                    gr.update(), gr.update(),
+                    "Enter a value for p.",
+                    gr.update(visible=False, value=""),
+                )
+                return
+            p = int(new_p)
+            log_lines = []
+            yield (
+                gr.update(), gr.update(),
+                f"**Running pipeline for p={p}...**",
+                gr.update(visible=True, value="Starting...\n"),
+            )
+            for line, is_err, is_done in run_pipeline_for_p_streaming(p):
+                log_lines.append(line)
+                # Keep last 200 lines to avoid memory bloat
+                display = "\n".join(log_lines[-200:])
+                if is_err:
+                    yield (
+                        gr.update(), gr.update(),
+                        f"**Error:** {line}",
+                        gr.update(value=display),
+                    )
+                    return
+                if is_done:
+                    new_moduli = get_available_moduli()
+                    new_choices = [str(v) for v in new_moduli]
+                    yield (
+                        gr.update(choices=new_choices, value=str(p)),
+                        gr.update(),
+                        f"**Success:** {line}",
+                        gr.update(value=display),
+                    )
+                    return
+                yield (
+                    gr.update(), gr.update(),
+                    f"**Running pipeline for p={p}...**",
+                    gr.update(value=display),
+                )
+
+        generate_btn.click(
+            fn=on_generate_click,
+            inputs=[new_p_input],
+            outputs=[p_dropdown, current_p, generate_status, generate_log],
+        )
+
+    return app
+
+
+if __name__ == "__main__":
+    app = create_app()
+    app.launch(theme=gr.themes.Soft(), css=CUSTOM_CSS)

hf_app/requirements.txt

@@ -0,0 +1,9 @@
+gradio>=4.0
+torch>=2.0.0
+--extra-index-url https://download.pytorch.org/whl/cpu
+numpy>=1.24
+matplotlib>=3.7
+Pillow>=9.0
+plotly>=5.0
+einops>=0.6
+scipy>=1.10

precompute/README.md

@@ -0,0 +1,337 @@
+# Pre-computation Pipeline
+
+Batch training and plot generation for all odd moduli $p$ in [3, 199]. Trains 5 model configurations per $p$ and generates publication-quality figures plus interactive JSON data files covering the paper's core results.
+
+All commands are run from the **project root directory**.
+
+## Quick Start (Shell Script)
+
+The easiest way to run the full pipeline for a single modulus:
+
+```bash
+# Run the complete pipeline for p=23
+bash precompute/run_pipeline.sh 23
+
+# Or using an environment variable
+P=23 bash precompute/run_pipeline.sh
+```
+
+This runs training, plot generation, analytical simulation, and verification in sequence.
+
+## Complete Pipeline (Single Modulus, Manual Steps)
+
+```bash
+# Step 1: Train all 5 model configurations
+python precompute/train_all.py --p 23 --output ./trained_models
+
+# Step 2: Generate model-based plots (21 PNGs + 6 JSONs + metadata)
+python precompute/generate_plots.py --p 23 --input ./trained_models --output ./precomputed_results
+
+# Step 3: Generate analytical simulation plots (2 PNGs, no model needed)
+python precompute/generate_analytical.py --p 23 --output ./precomputed_results
+
+# Step 4: Verify
+ls precomputed_results/p_023/
+```
+
+## Complete Pipeline (All Odd p)
+
+```bash
+# Train everything (225 runs total). Use --resume to skip completed runs.
+python precompute/train_all.py --all --output ./trained_models --resume
+
+# Generate all plots
+python precompute/generate_plots.py --all --input ./trained_models --output ./precomputed_results
+
+# Generate all analytical plots
+python precompute/generate_analytical.py --all --output ./precomputed_results
+```
+
+---
+
+## The 5 Model Configurations
+
+Each modulus is trained with 5 configurations that correspond to different sections of the paper:
+
+### 1. Standard Training (`standard`)
+
+The baseline experiment for Parts I--II (Mechanism & Dynamics). Demonstrates Fourier feature learning: neurons decompose modular addition into sparse frequency components with phase alignment (ψ ≈ 2φ).
+
+| Parameter | Value |
+|-----------|-------|
+| Activation | ReLU |
+| Initialization | random |
+| Optimizer | AdamW |
+| Learning rate | 5e-5 |
+| Weight decay | 0 |
+| Train fraction | 1.0 (all p² pairs) |
+| Epochs | 5,000 |
+| Init scale | 0.1 |
+
+**Used by:** Tab 1 (Overview), Tab 2 (Fourier Weights), Tab 3 (Phase Analysis), Tab 4 (Output Logits)
+
+### 2. Grokking (`grokking`)
+
+Reproduces the grokking phenomenon (Part III). The model memorizes training data first, then suddenly generalizes. Requires partial training data + weight decay.
+
+| Parameter | Value |
+|-----------|-------|
+| Activation | ReLU |
+| Initialization | random |
+| Optimizer | AdamW |
+| Learning rate | 1e-4 |
+| Weight decay | **2.0** |
+| Train fraction | **0.75** |
+| Epochs | **50,000** |
+| Init scale | 0.1 |
+
+**Used by:** Tab 1 (Overview, grokking curves), Tab 6 (Grokking)
+**Note:** Only runs for p ≥ 19 (smaller $p$ have too few test points for meaningful grokking).
+
+### 3. Quadratic Activation (`quad_random`)
+
+Uses σ(x) = x² activation. The quadratic nonlinearity directly implements the frequency factorization mechanism from the theory, enabling clean analysis of the lottery ticket mechanism.
+
+| Parameter | Value |
+|-----------|-------|
+| Activation | **Quad** |
+| Initialization | random |
+| Optimizer | AdamW |
+| Learning rate | 5e-5 |
+| Weight decay | 0 |
+| Train fraction | 1.0 |
+| Epochs | 5,000 |
+| Init scale | 0.1 |
+
+**Used by:** Tab 5 (Lottery Mechanism)
+
+### 4. Single-Frequency Quad (`quad_single_freq`)
+
+Initializes neurons at specific frequencies to study gradient dynamics under controlled conditions. Validates the phase alignment theorem and single-frequency preservation theorem.
+
+| Parameter | Value |
+|-----------|-------|
+| Activation | **Quad** |
+| Initialization | **single-freq** |
+| Optimizer | **SGD** |
+| Learning rate | **0.1** |
+| Weight decay | 0 |
+| Train fraction | 1.0 |
+| Epochs | 5,000 |
+| Init scale | **0.02** |
+
+**Used by:** Tab 7 (Gradient Dynamics, quadratic panels)
+
+### 5. Single-Frequency ReLU (`relu_single_freq`)
+
+Same as above but with ReLU activation. Shows that the theoretical results (proved for quadratic) hold approximately for ReLU with small harmonic leakage.
+
+| Parameter | Value |
+|-----------|-------|
+| Activation | **ReLU** |
+| Initialization | **single-freq** |
+| Optimizer | **SGD** |
+| Learning rate | **0.01** |
+| Weight decay | 0 |
+| Train fraction | 1.0 |
+| Epochs | 5,000 |
+| Init scale | **0.002** |
+
+**Used by:** Tab 7 (Gradient Dynamics, ReLU panels)
+
+---
+
+## Neuron Sizing
+
+The number of hidden neurons scales with $p$ to maintain the ratio from the baseline experiment ($p=23$, $d_\text{mlp}=512$):
+
+```
+d_mlp = max(512, ceil(512/529 * p²))
+```
+
+Examples: $p=3 \to 512$, $p=23 \to 512$, $p=53 \to 2720$, $p=97 \to 9108$, $p=199 \to 38329$.
+
+---
+
+## Blog Figure → Pipeline Output Mapping
+
+The table below maps every figure in the blog post to the corresponding file generated by the pipeline. Each figure is reproduced for every $p$, allowing users to verify the paper's claims across different moduli.
+
+### Part I: Mechanism (Tabs 2--4, standard run)
+
+| Blog Figure | Description | Pipeline Output | Tab | Verified? |
+|-------------|-------------|----------------|-----|-----------|
+| **Fig. 2** — Fourier sparsity of learned weights | DFT heatmap: each row is a neuron, each column is a Fourier mode (cos k, sin k). Sparse = one bright cell per row, confirming single-frequency specialization. | `pXXX_full_training_para_origin.png` | 2 | The heatmap applies `W_in @ fourier_basis.T` and `W_out.T @ fourier_basis.T` to show DFT coefficients. X-axis labels are Fourier mode names (Const, cos 1, sin 1, ...). Sparsity is visible as one dominant pair per neuron row. |
+| **Fig. 3** — Cosine fits to individual neurons | Raw learned weight values (dots) vs. best-fit cosine (dashed) for 3 representative neurons. Left: input weights θ_m. Right: output weights ξ_m. | `pXXX_lineplot_in.png`, `pXXX_lineplot_out.png` | 2 | Projects raw weights into Fourier space, keeps top-2 components, projects back. The fit quality demonstrates that each neuron is well-described by a single cosine. |
+| **Fig. 4** — Phase alignment ψ = 2φ | Scatter plot of (2φ_m mod 2π) vs (ψ_m mod 2π). All points lie on the diagonal y = x. | `pXXX_phase_relationship.png` | 3 | Computed via `compute_neuron()` for every neuron. The diagonal pattern is Observation 2 from the paper. |
+| **Fig. 5** — Higher-order phase symmetry | Polar plot: phase angles ι·φ_m on concentric rings for ι = 1, 2, 3, 4. Uniform spread confirms the cancellation condition Σ exp(i·ι·φ_m) ≈ 0. | `pXXX_phase_distribution.png` | 3 | Shows phases for the most common frequency group. For large p with many neurons, the uniform spread is clearly visible. |
+| **Fig. 6** — Magnitude homogeneity | Violin plots of α_m (input) and β_m (output) across all neurons. Tight concentration confirms magnitude homogeneity (Observation 3c). | `pXXX_magnitude_distribution.png` | 3 | Uses `compute_neuron()` to extract scale for every neuron. |
+| **Fig. 7** — Output logits (flawed indicator) | Heatmap of f(x,y)[j] for x=0. Bright red diagonal at j=(x+y) mod p (correct answer, coefficient p/2). Faint pink at j=2x mod p and j=2y mod p (spurious peaks, coefficient p/4). | `pXXX_output_logits.png` | 4 | Forward pass through the trained model with **matching activation** (ReLU for standard run). Rectangles highlight the correct answer and spurious peak positions. |
+
+### Part II: Dynamics (Tabs 5, 7, 8)
+
+| Blog Figure | Description | Pipeline Output | Tab | Verified? |
+|-------------|-------------|----------------|-----|-----------|
+| **Fig. 8** — Phase alignment dynamics | Phase trajectories (φ, ψ, 2φ) and magnitude growth (α, β) over training. Left: Quad activation. Right: ReLU. Shows ψ → 2φ convergence. | `pXXX_phase_align_quad.png`, `pXXX_phase_align_relu.png` | 7 | Tracks the neuron with largest final scale across all checkpoints. Shows phases converging and magnitudes growing. |
+| **Fig. 9** — Lottery ticket race | Left: phase misalignment D_m^k(t) for all frequencies within one neuron. The winner (smallest initial D) converges first. Right: magnitude β_m^k(t). Winner grows explosively. | `pXXX_lottery_mech_phase.png`, `pXXX_lottery_mech_magnitude.png` | 5 | Tracks all frequency components of a single neuron via `decode_scales_phis()` across checkpoints from the `quad_random` run. The winning frequency (highlighted in red) has the smallest initial misalignment. |
+| **Fig. 10** — Lottery outcome contour | Final magnitude β as a function of (initial magnitude, initial phase difference 2φ₀). Largest values at small D, symmetric about D = π. | `pXXX_lottery_beta_contour.png` | 5 | Simulates gradient flow on a 30×30 grid of initial conditions. Each point runs 100 steps of the analytical ODE. |
+| **Fig. 11** — Single-frequency preservation (Quad) | DFT heatmap at multiple training timepoints. The initialized frequency retains all energy; no cross-frequency leakage. | `pXXX_single_freq_quad.png` | 7 | Shows DFT of weights at 3 timepoints (step 0, mid, final) for the `quad_single_freq` run. Each column is a Fourier mode; sparsity confirms preservation. |
+| **Fig. 12a** — Single-frequency preservation (ReLU) | Same as Fig. 11 but with ReLU. Small harmonic leakage visible at 3k*, 5k* (input) and 2k*, 3k* (output), decaying as O(r⁻²). | `pXXX_single_freq_relu.png` | 7 | Shows DFT at 2 timepoints (step 0, final) for the `relu_single_freq` run. Dominant frequency overwhelms harmonics. |
+| **Fig. 12b** — Phase alignment under ReLU | Phase and magnitude trajectories for ReLU single-frequency init. Same zero-attractor behavior as Quad. | `pXXX_phase_align_relu.png` | 7 | Same as Fig. 8 right panel. |
+| — Decoupled ODE simulation | Pure ODE integration (no neural network) showing phase convergence and magnitude competition for all frequencies within one neuron. Two cases with different initial conditions. | `pXXX_phase_align_approx1.png`, `pXXX_phase_align_approx2.png` | 8 | Generated by `generate_analytical.py`, not `generate_plots.py`. Validates the theory without any training. |
+
+### Part III: Grokking (Tab 6)
+
+| Blog Figure | Description | Pipeline Output | Tab | Verified? |
+|-------------|-------------|----------------|-----|-----------|
+| **Fig. 13a** — Grokking loss curves | Training and test loss over 50k epochs. Three stages: (I) train loss drops, (II) test loss drops, (III) both near zero. Stage boundaries marked. | `pXXX_grokk_loss.json` → interactive Plotly chart | 6 | From `training_curves.json`. Stage boundaries detected by `grokking_stage_detector.py`. Shaded regions distinguish the three stages. |
+| **Fig. 13b** — Grokking accuracy curves | Training and test accuracy. Train → 100% in Stage I, test jumps in Stage II. | `pXXX_grokk_acc.json` → interactive Plotly chart | 6 | Computed by running forward pass on train/test data at each checkpoint. |
+| **Fig. 13c** — Phase alignment progress | Average |sin(D_m*)| over training. Decreases throughout, steepest in Stage II. | `pXXX_grokk_abs_phase_diff.png` | 6 | Computed via `decode_weights()` + `compute_neuron()` at each grokking checkpoint. |
+| **Fig. 13d** — IPR and parameter norm | Dual-axis: IPR (Fourier sparsity) increases sharply in Stage II; parameter norm shrinks in Stage III. | `pXXX_grokk_avg_ipr.png` | 6 | IPR uses the corrected per-frequency magnitude formula: A_k = sqrt(c_k² + s_k²), IPR = Σ A_k⁴ / (Σ A_k²)². Parameter norms from `training_curves.json`. |
+| **Fig. 14** — Memorization accuracy heatmap | Three panels at end of Stage I: (1) training data distribution under symmetry, (2) accuracy grid, (3) softmax probability at ground truth. Red rectangles = true test pairs. | `pXXX_grokk_memorization_accuracy.png` | 6 | Forward pass at the checkpoint closest to stage1_end. The symmetric architecture guarantees ~70% test accuracy during memorization. |
+| **Fig. 15** — Common-to-rare memorization | Four panels: training data distribution + accuracy at 3 timepoints during Stage I. Shows common pairs (both (i,j) and (j,i) in train) memorized before rare pairs (only one ordering). | `pXXX_grokk_memorization_common_to_rare.png` | 6 | Epochs selected at 0, stage1_end/2, stage1_end. Red rectangles mark asymmetric training pairs. |
+| **Fig. 16** — Weight evolution during grokking | 2×3 DFT heatmap grid showing θ_m and ξ_m at Step 0 (random init), end of Stage I (noisy multi-frequency), and end of Stage II (clean single-frequency). | `pXXX_grokk_decoded_weights_dynamic.png` | 6 | DFT coefficients `W @ fourier_basis.T` at 3 key epochs. The transition from diffuse to sparse confirms the sparsification narrative. |
+
+### Overview (Tab 1)
+
+| Blog Figure | Description | Pipeline Output | Tab | Verified? |
+|-------------|-------------|----------------|-----|-----------|
+| — Overview dashboard | 2×2 grid: standard loss + grokking loss (top), standard IPR + grokking IPR (bottom). Plus phase scatter from standard final checkpoint. | `pXXX_overview_loss_ipr.png`, `pXXX_overview_phase_scatter.png`, `pXXX_overview.json` | 1 | Combines data from standard and grokking runs. The phase scatter uses the same computation as Fig. 4. For $p < 19$, only the standard column is shown. |
+
+### Not Currently Generated (Blog-Only Figures)
+
+| Blog Figure | Description | Status |
+|-------------|-------------|--------|
+| **Fig. 1** — Architecture illustration | Schematic of the two-layer network, DFT decomposition, and mechanism overview. | Static illustration, not dependent on p. Could be included as a fixed image in the app. |
+
+---
+
+## Interactive JSON Data Files
+
+In addition to static PNG plots, the pipeline generates JSON files for interactive Plotly charts in the Gradio app:
+
+| File | Content | Used In |
+|------|---------|---------|
+| `pXXX_overview.json` | Standard loss/IPR + grokking loss/IPR time series | Tab 1: Interactive loss and IPR charts |
+| `pXXX_neuron_spectra.json` | Per-neuron Fourier magnitudes (W_in and W_out) for top-20 neurons, sorted by frequency | Tab 2: Neuron Inspector dropdown → bar chart of frequency decomposition |
+| `pXXX_logits_interactive.json` | Output logits for p representative (a,b) pairs, plus correct answers | Tab 4: Logit Explorer dropdown → bar chart with correct answer highlighted |
+| `pXXX_grokk_loss.json` | Full training/test loss curves + stage boundaries | Tab 6: Interactive loss chart with stage shading |
+| `pXXX_grokk_acc.json` | Accuracy at each checkpoint epoch + stage boundaries | Tab 6: Interactive accuracy chart with stage shading |
+| `pXXX_grokk_epoch_data.json` | p×p accuracy grids at ~10 evenly-spaced grokking epochs | Tab 6: Epoch Slider → heatmap animation across training |
+| `pXXX_metadata.json` | Config for all 5 runs + final metrics (loss, accuracy) | Displayed in the app's info panel for the selected $p$ |
+
+---
+
+## Output Structure
+
+All plots for a modulus are saved in a single flat directory. Each file is prefixed with `pXXX_` so the folder is self-contained and browsable:
+
+```
+precomputed_results/p_023/
+  # Metadata
+  p023_metadata.json
+
+  # Tab 1: Overview (Blog: summary of standard + grokking)
+  p023_overview_loss_ipr.png          # 2×2 grid: loss + IPR for both setups
+  p023_overview_phase_scatter.png     # Phase alignment scatter (same as Fig. 4)
+  p023_overview.json                  # Interactive data
+
+  # Tab 2: Fourier Weights (Blog: Figures 2, 3)
+  p023_full_training_para_origin.png  # DFT heatmap (Fig. 2)
+  p023_lineplot_in.png                # Cosine fits, input layer (Fig. 3 left)
+  p023_lineplot_out.png               # Cosine fits, output layer (Fig. 3 right)
+  p023_neuron_spectra.json            # Interactive: neuron inspector
+
+  # Tab 3: Phase Analysis (Blog: Figures 4, 5, 6)
+  p023_phase_distribution.png         # Polar phase plot (Fig. 5)
+  p023_phase_relationship.png         # 2φ vs ψ scatter (Fig. 4)
+  p023_magnitude_distribution.png     # Violin plots (Fig. 6)
+
+  # Tab 4: Output Logits (Blog: Figure 7)
+  p023_output_logits.png              # Logit heatmap (Fig. 7)
+  p023_logits_interactive.json        # Interactive: logit explorer
+
+  # Tab 5: Lottery Mechanism (Blog: Figures 9, 10)
+  p023_lottery_mech_magnitude.png     # Magnitude race (Fig. 9 right)
+  p023_lottery_mech_phase.png         # Phase misalignment race (Fig. 9 left)
+  p023_lottery_beta_contour.png       # Contour plot (Fig. 10)
+
+  # Tab 6: Grokking (Blog: Figures 13, 14, 15, 16)
+  p023_grokk_loss.json                # Interactive loss curves (Fig. 13a)
+  p023_grokk_acc.json                 # Interactive accuracy curves (Fig. 13b)
+  p023_grokk_abs_phase_diff.png       # Phase alignment progress (Fig. 13c)
+  p023_grokk_avg_ipr.png              # IPR + param norms (Fig. 13d)
+  p023_grokk_memorization_accuracy.png          # 3-panel heatmap (Fig. 14)
+  p023_grokk_memorization_common_to_rare.png    # 4-panel sequence (Fig. 15)
+  p023_grokk_decoded_weights_dynamic.png        # DFT evolution (Fig. 16)
+  p023_grokk_epoch_data.json          # Interactive: epoch slider
+
+  # Tab 7: Gradient Dynamics (Blog: Figures 8, 11, 12)
+  p023_phase_align_quad.png           # Phase + magnitude, Quad (Fig. 8 left)
+  p023_single_freq_quad.png           # DFT heatmap over time, Quad (Fig. 11)
+  p023_phase_align_relu.png           # Phase + magnitude, ReLU (Fig. 8 right / 12b)
+  p023_single_freq_relu.png           # DFT heatmap over time, ReLU (Fig. 12a)
+
+  # Tab 8: Decoupled Simulation (no blog figure number)
+  p023_phase_align_approx1.png        # ODE simulation case 1
+  p023_phase_align_approx2.png        # ODE simulation case 2
+```
+
+**29 files per $p$:** 21 PNGs + 6 JSONs from trained models, 2 PNGs from analytical simulation.
+
+---
+
+## Correctness Verification
+
+### How each computation matches the paper
+
+1. **DFT Decomposition (Figs. 2, 11, 12a, 16):** We compute `W @ fourier_basis.T` where `fourier_basis` is the orthonormal DFT basis from `get_fourier_basis(p)`. The basis has rows: [Const, cos 1, sin 1, cos 2, sin 2, ..., cos K, sin K] with K = (p-1)/2 for odd $p$. Each row is L2-normalized. This matches the standard real DFT on Z_p.
+
+2. **Phase extraction (Figs. 4, 8, 9, 13c):** For frequency k, the DFT gives coefficients (c_k, s_k) at indices (2k-1, 2k). The magnitude is α = sqrt(2/p) · sqrt(c_k² + s_k²), and the phase is φ = arctan2(-s_k, c_k). This convention matches the paper's θ_m[j] = α cos(ω_k j + φ) representation.
+
+3. **IPR (Figs. 13d, Overview):** Uses the corrected per-frequency magnitude formula: A_k = sqrt(c_k² + s_k²) (combining cos/sin pairs), then IPR = Σ A_k⁴ / (Σ A_k²)². This gives IPR → 1 for perfect single-frequency neurons, matching the paper's definition.
+
+4. **Phase alignment (Fig. 4):** The doubled-phase relationship ψ_m = 2φ_m is verified by extracting φ from W_in and ψ from W_out using the same `compute_neuron()` function, then plotting (2φ mod 2π) vs (ψ mod 2π).
+
+5. **Output logits (Fig. 7):** Forward pass uses the **same activation function** as training (ReLU for standard run). The flawed indicator structure (main diagonal + two ghost diagonals) is visible because the standard run trains to 100% accuracy with clean Fourier features.
+
+6. **Lottery mechanism (Figs. 9, 10):** Uses the `quad_random` run (quadratic activation, random init) which matches the theoretical setting. `decode_scales_phis()` extracts per-frequency magnitudes and phases at each checkpoint. The winning frequency is the one with smallest initial |D| = |2φ - ψ|.
+
+7. **Grokking stages (Figs. 13--16):** `grokking_stage_detector.py` identifies stage boundaries from training curves. Stage I ends when train accuracy ≈ 1.0, Stage II ends when test accuracy ≈ 1.0. Memorization heatmaps use forward pass at the closest checkpoint to stage1_end.
+
+8. **Analytical simulation (Tab 8):** Numerically integrates the four-variable ODE system from Section 5.3 of the paper. No neural network is involved — this validates the theoretical dynamics directly.
+
+### Why results generalize across $p$
+
+The paper's theory is stated for general odd $p$. Key properties that scale:
+
+- **Fourier basis:** Always has (p-1)/2 non-DC frequencies for any odd $p$.
+- **Phase alignment:** The ψ = 2φ relationship is a consequence of the gradient dynamics, independent of p.
+- **Lottery mechanism:** Random initial misalignments are uniform on [0, 2π) for any p.
+- **Grokking three stages:** The stage structure depends on the balance of loss gradient vs. weight decay, not on p specifically (though the stage durations and test accuracy during memorization may vary).
+- **Network width:** d_mlp scales as O(p²) to maintain the neuron-to-frequency ratio, ensuring enough neurons per frequency for diversification.
+
+---
+
+## Scripts
+
+| Script | Purpose |
+|--------|---------|
+| `run_pipeline.sh` | Runs the complete pipeline (train + plots + analytical + verify) for a single modulus. |
+| `train_all.py` | Trains all 5 model configurations. Saves checkpoints + `training_curves.json`. |
+| `generate_plots.py` | Loads trained models and generates all model-dependent plots (Tabs 1--7) plus interactive JSONs and metadata. |
+| `generate_analytical.py` | Runs gradient flow simulations to generate theory plots (Tab 8). No model needed. |
+| `prime_config.py` | Configuration: moduli list, d_mlp formula, training run parameters. |
+| `neuron_selector.py` | Automated neuron selection for plots (replaces hardcoded indices from notebooks). |
+| `grokking_stage_detector.py` | Detects memorization/transition/generalization stage boundaries from training curves. |
+
+---
+
+## Analytical Simulations (No Model Needed)
+
+`generate_analytical.py` produces 2 plots per $p$ by simulating gradient flow on decoupled frequency dynamics. These validate the theoretical analysis without training any model.
+
+- **Case 1**: Shows phase difference D* converging from initial conditions (φ₀=1.5, ψ₀=0.18)
+- **Case 2**: Different initial conditions (φ₀=-0.72, ψ₀=-2.91) showing convergence from the other side
+
+Both cases confirm the phase alignment theorem: D → 0 is the stable attractor, D → π is unstable.

precompute/generate_analytical.py

@@ -0,0 +1,358 @@
+#!/usr/bin/env python3
+"""
+Generate "Decoupled Simulation" plots -- analytical gradient flow simulations
+that don't require trained models.
+
+Produces 2 plots per p, saved to {output_dir}/p_{p:03d}/:
+  1. p{p:03d}_phase_align_approx1.png  -- case 1: longer simulation with annotations
+  2. p{p:03d}_phase_align_approx2.png  -- case 2: shorter simulation
+
+Usage:
+    python generate_analytical.py --all
+    python generate_analytical.py --p 23
+    python generate_analytical.py --p 23 --output ./my_output
+"""
+import argparse
+import os
+import sys
+
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+# Add project root to path so we can import src modules
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+from mechanism_base import get_fourier_basis, normalize_to_pi
+from prime_config import get_moduli, ANALYTICAL_CONFIGS
+
+# ---------------------------------------------------------------------------
+# Style constants
+# ---------------------------------------------------------------------------
+COLORS = ['#0D2758', '#60656F', '#DEA54B', '#A32015', '#347186']
+DPI = 150
+
+
+# ===========================================================================
+# Decouple dynamics simulation
+# ===========================================================================
+
+def gradient_update(theta, xi, p, device):
+    """
+    Compute the sum of gradients over all frequency modes k.
+
+    For each frequency k from 1 to (p-1)//2, project theta and xi onto the
+    Fourier basis to obtain 2-coefficient vectors, then compute alpha, phi,
+    beta, psi and the corresponding gradient contributions.
+    """
+    fourier_basis, _ = get_fourier_basis(p, device)
+    fourier_basis = fourier_basis.to(theta.dtype)
+    theta_coeff = fourier_basis @ theta
+    xi_coeff = fourier_basis @ xi
+
+    total_grad_theta = torch.zeros_like(theta)
+    total_grad_xi = torch.zeros_like(xi)
+
+    j_values = torch.arange(p, device=device, dtype=theta.dtype)
+    factor = np.sqrt(2.0 / p)
+
+    for k in range(1, p // 2 + 1):
+        coeff_indices = [k * 2 - 1, k * 2]
+        neuron_coeff_theta = theta_coeff[coeff_indices]
+        neuron_coeff_xi = xi_coeff[coeff_indices]
+
+        alpha = factor * torch.norm(neuron_coeff_theta, dim=0)
+        phi = torch.arctan2(-neuron_coeff_theta[1], neuron_coeff_theta[0])
+
+        beta = factor * torch.norm(neuron_coeff_xi, dim=0)
+        psi = torch.arctan2(-neuron_coeff_xi[1], neuron_coeff_xi[0])
+
+        w_k = 2 * np.pi * k / p
+        grad_theta_k = 2 * p * alpha * beta * torch.cos(w_k * j_values + psi - phi)
+        grad_xi_k = p * alpha.pow(2) * torch.cos(w_k * j_values + 2 * phi)
+
+        total_grad_theta += grad_theta_k / p ** 2
+        total_grad_xi += grad_xi_k / p ** 2
+
+    return total_grad_theta, total_grad_xi
+
+
+def simulate_gradient_flow(theta, xi, p, num_steps, learning_rate, device):
+    """Euler integration of the coupled gradient-flow ODEs."""
+    theta_history = [theta.clone()]
+    xi_history = [xi.clone()]
+
+    for _ in range(num_steps):
+        grad_theta, grad_xi = gradient_update(theta, xi, p, device)
+        theta = theta + learning_rate * grad_theta
+        xi = xi + learning_rate * grad_xi
+        theta_history.append(theta.clone())
+        xi_history.append(xi.clone())
+
+    return theta_history, xi_history
+
+
+def analyze_history(theta_history, xi_history, p, fourier_basis):
+    """
+    Extract time series of alpha, phi, beta, psi, delta for every frequency k.
+    """
+    theta_hist_tensor = torch.stack(theta_history)
+    xi_hist_tensor = torch.stack(xi_history)
+
+    theta_coeffs_hist = fourier_basis @ theta_hist_tensor.T
+    xi_coeffs_hist = fourier_basis @ xi_hist_tensor.T
+
+    results = {
+        'alphas': {}, 'phis': {}, 'betas': {}, 'psis': {}, 'deltas': {}
+    }
+    factor = np.sqrt(2.0 / p)
+
+    for k in range(1, p // 2 + 1):
+        idx = [k * 2 - 1, k * 2]
+        neuron_theta_hist = theta_coeffs_hist[idx, :]
+        neuron_xi_hist = xi_coeffs_hist[idx, :]
+
+        alphas_k = factor * torch.norm(neuron_theta_hist, dim=0)
+        phis_k = torch.atan2(-neuron_theta_hist[1, :], neuron_theta_hist[0, :])
+
+        betas_k = factor * torch.norm(neuron_xi_hist, dim=0)
+        psis_k = torch.atan2(-neuron_xi_hist[1, :], neuron_xi_hist[0, :])
+
+        deltas_k = normalize_to_pi(2 * phis_k - psis_k)
+
+        results['alphas'][k] = alphas_k.numpy()
+        results['phis'][k] = phis_k.numpy()
+        results['betas'][k] = betas_k.numpy()
+        results['psis'][k] = psis_k.numpy()
+        results['deltas'][k] = deltas_k.numpy()
+
+    return results
+
+
+def _run_decouple_simulation(p, init_k, num_steps, lr, init_phi, init_psi,
+                             amplitude, device):
+    """Initialize and run a single decouple-dynamics simulation."""
+    fourier_basis, _ = get_fourier_basis(p, device)
+    fourier_basis = fourier_basis.to(torch.float64)
+    w_k = 2 * np.pi * init_k / p
+
+    theta_init = amplitude * torch.tensor(
+        [np.cos(w_k * j + init_phi) for j in range(p)],
+        device=device, dtype=torch.float64,
+    )
+    xi_init = amplitude * torch.tensor(
+        [np.cos(w_k * j + init_psi) for j in range(p)],
+        device=device, dtype=torch.float64,
+    )
+
+    theta_history, xi_history = simulate_gradient_flow(
+        theta_init, xi_init, p, num_steps, lr, device,
+    )
+    results = analyze_history(theta_history, xi_history, p, fourier_basis)
+    return results
+
+
+def _plot_decouple(results, p, num_steps, lr, init_k, save_path,
+                   show_vline=True, vline_x=500):
+    """
+    Publication-quality 3-panel figure:
+      Top:    psi_k* and 2*phi_k* vs time
+      Middle: D_k* (phase difference) vs time, horizontal line at pi/2
+      Bottom: alpha_k* and beta_k* vs time
+    """
+    plt.rcParams['mathtext.fontset'] = 'cm'
+
+    alphas = np.array(results['alphas'][init_k])
+    betas = np.array(results['betas'][init_k])
+    deltas = np.array(results['deltas'][init_k])
+    phis = np.array(results['phis'][init_k])
+    psis = np.array(results['psis'][init_k])
+
+    # Phase wrapping fix: normalize 2*phi to [-pi,pi], adjust psi to
+    # stay within pi of 2*phi, then unwrap the time series so there
+    # are no discontinuous jumps at +-pi boundaries.
+    def _fix_phase_pair(two_phi_raw, psi_raw):
+        two_phi = normalize_to_pi(two_phi_raw)
+        psi_fixed = normalize_to_pi(psi_raw).copy()
+        diff = psi_fixed - two_phi
+        psi_fixed[diff > np.pi] -= 2 * np.pi
+        psi_fixed[diff < -np.pi] += 2 * np.pi
+        return np.unwrap(two_phi), np.unwrap(psi_fixed)
+
+    phis2_plot, psis_plot = _fix_phase_pair(2 * phis, psis)
+
+    x = np.arange(num_steps + 1) * lr
+    vline_kwargs = dict(color='gray', linestyle='--', linewidth=1.5)
+
+    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(7, 9), sharex=True)
+    fig.suptitle(f'Decoupled Gradient Flow (p={p})', fontsize=20, y=1.01)
+
+    # --- Top: phase alignment ---
+    for k in range(1, (p - 1) // 2 + 1):
+        if k != init_k:
+            bg_2phi, bg_psi = _fix_phase_pair(
+                2 * np.array(results['phis'][k]),
+                np.array(results['psis'][k]),
+            )
+            ax1.plot(x, bg_psi, lw=1.5, alpha=0.4, color='gray')
+            ax1.plot(x, bg_2phi, lw=1.5, alpha=0.4, color='gray',
+                     linestyle='--')
+    ax1.plot(x, psis_plot, color=COLORS[3], linewidth=2.5,
+             label=r"$\psi_{k^\star}$")
+    ax1.plot(x, phis2_plot, linewidth=2.5, color=COLORS[0],
+             label=r"$2\phi_{k^\star}$")
+    if show_vline:
+        ax1.axvline(x=vline_x, **vline_kwargs)
+    ax1.set_title('Dynamics of Phase Alignment', fontsize=18)
+    ax1.set_ylabel('Phase (radians)', fontsize=14)
+    ax1.legend(fontsize=18)
+    ax1.grid(True)
+
+    # --- Middle: phase difference ---
+    for k in range(1, (p - 1) // 2 + 1):
+        if k != init_k:
+            ax2.plot(x, np.array(results['deltas'][k]),
+                     lw=1.5, alpha=0.4, color='gray')
+    ax2.plot(x, deltas, color=COLORS[0], linewidth=2.5,
+             label=r"$D_{k^\star}$")
+    if show_vline:
+        ax2.axvline(x=vline_x, **vline_kwargs)
+        ax2.axhline(y=np.pi / 2, **vline_kwargs)
+        ax2.text(x=max(x) * 0.05, y=np.pi / 2 - 0.45,
+                 s=r"$D^\star_{k^\star}=\pi/2$", fontsize=16, color='black')
+    ax2.set_title('Dynamics of Phase Difference', fontsize=18)
+    ax2.set_ylabel('Phase (radians)', fontsize=14)
+    ax2.legend(fontsize=18)
+    ax2.grid(True)
+
+    # --- Bottom: magnitude evolution ---
+    for k in range(1, (p - 1) // 2 + 1):
+        if k != init_k:
+            ax3.plot(x, np.array(results['alphas'][k]),
+                     lw=1.5, alpha=0.4, color='gray')
+            ax3.plot(x, np.array(results['betas'][k]),
+                     lw=1.5, alpha=0.4, color='gray', linestyle='--')
+    ax3.plot(x, alphas, linewidth=2.5, color=COLORS[0],
+             label=r"$\alpha_{k^\star}$")
+    ax3.plot(x, betas, linewidth=2.5, color=COLORS[3],
+             label=r"$\beta_{k^\star}$")
+    if show_vline:
+        ax3.axvline(x=vline_x, **vline_kwargs)
+    ax3.set_title('Magnitude Evolution', fontsize=18)
+    ax3.set_xlabel('Time', fontsize=18)
+    ax3.set_ylabel('Magnitude', fontsize=14)
+    ax3.legend(fontsize=18)
+    ax3.grid(True)
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=DPI, bbox_inches='tight')
+    plt.close(fig)
+    print(f"  Saved {save_path}")
+
+
+def generate_decouple_dynamics(p, output_dir):
+    """Generate the two decouple-dynamics phase-alignment plots."""
+    max_freq = (p - 1) // 2
+    if max_freq < 1:
+        print(f"  SKIP: p={p} has no non-DC frequencies for analytical simulation")
+        return
+
+    cfg = ANALYTICAL_CONFIGS["decouple_dynamics"]
+    device = torch.device("cpu")
+    init_k = min(cfg["init_k"], max_freq)
+    amplitude = cfg["amplitude"]
+
+    # Case 1: longer simulation with vline annotations
+    print(f"  Running decouple dynamics case 1 (p={p}) ...")
+    results1 = _run_decouple_simulation(
+        p, init_k,
+        num_steps=cfg["num_steps_case1"],
+        lr=cfg["learning_rate_case1"],
+        init_phi=cfg["init_phi_case1"],
+        init_psi=cfg["init_psi_case1"],
+        amplitude=amplitude,
+        device=device,
+    )
+    _plot_decouple(
+        results1, p,
+        num_steps=cfg["num_steps_case1"],
+        lr=cfg["learning_rate_case1"],
+        init_k=init_k,
+        save_path=os.path.join(output_dir, f"p{p:03d}_phase_align_approx1.png"),
+        show_vline=True,
+        vline_x=cfg["num_steps_case1"] * cfg["learning_rate_case1"] * 0.36,
+    )
+
+    # Case 2: shorter simulation without vline annotations
+    print(f"  Running decouple dynamics case 2 (p={p}) ...")
+    results2 = _run_decouple_simulation(
+        p, init_k,
+        num_steps=cfg["num_steps_case2"],
+        lr=cfg["learning_rate_case2"],
+        init_phi=cfg["init_phi_case2"],
+        init_psi=cfg["init_psi_case2"],
+        amplitude=amplitude,
+        device=device,
+    )
+    _plot_decouple(
+        results2, p,
+        num_steps=cfg["num_steps_case2"],
+        lr=cfg["learning_rate_case2"],
+        init_k=init_k,
+        save_path=os.path.join(output_dir, f"p{p:03d}_phase_align_approx2.png"),
+        show_vline=False,
+    )
+
+
+# ===========================================================================
+# Entry point
+# ===========================================================================
+
+def generate_all_for_prime(p, output_base):
+    """Generate the 2 decoupled simulation plots for a single prime."""
+    output_dir = os.path.join(output_base, f"p_{p:03d}")
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(f"\n{'='*60}")
+    print(f"Generating decoupled simulation plots for p={p}")
+    print(f"Output: {output_dir}")
+    print(f"{'='*60}")
+
+    # Use float64 globally for numerical precision in simulations
+    prev_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(torch.float64)
+
+    try:
+        generate_decouple_dynamics(p, output_dir)
+    finally:
+        torch.set_default_dtype(prev_dtype)
+
+    print(f"[DONE] p={p}: 2 plots written to {output_dir}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Generate decoupled simulation plots (analytical, no model needed)'
+    )
+    parser.add_argument('--all', action='store_true',
+                        help='Generate plots for all odd p in [3, 199]')
+    parser.add_argument('--p', type=int,
+                        help='Generate plots for a specific p')
+    parser.add_argument('--output', type=str, default='./precomputed_results',
+                        help='Base output directory (default: ./precomputed_results)')
+    args = parser.parse_args()
+
+    if not args.all and args.p is None:
+        parser.error("Specify --all or --p P")
+
+    moduli = [args.p] if args.p else get_moduli()
+
+    for p in moduli:
+        generate_all_for_prime(p, args.output)
+
+    print(f"\nAll done. Processed {len(moduli)} value(s) of p.")
+
+
+if __name__ == "__main__":
+    main()

precompute/generate_plots.py

@@ -0,0 +1,2192 @@
+#!/usr/bin/env python3
+"""
+Main plot generation script for the HF app.
+Creates all model-dependent plots (Tabs 1-7) from trained checkpoints.
+
+Usage:
+    python generate_plots.py --all               # Generate for all primes
+    python generate_plots.py --p 23           # Generate for a specific p
+    python generate_plots.py --p 23 --input ./trained_models --output ./hf_app/precomputed_results
+"""
+import matplotlib
+matplotlib.use('Agg')
+
+import argparse
+import json
+import math
+import os
+import sys
+import traceback
+
+import matplotlib.colors as mcolors
+import matplotlib.patches as patches
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+from matplotlib.colors import LinearSegmentedColormap
+from matplotlib.ticker import FuncFormatter
+
+# Add project root to path so we can import src modules
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+sys.path.insert(0, PROJECT_ROOT)
+sys.path.insert(0, os.path.join(PROJECT_ROOT, 'src'))
+sys.path.insert(0, os.path.dirname(__file__))
+
+from src.mechanism_base import (
+    get_fourier_basis,
+    decode_weights,
+    compute_neuron,
+    decode_scales_phis,
+    normalize_to_pi,
+)
+from src.model_base import EmbedMLP
+from src.utils import cross_entropy_high_precision, acc_rate
+from precompute.neuron_selector import (
+    select_top_neurons_by_frequency,
+    select_lineplot_neurons,
+    select_phase_frequency,
+    select_lottery_neuron,
+)
+from precompute.grokking_stage_detector import detect_grokking_stages
+from precompute.prime_config import compute_d_mlp, TRAINING_RUNS, MIN_P_GROKKING
+
+# ---------- Lightweight train/test data regeneration ----------
+
+def _gen_train_test(p, frac_train=0.75, seed=42):
+    """
+    Regenerate train/test split deterministically without needing a Config object.
+    Mirrors the logic in utils.gen_train_test for the 'add' function.
+    Returns (train_data, test_data) where each is a tensor of shape (N, 2).
+    """
+    import random as _random
+    all_pairs = []
+    for i in range(p):
+        for j in range(p):
+            all_pairs.append((i, j))
+    data_tensor = torch.tensor(all_pairs, dtype=torch.long)
+    _random.seed(seed)
+    indices = torch.randperm(len(all_pairs))
+    data_tensor = data_tensor[indices]
+    if frac_train >= 1.0:
+        return data_tensor, data_tensor
+    div = int(frac_train * len(all_pairs))
+    return data_tensor[:div], data_tensor[div:]
+
+
+# ---------- Style constants ----------
+COLORS = ['#0D2758', '#60656F', '#DEA54B', '#A32015', '#347186']
+CMAP_DIVERGING = LinearSegmentedColormap.from_list(
+    'cividis_white_center', ['#0D2758', 'white', '#A32015'], N=256
+)
+CMAP_SEQUENTIAL = LinearSegmentedColormap.from_list(
+    'cividis_white_seq', ['white', '#0D2758'], N=256
+)
+DPI = 150
+plt.rcParams['mathtext.fontset'] = 'cm'
+
+
+def _save_fig(fig, path):
+    """Save a figure and close it."""
+    fig.savefig(path, dpi=DPI, bbox_inches='tight', format='png')
+    plt.close(fig)
+
+
+# ======================================================================
+# Helpers for loading checkpoints
+# ======================================================================
+
+def _find_run_dir(base_dir):
+    """
+    Given a run type directory (e.g. trained_models/p_023/standard/),
+    find the actual checkpoint directory.  It may be a timestamped
+    subdirectory, or the checkpoints may live directly in base_dir.
+    Returns the path that contains the .pth checkpoint files.
+    """
+    if not os.path.isdir(base_dir):
+        return None
+
+    # Check if .pth files live directly here
+    pth_files = [f for f in os.listdir(base_dir)
+                 if f.endswith('.pth') and f not in ('train_data.pth', 'test_data.pth')]
+    if pth_files:
+        return base_dir
+
+    # Otherwise look for a single timestamped subdirectory
+    subdirs = [d for d in os.listdir(base_dir)
+               if os.path.isdir(os.path.join(base_dir, d))]
+    for sd in sorted(subdirs):
+        candidate = os.path.join(base_dir, sd)
+        files = os.listdir(candidate)
+        if any(f.endswith('.pth') for f in files):
+            return candidate
+    return None
+
+
+def _load_checkpoints(run_dir, device='cpu'):
+    """
+    Load all numbered checkpoints from run_dir.
+    Returns dict {epoch_int: state_dict} sorted by epoch.
+    """
+    loaded = {}
+    exclude = {'final.pth', 'test_data.pth', 'train_data.pth'}
+    for fname in os.listdir(run_dir):
+        fpath = os.path.join(run_dir, fname)
+        if (os.path.isfile(fpath) and fname.endswith('.pth')
+                and fname not in exclude):
+            try:
+                epoch = int(os.path.splitext(fname)[0])
+            except ValueError:
+                continue
+            data = torch.load(fpath, weights_only=True, map_location=device)
+            if isinstance(data, dict) and 'model' in data:
+                loaded[epoch] = data['model']
+            else:
+                loaded[epoch] = data
+    return {k: loaded[k] for k in sorted(loaded)}
+
+
+def _load_final(run_dir, device='cpu'):
+    """Load the final.pth model data dict."""
+    fpath = os.path.join(run_dir, 'final.pth')
+    if not os.path.exists(fpath):
+        # Fall back to largest epoch checkpoint
+        ckpts = _load_checkpoints(run_dir, device)
+        if ckpts:
+            max_epoch = max(ckpts.keys())
+            return {'model': ckpts[max_epoch]}
+        return None
+    return torch.load(fpath, weights_only=True, map_location=device)
+
+
+def _load_training_curves(run_type_dir):
+    """Load training_curves.json from the run type directory."""
+    path = os.path.join(run_type_dir, 'training_curves.json')
+    if os.path.exists(path):
+        with open(path) as f:
+            return json.load(f)
+    # Fall back: check inside the checkpoint subdirectory
+    run_dir = _find_run_dir(run_type_dir)
+    if run_dir and run_dir != run_type_dir:
+        path = os.path.join(run_dir, 'training_curves.json')
+        if os.path.exists(path):
+            with open(path) as f:
+                return json.load(f)
+    # Fall back: try loading from final.pth
+    if run_dir:
+        final_path = os.path.join(run_dir, 'final.pth')
+        if os.path.exists(final_path):
+            data = torch.load(final_path, weights_only=True, map_location='cpu')
+            if isinstance(data, dict):
+                curves = {}
+                for key in ('train_losses', 'test_losses', 'train_accs', 'test_accs',
+                            'grad_norms', 'param_norms'):
+                    if key in data:
+                        val = data[key]
+                        if isinstance(val, torch.Tensor):
+                            val = val.cpu().tolist()
+                        curves[key] = val
+                if curves:
+                    return curves
+    return None
+
+
+# ======================================================================
+# PlotGenerator
+# ======================================================================
+
+class PlotGenerator:
+    """
+    Generates all model-dependent plots for a single prime p.
+
+    Parameters
+    ----------
+    p : int
+        The prime modulus.
+    input_dir : str
+        Path to trained_models/p_PPP/ containing run-type subdirectories.
+    output_dir : str
+        Path to hf_app/precomputed_results/p_PPP/ where plots are saved.
+    """
+
+    def __init__(self, p, input_dir, output_dir):
+        self.p = p
+        self.input_dir = input_dir
+        self.output_dir = output_dir
+        self.device = 'cpu'
+        self.d_vocab = p
+        self.d_model = p
+
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Infer d_mlp from checkpoint weights; fall back to formula
+        self.d_mlp = self._infer_d_mlp() or compute_d_mlp(p)
+
+        # Fourier basis (mechanism_base version with device arg)
+        self.fourier_basis, self.fourier_basis_names = get_fourier_basis(p, self.device)
+
+        # All (a,b) pairs and labels
+        self.all_data = torch.tensor(
+            [(i, j) for i in range(p) for j in range(p)], dtype=torch.long
+        )
+        self.all_labels = torch.tensor(
+            [(i + j) % p for i in range(p) for j in range(p)], dtype=torch.long
+        )
+
+    def _infer_d_mlp(self):
+        """Infer d_mlp from the first available checkpoint's weight shape."""
+        for run_name in TRAINING_RUNS:
+            run_type_dir = os.path.join(self.input_dir, run_name)
+            run_dir = _find_run_dir(run_type_dir)
+            if run_dir is None:
+                continue
+            final = _load_final(run_dir, 'cpu')
+            if final and 'model' in final and 'mlp.W_in' in final['model']:
+                d_mlp = final['model']['mlp.W_in'].shape[0]
+                print(f"  Inferred d_mlp={d_mlp} from {run_name} checkpoint")
+                return d_mlp
+        return None
+
+    # ------------------------------------------------------------------
+    # Path helpers
+    # ------------------------------------------------------------------
+
+    def _run_type_dir(self, run_name):
+        return os.path.join(self.input_dir, run_name)
+
+    def _run_dir(self, run_name):
+        return _find_run_dir(self._run_type_dir(run_name))
+
+    def _out(self, filename):
+        # Prefix every file with pXXX_ so folders are self-contained and browsable
+        return os.path.join(self.output_dir, f"p{self.p:03d}_{filename}")
+
+    # ------------------------------------------------------------------
+    # ------------------------------------------------------------------
+    # Shared IPR helper
+    # ------------------------------------------------------------------
+
+    def _compute_freq_ipr(self, W_dec):
+        """IPR over per-frequency magnitudes (combines cos+sin pairs).
+
+        IPR = sum_k A_k^4 / (sum_k A_k^2)^2, where A_k = sqrt(c_k^2 + s_k^2).
+        IPR → 1 means all energy at a single frequency.
+        Returns mean IPR across neurons.
+        """
+        K = (self.p - 1) // 2
+        A2 = torch.zeros(W_dec.shape[0], K)
+        for k in range(1, K + 1):
+            A2[:, k - 1] = W_dec[:, 2 * k - 1].pow(2) + W_dec[:, 2 * k].pow(2)
+        A4 = A2.pow(2)
+        denom = A2.sum(dim=1).pow(2)
+        valid = denom > 0
+        ipr = torch.zeros(W_dec.shape[0])
+        ipr[valid] = A4[valid].sum(dim=1) / denom[valid]
+        return ipr.mean()
+
+    def _ipr_at_checkpoint(self, model_sd):
+        """Compute average IPR (across both layers) for a single checkpoint."""
+        W_in_d, W_out_d, _ = decode_weights(model_sd, self.fourier_basis)
+        return ((self._compute_freq_ipr(W_in_d)
+                 + self._compute_freq_ipr(W_out_d)) / 2).item()
+
+    # ------------------------------------------------------------------
+    # Tab 1: Overview (standard loss+IPR, grokking loss+IPR, phase plot)
+    # ------------------------------------------------------------------
+
+    def generate_tab1(self):
+        """Generate overview plots: standard + grokking loss/IPR, plus phase scatter."""
+        print(f"  [Tab 1] Overview for p={self.p}")
+
+        # ---- Standard run: loss + IPR ----
+        std_dir = self._run_dir('standard')
+        std_epochs, std_loss, std_ipr = [], [], []
+        if std_dir is not None:
+            std_curves = _load_training_curves(self._run_type_dir('standard'))
+            std_ckpts = _load_checkpoints(std_dir, self.device)
+            if std_ckpts:
+                std_epochs = sorted(std_ckpts.keys())
+                std_ipr = [self._ipr_at_checkpoint(std_ckpts[ep]) for ep in std_epochs]
+                if std_curves and 'train_losses' in std_curves:
+                    se = std_epochs[1] - std_epochs[0] if len(std_epochs) > 1 else 200
+                    std_loss = std_curves['train_losses'][::se][:len(std_epochs)]
+
+        # ---- Grokking run: train/test loss + IPR ----
+        grokk_epochs, grokk_train_loss, grokk_test_loss, grokk_ipr = [], [], [], []
+        has_grokk = self.p >= MIN_P_GROKKING
+        if has_grokk:
+            grokk_dir = self._run_dir('grokking')
+            if grokk_dir is not None:
+                grokk_curves = _load_training_curves(self._run_type_dir('grokking'))
+                grokk_ckpts = _load_checkpoints(grokk_dir, self.device)
+                if grokk_ckpts:
+                    grokk_epochs = sorted(grokk_ckpts.keys())
+                    grokk_ipr = [self._ipr_at_checkpoint(grokk_ckpts[ep])
+                                 for ep in grokk_epochs]
+                    if grokk_curves:
+                        se = grokk_epochs[1] - grokk_epochs[0] if len(grokk_epochs) > 1 else 200
+                        if 'train_losses' in grokk_curves:
+                            grokk_train_loss = grokk_curves['train_losses'][::se][:len(grokk_epochs)]
+                        if 'test_losses' in grokk_curves:
+                            grokk_test_loss = grokk_curves['test_losses'][::se][:len(grokk_epochs)]
+
+        if not std_epochs and not grokk_epochs:
+            print("    SKIP: no checkpoints found for standard or grokking run")
+            return
+
+        # ---- Static plot: 2×2 grid (std loss, grokk loss, std IPR, grokk IPR) ----
+        n_cols = 2 if has_grokk and grokk_epochs else 1
+        fig, axes = plt.subplots(2, n_cols, figsize=(5 * n_cols, 7),
+                                 constrained_layout=True)
+        if n_cols == 1:
+            axes = axes.reshape(2, 1)
+
+        # Standard loss (top-left)
+        ax = axes[0, 0]
+        if std_loss:
+            ax.plot(std_epochs[:len(std_loss)], std_loss,
+                    color=COLORS[0], linewidth=1.5, label="Train Loss")
+        ax.set_title('Standard (ReLU, full data)', fontsize=14)
+        ax.set_ylabel('Loss', fontsize=13)
+        ax.legend(fontsize=11)
+        ax.grid(True, alpha=0.4)
+
+        # Standard IPR (bottom-left)
+        ax = axes[1, 0]
+        if std_ipr:
+            ax.plot(std_epochs[:len(std_ipr)], std_ipr,
+                    color=COLORS[3], linewidth=1.5, label="Avg. IPR")
+        ax.set_title('Standard IPR', fontsize=14)
+        ax.set_xlabel('Step', fontsize=13)
+        ax.set_ylabel('IPR', fontsize=13)
+        ax.set_ylim([0, 1.05])
+        ax.legend(fontsize=11)
+        ax.grid(True, alpha=0.4)
+
+        if n_cols == 2:
+            # Grokking loss (top-right)
+            ax = axes[0, 1]
+            gx = grokk_epochs
+            if grokk_train_loss:
+                ax.plot(gx[:len(grokk_train_loss)], grokk_train_loss,
+                        color=COLORS[0], linewidth=1.5, label="Train Loss")
+            if grokk_test_loss:
+                ax.plot(gx[:len(grokk_test_loss)], grokk_test_loss,
+                        color=COLORS[3], linewidth=1.5, label="Test Loss")
+            ax.set_title('Grokking (ReLU, 75% data, WD)', fontsize=14)
+            ax.legend(fontsize=11)
+            ax.grid(True, alpha=0.4)
+
+            # Grokking IPR (bottom-right)
+            ax = axes[1, 1]
+            if grokk_ipr:
+                ax.plot(gx[:len(grokk_ipr)], grokk_ipr,
+                        color=COLORS[3], linewidth=1.5, label="Avg. IPR")
+            ax.set_title('Grokking IPR', fontsize=14)
+            ax.set_xlabel('Step', fontsize=13)
+            ax.set_ylim([0, 1.05])
+            ax.legend(fontsize=11)
+            ax.grid(True, alpha=0.4)
+
+        _save_fig(fig, self._out('overview_loss_ipr.png'))
+
+        # ---- Phase relationship scatter from standard final checkpoint ----
+        if std_ckpts:
+            final_ep = max(std_ckpts.keys())
+            model_sd = std_ckpts[final_ep]
+            W_in_d, W_out_d, mfl = decode_weights(model_sd, self.fourier_basis)
+            n_neurons = W_in_d.shape[0]
+            phis_2, psis = [], []
+            for neuron in range(n_neurons):
+                _, phi = compute_neuron(neuron, mfl, W_in_d)
+                _, psi = compute_neuron(neuron, mfl, W_out_d)
+                two_phi = normalize_to_pi(2 * phi)
+                psi_n = normalize_to_pi(psi)
+                # Fix ±π wrap: keep ψ within π of 2φ
+                if psi_n - two_phi > np.pi:
+                    psi_n -= 2 * np.pi
+                elif psi_n - two_phi < -np.pi:
+                    psi_n += 2 * np.pi
+                phis_2.append(two_phi)
+                psis.append(psi_n)
+
+            fig, ax = plt.subplots(figsize=(5, 5))
+            ax.plot([-np.pi, np.pi], [-np.pi, np.pi], 'r-',
+                    linewidth=3, alpha=0.8,
+                    label=r'$\psi_m = 2\phi_m$', zorder=1)
+            ax.scatter(phis_2, psis, s=12, alpha=0.6, color=COLORS[0], zorder=2)
+            ax.legend(fontsize=12, loc='upper left')
+            ax.set_xlabel(r'$2\phi_m$', fontsize=14)
+            ax.set_ylabel(r'$\psi_m$', fontsize=14)
+            ax.set_title(r'Phase Alignment: $\psi_m = 2\phi_m$', fontsize=14)
+            ax.set_xlim([-np.pi, np.pi])
+            ax.set_ylim([-np.pi, np.pi])
+            ax.set_aspect('equal')
+            ax.grid(True, alpha=0.3)
+            _save_fig(fig, self._out('overview_phase_scatter.png'))
+
+        # ---- JSON for interactive Plotly charts ----
+        payload = {
+            'std_epochs': [int(e) for e in std_epochs],
+            'std_ipr': std_ipr,
+        }
+        if std_loss:
+            payload['std_train_loss'] = [float(v) for v in std_loss]
+
+        if has_grokk and grokk_epochs:
+            payload['grokk_epochs'] = [int(e) for e in grokk_epochs]
+            payload['grokk_ipr'] = grokk_ipr
+            if grokk_train_loss:
+                payload['grokk_train_loss'] = [float(v) for v in grokk_train_loss]
+            if grokk_test_loss:
+                payload['grokk_test_loss'] = [float(v) for v in grokk_test_loss]
+
+        with open(self._out('overview.json'), 'w') as f:
+            json.dump(payload, f)
+
+        files = ['overview_loss_ipr.png', 'overview.json']
+        if std_ckpts:
+            files.append('overview_phase_scatter.png')
+        print(f"    Saved {', '.join(files)}")
+
+    # ------------------------------------------------------------------
+    # Tab 2: Fourier Weights (heatmap + lineplots)
+    # ------------------------------------------------------------------
+
+    def generate_tab2(self):
+        """Generate full_training_para_origin.png, lineplot_in.png, lineplot_out.png."""
+        print(f"  [Tab 2] Fourier Weights for p={self.p}")
+        run_dir = self._run_dir('standard')
+        if run_dir is None:
+            print("    SKIP: standard run directory not found")
+            return
+
+        final_data = _load_final(run_dir, self.device)
+        if final_data is None:
+            print("    SKIP: no final checkpoint")
+            return
+        model_load = final_data['model']
+
+        W_in_decode, W_out_decode, max_freq_ls = decode_weights(
+            model_load, self.fourier_basis
+        )
+        d_mlp = W_in_decode.shape[0]
+        num_neurons = min(20, d_mlp)
+
+        # Sort neurons by frequency
+        sorted_indices = select_top_neurons_by_frequency(
+            max_freq_ls, W_in_decode, n=num_neurons
+        )
+
+        freq_ls = np.array([max_freq_ls[i] for i in sorted_indices])
+
+        # DFT coefficients for heatmap (matches blog Figure 2)
+        W_in_dft = W_in_decode[sorted_indices, :]
+        W_out_dft = W_out_decode[sorted_indices, :]
+        # Raw weights for line plots (matches blog Figure 3)
+        W_in_raw = model_load['mlp.W_in'][sorted_indices, :]
+        W_out_raw = model_load['mlp.W_out'].T[sorted_indices, :]
+
+        # Sort within selected set by frequency
+        sort_order = np.argsort(freq_ls)
+        ranked_W_in_dft = W_in_dft[sort_order, :]
+        ranked_W_out_dft = W_out_dft[sort_order, :]
+        ranked_W_in_raw = W_in_raw[sort_order, :]
+        ranked_W_out_raw = W_out_raw[sort_order, :]
+
+        # ---- Heatmap plot (DFT coefficients, matching blog Figure 2) ----
+        fb_names = self.fourier_basis_names
+        n_modes = len(fb_names)
+        fig_w = max(8, n_modes * 0.4)
+        fig_h = max(8, num_neurons * 0.35 + 3)
+        fig, axes = plt.subplots(
+            2, 1, figsize=(fig_w, fig_h), constrained_layout=True,
+            gridspec_kw={"hspace": 0.15}
+        )
+
+        # W_in DFT
+        ax_in = axes[0]
+        W_in_np = ranked_W_in_dft.detach().cpu().numpy()
+        abs_max_in = np.abs(W_in_np).max()
+        im_in = ax_in.imshow(
+            W_in_np,
+            cmap=CMAP_DIVERGING, vmin=-abs_max_in, vmax=abs_max_in,
+            aspect='auto'
+        )
+        ax_in.set_title(r'First-Layer $\theta_m$ after DFT', fontsize=18)
+        fig.colorbar(im_in, ax=ax_in, shrink=0.8)
+        y_locs = np.arange(num_neurons)
+        ax_in.set_yticks(y_locs)
+        ax_in.set_yticklabels(y_locs, fontsize=10)
+        ax_in.set_ylabel('Neuron #', fontsize=14)
+        x_locs = np.arange(n_modes)
+        ax_in.set_xticks(x_locs)
+        ax_in.set_xticklabels(fb_names, rotation=90, fontsize=10)
+
+        # W_out DFT
+        ax_out = axes[1]
+        W_out_np = ranked_W_out_dft.detach().cpu().numpy()
+        abs_max_out = np.abs(W_out_np).max()
+        im_out = ax_out.imshow(
+            W_out_np,
+            cmap=CMAP_DIVERGING, vmin=-abs_max_out, vmax=abs_max_out,
+            aspect='auto'
+        )
+        ax_out.set_title(r'Second-Layer $\xi_m$ after DFT', fontsize=18)
+        fig.colorbar(im_out, ax=ax_out, shrink=0.8)
+        ax_out.set_yticks(y_locs)
+        ax_out.set_yticklabels(y_locs, fontsize=10)
+        ax_out.set_ylabel('Neuron #', fontsize=14)
+        ax_out.set_xticks(x_locs)
+        ax_out.set_xticklabels(fb_names, rotation=90, fontsize=10)
+        ax_out.set_xlabel('Fourier Component', fontsize=14)
+
+        _save_fig(fig, self._out('full_training_para_origin.png'))
+
+        # ---- Line plots (raw weights + cosine fits, matching blog Figure 3) ----
+        lineplot_idx = select_lineplot_neurons(list(range(num_neurons)), n=3)
+        fb = self.fourier_basis
+        positions = np.arange(ranked_W_in_raw.shape[1])
+
+        for tag, weight_data, title_tex in [
+            ('lineplot_in', ranked_W_in_raw, r'First-Layer Parameters $\theta_m$'),
+            ('lineplot_out', ranked_W_out_raw, r'Second-Layer Parameters $\xi_m$'),
+        ]:
+            if hasattr(weight_data, 'detach'):
+                weight_np = weight_data.detach().cpu()
+            else:
+                weight_np = weight_data
+
+            top3 = weight_np[lineplot_idx]
+
+            lp_w = max(8, self.p * 0.35)
+            fig, axes_lp = plt.subplots(
+                nrows=3, ncols=1, figsize=(lp_w, 8),
+                constrained_layout=True,
+                gridspec_kw={'hspace': 0.08}
+            )
+
+            for i, ax in enumerate(axes_lp):
+                data = top3[i]
+                if isinstance(data, torch.Tensor):
+                    data_t = data.float()
+                else:
+                    data_t = torch.tensor(data, dtype=torch.float32)
+                # Project into Fourier space, keep top 2 components, project back
+                proj = data_t @ fb.T
+                abs_proj = torch.abs(proj)
+                _, top2_idx = torch.topk(abs_proj, 2)
+                mask = torch.zeros_like(proj)
+                mask[top2_idx] = proj[top2_idx]
+                data_est = mask @ fb
+                data_np = data_t.numpy()
+                data_est_np = data_est.numpy()
+
+                ax.plot(data_np, marker='o', markersize=5,
+                        color=COLORS[0], linewidth=1.5, linestyle=':',
+                        label="Actual")
+                ax.plot(data_est_np, marker='o', markersize=5,
+                        color=COLORS[3], linewidth=1.5, linestyle=':',
+                        alpha=0.7, label="Fitted")
+                ax.set_ylim(-0.9, 0.9)
+                ax.set_ylabel(f'Neuron #{i+1}', fontsize=14)
+                ax.set_xticks(positions)
+                ax.grid(True, which='major', axis='both',
+                        linestyle='--', linewidth=0.5, alpha=0.6)
+                if i < len(axes_lp) - 1:
+                    ax.set_xticklabels([])
+                ax.legend(fontsize=12, loc="upper right")
+
+            axes_lp[-1].set_xlabel('Input Dimension', fontsize=14)
+            axes_lp[-1].set_xticks(positions)
+            axes_lp[-1].set_xticklabels(
+                np.arange(ranked_W_in_raw.shape[1]), rotation=0, fontsize=10
+            )
+            axes_lp[0].set_title(title_tex, fontsize=18)
+
+            _save_fig(fig, self._out(f'{tag}.png'))
+
+        print("    Saved full_training_para_origin.png, lineplot_in.png, lineplot_out.png")
+
+    # ------------------------------------------------------------------
+    # Tab 3: Phase Analysis
+    # ------------------------------------------------------------------
+
+    def generate_tab3(self):
+        """Generate phase_distribution.png, phase_relationship.png, magnitude_distribution.png."""
+        print(f"  [Tab 3] Phase Analysis for p={self.p}")
+        run_dir = self._run_dir('standard')
+        if run_dir is None:
+            print("    SKIP: standard run directory not found")
+            return
+
+        final_data = _load_final(run_dir, self.device)
+        if final_data is None:
+            print("    SKIP: no final checkpoint")
+            return
+        model_load = final_data['model']
+
+        W_in_decode, W_out_decode, max_freq_ls = decode_weights(
+            model_load, self.fourier_basis
+        )
+        d_mlp = W_in_decode.shape[0]
+
+        # Compute all neuron phases and magnitudes
+        coeff_in_scale_ls = []
+        coeff_out_scale_ls = []
+        coeff_phi_ls = []
+        coeff_psi_ls = []
+
+        for neuron in range(d_mlp):
+            s_in, phi_in = compute_neuron(neuron, max_freq_ls, W_in_decode)
+            s_out, phi_out = compute_neuron(neuron, max_freq_ls, W_out_decode)
+            coeff_in_scale_ls.append(s_in)
+            coeff_out_scale_ls.append(s_out)
+            coeff_phi_ls.append(phi_in)
+            coeff_psi_ls.append(phi_out)
+
+        coeff_phi_arr = np.array(coeff_phi_ls)
+        coeff_psi_arr = np.array(coeff_psi_ls)
+
+        # ---- Phase distribution on concentric circles ----
+        # Select the most common non-zero frequency for phase analysis
+        target_freq = select_phase_frequency(max_freq_ls, self.p)
+        freq_neurons = [i for i, f in enumerate(max_freq_ls) if f == target_freq]
+        phi_subset = np.array([coeff_phi_ls[i] for i in freq_neurons])
+
+        theta = np.linspace(0, 2 * np.pi, 300)
+        multipliers = [1, 2, 3, 4]
+        radii = [1.0, 0.88, 0.76, 0.64]
+
+        fig, ax = plt.subplots(figsize=(4, 4))
+        for m, r in zip(multipliers, radii):
+            x_c, y_c = r * np.cos(theta), r * np.sin(theta)
+            ax.plot(x_c, y_c, linewidth=0.8, color='gray', alpha=0.6)
+
+            x_pts = r * np.cos(m * phi_subset)
+            y_pts = r * np.sin(m * phi_subset)
+            label = fr'$\phi_m$' if m == 1 else fr'${m}\phi_m$'
+            ax.scatter(x_pts, y_pts, s=20, marker='o',
+                       color=COLORS[m - 1], label=label)
+
+        ax.legend(
+            fontsize=15, loc='upper center', columnspacing=0.2,
+            handletextpad=0.1, bbox_to_anchor=(0.5, 1.15), ncol=4, frameon=False
+        )
+        ax.set_xlabel(r'$\cos(\phi_m)$', fontsize=19)
+        ax.set_ylabel(r'$\sin(\phi_m)$', fontsize=19)
+        ax.set_xticks([])
+        ax.set_yticks([])
+        for spine in ax.spines.values():
+            spine.set_visible(False)
+
+        _save_fig(fig, self._out('phase_distribution.png'))
+
+        # ---- Phase relationship: 2*phi vs psi ----
+        coeff_2phi_arr = np.array([normalize_to_pi(2 * phi) for phi in coeff_phi_arr])
+        coeff_psi_plot = coeff_psi_arr.copy()
+        # Fix ±π wrap: keep ψ within π of 2φ so boundary points stay on diagonal
+        diff = coeff_psi_plot - coeff_2phi_arr
+        coeff_psi_plot[diff > np.pi] -= 2 * np.pi
+        coeff_psi_plot[diff < -np.pi] += 2 * np.pi
+
+        fig, ax = plt.subplots(figsize=(5, 5))
+        ax.plot([-np.pi, np.pi], [-np.pi, np.pi], 'r-', linewidth=3, alpha=0.8,
+                label=r'$\psi_m = 2\phi_m$', zorder=1)
+        ax.scatter(
+            coeff_2phi_arr, coeff_psi_plot,
+            marker='.', color=COLORS[0], s=20, zorder=2
+        )
+        ax.legend(fontsize=12, loc='upper left')
+        ax.set_xlabel(r'$2\phi_m$', fontsize=14)
+        ax.set_ylabel(r'$\psi_m$', fontsize=14)
+        ax.set_title(r'Phase Alignment: $\psi_m = 2\phi_m$', fontsize=14)
+        ax.set_xlim(-np.pi * 1.1, np.pi * 1.1)
+        ax.set_ylim(-np.pi * 1.1, np.pi * 1.1)
+        ax.set_aspect('equal')
+        ax.grid(True, alpha=0.3)
+
+        _save_fig(fig, self._out('phase_relationship.png'))
+
+        # ---- Magnitude distribution (violin) ----
+        fig, ax = plt.subplots(figsize=(4, 4))
+        data_for_plot = [coeff_in_scale_ls, coeff_out_scale_ls]
+        positions = [1, 2]
+
+        parts = ax.violinplot(
+            data_for_plot, positions=positions, widths=0.6,
+            showmeans=True, showmedians=True, showextrema=True
+        )
+        for pc in parts['bodies']:
+            pc.set_facecolor(COLORS[0])
+            pc.set_alpha(0.7)
+        parts['cmedians'].set_color(COLORS[2])
+        parts['cmedians'].set_linewidth(2)
+        parts['cmeans'].set_color(COLORS[2])
+        parts['cmeans'].set_linewidth(2)
+        parts['cbars'].set_color(COLORS[0])
+        parts['cbars'].set_linewidth(1.5)
+        parts['cmaxes'].set_color(COLORS[0])
+        parts['cmins'].set_color(COLORS[0])
+
+        ax.set_xticks(positions)
+        ax.set_xticklabels(['First-Layer', 'Second-Layer'], fontsize=14)
+        ax.set_ylabel('Magnitude', fontsize=19)
+        ax.grid(True, alpha=0.3)
+        plt.tight_layout()
+
+        _save_fig(fig, self._out('magnitude_distribution.png'))
+
+        print("    Saved phase_distribution.png, phase_relationship.png, magnitude_distribution.png")
+
+    # ------------------------------------------------------------------
+    # Tab 4: Output Logits
+    # ------------------------------------------------------------------
+
+    def generate_tab4(self):
+        """Generate output_logits.png."""
+        print(f"  [Tab 4] Output Logits for p={self.p}")
+        run_dir = self._run_dir('standard')
+        if run_dir is None:
+            print("    SKIP: standard run directory not found")
+            return
+
+        final_data = _load_final(run_dir, self.device)
+        if final_data is None:
+            print("    SKIP: no final checkpoint")
+            return
+        model_load = final_data['model']
+
+        p = self.p
+        act_type = TRAINING_RUNS['standard']['act_type']
+        model = EmbedMLP(
+            d_vocab=self.d_vocab,
+            d_model=self.d_model,
+            d_mlp=self.d_mlp,
+            act_type=act_type,
+            use_cache=False
+        )
+        model.to(self.device)
+        model.load_state_dict(model_load)
+        model.eval()
+
+        with torch.no_grad():
+            logits = model(self.all_data).squeeze(1)
+
+        logits_np = logits.cpu().numpy()
+
+        # Show first p pairs (first row of the input grid)
+        interval_start = 0
+        interval_end = p
+        logits_interval = logits_np[interval_start:interval_end]
+        selected_pairs = self.all_data[interval_start:interval_end]
+
+        fig, ax = plt.subplots(figsize=(7, 6))
+        abs_max = np.abs(logits_np).max() * 0.8
+        im = ax.imshow(
+            logits_interval.T, cmap=CMAP_DIVERGING, aspect='auto',
+            vmin=-abs_max, vmax=abs_max
+        )
+
+        # Highlight target positions with rectangles
+        for i, (x_val_t, y_val_t) in enumerate(selected_pairs):
+            x_val = x_val_t.item()
+            y_val = y_val_t.item()
+            target_2x = (2 * x_val) % p
+            target_2y = (2 * y_val) % p
+            target_sum = (x_val + y_val) % p
+
+            rect_2x = patches.Rectangle(
+                (i - 0.5, target_2x - 0.5), 1, 1,
+                linewidth=1.6, edgecolor='#0D2758', facecolor='none', alpha=0.9
+            )
+            ax.add_patch(rect_2x)
+            if target_2y != target_2x:
+                rect_2y = patches.Rectangle(
+                    (i - 0.5, target_2y - 0.5), 1, 1,
+                    linewidth=1.6, edgecolor='#0D2758', facecolor='none', alpha=0.9
+                )
+                ax.add_patch(rect_2y)
+            rect_sum = patches.Rectangle(
+                (i - 0.5, target_sum - 0.5), 1, 1,
+                linewidth=1.6, edgecolor='#0D2758', facecolor='none', alpha=0.9
+            )
+            ax.add_patch(rect_sum)
+
+        n_pairs = interval_end - interval_start
+        if n_pairs <= 50:
+            x_positions = np.arange(n_pairs)
+            x_labels = [f"({selected_pairs[i][0].item()},{selected_pairs[i][1].item()})"
+                        for i in range(n_pairs)]
+            ax.set_xticks(x_positions)
+            ax.set_xticklabels(x_labels, rotation=90, ha='right', fontsize=14)
+        else:
+            n_labels = min(25, n_pairs)
+            step = n_pairs // n_labels
+            x_positions = np.arange(0, n_pairs, step)
+            x_labels = [f"({selected_pairs[i][0].item()},{selected_pairs[i][1].item()})"
+                        for i in x_positions]
+            ax.set_xticks(x_positions)
+            ax.set_xticklabels(x_labels, rotation=90, ha='right', fontsize=14)
+
+        ax.set_yticks(np.arange(p))
+        ax.set_yticklabels(np.arange(p), fontsize=14)
+        ax.set_xlabel("Input Pair", fontsize=18)
+        ax.set_ylabel("Output", fontsize=18)
+        plt.colorbar(im, ax=ax)
+        ax.grid(True, alpha=0.2, linestyle=':', linewidth=0.5, axis='x')
+        plt.tight_layout()
+
+        _save_fig(fig, self._out('output_logits.png'))
+        print("    Saved output_logits.png")
+
+    # ------------------------------------------------------------------
+    # Tab 5: Grokking
+    # ------------------------------------------------------------------
+
+    def generate_tab5(self):
+        """Generate grokking-related plots."""
+        print(f"  [Tab 5] Grokking for p={self.p}")
+        if self.p < MIN_P_GROKKING:
+            print(f"    SKIP: p={self.p} < {MIN_P_GROKKING} (too few test points for grokking)")
+            return
+        run_dir = self._run_dir('grokking')
+        if run_dir is None:
+            print("    SKIP: grokking run directory not found")
+            return
+
+        curves = _load_training_curves(self._run_type_dir('grokking'))
+        checkpoints = _load_checkpoints(run_dir, self.device)
+
+        if not checkpoints:
+            print("    SKIP: no grokking checkpoints")
+            return
+
+        epochs = sorted(checkpoints.keys())
+        p = self.p
+        d_mlp = self.d_mlp
+        act_type = TRAINING_RUNS['grokking']['act_type']
+
+        # Load train/test data
+        train_data_path = os.path.join(run_dir, 'train_data.pth')
+        test_data_path = os.path.join(run_dir, 'test_data.pth')
+        train_data = None
+        test_data = None
+        train_labels = None
+        test_labels = None
+        if os.path.exists(train_data_path):
+            raw = torch.load(train_data_path, weights_only=False,
+                             map_location=self.device)
+            # Handle both formats: plain tensor or (pairs, labels) tuple
+            if isinstance(raw, (tuple, list)):
+                train_data, train_labels = raw[0], raw[1]
+            else:
+                train_data = raw
+        if os.path.exists(test_data_path):
+            raw = torch.load(test_data_path, weights_only=False,
+                             map_location=self.device)
+            if isinstance(raw, (tuple, list)):
+                test_data, test_labels = raw[0], raw[1]
+            else:
+                test_data = raw
+
+        # Fallback: regenerate data deterministically if files are missing
+        if train_data is None or test_data is None:
+            grokk_cfg = TRAINING_RUNS['grokking']
+            frac = grokk_cfg['frac_train']
+            seed = grokk_cfg['seed']
+            print(f"    Regenerating train/test data (frac={frac}, seed={seed})")
+            train_data, test_data = _gen_train_test(p, frac_train=frac, seed=seed)
+
+        # Compute labels from pairs if not loaded directly
+        if train_labels is None and train_data is not None:
+            train_labels = torch.tensor(
+                [(train_data[i, 0].item() + train_data[i, 1].item()) % p
+                 for i in range(train_data.shape[0])],
+                dtype=torch.long
+            )
+        if test_labels is None and test_data is not None:
+            test_labels = torch.tensor(
+                [(test_data[i, 0].item() + test_data[i, 1].item()) % p
+                 for i in range(test_data.shape[0])],
+                dtype=torch.long
+            )
+
+        # Detect stage boundaries
+        train_losses = curves.get('train_losses', []) if curves else []
+        test_losses = curves.get('test_losses', []) if curves else []
+        train_accs_curve = curves.get('train_accs', None) if curves else None
+        test_accs_curve = curves.get('test_accs', None) if curves else None
+
+        stage1_end, stage2_end = detect_grokking_stages(
+            train_losses, test_losses, train_accs_curve, test_accs_curve
+        )
+        if stage1_end is None:
+            stage1_end = len(epochs) // 5
+        if stage2_end is None:
+            stage2_end = len(epochs) * 3 // 5
+
+        # ---- Loss JSON + static PNG ----
+        if train_losses:
+            loss_data = {
+                'train_losses': train_losses,
+                'test_losses': test_losses,
+                'stage1_end': stage1_end,
+                'stage2_end': stage2_end,
+            }
+            with open(self._out('grokk_loss.json'), 'w') as f:
+                json.dump(loss_data, f)
+
+            # Static loss PNG (matches blog Figure 13a)
+            max_step = min(len(train_losses), len(test_losses)) if test_losses else len(train_losses)
+            fig, ax = plt.subplots(figsize=(4, 4))
+            ax.plot(train_losses[:max_step], color='#0D2758', linewidth=2, label='Train')
+            if test_losses:
+                ax.plot(test_losses[:max_step], color='#A32015', linewidth=2, label='Test')
+            ax.axvspan(0, stage1_end, alpha=0.15, color='#D4AF37')
+            ax.axvspan(stage1_end, stage2_end, alpha=0.15, color='#8B7355')
+            ax.axvspan(stage2_end, max_step, alpha=0.15, color='#60656F')
+            ax.axvline(x=stage1_end, color='black', linestyle='--', linewidth=1)
+            ax.axvline(x=stage2_end, color='black', linestyle='--', linewidth=1)
+            ax.set_xlabel('Step', fontsize=16)
+            ax.set_ylabel('Loss', fontsize=16)
+            ax.legend(fontsize=16, loc='upper right')
+            ax.grid(True, linestyle='--', alpha=0.5)
+            plt.tight_layout()
+            _save_fig(fig, self._out('grokk_loss.png'))
+
+        # ---- Accuracy: compute from checkpoints if not in curves ----
+        train_accs = []
+        test_accs = []
+        if train_data is not None and test_data is not None:
+            for ep in epochs:
+                model = EmbedMLP(
+                    d_vocab=self.d_vocab, d_model=self.d_model,
+                    d_mlp=d_mlp, act_type=act_type, use_cache=False
+                ).to(self.device)
+                model.load_state_dict(checkpoints[ep])
+                model.eval()
+                with torch.no_grad():
+                    tr_logits = model(train_data)
+                    te_logits = model(test_data)
+                train_accs.append(acc_rate(tr_logits, train_labels))
+                test_accs.append(acc_rate(te_logits, test_labels))
+        elif train_accs_curve is not None:
+            # Use curves data, subsample to match checkpoint epochs
+            save_every = epochs[1] - epochs[0] if len(epochs) > 1 else 200
+            train_accs = train_accs_curve[::save_every][:len(epochs)]
+            test_accs = test_accs_curve[::save_every][:len(epochs)]
+
+        acc_data = {
+            'epochs': epochs,
+            'train_accs': train_accs,
+            'test_accs': test_accs,
+            'stage1_end': stage1_end,
+            'stage2_end': stage2_end,
+        }
+        with open(self._out('grokk_acc.json'), 'w') as f:
+            json.dump(acc_data, f)
+
+        # Static accuracy PNG (matches blog Figure 13b)
+        if train_accs and test_accs:
+            fig, ax = plt.subplots(figsize=(4, 4))
+            ax.axvspan(0, stage1_end, alpha=0.15, color='#D4AF37')
+            ax.axvspan(stage1_end, stage2_end, alpha=0.15, color='#8B7355')
+            ax.axvspan(stage2_end, epochs[-1] if epochs else stage2_end,
+                       alpha=0.15, color='#60656F')
+            ax.axvline(x=stage1_end, color='black', linestyle='--', linewidth=1)
+            ax.axvline(x=stage2_end, color='black', linestyle='--', linewidth=1)
+            ax.plot(epochs[:len(train_accs)], train_accs,
+                    label='Train', color='#0D2758', linewidth=2.5)
+            ax.plot(epochs[:len(test_accs)], test_accs,
+                    label='Test', color='#A32015', linewidth=2.5)
+            ax.set_xlabel('Step', fontsize=16)
+            ax.set_ylabel('Accuracy', fontsize=16)
+            ax.legend(fontsize=16, loc='lower right')
+            ax.grid(True, linestyle='--', alpha=0.5)
+            plt.tight_layout()
+            _save_fig(fig, self._out('grokk_acc.png'))
+
+        # ---- Phase difference |sin(D*)| ----
+        abs_phase_diff = []
+        sparse_level = []
+
+        for ep in epochs:
+            model_sd = checkpoints[ep]
+            W_in_d, W_out_d, mfl = decode_weights(model_sd, self.fourier_basis)
+
+            sparse_level.append(self._ipr_at_checkpoint(model_sd))
+
+            phase_diffs = []
+            for neuron in range(W_in_d.shape[0]):
+                _, phi_in = compute_neuron(neuron, mfl, W_in_d)
+                _, phi_out = compute_neuron(neuron, mfl, W_out_d)
+                phase_diffs.append(normalize_to_pi(phi_out - 2 * phi_in))
+            phase_diffs = np.array(phase_diffs)
+            abs_phase_diff.append(np.mean(np.abs(np.sin(phase_diffs))))
+
+        # Limit to reasonable number of points for plotting
+        n_plot = min(len(epochs), 100)
+        x_phase = np.array(epochs[:n_plot])
+
+        fig, ax = plt.subplots(figsize=(4, 4))
+        ax.axvspan(0, stage1_end, alpha=0.15, color='#D4AF37')
+        ax.axvspan(stage1_end, min(stage2_end, x_phase[-1] if len(x_phase) else stage2_end),
+                   alpha=0.15, color='#8B7355')
+        if len(x_phase):
+            ax.axvspan(stage2_end, x_phase[-1], alpha=0.15, color='#60656F')
+        ax.axvline(x=stage1_end, color='black', linestyle='--', linewidth=1)
+        ax.axvline(x=stage2_end, color='black', linestyle='--', linewidth=1)
+        ax.plot(x_phase, abs_phase_diff[:n_plot], marker='x', markersize=5,
+                color='#986d56', label=r"Avg. $|\sin(D_m^\star)|$", linewidth=1.5)
+        ax.set_xlabel('Step', fontsize=16)
+        ax.set_ylabel('Average Value', fontsize=16)
+        ax.set_ylim([0, 0.65])
+        ax.legend(fontsize=16, loc="upper right")
+        ax.grid(True, alpha=0.5, linestyle='--')
+        plt.tight_layout()
+        _save_fig(fig, self._out('grokk_abs_phase_diff.png'))
+
+        # ---- IPR + param norms (dual axis) ----
+        x_all = np.array(epochs)
+        param_norms = []
+        if curves and 'param_norms' in curves:
+            save_every = epochs[1] - epochs[0] if len(epochs) > 1 else 200
+            param_norms = curves['param_norms'][::save_every][:len(epochs)]
+
+        fig, ax1 = plt.subplots(figsize=(4, 4))
+        ax1.axvspan(0, stage1_end, alpha=0.15, color='#D4AF37')
+        ax1.axvspan(stage1_end, min(stage2_end, x_all[-1] if len(x_all) else stage2_end),
+                    alpha=0.15, color='#8B7355')
+        if len(x_all):
+            ax1.axvspan(stage2_end, x_all[-1], alpha=0.15, color='#60656F')
+        ax1.axvline(x=stage1_end, color='black', linestyle='--', linewidth=1)
+        ax1.axvline(x=stage2_end, color='black', linestyle='--', linewidth=1)
+
+        line1 = ax1.plot(x_all, sparse_level, marker='x', markersize=3,
+                         color='#986d56', label=r"Avg. IPR", linewidth=1.5)
+        ax1.set_xlabel('Step', fontsize=16)
+        ax1.tick_params(axis='y')
+        ax1.set_ylim([0, 1.05])
+
+        if param_norms:
+            ax2 = ax1.twinx()
+            line2 = ax2.plot(x_all[:len(param_norms)], param_norms,
+                             marker='o', markersize=3, color='#2E5266',
+                             label=r"Param. Norm", linewidth=1.5)
+            ax2.tick_params(axis='y')
+            lines = line1 + line2
+            labels = [l.get_label() for l in lines]
+            ax1.legend(lines, labels, fontsize=16, loc="lower right")
+        else:
+            ax1.legend(fontsize=16, loc="lower right")
+
+        ax1.grid(True, alpha=0.5, linestyle='--')
+        plt.tight_layout()
+        _save_fig(fig, self._out('grokk_avg_ipr.png'))
+
+        # ---- Memorization accuracy (3-panel) ----
+        if train_data is not None:
+            # Find a checkpoint near stage1_end
+            closest_epoch = min(epochs, key=lambda e: abs(e - stage1_end))
+            model_sd = checkpoints[closest_epoch]
+
+            model = EmbedMLP(
+                d_vocab=self.d_vocab, d_model=self.d_model,
+                d_mlp=d_mlp, act_type=act_type, use_cache=False
+            ).to(self.device)
+            model.load_state_dict(model_sd)
+            model.eval()
+
+            with torch.no_grad():
+                logits = model(self.all_data).squeeze(1)
+
+            train_set = set([(int(i), int(j)) for i, j in train_data])
+            true_test_points = []
+
+            train_mask = torch.zeros(p, p)
+            for i in range(p):
+                for j in range(p):
+                    if (i, j) in train_set:
+                        train_mask[i, j] = 1.0
+                    elif (j, i) in train_set:
+                        train_mask[i, j] = 0.65
+                    else:
+                        train_mask[i, j] = 0.0
+                        true_test_points.append((i, j))
+
+            predicted = torch.argmax(logits, dim=1).view(p, p)
+            gt_grid = self.all_labels.view(p, p)
+            accuracy_mask = (predicted == gt_grid).float()
+
+            probs = torch.softmax(logits, dim=1)
+            gt_probs = torch.zeros(p * p)
+            for idx in range(p * p):
+                i_val = self.all_data[idx, 0].item()
+                j_val = self.all_data[idx, 1].item()
+                correct = (i_val + j_val) % p
+                gt_probs[idx] = probs[idx, correct]
+            gt_probs_grid = gt_probs.view(p, p)
+
+            fig = plt.figure(figsize=(20, 6))
+            gs = fig.add_gridspec(1, 3, width_ratios=[1, 1, 1.1], wspace=0.15)
+
+            ax1 = fig.add_subplot(gs[0])
+            ax2 = fig.add_subplot(gs[1])
+            ax3 = fig.add_subplot(gs[2])
+
+            # Train mask
+            im1 = ax1.imshow(train_mask.numpy(), cmap=CMAP_SEQUENTIAL,
+                             vmin=0, vmax=1, aspect='equal')
+            ax1.set_title('Training Data under Symmetry', fontsize=21)
+            ax1.set_ylabel('First Input', fontsize=18)
+            ax1.set_xlabel('Second Input', fontsize=18)
+            locs = np.arange(p)
+            ax1.set_xticks(locs)
+            ax1.set_yticks(locs)
+            ax1.set_xticklabels(locs, fontsize=11)
+            ax1.set_yticklabels(locs, fontsize=11)
+            for ti, tj in true_test_points:
+                rect = plt.Rectangle((tj - 0.5, ti - 0.5), 1, 1,
+                                     linewidth=2.5, edgecolor='red', facecolor='none')
+                ax1.add_patch(rect)
+
+            # Accuracy mask
+            im2 = ax2.imshow(accuracy_mask.numpy(), cmap=CMAP_SEQUENTIAL,
+                             vmin=0, vmax=1, aspect='equal')
+            ax2.set_title('Accuracy before Grokking', fontsize=21)
+            ax2.set_xlabel('Second Input', fontsize=18)
+            ax2.set_xticks(locs)
+            ax2.set_yticks(locs)
+            ax2.set_xticklabels(locs, fontsize=11)
+            ax2.set_yticklabels(locs, fontsize=11)
+            for ti, tj in true_test_points:
+                rect = plt.Rectangle((tj - 0.5, ti - 0.5), 1, 1,
+                                     linewidth=2.5, edgecolor='red', facecolor='none')
+                ax2.add_patch(rect)
+
+            # Softmax probability
+            prob_max = gt_probs_grid.max().item()
+            im3 = ax3.imshow(gt_probs_grid.detach().numpy(), cmap=CMAP_SEQUENTIAL,
+                             vmin=0, vmax=prob_max, aspect='equal')
+            ax3.set_title('Softmax Weight at Ground-Truth', fontsize=21)
+            ax3.set_xlabel('Second Input', fontsize=18)
+            ax3.set_xticks(locs)
+            ax3.set_yticks(locs)
+            ax3.set_xticklabels(locs, fontsize=11)
+            ax3.set_yticklabels(locs, fontsize=11)
+            for ti, tj in true_test_points:
+                rect = plt.Rectangle((tj - 0.5, ti - 0.5), 1, 1,
+                                     linewidth=2.5, edgecolor='red', facecolor='none')
+                ax3.add_patch(rect)
+            cbar3 = fig.colorbar(im3, ax=ax3, fraction=0.046, pad=0.04)
+            cbar3.ax.tick_params(labelsize=12)
+            plt.tight_layout()
+            _save_fig(fig, self._out('grokk_memorization_accuracy.png'))
+
+        # ---- Memorization common-to-rare (4-panel) ----
+        if train_data is not None:
+            train_set = set([(int(i), int(j)) for i, j in train_data])
+            asymmetric_train_points = []
+            train_mask_dist = torch.zeros(p, p)
+            for i in range(p):
+                for j in range(p):
+                    if (i, j) in train_set and (j, i) in train_set:
+                        train_mask_dist[i, j] = 1.0
+                    elif (i, j) in train_set and (j, i) not in train_set:
+                        train_mask_dist[i, j] = 0.5
+                        asymmetric_train_points.append((i, j))
+                    else:
+                        train_mask_dist[i, j] = 0.0
+
+            # Pick 3 epochs: 0, ~stage1/2, ~stage1
+            selected_epochs = [0]
+            mid_epoch = min(epochs, key=lambda e: abs(e - stage1_end // 2))
+            end_epoch = min(epochs, key=lambda e: abs(e - stage1_end))
+            if mid_epoch not in selected_epochs:
+                selected_epochs.append(mid_epoch)
+            if end_epoch not in selected_epochs:
+                selected_epochs.append(end_epoch)
+            # Ensure we have exactly 3 + distribution = 4 panels
+            while len(selected_epochs) < 3:
+                selected_epochs.append(epochs[min(len(epochs) - 1, 2)])
+
+            fig = plt.figure(figsize=(26, 6))
+            gs = fig.add_gridspec(
+                1, 4, width_ratios=[1, 1, 1, 1.1], wspace=0.15
+            )
+
+            # Panel 1: training data distribution
+            ax_d = fig.add_subplot(gs[0])
+            ax_d.imshow(train_mask_dist.numpy(), cmap=CMAP_SEQUENTIAL,
+                        vmin=0, vmax=1, aspect='equal')
+            ax_d.set_title('Training Data Distribution', fontsize=21)
+            ax_d.set_ylabel('First Input', fontsize=18)
+            ax_d.set_xlabel('Second Input', fontsize=18)
+            locs = np.arange(p)
+            ax_d.set_xticks(locs)
+            ax_d.set_yticks(locs)
+            ax_d.set_xticklabels(locs, fontsize=11)
+            ax_d.set_yticklabels(locs, fontsize=11)
+            for ti, tj in asymmetric_train_points:
+                rect = plt.Rectangle((tj - 0.5, ti - 0.5), 1, 1,
+                                     linewidth=2.0, edgecolor='red', facecolor='none')
+                ax_d.add_patch(rect)
+
+            # Panels 2-4: accuracy at selected epochs
+            for panel_idx, sel_ep in enumerate(selected_epochs):
+                ax_p = fig.add_subplot(gs[panel_idx + 1])
+                model_p = EmbedMLP(
+                    d_vocab=self.d_vocab, d_model=self.d_model,
+                    d_mlp=d_mlp, act_type=act_type, use_cache=False
+                ).to(self.device)
+                model_p.load_state_dict(checkpoints[sel_ep])
+                model_p.eval()
+                with torch.no_grad():
+                    logits_p = model_p(self.all_data).squeeze(1)
+                pred_p = torch.argmax(logits_p, dim=1).view(p, p)
+                acc_p = (pred_p == gt_grid).float()
+
+                ax_p.imshow(acc_p.numpy(), cmap=CMAP_SEQUENTIAL,
+                            vmin=0, vmax=1, aspect='equal')
+                ax_p.set_title(f'Accuracy at Step {sel_ep}', fontsize=21)
+                ax_p.set_xlabel('Second Input', fontsize=18)
+                ax_p.set_xticks(locs)
+                ax_p.set_yticks(locs)
+                ax_p.set_xticklabels(locs, fontsize=11)
+                ax_p.set_yticklabels(locs, fontsize=11)
+                for ti, tj in asymmetric_train_points:
+                    rect = plt.Rectangle((tj - 0.5, ti - 0.5), 1, 1,
+                                         linewidth=2.0, edgecolor='red', facecolor='none')
+                    ax_p.add_patch(rect)
+
+            plt.tight_layout()
+            _save_fig(fig, self._out('grokk_memorization_common_to_rare.png'))
+
+        # ---- Decoded weights dynamic (3 timepoints) ----
+        # Pick 3 representative epochs: 0, stage1, stage2
+        key_epochs = [0]
+        ep_s1 = min(epochs, key=lambda e: abs(e - stage1_end))
+        ep_s2 = min(epochs, key=lambda e: abs(e - stage2_end))
+        if ep_s1 not in key_epochs:
+            key_epochs.append(ep_s1)
+        if ep_s2 not in key_epochs:
+            key_epochs.append(ep_s2)
+        while len(key_epochs) < 3:
+            key_epochs.append(epochs[-1])
+
+        num_components = min(20, d_mlp)
+        n = len(key_epochs)
+        fig, axes = plt.subplots(
+            2, n, figsize=(18, 3.3 * n),
+            gridspec_kw={"hspace": 0.05}, constrained_layout=True
+        )
+        if n == 1:
+            axes = axes.reshape(2, 1)
+
+        x_locs = np.arange(len(self.fourier_basis_names))
+        y_locs = np.arange(num_components)
+
+        for col, key in enumerate(key_epochs):
+            W_in = checkpoints[key]['mlp.W_in']
+            W_out = checkpoints[key]['mlp.W_out']
+
+            data_in = (W_in @ self.fourier_basis.T)[:num_components]
+            data_in_np = data_in.detach().cpu().numpy()
+            abs_max_in = np.abs(data_in_np).max()
+            ax_in = axes[0, col]
+            im_in = ax_in.imshow(
+                data_in_np, cmap=CMAP_DIVERGING,
+                vmin=-abs_max_in, vmax=abs_max_in, aspect='auto'
+            )
+            ax_in.set_title(rf'Step {key}, $\theta_m$ after DFT', fontsize=18)
+            ax_in.set_xticks(x_locs)
+            ax_in.set_xticklabels(self.fourier_basis_names, rotation=90, fontsize=11)
+            ax_in.set_yticks(y_locs)
+            ax_in.set_yticklabels(y_locs)
+            if col == 0:
+                ax_in.set_ylabel('Neuron #', fontsize=16)
+            fig.colorbar(im_in, ax=ax_in)
+
+            data_out = (W_out.T @ self.fourier_basis.T)[:num_components]
+            data_out_np = data_out.detach().cpu().numpy()
+            abs_max_out = np.abs(data_out_np).max() * 0.85
+            ax_out = axes[1, col]
+            im_out = ax_out.imshow(
+                data_out_np, cmap=CMAP_DIVERGING,
+                vmin=-abs_max_out, vmax=abs_max_out, aspect='auto'
+            )
+            ax_out.set_title(rf'Step {key}, $\xi_m$ after DFT', fontsize=18)
+            ax_out.set_xticks(x_locs)
+            ax_out.set_xticklabels(self.fourier_basis_names, rotation=90, fontsize=11)
+            ax_out.set_yticks(y_locs)
+            ax_out.set_yticklabels(y_locs)
+            if col == 0:
+                ax_out.set_ylabel('Neuron #', fontsize=16)
+            fig.colorbar(im_out, ax=ax_out)
+
+        _save_fig(fig, self._out('grokk_decoded_weights_dynamic.png'))
+
+        print("    Saved grokk_loss.json, grokk_loss.png, grokk_acc.json, grokk_acc.png, "
+              "grokk_abs_phase_diff.png, grokk_avg_ipr.png, "
+              "grokk_memorization_accuracy.png, "
+              "grokk_memorization_common_to_rare.png, grokk_decoded_weights_dynamic.png")
+
+    # ------------------------------------------------------------------
+    # Tab 6: Lottery Mechanism
+    # ------------------------------------------------------------------
+
+    def generate_tab6(self):
+        """Generate lottery mechanism plots."""
+        print(f"  [Tab 6] Lottery Mechanism for p={self.p}")
+        run_dir = self._run_dir('quad_random')
+        if run_dir is None:
+            print("    SKIP: quad_random run directory not found")
+            return
+
+        checkpoints = _load_checkpoints(run_dir, self.device)
+        if not checkpoints:
+            print("    SKIP: no quad_random checkpoints")
+            return
+
+        final_data = _load_final(run_dir, self.device)
+        if final_data is None:
+            print("    SKIP: no final quad_random checkpoint")
+            return
+        model_load_final = final_data['model']
+
+        # Select best neuron
+        neuron_id = select_lottery_neuron(
+            model_load_final, self.fourier_basis, decode_scales_phis
+        )
+
+        epochs = sorted(checkpoints.keys())
+        p = self.p
+
+        # Collect per-checkpoint scales and phase diffs for the selected neuron
+        scales_list = []
+        diff_list = []
+        for ep in epochs:
+            scales, phis, psis = decode_scales_phis(
+                checkpoints[ep], self.fourier_basis
+            )
+            scales_list.append(scales[neuron_id])
+            diff_list.append(normalize_to_pi(
+                psis[neuron_id] - 2 * phis[neuron_id]
+            ))
+
+        # Stack: [num_checkpoints, K+1], skip DC
+        scales_all = torch.stack(scales_list, dim=0)[:, 1:]
+        diff_all = torch.stack(diff_list, dim=0)[:, 1:]
+
+        # Determine which frequency this neuron specializes in
+        _, _, max_freq_ls = decode_weights(model_load_final, self.fourier_basis)
+        max_freq = max_freq_ls[neuron_id] - 1  # 0-indexed into scales_all
+
+        scales_np = scales_all.cpu().numpy()
+        diff_np = diff_all.cpu().numpy()
+        num_models, num_freqs = scales_np.shape
+        n_plot = min(num_models, 160)
+        scales_np = scales_np[:n_plot]
+        diff_np = diff_np[:n_plot]
+        x_idx = np.arange(n_plot)
+
+        # Color gradient for non-highlighted frequencies
+        base_rgb = np.array(mcolors.to_rgb(COLORS[0]))
+        gray_rgb = np.array(mcolors.to_rgb('white'))
+        highlight_color = COLORS[3]
+
+        nonmax = [f for f in range(num_freqs) if f != max_freq]
+        final_scales = scales_np[-1]
+        sorted_nonmax = sorted(nonmax, key=lambda f: final_scales[f])
+        M = len(sorted_nonmax)
+
+        # Compute save_every for x-axis formatter
+        save_every = epochs[1] - epochs[0] if len(epochs) > 1 else 200
+
+        # ---- Magnitude plot ----
+        fig, ax = plt.subplots(figsize=(4, 4))
+        for idx, f in enumerate(sorted_nonmax):
+            blend = idx / (M - 1) if M > 1 else 0.0
+            col_rgb = (1 - blend - 0.05) * gray_rgb + (blend + 0.05) * base_rgb
+            ax.plot(x_idx, scales_np[:, f], color=col_rgb, linestyle=':',
+                    marker='x', linewidth=3.5, markersize=1.5,
+                    label=f"Freq. {f + 1}")
+
+        ax.plot(x_idx, scales_np[:, max_freq], color=highlight_color,
+                linestyle=':', marker='x', linewidth=3.5, markersize=1.5,
+                label=f"Freq. {max_freq + 1}")
+
+        ax.xaxis.set_major_formatter(
+            FuncFormatter(lambda val, pos: f"{int(val * save_every)}")
+        )
+        ax.legend(loc='upper left', bbox_to_anchor=(1.02, 1),
+                  borderaxespad=0.2, frameon=False, fontsize=13)
+        ax.set_xlabel("Step", fontsize=16)
+        ax.set_ylabel("Magnitude", fontsize=16)
+        ax.grid(True, linestyle='--', linewidth=0.5, alpha=0.7)
+        _save_fig(fig, self._out('lottery_mech_magnitude.png'))
+
+        # ---- Phase misalignment plot ----
+        fig, ax = plt.subplots(figsize=(4, 4))
+        ax.axhline(y=0, color='black', linewidth=1, linestyle='dotted')
+        for idx, f in enumerate(sorted_nonmax):
+            blend = idx / (M - 1) if M > 1 else 0.0
+            col_rgb = (1 - blend - 0.05) * gray_rgb + (blend + 0.05) * base_rgb
+            ax.plot(x_idx, diff_np[:, f], linestyle=':', marker='x',
+                    linewidth=3.5, markersize=1.5, color=col_rgb,
+                    label=f"Freq. {f}")
+
+        ax.plot(x_idx, diff_np[:, max_freq], linestyle=':', marker='x',
+                linewidth=3.5, markersize=1.5, color=highlight_color,
+                label=f"Freq. {max_freq}")
+
+        ax.xaxis.set_major_formatter(
+            FuncFormatter(lambda val, pos: f"{int(val * save_every)}")
+        )
+        ax.set_xlabel("Step", fontsize=16)
+        ax.set_ylabel("Misalignment", fontsize=16)
+        ax.grid(True, linestyle='--', linewidth=0.5, alpha=0.7)
+        _save_fig(fig, self._out('lottery_mech_phase.png'))
+
+        # ---- Beta contour: simulate gradient flow ----
+        self._generate_lottery_contour()
+
+        print("    Saved lottery_mech_magnitude.png, lottery_mech_phase.png, "
+              "lottery_beta_contour.png")
+
+    def _generate_lottery_contour(self):
+        """Simulate gradient flow for a grid of (init_magnitude, init_phase_diff)."""
+        p = self.p
+        device = self.device
+        init_k = 1
+        init_psi = 0.0
+        num_steps = 100
+        learning_rate = 0.01
+
+        fourier_basis, _ = get_fourier_basis(p, device)
+        fourier_basis = fourier_basis.to(torch.get_default_dtype())
+
+        initial_scales = np.linspace(0.01, 0.02, num=30)
+        phi0_vals = np.linspace(0, np.pi, num=30)
+
+        results = []
+        for scale in initial_scales:
+            for phi0 in phi0_vals:
+                w_k = 2 * np.pi * init_k / p
+                theta = scale * torch.tensor(
+                    [np.cos(w_k * j + phi0) for j in range(p)],
+                    device=device
+                )
+                xi = scale * torch.tensor(
+                    [np.cos(w_k * j + init_psi) for j in range(p)],
+                    device=device
+                )
+
+                # Run gradient flow simulation
+                for _ in range(num_steps):
+                    theta, xi = self._gradient_flow_step(
+                        theta, xi, init_k, p, learning_rate, fourier_basis
+                    )
+
+                # Compute final beta
+                coeffs_xi = fourier_basis.to(xi.dtype) @ xi
+                idx = [init_k * 2 - 1, init_k * 2]
+                xi_n = coeffs_xi[idx]
+                beta_f = torch.norm(xi_n).item() * np.sqrt(2 / p)
+
+                results.append({
+                    "init_scale": scale,
+                    "init_diff": 2 * phi0,
+                    "beta_f": beta_f,
+                })
+
+        # Pivot into grid
+        n_scales = len(initial_scales)
+        n_phis = len(phi0_vals)
+        Z = np.zeros((n_phis, n_scales))
+        for i, r in enumerate(results):
+            row = i % n_phis
+            col = i // n_phis
+            Z[row, col] = r['beta_f']
+
+        X, Y = np.meshgrid(initial_scales, 2 * phi0_vals)
+
+        fig = plt.figure(figsize=(4.5, 4))
+        cf = plt.contourf(X, Y, Z, levels=12, cmap=CMAP_DIVERGING, extend='both')
+        plt.axhline(y=np.pi, color='white', linewidth=1, linestyle=':')
+        plt.xlabel("Initial Magnitude", fontsize=16)
+        plt.ylabel("Initial Phase Difference", fontsize=16)
+        plt.title("Contour of Final Magnitude", fontsize=16)
+        plt.colorbar(cf)
+        plt.tight_layout()
+        _save_fig(fig, self._out('lottery_beta_contour.png'))
+
+    @staticmethod
+    def _gradient_flow_step(theta, xi, init_k, p, lr, fourier_basis):
+        """One step of analytical gradient flow."""
+        fb = fourier_basis.to(theta.dtype)
+        theta_coeff = fb @ theta
+        xi_coeff = fb @ xi
+
+        neuron_coeff_theta = theta_coeff[[init_k * 2 - 1, init_k * 2]]
+        alpha = np.sqrt(2 / p) * torch.sqrt(
+            torch.sum(neuron_coeff_theta.pow(2))
+        ).item()
+        phi = np.arctan2(
+            -neuron_coeff_theta[1].item(), neuron_coeff_theta[0].item()
+        )
+
+        neuron_coeff_xi = xi_coeff[[init_k * 2 - 1, init_k * 2]]
+        beta = np.sqrt(2 / p) * torch.sqrt(
+            torch.sum(neuron_coeff_xi.pow(2))
+        ).item()
+        psi = np.arctan2(
+            -neuron_coeff_xi[1].item(), neuron_coeff_xi[0].item()
+        )
+
+        w_k = 2 * np.pi * init_k / p
+        grad_theta = torch.tensor(
+            [2 * p * alpha * beta * np.cos(w_k * j + psi - phi)
+             for j in range(p)],
+            device=theta.device
+        )
+        grad_xi = torch.tensor(
+            [p * alpha ** 2 * np.cos(w_k * j + 2 * phi)
+             for j in range(p)],
+            device=theta.device
+        )
+
+        theta = theta + lr * grad_theta
+        xi = xi + lr * grad_xi
+        return theta, xi
+
+    # ------------------------------------------------------------------
+    # Tab 7: Gradient Dynamics
+    # ------------------------------------------------------------------
+
+    def generate_tab7(self):
+        """Generate gradient dynamics plots for quad_single_freq and relu_single_freq."""
+        print(f"  [Tab 7] Gradient Dynamics for p={self.p}")
+
+        for run_name, act_name, prefix in [
+            ('quad_single_freq', 'Quad', 'quad'),
+            ('relu_single_freq', 'ReLU', 'relu'),
+        ]:
+            run_dir = self._run_dir(run_name)
+            if run_dir is None:
+                print(f"    SKIP: {run_name} run directory not found")
+                continue
+
+            checkpoints = _load_checkpoints(run_dir, self.device)
+            if not checkpoints:
+                print(f"    SKIP: no {run_name} checkpoints")
+                continue
+
+            epochs = sorted(checkpoints.keys())
+            d_mlp = self.d_mlp
+
+            # Build all neuron records across epochs
+            all_neuron_records = []
+            for ep in epochs:
+                model_sd = checkpoints[ep]
+                W_in_d, W_out_d, mfl = decode_weights(model_sd, self.fourier_basis)
+                for neuron in range(W_in_d.shape[0]):
+                    s_in, phi_in = compute_neuron(neuron, mfl, W_in_d)
+                    s_out, phi_out = compute_neuron(neuron, mfl, W_out_d)
+                    all_neuron_records.append({
+                        'epoch': ep,
+                        'neuron': neuron,
+                        'scale_in': s_in,
+                        'phi_in': phi_in,
+                        'scale_out': s_out,
+                        'phi_out': phi_out,
+                    })
+
+            # Select a neuron that shows clear phase alignment
+            # Pick neuron with largest final scale
+            final_records = [r for r in all_neuron_records if r['epoch'] == epochs[-1]]
+            if not final_records:
+                continue
+            best_neuron = max(final_records, key=lambda r: r['scale_in'])['neuron']
+
+            # Extract trajectory for this neuron
+            neuron_records = [r for r in all_neuron_records if r['neuron'] == best_neuron]
+            # Remove last few points if noisy (as notebooks do)
+            trim = max(0, len(neuron_records) - 4) if prefix == 'relu' else max(0, len(neuron_records) - 14)
+            neuron_records = neuron_records[:trim] if trim > 0 else neuron_records
+
+            phi_in_raw = [r['phi_in'] for r in neuron_records]
+            phi_out_raw = [r['phi_out'] for r in neuron_records]
+            scale_in_list = [r['scale_in'] for r in neuron_records]
+            scale_out_list = [r['scale_out'] for r in neuron_records]
+
+            # Phase wrapping fix: normalize 2*phi to [-pi, pi], then adjust
+            # psi to stay within pi of 2*phi (same fix as Tab 3 scatter).
+            phi2_in_list = [normalize_to_pi(2 * v) for v in phi_in_raw]
+            phi_out_list = []
+            for two_phi, psi in zip(phi2_in_list, phi_out_raw):
+                psi_n = normalize_to_pi(psi)
+                if psi_n - two_phi > np.pi:
+                    psi_n -= 2 * np.pi
+                elif psi_n - two_phi < -np.pi:
+                    psi_n += 2 * np.pi
+                phi_out_list.append(psi_n)
+
+            # Unwrap time series to remove remaining jumps at +-pi boundary
+            phi_in_list = list(np.unwrap(phi_in_raw))
+            phi2_in_list = list(np.unwrap(phi2_in_list))
+            phi_out_list = list(np.unwrap(phi_out_list))
+
+            x = np.arange(len(phi_in_list)) * (epochs[1] - epochs[0] if len(epochs) > 1 else 200)
+
+            # ---- Phase alignment + magnitude plot ----
+            fig_width = 8 if prefix == 'quad' else 5
+            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(fig_width, 6), sharex=True)
+
+            ax1.plot(x, phi_in_list, marker='o', markersize=4,
+                     color=COLORS[1], label=r"$\phi_m^\star$")
+            ax1.plot(x, phi_out_list, marker='x', markersize=4,
+                     color=COLORS[3], label=r"$\psi_m^\star$")
+            ax1.plot(x, phi2_in_list, marker='^', markersize=4,
+                     color=COLORS[0], label=r"$2\phi_m^\star$")
+            ax1.set_title('Phase Alignment of Neuron $m$', fontsize=16)
+            ax1.legend(fontsize=18, loc="upper right")
+            ax1.grid(True)
+
+            ax2.plot(x, scale_in_list, marker='o', markersize=4,
+                     color=COLORS[0], label=r"$\alpha_m^\star$")
+            ax2.plot(x, scale_out_list, marker='x', markersize=4,
+                     color=COLORS[3], label=r"$\beta_m^\star$")
+            ax2.set_title('Magnitude Growth of Neuron $m$', fontsize=16)
+            ax2.set_xlabel('Step', fontsize=16)
+            ax2.legend(fontsize=18, loc="upper left")
+            ax2.grid(True)
+
+            plt.tight_layout()
+            _save_fig(fig, self._out(f'phase_align_{prefix}.png'))
+
+            # ---- Decoded weights at timepoints ----
+            if prefix == 'quad':
+                keys = [0]
+                mid = min(epochs, key=lambda e: abs(e - 1000))
+                end = epochs[-1]
+                if mid not in keys:
+                    keys.append(mid)
+                if end not in keys:
+                    keys.append(end)
+            else:
+                keys = [0, epochs[-1]]
+
+            num_components = min(20, d_mlp)
+            n = len(keys)
+            fig, axes = plt.subplots(
+                2, n, figsize=(12 if n <= 2 else 18, 4 * n if n <= 2 else 3.3 * n),
+                gridspec_kw={"hspace": 0.05}, constrained_layout=True
+            )
+            if n == 1:
+                axes = axes.reshape(2, 1)
+
+            x_locs = np.arange(len(self.fourier_basis_names))
+            y_locs = np.arange(num_components)
+
+            for col, key in enumerate(keys):
+                if key not in checkpoints:
+                    key = min(checkpoints.keys(), key=lambda e: abs(e - key))
+                W_in = checkpoints[key]['mlp.W_in']
+                W_out = checkpoints[key]['mlp.W_out']
+
+                data_in = (W_in @ self.fourier_basis.T)[:num_components]
+                data_in_np = data_in.detach().cpu().numpy()
+                abs_max_in = np.abs(data_in_np).max()
+                ax_in = axes[0, col]
+                im_in = ax_in.imshow(
+                    data_in_np, cmap=CMAP_DIVERGING,
+                    vmin=-abs_max_in, vmax=abs_max_in, aspect='auto'
+                )
+                ax_in.set_title(rf'Step {key}, $\theta_m$ after DFT', fontsize=18)
+                ax_in.set_xticks(x_locs)
+                ax_in.set_xticklabels(self.fourier_basis_names, rotation=90, fontsize=11)
+                ax_in.set_yticks(y_locs)
+                ax_in.set_yticklabels(y_locs)
+                if col == 0:
+                    ax_in.set_ylabel('Neuron #', fontsize=16)
+                fig.colorbar(im_in, ax=ax_in)
+
+                data_out = (W_out.T @ self.fourier_basis.T)[:num_components]
+                data_out_np = data_out.detach().cpu().numpy()
+                abs_max_out = np.abs(data_out_np).max()
+                ax_out = axes[1, col]
+                im_out = ax_out.imshow(
+                    data_out_np, cmap=CMAP_DIVERGING,
+                    vmin=-abs_max_out, vmax=abs_max_out, aspect='auto'
+                )
+                ax_out.set_title(rf'Step {key}, $\xi_m$ after DFT', fontsize=18)
+                ax_out.set_xticks(x_locs)
+                ax_out.set_xticklabels(self.fourier_basis_names, rotation=90, fontsize=11)
+                ax_out.set_yticks(y_locs)
+                ax_out.set_yticklabels(y_locs)
+                if col == 0:
+                    ax_out.set_ylabel('Neuron #', fontsize=16)
+                fig.colorbar(im_out, ax=ax_out)
+
+            _save_fig(fig, self._out(f'single_freq_{prefix}.png'))
+
+            print(f"    Saved phase_align_{prefix}.png, single_freq_{prefix}.png")
+
+    # ------------------------------------------------------------------
+    # Metadata JSON
+    # ------------------------------------------------------------------
+
+    def _save_metadata(self):
+        """Save a metadata JSON summarizing config and final metrics."""
+        print(f"  [Meta] Saving metadata for p={self.p}")
+        meta = {
+            'prime': self.p,
+            'd_mlp': self.d_mlp,
+            'training_runs': {},
+            'final_metrics': {},
+        }
+        for run_name, params in TRAINING_RUNS.items():
+            meta['training_runs'][run_name] = {
+                'act_type': params['act_type'],
+                'lr': params['lr'],
+                'weight_decay': params['weight_decay'],
+                'num_epochs': params['num_epochs'],
+                'frac_train': params['frac_train'],
+                'init_type': params['init_type'],
+                'init_scale': params['init_scale'],
+                'optimizer': params['optimizer'],
+            }
+            curves = _load_training_curves(self._run_type_dir(run_name))
+            if curves:
+                metrics = {}
+                if 'train_accs' in curves and curves['train_accs']:
+                    metrics['train_acc'] = curves['train_accs'][-1]
+                if 'test_accs' in curves and curves['test_accs']:
+                    metrics['test_acc'] = curves['test_accs'][-1]
+                if 'train_losses' in curves and curves['train_losses']:
+                    metrics['train_loss'] = curves['train_losses'][-1]
+                if 'test_losses' in curves and curves['test_losses']:
+                    metrics['test_loss'] = curves['test_losses'][-1]
+                if metrics:
+                    meta['final_metrics'][run_name] = metrics
+
+        with open(self._out('metadata.json'), 'w') as f:
+            json.dump(meta, f, indent=2)
+        print("    Saved metadata.json")
+
+    # ------------------------------------------------------------------
+    # Interactive JSON precomputation
+    # ------------------------------------------------------------------
+
+    def _precompute_neuron_spectra(self):
+        """Precompute per-neuron Fourier magnitude spectra for top-20 neurons."""
+        print(f"  [Interactive] Neuron spectra for p={self.p}")
+        run_dir = self._run_dir('standard')
+        if run_dir is None:
+            print("    SKIP: standard run directory not found")
+            return
+
+        final_data = _load_final(run_dir, self.device)
+        if final_data is None:
+            print("    SKIP: no final checkpoint")
+            return
+        model_load = final_data['model']
+
+        W_in_decode, W_out_decode, max_freq_ls = decode_weights(
+            model_load, self.fourier_basis
+        )
+        d_mlp = W_in_decode.shape[0]
+        num_neurons = min(20, d_mlp)
+
+        sorted_indices = select_top_neurons_by_frequency(
+            max_freq_ls, W_in_decode, n=num_neurons
+        )
+
+        fb_names = self.fourier_basis_names
+        spectra = {}
+        for rank, neuron_idx in enumerate(sorted_indices):
+            # Fourier magnitudes for W_in
+            magnitudes_in = W_in_decode[neuron_idx].abs().cpu().tolist()
+            magnitudes_out = W_out_decode[neuron_idx].abs().cpu().tolist()
+            spectra[f"neuron_{rank}"] = {
+                'global_index': int(neuron_idx),
+                'dominant_freq': int(max_freq_ls[neuron_idx]),
+                'fourier_magnitudes_in': magnitudes_in,
+                'fourier_magnitudes_out': magnitudes_out,
+            }
+
+        payload = {
+            'fourier_basis_names': fb_names,
+            'neurons': spectra,
+        }
+        with open(self._out('neuron_spectra.json'), 'w') as f:
+            json.dump(payload, f)
+        print("    Saved neuron_spectra.json")
+
+    def _precompute_logit_explorer(self):
+        """Precompute logits for representative (a,b) pairs."""
+        print(f"  [Interactive] Logit explorer for p={self.p}")
+        run_dir = self._run_dir('standard')
+        if run_dir is None:
+            print("    SKIP: standard run directory not found")
+            return
+
+        final_data = _load_final(run_dir, self.device)
+        if final_data is None:
+            print("    SKIP: no final checkpoint")
+            return
+        model_load = final_data['model']
+
+        p = self.p
+        act_type = TRAINING_RUNS['standard']['act_type']
+        model = EmbedMLP(
+            d_vocab=self.d_vocab, d_model=self.d_model,
+            d_mlp=self.d_mlp, act_type=act_type, use_cache=False
+        )
+        model.to(self.device)
+        model.load_state_dict(model_load)
+        model.eval()
+
+        # Select p representative pairs: (0,0), (1,2), (3,5), ... spread across inputs
+        pairs = []
+        step = max(1, (p * p) // p)
+        for idx in range(0, p * p, step):
+            a = idx // p
+            b = idx % p
+            pairs.append((a, b))
+            if len(pairs) >= p:
+                break
+
+        pair_tensor = torch.tensor(pairs, dtype=torch.long, device=self.device)
+        with torch.no_grad():
+            logits = model(pair_tensor).squeeze(1)  # [n_pairs, p]
+
+        payload = {
+            'pairs': pairs,
+            'correct_answers': [(a + b) % p for a, b in pairs],
+            'logits': logits.cpu().tolist(),
+            'output_classes': list(range(p)),
+        }
+        with open(self._out('logits_interactive.json'), 'w') as f:
+            json.dump(payload, f)
+        print("    Saved logits_interactive.json")
+
+    def _precompute_grokk_slider(self):
+        """Precompute accuracy grids at ~10 grokking checkpoints for epoch slider."""
+        print(f"  [Interactive] Grokking epoch slider for p={self.p}")
+        if self.p < MIN_P_GROKKING:
+            print(f"    SKIP: p={self.p} < {MIN_P_GROKKING}")
+            return
+        run_dir = self._run_dir('grokking')
+        if run_dir is None:
+            print("    SKIP: grokking run directory not found")
+            return
+
+        checkpoints = _load_checkpoints(run_dir, self.device)
+        if not checkpoints:
+            print("    SKIP: no grokking checkpoints")
+            return
+
+        epochs = sorted(checkpoints.keys())
+        p = self.p
+        d_mlp = self.d_mlp
+        act_type = TRAINING_RUNS['grokking']['act_type']
+        gt_grid = self.all_labels.view(p, p)
+
+        # Subsample ~10 epochs evenly
+        n_snapshots = min(10, len(epochs))
+        indices = np.linspace(0, len(epochs) - 1, n_snapshots, dtype=int)
+        selected_epochs = [epochs[i] for i in indices]
+
+        epoch_data = []
+        for ep in selected_epochs:
+            model = EmbedMLP(
+                d_vocab=self.d_vocab, d_model=self.d_model,
+                d_mlp=d_mlp, act_type=act_type, use_cache=False
+            ).to(self.device)
+            model.load_state_dict(checkpoints[ep])
+            model.eval()
+            with torch.no_grad():
+                logits = model(self.all_data).squeeze(1)
+            predicted = torch.argmax(logits, dim=1).view(p, p)
+            accuracy_grid = (predicted == gt_grid).float().cpu().tolist()
+            epoch_data.append({
+                'epoch': int(ep),
+                'accuracy_grid': accuracy_grid,
+            })
+
+        payload = {
+            'prime': p,
+            'epochs': [d['epoch'] for d in epoch_data],
+            'grids': [d['accuracy_grid'] for d in epoch_data],
+        }
+        with open(self._out('grokk_epoch_data.json'), 'w') as f:
+            json.dump(payload, f)
+        print("    Saved grokk_epoch_data.json")
+
+    # ------------------------------------------------------------------
+    # Training Log consolidation
+    # ------------------------------------------------------------------
+
+    def _save_training_log(self):
+        """Consolidate training logs from all runs into a precomputed JSON.
+
+        For each run, includes:
+        - config: hyperparameters
+        - log_text: human-readable formatted log
+        - table: subsampled per-epoch metrics for display
+        """
+        print(f"  [Log] Saving training log for p={self.p}")
+        all_runs = {}
+
+        for run_name, params in TRAINING_RUNS.items():
+            run_type_dir = self._run_type_dir(run_name)
+            curves = _load_training_curves(run_type_dir)
+            if curves is None:
+                continue
+
+            # Also check for a pre-saved training_log.txt
+            log_text_path = os.path.join(run_type_dir, "training_log.txt")
+            if os.path.exists(log_text_path):
+                with open(log_text_path) as f:
+                    log_text = f.read()
+            else:
+                # Reconstruct from curves data
+                log_text = self._reconstruct_log_text(
+                    run_name, params, curves
+                )
+
+            # Build a subsampled table (~100 rows max)
+            n_epochs = len(curves.get('train_losses', []))
+            step = max(1, n_epochs // 100)
+            indices = list(range(0, n_epochs, step))
+            if n_epochs > 0 and (n_epochs - 1) not in indices:
+                indices.append(n_epochs - 1)
+
+            table = []
+            for i in indices:
+                row = {'epoch': i}
+                for key in ('train_losses', 'test_losses', 'train_accs',
+                            'test_accs', 'grad_norms', 'param_norms'):
+                    vals = curves.get(key, [])
+                    row[key.replace('_', '_')] = (
+                        round(vals[i], 6) if i < len(vals) else None
+                    )
+                table.append(row)
+
+            all_runs[run_name] = {
+                'config': {
+                    'prime': self.p,
+                    'd_mlp': self.d_mlp,
+                    'act_type': params['act_type'],
+                    'init_type': params['init_type'],
+                    'init_scale': params['init_scale'],
+                    'optimizer': params['optimizer'],
+                    'lr': params['lr'],
+                    'weight_decay': params['weight_decay'],
+                    'frac_train': params['frac_train'],
+                    'num_epochs': params['num_epochs'],
+                    'seed': params['seed'],
+                },
+                'log_text': log_text,
+                'table': table,
+                'total_epochs': n_epochs,
+            }
+
+        if all_runs:
+            with open(self._out('training_log.json'), 'w') as f:
+                json.dump(all_runs, f)
+            print(f"    Saved training_log.json ({len(all_runs)} runs)")
+        else:
+            print("    SKIP: no training curves found")
+
+    def _reconstruct_log_text(self, run_name, params, curves):
+        """Reconstruct a human-readable training log from curves data."""
+        lines = []
+        lines.append(f"{'=' * 70}")
+        lines.append(f"Training Log: p={self.p}, run={run_name}")
+        lines.append(f"{'=' * 70}")
+        lines.append("")
+        lines.append("Configuration:")
+        lines.append(f"  prime (p)       = {self.p}")
+        lines.append(f"  d_mlp           = {self.d_mlp}")
+        lines.append(f"  activation      = {params['act_type']}")
+        lines.append(f"  init_type       = {params['init_type']}")
+        lines.append(f"  init_scale      = {params['init_scale']}")
+        lines.append(f"  optimizer       = {params['optimizer']}")
+        lines.append(f"  learning_rate   = {params['lr']}")
+        lines.append(f"  weight_decay    = {params['weight_decay']}")
+        lines.append(f"  frac_train      = {params['frac_train']}")
+        lines.append(f"  num_epochs      = {params['num_epochs']}")
+        lines.append(f"  seed            = {params['seed']}")
+        lines.append("")
+        lines.append(f"{'─' * 70}")
+        lines.append(
+            f"{'Epoch':>8s}  {'Train Loss':>12s}  {'Test Loss':>12s}  "
+            f"{'Train Acc':>10s}  {'Test Acc':>10s}  "
+            f"{'Grad Norm':>10s}  {'Param Norm':>11s}"
+        )
+        lines.append(f"{'─' * 70}")
+
+        train_losses = curves.get('train_losses', [])
+        test_losses = curves.get('test_losses', [])
+        train_accs = curves.get('train_accs', [])
+        test_accs = curves.get('test_accs', [])
+        grad_norms = curves.get('grad_norms', [])
+        param_norms = curves.get('param_norms', [])
+        n_epochs = len(train_losses)
+
+        step = max(1, n_epochs // 100)
+        indices = list(range(0, n_epochs, step))
+        if n_epochs > 0 and (n_epochs - 1) not in indices:
+            indices.append(n_epochs - 1)
+
+        for i in indices:
+            tl = f"{train_losses[i]:.6f}" if i < len(train_losses) else "N/A"
+            tel = f"{test_losses[i]:.6f}" if i < len(test_losses) else "N/A"
+            ta = f"{train_accs[i]:.4f}" if i < len(train_accs) else "N/A"
+            tea = f"{test_accs[i]:.4f}" if i < len(test_accs) else "N/A"
+            gn = f"{grad_norms[i]:.4f}" if i < len(grad_norms) else "N/A"
+            pn = f"{param_norms[i]:.4f}" if i < len(param_norms) else "N/A"
+            lines.append(
+                f"{i:>8d}  {tl:>12s}  {tel:>12s}  "
+                f"{ta:>10s}  {tea:>10s}  "
+                f"{gn:>10s}  {pn:>11s}"
+            )
+
+        lines.append(f"{'─' * 70}")
+        lines.append("")
+        lines.append("Final Results:")
+        if train_losses:
+            lines.append(f"  Train Loss  = {train_losses[-1]:.6f}")
+        if test_losses:
+            lines.append(f"  Test Loss   = {test_losses[-1]:.6f}")
+        if train_accs:
+            lines.append(f"  Train Acc   = {train_accs[-1]:.4f}")
+        if test_accs:
+            lines.append(f"  Test Acc    = {test_accs[-1]:.4f}")
+        if param_norms:
+            lines.append(f"  Param Norm  = {param_norms[-1]:.4f}")
+        lines.append(f"\nTotal epochs trained: {n_epochs}")
+        return "\n".join(lines)
+
+    # ------------------------------------------------------------------
+    # Generate all
+    # ------------------------------------------------------------------
+
+    def generate_all(self):
+        """Generate all tab plots with error handling."""
+        print(f"\n{'=' * 60}")
+        print(f"Generating plots for p={self.p}")
+        print(f"  Input:  {self.input_dir}")
+        print(f"  Output: {self.output_dir}")
+        print(f"{'=' * 60}")
+
+        # Save metadata and training logs first
+        try:
+            self._save_metadata()
+        except Exception as e:
+            print(f"  [ERROR] metadata failed: {e}")
+            traceback.print_exc()
+
+        try:
+            self._save_training_log()
+        except Exception as e:
+            print(f"  [ERROR] training log failed: {e}")
+            traceback.print_exc()
+
+        generators = [
+            ('Tab 1', self.generate_tab1),
+            ('Tab 2', self.generate_tab2),
+            ('Tab 3', self.generate_tab3),
+            ('Tab 4', self.generate_tab4),
+            ('Tab 5', self.generate_tab5),
+            ('Tab 6', self.generate_tab6),
+            ('Tab 7', self.generate_tab7),
+        ]
+
+        for name, gen_fn in generators:
+            try:
+                gen_fn()
+            except Exception as e:
+                print(f"  [ERROR] {name} failed: {e}")
+                traceback.print_exc()
+
+        # Precompute interactive JSON data
+        interactive = [
+            ('Neuron Spectra', self._precompute_neuron_spectra),
+            ('Logit Explorer', self._precompute_logit_explorer),
+            ('Grokking Slider', self._precompute_grokk_slider),
+        ]
+        for name, fn in interactive:
+            try:
+                fn()
+            except Exception as e:
+                print(f"  [ERROR] {name} failed: {e}")
+                traceback.print_exc()
+
+        print(f"\nDone generating plots for p={self.p}")
+
+
+# ======================================================================
+# CLI
+# ======================================================================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Generate all model-dependent plots for the HF app.'
+    )
+    parser.add_argument('--all', action='store_true',
+                        help='Generate plots for all p found in input dir')
+    parser.add_argument('--p', type=int,
+                        help='Generate plots for a specific p')
+    parser.add_argument('--input', type=str, default='./trained_models',
+                        help='Base input directory containing p_PPP subdirs')
+    parser.add_argument('--output', type=str,
+                        default='./precomputed_results',
+                        help='Base output directory for precomputed results')
+    args = parser.parse_args()
+
+    if not args.all and args.p is None:
+        parser.error("Specify --all or --p P")
+
+    if args.p:
+        moduli = [args.p]
+    else:
+        # Discover moduli from input directory
+        moduli = []
+        if os.path.isdir(args.input):
+            for d in sorted(os.listdir(args.input)):
+                if d.startswith('p_'):
+                    try:
+                        p = int(d.split('_')[1])
+                        moduli.append(p)
+                    except (ValueError, IndexError):
+                        pass
+        if not moduli:
+            print(f"No p_PPP directories found in {args.input}")
+            sys.exit(1)
+
+    total = len(moduli)
+    for i, p in enumerate(moduli):
+        print(f"\n[{i + 1}/{total}] Processing p={p}")
+        # Handle both p_23 and p_023 naming conventions
+        input_dir = os.path.join(args.input, f'p_{p:03d}')
+        if not os.path.isdir(input_dir):
+            input_dir = os.path.join(args.input, f'p_{p}')
+        if not os.path.isdir(input_dir):
+            print(f"  Input directory not found: {input_dir}")
+            continue
+
+        output_dir = os.path.join(args.output, f'p_{p:03d}')
+
+        gen = PlotGenerator(p=p, input_dir=input_dir, output_dir=output_dir)
+        gen.generate_all()
+
+    print(f"\nAll done. Processed {total} prime(s).")
+
+
+if __name__ == '__main__':
+    main()

precompute/grokking_stage_detector.py

@@ -0,0 +1,55 @@
+"""
+Automatic detection of grokking stage boundaries from training curves.
+
+Three stages:
+  Stage 1 (Memorization): Train accuracy rises, test accuracy stays low
+  Stage 2 (Transition): Test accuracy starts climbing
+  Stage 3 (Generalization): Test accuracy near 1.0
+
+Returns (stage1_end, stage2_end) as epoch indices, or (None, None) if
+grokking is not detected.
+"""
+
+
+def detect_grokking_stages(train_losses, test_losses, train_accs=None, test_accs=None):
+    """
+    Detect memorization -> transition -> generalization boundaries.
+
+    Heuristic:
+    - stage1_end: first epoch where train accuracy >= 0.95 (memorization complete)
+    - stage2_end: first epoch where test accuracy >= 0.95 (generalization reached)
+
+    Fallback (if accuracy curves not available):
+    - stage1_end: first epoch where train loss < 0.1
+    - stage2_end: first epoch where test loss < 0.1
+    """
+    if train_accs is not None and test_accs is not None:
+        stage1_end = None
+        for i, a in enumerate(train_accs):
+            if a >= 0.95:
+                stage1_end = i
+                break
+
+        stage2_end = None
+        for i, a in enumerate(test_accs):
+            if a >= 0.95:
+                stage2_end = i
+                break
+
+        # Sanity: stage1 should come before stage2
+        if stage1_end is not None and stage2_end is not None and stage1_end >= stage2_end:
+            stage1_end = stage2_end // 3
+    else:
+        stage1_end = None
+        for i, loss in enumerate(train_losses):
+            if loss < 0.1:
+                stage1_end = i
+                break
+
+        stage2_end = None
+        for i, loss in enumerate(test_losses):
+            if loss < 0.1:
+                stage2_end = i
+                break
+
+    return stage1_end, stage2_end

precompute/neuron_selector.py

@@ -0,0 +1,98 @@
+"""
+Automated neuron selection strategies for all primes.
+Replaces hard-coded neuron indices from the analysis notebooks.
+"""
+import torch
+import numpy as np
+from collections import Counter
+
+
+def select_top_neurons_by_frequency(max_freq_ls, W_in_decode, n=20):
+    """
+    Select top N neurons covering all frequencies (round-robin).
+    Used for heatmap plots (Tab 2).
+
+    Picks the highest-magnitude neuron from each frequency in turn,
+    cycling through frequencies until n neurons are selected. This ensures
+    the heatmap shows diversification across all frequencies, matching
+    the blog's Figure 2.
+
+    Returns list of neuron indices into the original d_mlp-sized arrays.
+    """
+    d_mlp = W_in_decode.shape[0]
+    magnitudes = W_in_decode.abs().max(dim=1).values
+
+    # Group neurons by their dominant frequency, sorted by magnitude (descending)
+    from collections import defaultdict
+    freq_groups = defaultdict(list)
+    for i in range(d_mlp):
+        f = max_freq_ls[i]
+        if f > 0:  # skip DC neurons
+            freq_groups[f].append((magnitudes[i].item(), i))
+
+    # Sort each group by magnitude descending
+    for f in freq_groups:
+        freq_groups[f].sort(key=lambda x: -x[0])
+
+    # Round-robin across frequencies (ascending order)
+    freqs_sorted = sorted(freq_groups.keys())
+    selected = []
+    pointers = {f: 0 for f in freqs_sorted}
+
+    while len(selected) < min(n, d_mlp) and freqs_sorted:
+        exhausted = []
+        for f in freqs_sorted:
+            if len(selected) >= n:
+                break
+            if pointers[f] < len(freq_groups[f]):
+                _, idx = freq_groups[f][pointers[f]]
+                selected.append(idx)
+                pointers[f] += 1
+            else:
+                exhausted.append(f)
+        for f in exhausted:
+            freqs_sorted.remove(f)
+
+    return selected
+
+
+def select_lineplot_neurons(sorted_indices, n=3):
+    """
+    Select first N neurons from the frequency-sorted set for line plots (Tab 2).
+    Picks neurons evenly spaced through the sorted list to show diverse frequencies.
+    """
+    if len(sorted_indices) <= n:
+        return list(range(len(sorted_indices)))
+    step = len(sorted_indices) // n
+    return [i * step for i in range(n)]
+
+
+def select_phase_frequency(max_freq_ls, p):
+    """
+    Choose the frequency for phase distribution analysis (Tab 3).
+    Picks the frequency with the most neurons assigned to it (mode),
+    excluding frequency 0 (DC component).
+    """
+    freq_counts = Counter(f for f in max_freq_ls if f > 0)
+    if not freq_counts:
+        return 1
+    return freq_counts.most_common(1)[0][0]
+
+
+def select_lottery_neuron(model_load, fourier_basis, decode_scales_phis_fn):
+    """
+    Find the neuron with the clearest frequency specialization (Tab 6).
+    Picks the neuron with the highest ratio of dominant frequency scale
+    to second-highest frequency scale.
+    """
+    scales, _, _ = decode_scales_phis_fn(model_load, fourier_basis)
+    # scales: [n_neurons, K+1], skip DC at index 0
+    scales_no_dc = scales[:, 1:]
+
+    if scales_no_dc.shape[1] < 2:
+        return 0
+
+    sorted_scales, _ = torch.sort(scales_no_dc, dim=1, descending=True)
+    ratio = sorted_scales[:, 0] / (sorted_scales[:, 1] + 1e-10)
+
+    return ratio.argmax().item()

precompute/prime_config.py

@@ -0,0 +1,135 @@
+"""
+Configuration for all moduli and training runs.
+Defines the d_mlp sizing formula and the 5 training run configurations.
+p can be any odd number >= 3 (not restricted to primes).
+"""
+import math
+
+
+def get_moduli(low=3, high=199):
+    """Return all odd numbers in [low, high]."""
+    moduli = []
+    for n in range(low, high + 1):
+        if n >= 3 and n % 2 == 1:
+            moduli.append(n)
+    return moduli
+
+
+# Keep old name as alias for backward compatibility
+get_primes = get_moduli
+
+
+def compute_d_mlp(p: int) -> int:
+    """
+    Compute d_mlp maintaining the ratio from p=23, d_mlp=512.
+    Formula: d_mlp = max(512, ceil(512/529 * p^2))
+    Can have more neurons but not less than the ratio dictates.
+    """
+    ratio = 512 / (23 ** 2)  # 512/529 ≈ 0.9679
+    return max(512, math.ceil(ratio * p * p))
+
+
+# Minimum p overall (p=2 has 0 non-DC frequencies, making Fourier analysis degenerate)
+MIN_P = 3
+
+# Minimum p for grokking experiments (need enough test data for meaningful split)
+MIN_P_GROKKING = 19
+
+# Backward-compatible aliases
+MIN_PRIME = MIN_P
+MIN_PRIME_GROKKING = MIN_P_GROKKING
+
+# 5 training run configurations per p
+TRAINING_RUNS = {
+    "standard": {
+        "embed_type": "one_hot",
+        "init_type": "random",
+        "optimizer": "AdamW",
+        "act_type": "ReLU",
+        "lr": 5e-5,
+        "weight_decay": 0,
+        "frac_train": 1.0,
+        "num_epochs": 5000,
+        "save_every": 200,
+        "init_scale": 0.1,
+        "save_models": True,
+        "batch_style": "full",
+        "seed": 42,
+    },
+    "grokking": {
+        "embed_type": "one_hot",
+        "init_type": "random",
+        "optimizer": "AdamW",
+        "act_type": "ReLU",
+        "lr": 1e-4,
+        "weight_decay": 2.0,
+        "frac_train": 0.75,
+        "num_epochs": 50000,
+        "save_every": 200,
+        "init_scale": 0.1,
+        "save_models": True,
+        "batch_style": "full",
+        "seed": 42,
+    },
+    "quad_random": {
+        "embed_type": "one_hot",
+        "init_type": "random",
+        "optimizer": "AdamW",
+        "act_type": "Quad",
+        "lr": 5e-5,
+        "weight_decay": 0,
+        "frac_train": 1.0,
+        "num_epochs": 5000,
+        "save_every": 200,
+        "init_scale": 0.1,
+        "save_models": True,
+        "batch_style": "full",
+        "seed": 42,
+    },
+    "quad_single_freq": {
+        "embed_type": "one_hot",
+        "init_type": "single-freq",
+        "optimizer": "SGD",
+        "act_type": "Quad",
+        "lr": 0.1,
+        "weight_decay": 0,
+        "frac_train": 1.0,
+        "num_epochs": 5000,
+        "save_every": 200,
+        "init_scale": 0.02,
+        "save_models": True,
+        "batch_style": "full",
+        "seed": 42,
+    },
+    "relu_single_freq": {
+        "embed_type": "one_hot",
+        "init_type": "single-freq",
+        "optimizer": "SGD",
+        "act_type": "ReLU",
+        "lr": 0.01,
+        "weight_decay": 0,
+        "frac_train": 1.0,
+        "num_epochs": 5000,
+        "save_every": 200,
+        "init_scale": 0.002,
+        "save_models": True,
+        "batch_style": "full",
+        "seed": 42,
+    },
+}
+
+# Analytical computation configs (no training needed)
+ANALYTICAL_CONFIGS = {
+    "decouple_dynamics": {
+        "init_k": 2,
+        "num_steps_case1": 1400,
+        "learning_rate_case1": 1,
+        "init_phi_case1": 1.5,
+        "init_psi_case1": 0.18,
+        "num_steps_case2": 700,
+        "learning_rate_case2": 1,
+        "init_phi_case2": -0.72,
+        "init_psi_case2": -2.91,
+        "amplitude": 0.02,
+    },
+}

precompute/run_all.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+# Pre-compute results for all odd p in [3, MAX_P].
+# Deletes checkpoints after each p to save disk space.
+#
+# Usage:
+#   bash precompute/run_all.sh          # p = 3, 5, 7, ..., 99
+#   MAX_P=199 bash precompute/run_all.sh  # p = 3, 5, 7, ..., 199
+#
+# Run from the project root directory.
+
+MAX_P=${MAX_P:-99}
+
+set -e
+echo "=== Pre-computing all odd p in [3, $MAX_P] ==="
+
+COMPLETED=0
+FAILED=0
+
+for P in $(seq 3 2 "$MAX_P"); do
+    echo ""
+    echo "========================================"
+    echo "  Processing p=$P"
+    echo "========================================"
+    if CLEANUP=1 bash precompute/run_pipeline.sh "$P"; then
+        COMPLETED=$((COMPLETED + 1))
+    else
+        echo "[FAIL] p=$P failed"
+        FAILED=$((FAILED + 1))
+    fi
+done
+
+echo ""
+echo "=== All done. Completed: $COMPLETED, Failed: $FAILED ==="
+echo "=== Precomputed results size: ==="
+du -sh precomputed_results/

precompute/run_pipeline.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+# Full pre-computation pipeline for a single modulus p (any odd number >= 3).
+#
+# Usage:
+#   bash precompute/run_pipeline.sh 23
+#   bash precompute/run_pipeline.sh 9 --d_mlp 128
+#   P=23 bash precompute/run_pipeline.sh
+#
+#   # Delete checkpoints after generating plots (saves disk space):
+#   CLEANUP=1 bash precompute/run_pipeline.sh 97
+#
+# Run from the project root directory.
+
+P=${1:-${P:-23}}
+shift 2>/dev/null || true  # consume the p arg
+
+# CLEANUP=1 to delete model checkpoints after plot generation
+CLEANUP=${CLEANUP:-0}
+
+# Collect remaining args (e.g. --d_mlp 128) to pass to train_all.py
+EXTRA_ARGS="$@"
+
+set -e
+echo "=== Running full pipeline for p=$P $EXTRA_ARGS ==="
+
+# Step 1: Train all 5 configurations
+echo ""
+echo "--- Step 1/4: Training ---"
+python precompute/train_all.py --p "$P" --output ./trained_models --resume $EXTRA_ARGS
+
+# Step 2: Generate model-based plots (d_mlp inferred from checkpoint)
+echo ""
+echo "--- Step 2/4: Generating model-based plots ---"
+python precompute/generate_plots.py --p "$P" --input ./trained_models --output ./precomputed_results
+
+# Step 3: Generate analytical simulation plots
+echo ""
+echo "--- Step 3/4: Generating analytical plots ---"
+python precompute/generate_analytical.py --p "$P" --output ./precomputed_results
+
+# Step 4: Cleanup checkpoints if requested
+PADDED=$(printf '%03d' "$P")
+MODEL_DIR="trained_models/p_${PADDED}"
+if [ "$CLEANUP" = "1" ] && [ -d "$MODEL_DIR" ]; then
+    echo ""
+    echo "--- Cleanup: Deleting checkpoints for p=$P ---"
+    SIZE=$(du -sh "$MODEL_DIR" | cut -f1)
+    rm -rf "$MODEL_DIR"
+    echo "    Freed $SIZE from $MODEL_DIR"
+fi
+
+# Step 5: Verify
+echo ""
+echo "--- Verification ---"
+RESULT_DIR="precomputed_results/p_${PADDED}"
+echo "=== Results in ${RESULT_DIR}/ ==="
+ls -la "${RESULT_DIR}/"
+FILE_COUNT=$(ls -1 "${RESULT_DIR}/" | wc -l | tr -d ' ')
+echo "=== Total files: ${FILE_COUNT} ==="
+echo "=== Pipeline complete for p=$P ==="

precompute/train_all.py

@@ -0,0 +1,290 @@
+#!/usr/bin/env python3
+"""
+Batch training script for all odd moduli p in [3, 199].
+
+Usage:
+    # Train all runs for all odd p
+    python train_all.py --all
+
+    # Train specific p
+    python train_all.py --p 23
+
+    # Train specific run type for a p
+    python train_all.py --p 23 --run standard
+
+    # Resume (skips completed runs)
+    python train_all.py --all --resume
+
+    # Custom output directory
+    python train_all.py --all --output ./my_models
+"""
+import argparse
+import json
+import os
+import sys
+
+# Add src to path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+
+import torch
+from prime_config import get_moduli, compute_d_mlp, TRAINING_RUNS, MIN_P, MIN_P_GROKKING
+from utils import Config
+from nnTrainer import Trainer
+
+
+def build_config_dict(p, run_params, d_mlp_override=None):
+    """Build a nested config dict compatible with the Config class."""
+    d_mlp = d_mlp_override if d_mlp_override is not None else compute_d_mlp(p)
+    return {
+        'data': {
+            'p': p,
+            'd_vocab': None,
+            'fn_name': 'add',
+            'frac_train': run_params['frac_train'],
+            'batch_style': run_params['batch_style'],
+        },
+        'model': {
+            'd_model': None,
+            'd_mlp': d_mlp,
+            'act_type': run_params['act_type'],
+            'embed_type': run_params['embed_type'],
+            'init_type': run_params['init_type'],
+            'init_scale': run_params['init_scale'],
+        },
+        'training': {
+            'num_epochs': run_params['num_epochs'],
+            'lr': run_params['lr'],
+            'weight_decay': run_params['weight_decay'],
+            'optimizer': run_params['optimizer'],
+            'stopping_thresh': -1,
+            'save_models': run_params['save_models'],
+            'save_every': run_params['save_every'],
+            'seed': run_params['seed'],
+        },
+    }
+
+
+def _save_training_log(output_dir, p, run_name, run_params, d_mlp, curves):
+    """Save a human-readable training_log.txt summarizing the run."""
+    log_path = os.path.join(output_dir, "training_log.txt")
+    n_epochs = len(curves.get('train_losses', []))
+    with open(log_path, 'w') as f:
+        f.write(f"{'=' * 70}\n")
+        f.write(f"Training Log: p={p}, run={run_name}\n")
+        f.write(f"{'=' * 70}\n\n")
+        f.write(f"Configuration:\n")
+        f.write(f"  prime (p)       = {p}\n")
+        f.write(f"  d_mlp           = {d_mlp}\n")
+        f.write(f"  activation      = {run_params['act_type']}\n")
+        f.write(f"  init_type       = {run_params['init_type']}\n")
+        f.write(f"  init_scale      = {run_params['init_scale']}\n")
+        f.write(f"  optimizer       = {run_params['optimizer']}\n")
+        f.write(f"  learning_rate   = {run_params['lr']}\n")
+        f.write(f"  weight_decay    = {run_params['weight_decay']}\n")
+        f.write(f"  frac_train      = {run_params['frac_train']}\n")
+        f.write(f"  num_epochs      = {run_params['num_epochs']}\n")
+        f.write(f"  batch_style     = {run_params['batch_style']}\n")
+        f.write(f"  seed            = {run_params['seed']}\n")
+        f.write(f"\n{'─' * 70}\n")
+        f.write(f"{'Epoch':>8s}  {'Train Loss':>12s}  {'Test Loss':>12s}  "
+                f"{'Train Acc':>10s}  {'Test Acc':>10s}  "
+                f"{'Grad Norm':>10s}  {'Param Norm':>11s}\n")
+        f.write(f"{'─' * 70}\n")
+
+        # Print every 100 epochs + the last epoch
+        train_losses = curves.get('train_losses', [])
+        test_losses = curves.get('test_losses', [])
+        train_accs = curves.get('train_accs', [])
+        test_accs = curves.get('test_accs', [])
+        grad_norms = curves.get('grad_norms', [])
+        param_norms = curves.get('param_norms', [])
+
+        step = max(1, n_epochs // 100)  # ~100 lines
+        indices = list(range(0, n_epochs, step))
+        if n_epochs > 0 and (n_epochs - 1) not in indices:
+            indices.append(n_epochs - 1)
+
+        for i in indices:
+            tl = f"{train_losses[i]:.6f}" if i < len(train_losses) else "N/A"
+            tel = f"{test_losses[i]:.6f}" if i < len(test_losses) else "N/A"
+            ta = f"{train_accs[i]:.4f}" if i < len(train_accs) else "N/A"
+            tea = f"{test_accs[i]:.4f}" if i < len(test_accs) else "N/A"
+            gn = f"{grad_norms[i]:.4f}" if i < len(grad_norms) else "N/A"
+            pn = f"{param_norms[i]:.4f}" if i < len(param_norms) else "N/A"
+            f.write(f"{i:>8d}  {tl:>12s}  {tel:>12s}  "
+                    f"{ta:>10s}  {tea:>10s}  "
+                    f"{gn:>10s}  {pn:>11s}\n")
+
+        f.write(f"{'─' * 70}\n\n")
+        f.write(f"Final Results:\n")
+        if train_losses:
+            f.write(f"  Train Loss  = {train_losses[-1]:.6f}\n")
+        if test_losses:
+            f.write(f"  Test Loss   = {test_losses[-1]:.6f}\n")
+        if train_accs:
+            f.write(f"  Train Acc   = {train_accs[-1]:.4f}\n")
+        if test_accs:
+            f.write(f"  Test Acc    = {test_accs[-1]:.4f}\n")
+        if param_norms:
+            f.write(f"  Param Norm  = {param_norms[-1]:.4f}\n")
+        f.write(f"\nTotal epochs trained: {n_epochs}\n")
+
+
+def run_training(p, run_name, output_base, d_mlp_override=None):
+    """Train a single run for a single prime."""
+    if p < MIN_P:
+        print(f"[SKIP] p={p}, run={run_name}: p < {MIN_P} (too few Fourier frequencies)")
+        return
+
+    # Single-freq init needs at least 1 non-DC frequency: (p-1)//2 >= 1 → p >= 3
+    if run_name in ('quad_single_freq', 'relu_single_freq') and (p - 1) // 2 < 1:
+        print(f"[SKIP] p={p}, run={run_name}: no non-DC frequencies for single-freq init")
+        return
+
+    if run_name == 'grokking' and p < MIN_P_GROKKING:
+        print(f"[SKIP] p={p}, run={run_name}: p < {MIN_P_GROKKING} (too few test points)")
+        return
+
+    run_params = TRAINING_RUNS[run_name]
+    config_dict = build_config_dict(p, run_params, d_mlp_override)
+    d_mlp = d_mlp_override if d_mlp_override is not None else compute_d_mlp(p)
+
+    output_dir = os.path.join(output_base, f"p_{p:03d}", run_name)
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Check if already completed
+    marker = os.path.join(output_dir, "DONE")
+    if os.path.exists(marker):
+        print(f"[SKIP] p={p}, run={run_name} already completed")
+        return
+
+    print(f"[TRAIN] p={p}, d_mlp={d_mlp}, run={run_name}, "
+          f"epochs={run_params['num_epochs']}")
+
+    config = Config(config_dict)
+    trainer = Trainer(config=config, use_wandb=False)
+
+    # Override save directory so checkpoints go into our output structure
+    trainer.save_dir = output_dir
+    run_subdir = os.path.join(output_dir, trainer.run_name)
+    os.makedirs(run_subdir, exist_ok=True)
+
+    # Re-save train/test data to the overridden location so generate_plots.py
+    # can find them (Trainer.__init__ saves to the original save_dir)
+    torch.save(trainer.train, os.path.join(run_subdir, 'train_data.pth'))
+    torch.save(trainer.test, os.path.join(run_subdir, 'test_data.pth'))
+
+    trainer.initial_save_if_appropriate()
+
+    # Plateau early-stopping for grokking: after 10K epochs, if curves
+    # haven't changed in the last 1000 epochs, stop training.
+    plateau_check = (run_name == 'grokking')
+    plateau_min_epoch = 10000
+    plateau_window = 1000
+    plateau_loss_tol = 1e-3   # absolute change in loss
+    plateau_acc_tol = 0.005   # absolute change in accuracy
+
+    for epoch in range(config.num_epochs):
+        train_loss, test_loss = trainer.do_a_training_step(epoch)
+
+        if test_loss.item() < config.stopping_thresh:
+            print(f"  Early stopping at epoch {epoch}: "
+                  f"test loss {test_loss.item():.6f}")
+            break
+
+        # Plateau detection for grokking
+        if (plateau_check and epoch >= plateau_min_epoch
+                and epoch % plateau_window == 0):
+            tl = trainer.train_losses
+            tel = trainer.test_losses
+            ta = trainer.train_accs
+            tea = trainer.test_accs
+            w = plateau_window
+            if len(tl) >= w and len(tel) >= w:
+                tl_flat = (max(tl[-w:]) - min(tl[-w:])) < plateau_loss_tol
+                tel_flat = (max(tel[-w:]) - min(tel[-w:])) < plateau_loss_tol
+                ta_flat = (not ta) or (max(ta[-w:]) - min(ta[-w:])) < plateau_acc_tol
+                tea_flat = (not tea) or (max(tea[-w:]) - min(tea[-w:])) < plateau_acc_tol
+                if tl_flat and tel_flat and ta_flat and tea_flat:
+                    print(f"  Plateau early stopping at epoch {epoch}: "
+                          f"no change in last {w} epochs")
+                    break
+
+        if config.is_it_time_to_save(epoch=epoch):
+            trainer.save_epoch(epoch=epoch, save_to_wandb=False, local_save=True)
+
+    trainer.post_training_save(
+        save_optimizer_and_scheduler=False, log_to_wandb=False
+    )
+
+    # Save training curves as JSON for plot generation
+    curves = {
+        'train_losses': trainer.train_losses,
+        'test_losses': trainer.test_losses,
+        'train_accs': trainer.train_accs,
+        'test_accs': trainer.test_accs,
+        'grad_norms': trainer.grad_norms,
+        'param_norms': trainer.param_norms,
+    }
+    curves_path = os.path.join(output_dir, "training_curves.json")
+    with open(curves_path, 'w') as f:
+        json.dump(curves, f)
+
+    # Save a human-readable training log
+    _save_training_log(output_dir, p, run_name, run_params, d_mlp, curves)
+
+    # Write completion marker
+    with open(marker, 'w') as f:
+        f.write(f"p={p} run={run_name} completed\n")
+
+    print(f"[DONE] p={p}, run={run_name}, "
+          f"train_acc={trainer.train_accs[-1]:.4f}, "
+          f"test_acc={trainer.test_accs[-1]:.4f}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Batch training for modular addition experiments'
+    )
+    parser.add_argument('--all', action='store_true',
+                        help='Train all odd p in [3, 199]')
+    parser.add_argument('--p', type=int,
+                        help='Train a specific odd modulus p')
+    parser.add_argument('--run', type=str, choices=list(TRAINING_RUNS.keys()),
+                        help='Train a specific run type')
+    parser.add_argument('--output', type=str, default='./trained_models',
+                        help='Output directory for trained models')
+    parser.add_argument('--d_mlp', type=int, default=None,
+                        help='Override d_mlp (number of hidden neurons). '
+                             'Default: auto-computed from p.')
+    parser.add_argument('--resume', action='store_true',
+                        help='Skip already-completed runs (checks DONE marker)')
+    args = parser.parse_args()
+
+    if not args.all and args.p is None:
+        parser.error("Specify --all or --p P")
+
+    moduli = [args.p] if args.p else get_moduli()
+    runs = [args.run] if args.run else list(TRAINING_RUNS.keys())
+
+    total = len(moduli) * len(runs)
+    completed = 0
+
+    for p in moduli:
+        for run_name in runs:
+            completed += 1
+            print(f"\n{'='*60}")
+            print(f"[{completed}/{total}] p={p}, run={run_name}")
+            print(f"{'='*60}")
+            try:
+                run_training(p, run_name, args.output, d_mlp_override=args.d_mlp)
+            except Exception as e:
+                print(f"[FAIL] p={p}, run={run_name}: {e}")
+                import traceback
+                traceback.print_exc()
+
+    print(f"\nAll done. {completed} runs processed.")
+
+
+if __name__ == "__main__":
+    main()

precomputed_results/p_015/p015_logits_interactive.json

@@ -0,0 +1 @@
+{"pairs": [[0, 0], [1, 0], [2, 0], [3, 0], [4, 0], [5, 0], [6, 0], [7, 0], [8, 0], [9, 0], [10, 0], [11, 0], [12, 0], [13, 0], [14, 0]], "correct_answers": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], "logits": [[5.845241069793701, -0.2678382396697998, -0.27407053112983704, 0.5236045718193054, 0.1727389544248581, -0.5161969065666199, 0.32797300815582275, 0.03217152878642082, 0.9118141531944275, 0.1262667030096054, -0.48561030626296997, -0.08425326645374298, 0.7506414651870728, 0.26417604088783264, 0.06685394793748856], [1.2039813995361328, 6.88271427154541, 1.5816727876663208, -1.3607698678970337, 0.2838805019855499, -0.1305035948753357, -0.1361088901758194, 0.1741233468055725, -0.20070961117744446, -0.19441378116607666, -0.07290667295455933, -0.036422792822122574, -0.045937348157167435, -0.4817635118961334, -1.3201030492782593], [1.021496057510376, 0.460868775844574, 6.665925979614258, -0.31677037477493286, 0.5722809433937073, -0.4090900421142578, -1.4602270126342773, -1.284806728363037, 0.1630801558494568, -0.14396820962429047, 0.21148079633712769, 0.009713075123727322, -0.13466662168502808, -1.4001739025115967, -0.48616650700569153], [2.0029213428497314, -1.0664080381393433, 0.1110871359705925, 6.9900641441345215, -0.6577256321907043, 0.14176632463932037, 0.5043318271636963, -0.9338439702987671, -0.5924315452575684, 0.0853288546204567, -0.8137574791908264, 0.2345130294561386, -0.747736930847168, -0.7098338603973389, -0.09835103899240494], [1.1424025297164917, -0.11194394528865814, -0.35566678643226624, 0.07304318249225616, 6.601017951965332, -0.7027722001075745, -0.2666034996509552, -0.44442567229270935, 0.8163065314292908, -0.5371125340461731, 0.3994847536087036, -2.494434118270874, -0.7756778001785278, -0.05254651606082916, -0.19730210304260254], [-0.2894173860549927, -0.4359486699104309, -0.5648083686828613, 0.2765766680240631, -0.8157410025596619, 6.83429479598999, -0.35287508368492126, -0.13654330372810364, 0.4015010595321655, -0.5329406261444092, -0.8750499486923218, 0.27872195839881897, 0.017333246767520905, 0.40312278270721436, -0.36508798599243164], [0.8411731719970703, -0.7100308537483215, -0.03279941901564598, -1.725080132484436, 0.041413623839616776, -0.020579706877470016, 7.065398216247559, -0.6480932235717773, 0.3441147208213806, -0.6311541795730591, -0.5848420858383179, 0.23020878434181213, 1.6761562824249268, -0.4346846342086792, 0.1869385540485382], [1.0707398653030396, 0.15244746208190918, -0.8544912338256836, -0.15931977331638336, -0.30050942301750183, -0.14030054211616516, -0.8752976059913635, 6.867440223693848, -0.9738104343414307, -0.016967639327049255, 0.06942860782146454, -0.36363551020622253, -0.5302596092224121, -0.04578635096549988, 1.362797498703003], [1.3280404806137085, 0.42291590571403503, 0.2617959976196289, -0.14120015501976013, 0.16784554719924927, -0.3041268587112427, -0.9258386492729187, -1.557655930519104, 6.762355327606201, -1.1574827432632446, 0.6234251260757446, 0.20528843998908997, -0.12094193696975708, -0.20512300729751587, 0.5443016290664673], [1.1472804546356201, -0.6102383732795715, -0.025639446452260017, 1.0203014612197876, -0.7307616472244263, -0.12430916726589203, -1.7754124402999878, -0.18994221091270447, -0.2811968922615051, 7.137627601623535, -0.20736677944660187, -0.2558075189590454, -0.8490561842918396, -0.11468800902366638, -0.2236211597919464], [-0.4398433268070221, -0.21957175433635712, 0.3112000524997711, -0.31722167134284973, -0.045922309160232544, -0.8279529213905334, 0.09667133539915085, -0.1345088630914688, -0.4463891386985779, -0.33588770031929016, 6.491363525390625, -0.12330206483602524, 0.10725130885839462, 0.42743009328842163, -0.8439421057701111], [0.8614866733551025, -0.09104573726654053, -0.26044338941574097, -1.4104074239730835, -1.4556188583374023, -0.1587509959936142, -0.8687411546707153, 0.6269921064376831, -0.11933901160955429, -0.3620198369026184, 0.2740001678466797, 6.834388256072998, 0.9378089308738708, -0.8011277914047241, 0.10253019630908966], [0.8134461045265198, -0.5108746886253357, -0.11389578878879547, -0.7348212003707886, -0.3866046965122223, -0.30147287249565125, -1.0970127582550049, -0.5308148264884949, -0.25204434990882874, 0.6024379134178162, -0.5250549912452698, 0.2944880723953247, 6.6044793128967285, -0.06822022795677185, 0.009033107198774815], [1.3036129474639893, 0.3358013927936554, -1.516667366027832, -0.7404187321662903, 0.37273597717285156, 0.4566952586174011, -0.5649716854095459, -0.16146983206272125, -0.5403905510902405, -0.6627720594406128, 0.10609856992959976, 1.1525702476501465, 0.3066888749599457, 6.983869552612305, 0.6811028122901917], [1.5695332288742065, -1.1690753698349, 0.29273056983947754, -0.6167480945587158, -0.481355756521225, 0.3711283802986145, -0.3477175235748291, 0.06404877454042435, -0.34069737792015076, -0.5002283453941345, -0.22661544382572174, 0.29657235741615295, -1.5196651220321655, 1.143660306930542, 6.862580299377441]], "output_classes": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]}

precomputed_results/p_015/p015_metadata.json

@@ -0,0 +1,82 @@
+{
+  "prime": 15,
+  "d_mlp": 512,
+  "training_runs": {
+    "standard": {
+      "act_type": "ReLU",
+      "lr": 5e-05,
+      "weight_decay": 0,
+      "num_epochs": 5000,
+      "frac_train": 1.0,
+      "init_type": "random",
+      "init_scale": 0.1,
+      "optimizer": "AdamW"
+    },
+    "grokking": {
+      "act_type": "ReLU",
+      "lr": 0.0001,
+      "weight_decay": 2.0,
+      "num_epochs": 50000,
+      "frac_train": 0.75,
+      "init_type": "random",
+      "init_scale": 0.1,
+      "optimizer": "AdamW"
+    },
+    "quad_random": {
+      "act_type": "Quad",
+      "lr": 5e-05,
+      "weight_decay": 0,
+      "num_epochs": 5000,
+      "frac_train": 1.0,
+      "init_type": "random",
+      "init_scale": 0.1,
+      "optimizer": "AdamW"
+    },
+    "quad_single_freq": {
+      "act_type": "Quad",
+      "lr": 0.1,
+      "weight_decay": 0,
+      "num_epochs": 5000,
+      "frac_train": 1.0,
+      "init_type": "single-freq",
+      "init_scale": 0.02,
+      "optimizer": "SGD"
+    },
+    "relu_single_freq": {
+      "act_type": "ReLU",
+      "lr": 0.01,
+      "weight_decay": 0,
+      "num_epochs": 5000,
+      "frac_train": 1.0,
+      "init_type": "single-freq",
+      "init_scale": 0.002,
+      "optimizer": "SGD"
+    }
+  },
+  "final_metrics": {
+    "standard": {
+      "train_acc": 1.0,
+      "test_acc": 1.0,
+      "train_loss": 0.020928841084241867,
+      "test_loss": 0.020928841084241867
+    },
+    "quad_random": {
+      "train_acc": 1.0,
+      "test_acc": 1.0,
+      "train_loss": 0.0036203155759721994,
+      "test_loss": 0.0036203155759721994
+    },
+    "quad_single_freq": {
+      "train_acc": 1.0,
+      "test_acc": 1.0,
+      "train_loss": 0.04876862093806267,
+      "test_loss": 0.04876862093806267
+    },
+    "relu_single_freq": {
+      "train_acc": 1.0,
+      "test_acc": 1.0,
+      "train_loss": 2.7064406871795654,
+      "test_loss": 2.7064406871795654
+    }
+  }
+}

precomputed_results/p_015/p015_neuron_spectra.json

@@ -0,0 +1 @@
+{"fourier_basis_names": ["Const", "cos 1", "sin 1", "cos 2", "sin 2", "cos 3", "sin 3", "cos 4", "sin 4", "cos 5", "sin 5", "cos 6", "sin 6", "cos 7", "sin 7"], "neurons": {"neuron_0": {"global_index": 298, "dominant_freq": 1, "fourier_magnitudes_in": [0.009785229340195656, 0.835197389125824, 0.06538262963294983, 0.023014001548290253, 0.11825679987668991, 0.13374952971935272, 0.06052215397357941, 0.05648910254240036, 0.19347421824932098, 0.022295288741588593, 0.025461282581090927, 0.10164802521467209, 0.13591913878917694, 0.05420650914311409, 0.015495105646550655], "fourier_magnitudes_out": [0.05270613357424736, 0.8947794437408447, 0.023624544963240623, 0.0090867318212986, 0.013868873007595539, 0.23038825392723083, 0.01721891760826111, 0.03962008282542229, 0.026063552126288414, 0.13086184859275818, 0.030048014596104622, 0.0664755254983902, 0.05337228253483772, 0.13469304144382477, 0.026413951069116592]}, "neuron_1": {"global_index": 322, "dominant_freq": 2, "fourier_magnitudes_in": [0.0048967949114739895, 0.05831458047032356, 0.047817476093769073, 0.8282294273376465, 0.08354193717241287, 0.07255098968744278, 0.11165464669466019, 0.025725897401571274, 0.1043122410774231, 0.06207004189491272, 0.03182118758559227, 0.16142725944519043, 0.05974670872092247, 0.07021868228912354, 0.14918409287929535], "fourier_magnitudes_out": [0.0682767927646637, 0.04459874704480171, 0.048241592943668365, 0.8725450038909912, 0.12229588627815247, 0.00944548286497593, 0.04323046654462814, 0.04066399484872818, 0.016419490799307823, 0.08050032705068588, 0.10128692537546158, 0.2789347767829895, 0.11171453446149826, 0.0022370540536940098, 0.05760588496923447]}, "neuron_2": {"global_index": 506, "dominant_freq": 3, "fourier_magnitudes_in": [0.16119761765003204, 0.007301056291908026, 0.015491282567381859, 0.008759702555835247, 0.012521358206868172, 0.8759171366691589, 0.027130309492349625, 0.011930187232792377, 0.008996764197945595, 0.017064398154616356, 0.0032833898440003395, 0.2937963902950287, 0.03241246938705444, 0.01809331774711609, 0.01297019049525261], "fourier_magnitudes_out": [0.1367007941007614, 0.08054599165916443, 0.009201026521623135, 0.04534965381026268, 0.08026888966560364, 0.8633893728256226, 0.050462786108255386, 0.0044480981305241585, 0.03184027969837189, 0.05210161209106445, 0.04297168180346489, 0.32101351022720337, 0.06568200141191483, 0.035220254212617874, 0.014398200437426567]}, "neuron_3": {"global_index": 163, "dominant_freq": 4, "fourier_magnitudes_in": [0.031200144439935684, 0.04270334914326668, 0.08881780505180359, 0.030023256316781044, 0.08304275572299957, 0.04510613903403282, 0.1826746016740799, 0.03483400121331215, 0.8419885039329529, 0.013886661268770695, 0.08348363637924194, 0.050276018679142, 0.10324658453464508, 0.007152870763093233, 0.037329111248254776], "fourier_magnitudes_out": [0.0033031131606549025, 0.04313803091645241, 0.04859255626797676, 0.03949907422065735, 0.0202469564974308, 0.20968425273895264, 0.08351283520460129, 0.8787142038345337, 0.06666962057352066, 0.12211224436759949, 0.0748063325881958, 0.03903647139668465, 0.004882670473307371, 0.03963864594697952, 0.00997573509812355]}, "neuron_4": {"global_index": 56, "dominant_freq": 5, "fourier_magnitudes_in": [0.3008343279361725, 0.00417996384203434, 0.0018625137163326144, 0.0030638582538813353, 0.0034040315076708794, 0.0014134111115708947, 0.004357232246547937, 0.0004787891812156886, 0.004558315966278315, 0.8440173864364624, 0.003970965277403593, 0.0037021928001195192, 0.00269315205514431, 0.004479836206883192, 0.0009506550850346684], "fourier_magnitudes_out": [0.34606269001960754, 0.027366885915398598, 0.036413464695215225, 0.017639687284827232, 0.027893830090761185, 0.009557241573929787, 0.03348158299922943, 0.02666180580854416, 0.006313585676252842, 0.9367755651473999, 0.021462714299559593, 0.004096114542335272, 0.023737004026770592, 0.012917638756334782, 0.005492078140377998]}, "neuron_5": {"global_index": 382, "dominant_freq": 6, "fourier_magnitudes_in": [0.17134937644004822, 0.030019262805581093, 0.003407673677429557, 0.009750651195645332, 0.017143065109848976, 0.29516395926475525, 0.04455732926726341, 0.0013349098153412342, 0.022812439128756523, 0.007311187218874693, 0.0026531207840889692, 0.8911311626434326, 0.017707478255033493, 0.012044227682054043, 0.007314047310501337], "fourier_magnitudes_out": [0.15681979060173035, 0.02663942240178585, 0.015795163810253143, 0.03183992952108383, 0.019026663154363632, 0.34929507970809937, 0.021825529634952545, 0.0012088950024917722, 0.013926293700933456, 0.004897939506918192, 0.010765696875751019, 0.9581749439239502, 0.012587963603436947, 0.02056090161204338, 0.01620820164680481]}, "neuron_6": {"global_index": 44, "dominant_freq": 7, "fourier_magnitudes_in": [0.016063887625932693, 0.005506421905010939, 0.06481847167015076, 0.050152119249105453, 0.08404979109764099, 0.04145295172929764, 0.0641774982213974, 0.06344524770975113, 0.07214409857988358, 0.09170015156269073, 0.048928599804639816, 0.22082282602787018, 0.04532390087842941, 0.8467198014259338, 0.08364292979240417], "fourier_magnitudes_out": [0.015495719388127327, 0.0706552267074585, 0.03188098222017288, 0.022803593426942825, 0.05851001664996147, 0.006458980031311512, 0.07394170761108398, 0.06014903634786606, 0.08831571042537689, 0.10563234984874725, 0.11405274271965027, 0.26272422075271606, 0.13614341616630554, 0.9203148484230042, 0.14940021932125092]}, "neuron_7": {"global_index": 157, "dominant_freq": 1, "fourier_magnitudes_in": [0.023283498361706734, 0.14338412880897522, 0.8202316761016846, 0.05399405211210251, 0.058775369077920914, 0.12734520435333252, 0.1968384087085724, 0.070158950984478, 0.07580292969942093, 0.14696937799453735, 0.0658537819981575, 0.08757402002811432, 0.050403349101543427, 0.07286650687456131, 0.030624864622950554], "fourier_magnitudes_out": [0.10438423603773117, 0.8233321905136108, 0.2895442247390747, 0.11436242610216141, 0.043502938002347946, 0.08488017320632935, 0.21032141149044037, 0.03133478760719299, 0.147483691573143, 0.09128359705209732, 0.11333037912845612, 0.042169392108917236, 0.14409606158733368, 0.06334586441516876, 0.010641958564519882]}, "neuron_8": {"global_index": 436, "dominant_freq": 2, "fourier_magnitudes_in": [0.02816132642328739, 0.0366082526743412, 0.09432870894670486, 0.82112056016922, 0.2999240458011627, 0.07823537290096283, 0.07104597240686417, 0.010237633250653744, 0.05316127836704254, 0.01780366338789463, 0.10813448578119278, 0.10922082513570786, 0.2154001146554947, 0.06505948305130005, 0.005772633943706751], "fourier_magnitudes_out": [0.07935837656259537, 0.05681774765253067, 0.11203812062740326, 0.7207208275794983, 0.5456924438476562, 0.08229123800992966, 0.0757966861128807, 0.025318337604403496, 0.12979470193386078, 0.16904646158218384, 0.01633336953818798, 0.10492201149463654, 0.25266337394714355, 0.07896621525287628, 0.056506332010030746]}, "neuron_9": {"global_index": 339, "dominant_freq": 3, "fourier_magnitudes_in": [0.00793448369950056, 0.23576614260673523, 0.05173733830451965, 0.0808015987277031, 0.08904809504747391, 0.8467172980308533, 0.07009700685739517, 0.2346189022064209, 0.06041780859231949, 0.22861061990261078, 0.005940192844718695, 0.08250848948955536, 0.05041716247797012, 0.09760993719100952, 0.08870971947908401], "fourier_magnitudes_out": [0.07995986938476562, 0.058010708540678024, 0.0238045621663332, 0.044154077768325806, 0.06384548544883728, 0.8888619542121887, 0.1083059012889862, 0.05018500238656998, 0.06562003493309021, 0.032845836132764816, 0.06583794951438904, 0.19849520921707153, 0.08598171174526215, 0.048292793333530426, 0.056917209178209305]}, "neuron_10": {"global_index": 405, "dominant_freq": 4, "fourier_magnitudes_in": [0.04605866223573685, 0.10036957263946533, 0.1019270122051239, 0.08908756822347641, 0.05084078758955002, 0.09276082366704941, 0.1580628901720047, 0.11540821939706802, 0.8267641663551331, 0.07836127281188965, 0.007629718631505966, 0.10522190481424332, 0.08720945566892624, 0.07787368446588516, 0.07161819189786911], "fourier_magnitudes_out": [0.0646987333893776, 0.06891193985939026, 0.1397920846939087, 0.05930565670132637, 0.08342889696359634, 0.09649658203125, 0.19548238813877106, 0.8096501231193542, 0.276018351316452, 0.04780340939760208, 0.11156941205263138, 0.02261027880012989, 0.12174452841281891, 0.13799422979354858, 0.09665590524673462]}, "neuron_11": {"global_index": 363, "dominant_freq": 5, "fourier_magnitudes_in": [0.26700034737586975, 0.027114098891615868, 0.030112946406006813, 0.004235479515045881, 0.04029411822557449, 0.03277970105409622, 0.02381693758070469, 0.03963545709848404, 0.008424329571425915, 0.8159627914428711, 0.03508993238210678, 0.012520099990069866, 0.03853161633014679, 0.037013012915849686, 0.016481192782521248], "fourier_magnitudes_out": [0.3175621032714844, 0.022698314860463142, 0.029232148081064224, 0.009378736838698387, 0.04299170523881912, 0.021195683628320694, 0.006493568420410156, 0.02105909213423729, 0.005779783241450787, 0.9192318916320801, 0.025057973340153694, 0.015149646438658237, 0.004153982736170292, 0.01740305870771408, 0.028322339057922363]}, "neuron_12": {"global_index": 195, "dominant_freq": 6, "fourier_magnitudes_in": [0.09348957985639572, 0.09145021438598633, 0.019759872928261757, 0.14044030010700226, 0.010322250425815582, 0.09715981036424637, 0.15762653946876526, 0.08852937072515488, 0.012697099708020687, 0.10808465629816055, 0.005568717140704393, 0.11048246920108795, 0.8844196796417236, 0.06171249970793724, 0.021244505420327187], "fourier_magnitudes_out": [0.16980455815792084, 0.02732899598777294, 0.036924730986356735, 0.05185381695628166, 0.045569196343421936, 0.26525792479515076, 0.05154341459274292, 0.0016775119584053755, 0.005609582178294659, 0.06192351505160332, 0.05204838141798973, 0.878484845161438, 0.06501992046833038, 0.029970666393637657, 0.007208859547972679]}, "neuron_13": {"global_index": 400, "dominant_freq": 7, "fourier_magnitudes_in": [0.03170298784971237, 0.01636633276939392, 0.02570384368300438, 0.059365708380937576, 0.038365013897418976, 0.04416408762335777, 0.03580497205257416, 0.08423422276973724, 0.05347143113613129, 0.14569905400276184, 0.11170309782028198, 0.19443102180957794, 0.1770995706319809, 0.16263896226882935, 0.8427256941795349], "fourier_magnitudes_out": [0.001597732538357377, 0.03749779239296913, 0.021893106400966644, 0.021266808733344078, 0.003741121618077159, 0.06272387504577637, 0.031196942552924156, 0.06795763224363327, 0.09264159947633743, 0.009625964798033237, 0.16665540635585785, 0.20691494643688202, 0.2753947973251343, 0.8446058630943298, 0.2736137807369232]}, "neuron_14": {"global_index": 204, "dominant_freq": 1, "fourier_magnitudes_in": [0.009332116693258286, 0.18445806205272675, 0.8107064962387085, 0.0019745242316275835, 0.021862220019102097, 0.164852112531662, 0.19152098894119263, 0.0416061170399189, 0.006671736016869545, 0.12151989340782166, 0.0737357810139656, 0.07512296736240387, 0.00437825545668602, 0.09026234596967697, 0.010851189494132996], "fourier_magnitudes_out": [0.00363162811845541, 0.7924743294715881, 0.3171338140964508, 0.004241126589477062, 0.01878344640135765, 0.12165644019842148, 0.20536403357982635, 0.018597787246108055, 0.0810178741812706, 0.016301296651363373, 0.17405299842357635, 0.0980185940861702, 0.04579637944698334, 0.14084625244140625, 0.08989735692739487]}, "neuron_15": {"global_index": 200, "dominant_freq": 2, "fourier_magnitudes_in": [0.18793319165706635, 0.07831189781427383, 0.0546899177134037, 0.26695817708969116, 0.8170905709266663, 0.05296998843550682, 0.1943284422159195, 0.23756185173988342, 0.15034066140651703, 0.017131371423602104, 0.0028046099469065666, 0.11494778096675873, 0.10373453050851822, 0.12329491227865219, 0.18243563175201416], "fourier_magnitudes_out": [0.09242420643568039, 0.0712248757481575, 0.12614493072032928, 0.7048966288566589, 0.5352984666824341, 0.1033194437623024, 0.05411561205983162, 0.020737502723932266, 0.1372404545545578, 0.13828690350055695, 0.03525715693831444, 0.11022429913282394, 0.2024645209312439, 0.14993123710155487, 0.08065018057823181]}, "neuron_16": {"global_index": 29, "dominant_freq": 3, "fourier_magnitudes_in": [0.06590171158313751, 0.1875513195991516, 0.13940751552581787, 0.05203652381896973, 0.12810246646404266, 0.8321968913078308, 0.14579933881759644, 0.1637534499168396, 0.11946561932563782, 0.1717541664838791, 0.013540252111852169, 0.15768876671791077, 0.04509717971086502, 0.05775374174118042, 0.17158927023410797], "fourier_magnitudes_out": [0.004155661445111036, 0.03632631525397301, 0.11120419949293137, 0.07094360142946243, 0.0025094610173255205, 0.8083617091178894, 0.2703225910663605, 0.04779119789600372, 0.01505393534898758, 0.019837403669953346, 0.020111212506890297, 0.13071304559707642, 0.1915966272354126, 0.012666973285377026, 0.03241927549242973]}, "neuron_17": {"global_index": 263, "dominant_freq": 4, "fourier_magnitudes_in": [0.04006792977452278, 0.04470885172486305, 0.03644431754946709, 0.09875550121068954, 0.0844721719622612, 0.2842579483985901, 0.07786386460065842, 0.8181069493293762, 0.07545538246631622, 0.1041831225156784, 0.08087007701396942, 0.013009922578930855, 0.05374658852815628, 0.010971452109515667, 0.05345306172966957], "fourier_magnitudes_out": [0.0032453967723995447, 0.04532613605260849, 0.050931766629219055, 0.06243853271007538, 0.1471986621618271, 0.07246130704879761, 0.0618220679461956, 0.793653130531311, 0.21675735712051392, 0.10524935275316238, 0.002365024061873555, 0.05830618739128113, 0.12600253522396088, 0.015407783910632133, 0.04248036816716194]}, "neuron_18": {"global_index": 468, "dominant_freq": 5, "fourier_magnitudes_in": [0.27209052443504333, 0.01535376999527216, 0.024236787110567093, 0.014399196021258831, 0.026386309415102005, 0.03005525842308998, 0.0007289683562703431, 0.013318601995706558, 0.025410467758774757, 0.8113285303115845, 0.024080565199255943, 0.028674796223640442, 0.0011739643523469567, 0.015659669414162636, 0.025659453123807907], "fourier_magnitudes_out": [0.2273026406764984, 0.010362203232944012, 0.08463546633720398, 0.048502951860427856, 0.006748202722519636, 0.008361948654055595, 0.017721673473715782, 0.055468566715717316, 0.07133140414953232, 0.8983375430107117, 0.06657372415065765, 0.07770014554262161, 0.05767229199409485, 0.02166377194225788, 0.007756791543215513]}, "neuron_19": {"global_index": 502, "dominant_freq": 6, "fourier_magnitudes_in": [0.161783829331398, 0.011249735951423645, 0.011750025674700737, 0.006551133934408426, 0.005556976888328791, 0.3025686740875244, 0.004988331813365221, 0.01731167919933796, 0.009423403069376945, 0.001427029725164175, 0.0038674750830978155, 0.8558500409126282, 0.014558211900293827, 0.0111812399700284, 0.0006309517193585634], "fourier_magnitudes_out": [0.17063283920288086, 0.016314541921019554, 0.05021467059850693, 0.016806060448288918, 0.04298456758260727, 0.2806430757045746, 0.014211905188858509, 0.004449731670320034, 0.02185676619410515, 0.006607384420931339, 0.005355009343475103, 0.875554621219635, 0.016872091218829155, 0.022274665534496307, 0.03286416828632355]}}}

precomputed_results/p_015/p015_overview.json

@@ -0,0 +1 @@
+{"std_epochs": [0, 200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800, 2000, 2200, 2400, 2600, 2800, 3000, 3200, 3400, 3600, 3800, 4000, 4200, 4400, 4600, 4800], "std_ipr": [0.2509222626686096, 0.26378345489501953, 0.3110557198524475, 0.36821821331977844, 0.4191873371601105, 0.46056264638900757, 0.4932430386543274, 0.518873393535614, 0.5395100712776184, 0.5563467741012573, 0.5705686807632446, 0.5829041600227356, 0.5938962697982788, 0.6040970087051392, 0.6136700510978699, 0.6225748658180237, 0.6310203671455383, 0.6389821171760559, 0.6463097333908081, 0.6528617739677429, 0.6585475206375122, 0.6634621620178223, 0.6678526401519775, 0.6718972325325012, 0.6756421327590942], "std_train_loss": [2.7085084915161133, 2.677509307861328, 2.6341302394866943, 2.5730724334716797, 2.492614269256592, 2.3927817344665527, 2.2741761207580566, 2.137855052947998, 1.9853696823120117, 1.8187763690948486, 1.6410351991653442, 1.4555929899215698, 1.2665268182754517, 1.0787031650543213, 0.8973091840744019, 0.7275450825691223, 0.5740776062011719, 0.44004037976264954, 0.32742777466773987, 0.23663438856601715, 0.1664789468050003, 0.1142977625131607, 0.0767836645245552, 0.05059399828314781, 0.032766345888376236]}

precomputed_results/p_023/p023_grokk_acc.json

@@ -0,0 +1 @@
+{"epochs": [0, 200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800, 2000, 2200, 2400, 2600, 2800, 3000, 3200, 3400, 3600, 3800, 4000, 4200, 4400, 4600, 4800, 5000, 5200, 5400, 5600, 5800, 6000, 6200, 6400, 6600, 6800, 7000, 7200, 7400, 7600, 7800, 8000, 8200, 8400, 8600, 8800, 9000, 9200, 9400, 9600, 9800, 10000, 10200, 10400, 10600, 10800, 11000, 11200, 11400, 11600, 11800, 12000, 12200, 12400, 12600, 12800, 13000, 13200, 13400, 13600, 13800, 14000, 14200, 14400, 14600, 14800, 15000, 15200, 15400, 15600, 15800, 16000, 16200, 16400, 16600, 16800, 17000, 17200, 17400, 17600, 17800, 18000, 18200, 18400, 18600, 18800, 19000, 19200, 19400, 19600, 19800, 20000, 20200, 20400, 20600, 20800, 21000, 21200, 21400, 21600, 21800, 22000, 22200, 22400, 22600, 22800, 23000, 23200, 23400, 23600, 23800], "train_accs": [0.045454545454545456, 0.23737373737373738, 0.5404040404040404, 0.7272727272727273, 0.7348484848484849, 0.7424242424242424, 0.76010101010101, 0.8055555555555556, 0.8838383838383839, 0.9671717171717171, 0.9924242424242424, 0.9974747474747475, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], "test_accs": [0.03759398496240601, 0.007518796992481203, 0.007518796992481203, 0.007518796992481203, 0.015037593984962405, 0.03007518796992481, 0.06766917293233082, 0.18045112781954886, 0.39849624060150374, 0.6466165413533834, 0.7142857142857143, 0.7293233082706767, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7368421052631579, 0.7518796992481203, 0.7669172932330827, 0.7819548872180451, 0.7819548872180451, 0.7819548872180451, 0.7819548872180451, 0.7819548872180451, 0.7819548872180451, 0.7819548872180451, 0.7969924812030075, 0.8270676691729323, 0.8270676691729323, 0.8421052631578947, 0.9022556390977443, 0.9022556390977443, 0.9022556390977443, 0.9022556390977443, 0.9172932330827067, 0.9323308270676691, 0.9323308270676691, 0.9323308270676691, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939, 0.9624060150375939], "stage1_end": 1744, "stage2_end": 9431}

precomputed_results/p_023/p023_grokk_epoch_data.json

@@ -0,0 +1 @@
+{"prime": 23, "epochs": [0, 2600, 5200, 7800, 10400, 13200, 15800, 18400, 21000, 23800], "grids": [[[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], [[1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], [[1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], [[1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]]]}