Notebook 04 — Memory capacity vs chain parameters + Optuna hyperparameter sweeps

[rosspeili]

Context

Issue #6 adds memory.py with Dambre et al. (2012) memory capacity (MC) for QRC and ESN feature collectors. Notebook 04 is the first systematic parameter study in this repo: how MC scales with chain length N, evolution time tau, and disorder W.

OVERVIEW.md Workstream B3 lists hyperparameter sweeps (N, tau, W, Jz/Jxy ratio) via Optuna. The [tuning] extra in pyproject.toml already declares optuna>=3.0 — this notebook is its first consumer.

docs/notebooks.md specifies output: figures/memory_capacity.png

Goal

Create notebooks/04_memory_capacity.ipynb that:

1. Computes total memory capacity MC for QRC across a parameter grid (or Optuna study)
2. Optionally compares quantum MC vs ESN MC on identical protocol (Issue Classical Echo State Network baseline module + unit tests #5 + Memory capacity metric (Dambre et al. 2012) in package #6)
3. Visualizes MC vs N, tau, W (and optionally J_z/J_xy)
4. Saves figures/memory_capacity.png and data/memory_capacity.csv
5. Uses Optuna for at least one targeted sweep (e.g. maximize test_nmse on NARMA-10 or maximize MC)

Prerequisites

Install notebook extras from Spintronics/ root:

pip install -e "./spintronic-qrc[notebooks]"

Requires Issues #1, #5, #6 merged (QRC pipeline, ESN baseline, memory capacity module).

Notebook structure

Cell 1 — Title and context (markdown)

• Explain Dambre memory capacity: sum of per-delay squared correlations achievable by linear readout
• State impulse protocol (from memory.py) — do not rederive math in notebook
• Cite: Dambre et al. (2012), Inubushi & Yoshimura (2017), Fujii & Nakajima (2017)

Cell 2 — Setup

from __future__ import annotations

import os
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import optuna

from spintronic_qrc.pipeline import QRCConfig
from spintronic_qrc.esn import ESNConfig
from spintronic_qrc.memory import (
    MemoryCapacityConfig,
    qrc_memory_capacity,
    esn_memory_capacity,
)
from spintronic_qrc import utils

os.makedirs("../figures", exist_ok=True)
os.makedirs("../data", exist_ok=True)

optuna.logging.set_verbosity(optuna.logging.WARNING)

Cell 3 — Baseline MC comparison (QRC vs ESN)

Fixed small config for a fast headline number:

Parameter	QRC	ESN
Reservoir size	n_sites=4	n_reservoir=100
k_max	10	10
n_samples	500	500
washout	50	50
seed	42	42

mc_cfg = MemoryCapacityConfig(k_max=10, n_samples=500, washout=50, seed=42)
qrc_mc = qrc_memory_capacity(QRCConfig(n_sites=4, seed=42), mc_cfg)
esn_mc = esn_memory_capacity(ESNConfig(n_reservoir=100, seed=42), mc_cfg)
print(f"QRC total MC: {qrc_mc.total_mc:.2f}")
print(f"ESN total MC: {esn_mc.total_mc:.2f}")

Bar chart or table: total_mc QRC vs ESN.

Cell 4 — Parameter grid sweep (QRC memory capacity)

Sweep one parameter at a time (keep others fixed) to keep runtime manageable:

Sweep A — chain length N: n_sites in [3, 4, 5, 6] (or [4, 5, 6, 7] if fast enough)

Sweep B — evolution time tau: evolution_time in [0.2, 0.5, 0.8, 1.0, 1.5]

Sweep C — disorder W: disorder in [0.0, 0.25, 0.5, 0.75, 1.0]

For each point, call qrc_memory_capacity() with reduced k_max/n_samples if needed for speed. Store results in a list of dicts -> DataFrame.

Cell 5 — Optuna study (at least one)

Example objective — maximize QRC memory capacity:

def objective(trial: optuna.Trial) -> float:
    n_sites = trial.suggest_int("n_sites", 3, 6)
    evolution_time = trial.suggest_float("evolution_time", 0.2, 1.5)
    disorder = trial.suggest_float("disorder", 0.0, 1.0)
    j_z = trial.suggest_float("J_z", 0.5, 1.5)

    cfg = QRCConfig(
        n_sites=n_sites,
        evolution_time=evolution_time,
        disorder=disorder,
        J_z=j_z,
        seed=42,
    )
    mc = qrc_memory_capacity(cfg, MemoryCapacityConfig(k_max=8, n_samples=300, washout=50))
    return mc.total_mc

Run n_trials=20–30 (document count). Plot optimization history or parameter importances if time permits.

Alternative objective (optional second study): minimize test_nmse on short NARMA-10 via run_qrc() — ties MC sweep to prediction task.

Cell 6 — Visualization

Multi-panel figure for figures/memory_capacity.png:

1. MC vs n_sites
2. MC vs evolution_time
3. MC vs disorder W
4. Optional: QRC vs ESN total_mc bar chart

Use utils palette. 150 DPI, bbox_inches="tight".

Cell 7 — Export CSV

data/memory_capacity.csv with columns:

run_type, n_sites, evolution_time, disorder, J_xy, J_z, k_max, total_mc, esn_total_mc, seed, notes

Include grid sweep rows and best Optuna trial row.

Cell 8 — Sanity checks

• All total_mc values finite and >= 0
• QRC MC varies across at least one swept parameter (not flat line everywhere)
• Figure and CSV exist

Runtime guidance

Full grid + Optuna can be slow. Acceptable strategies (document which you used):

• Reduce k_max to 8 and n_samples to 300 during development
• Run grid sweeps sequentially, not full factorial
• Cap Optuna at 20 trials for committed notebook outputs

Target: committed notebook outputs reproducible in under ~30 min CPU at reduced settings.

Acceptance criteria

• notebooks/04_memory_capacity.ipynb runs top-to-bottom
• figures/memory_capacity.png committed
• data/memory_capacity.csv committed
• Uses memory.qrc_memory_capacity and esn_memory_capacity (no duplicated MC math)
• At least one Optuna study executed and best params recorded
• OVERVIEW.md deliverables: check off notebook 04
• Optional: 2–3 line MC summary in OVERVIEW Key results section

Files likely touched

• notebooks/04_memory_capacity.ipynb (new)
• figures/memory_capacity.png (new)
• data/memory_capacity.csv (new)
• OVERVIEW.md (optional results stub)

Depends on

• Issue Implement end-to-end QRC pipeline (ApproxTimeEvolution + feature matrix builder) #1 (QRC pipeline) — required
• Issue Classical Echo State Network baseline module + unit tests #5 (ESN) — required for QRC vs ESN comparison cells
• Issue Memory capacity metric (Dambre et al. 2012) in package #6 (memory capacity module) — required

Blocks

• benchmarks/ reproducible sweep scripts
• Paper results section on memory capacity scaling