results: identify state-dependent physical drift

1 files changed, 493 insertions, 0 deletions

diff --git a/experiments/analyze_physical_drift_state_dependence.py b/experiments/analyze_physical_drift_state_dependence.py
new file mode 100644
index 0000000..42ca420
--- /dev/null
+++ b/experiments/analyze_physical_drift_state_dependence.py
@@ -0,0 +1,493 @@
+#!/usr/bin/env python3
+"""Test whether released physical drift is state dependent.
+
+Dillavou et al. approximate the two-edge physical network with a constant
+bias vector.  Their appendix also notes that the experimental drift rate is
+not constant.  This script tests that statement directly from the four
+released DriftTests traces.
+
+Once a trace is close to a one-task solution line, the clean learning force is
+normal to that line.  Its measured tangential velocity therefore estimates
+the tangential projection of the physical bias.  We compare:
+
+    constant:       B(G) = b
+    local affine:   B_i(G_i) = b_i + k_i (G_i - G_ref_i)
+
+The second model uses exactly the per-edge local information available to an
+affine SDIL predictor.  It is a measured-data identifiability diagnostic, not
+yet evidence that SDIL improves task learning.
+"""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+import hashlib
+import json
+from pathlib import Path
+from typing import Dict, List, Sequence
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.io import loadmat
+
+
+ZENODO_RECORD = "15692914"
+ZENODO_DOI = "10.5281/zenodo.15692914"
+RELEASE = "v1.0.1"
+SOURCE_TREE = "maguzj-imperfect-learning-physical-systems-71b8d72"
+
+
+@dataclass(frozen=True)
+class TraceSpec:
+    filename: str
+    task: str
+    solution_slope: float
+    solution_intercept: float
+
+
+@dataclass(frozen=True)
+class PairSpec:
+    name: str
+    reference_gate: tuple[float, float]
+    published_bias_at_reference: tuple[float, float]
+    traces: tuple[TraceSpec, TraceSpec]
+
+
+PAIR_SPECS = (
+    PairSpec(
+        name="experiment_1",
+        reference_gate=(3.0, 3.5),
+        published_bias_at_reference=(2.9, 4.7),
+        traces=(
+            TraceSpec("exp1DriftTest_1.mat", "beta", 0.43914, 1.6061),
+            TraceSpec("exp1DriftTest_2.mat", "alpha", 2.7257, -3.3003),
+        ),
+    ),
+    PairSpec(
+        name="experiment_2",
+        reference_gate=(3.1, 4.3),
+        published_bias_at_reference=(-0.3, 2.0),
+        traces=(
+            TraceSpec("exp2DriftTest_1.mat", "beta", 0.74821, 1.9224),
+            TraceSpec("exp2DriftTest_2.mat", "alpha", 2.5678, -2.901),
+        ),
+    ),
+)
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def rmse(observed: np.ndarray, predicted: np.ndarray) -> float:
+    return float(np.sqrt(np.mean((observed - predicted) ** 2)))
+
+
+def r_squared(observed: np.ndarray, predicted: np.ndarray) -> float:
+    denominator = float(np.sum((observed - np.mean(observed)) ** 2))
+    if denominator == 0.0:
+        return float("nan")
+    return float(1.0 - np.sum((observed - predicted) ** 2) / denominator)
+
+
+def load_trace(
+    drift_dir: Path,
+    spec: TraceSpec,
+    reference_gate: Sequence[float],
+    *,
+    minimum_time: float,
+    terminal_headroom: float,
+    maximum_line_distance: float,
+    trim_selected_ends: int,
+) -> dict:
+    path = drift_dir / spec.filename
+    raw = loadmat(path, squeeze_me=True)
+    time = np.asarray(raw["time"], dtype=float)
+    gate_minus = np.asarray(raw["GMINUS"], dtype=float)
+    gate_plus = np.asarray(raw["GPLUS"], dtype=float)
+    if time.ndim != 1 or len(time) != 300:
+        raise ValueError(f"unexpected time array in {path}: {time.shape}")
+    if not (time.shape == gate_minus.shape == gate_plus.shape):
+        raise ValueError(f"array shape disagreement in {path}")
+    if not np.all(np.isfinite(np.column_stack((time, gate_minus, gate_plus)))):
+        raise ValueError(f"nonfinite source value in {path}")
+
+    # The first two entries straddle the experiment initialization and a
+    # roughly 0.9 s gap.  The remaining samples have a stable ~30 ms cadence.
+    time = time[2:]
+    gate_minus = gate_minus[2:]
+    gate_plus = gate_plus[2:]
+    sample_period = float(np.median(np.diff(time)))
+    velocity_minus = np.gradient(gate_minus, time)
+    velocity_plus = np.gradient(gate_plus, time)
+
+    slope = spec.solution_slope
+    intercept = spec.solution_intercept
+    tangent = np.asarray((1.0, slope), dtype=float)
+    tangent /= np.linalg.norm(tangent)
+    signed_distance = (
+        gate_plus - slope * gate_minus - intercept
+    ) / np.sqrt(1.0 + slope * slope)
+    # Every released trace terminates on a clear gate-voltage plateau.  Exclude
+    # the final headroom before either observed plateau so a hardware rail is
+    # not mislabeled as a state-dependent bias.  Trimming the selected ends
+    # then ensures that centered differences use only pre-plateau samples.
+    selected = (
+        (time >= minimum_time)
+        & (gate_minus < np.max(gate_minus) - terminal_headroom)
+        & (gate_plus < np.max(gate_plus) - terminal_headroom)
+        & (np.abs(signed_distance) < maximum_line_distance)
+    )
+    indices = np.flatnonzero(selected)
+    if len(indices) <= 2 * trim_selected_ends:
+        raise ValueError(f"too few usable solution-line samples in {path}")
+    if np.any(np.diff(indices) != 1):
+        raise ValueError(f"selected solution-line window is not contiguous in {path}")
+    indices = indices[trim_selected_ends:-trim_selected_ends]
+
+    gates = np.column_stack((gate_minus[indices], gate_plus[indices]))
+    velocities = np.column_stack(
+        (velocity_minus[indices], velocity_plus[indices]))
+    reference = np.asarray(reference_gate, dtype=float)
+    design = np.column_stack((
+        np.full(len(indices), tangent[0]),
+        np.full(len(indices), tangent[1]),
+        tangent[0] * (gates[:, 0] - reference[0]),
+        tangent[1] * (gates[:, 1] - reference[1]),
+    ))
+    projected_velocity = velocities @ tangent
+    path_coordinate = gates @ tangent
+    if len(indices) < 6:
+        raise ValueError(f"need at least six retained samples in {path}")
+
+    return {
+        "spec": spec,
+        "source_file": {"path": str(path), "sha256": sha256(path)},
+        "sample_period_seconds": sample_period,
+        "raw_gate_minus": gate_minus,
+        "raw_gate_plus": gate_plus,
+        "time": time,
+        "selected_indices": indices,
+        "selected_time": time[indices],
+        "selected_gates": gates,
+        "selected_line_distance": signed_distance[indices],
+        "tangent": tangent,
+        "path_coordinate": path_coordinate,
+        "projected_velocity": projected_velocity,
+        "design": design,
+    }
+
+
+def fit_columns(design: np.ndarray, target: np.ndarray, columns: List[int]) -> dict:
+    selected = design[:, columns]
+    parameters, _, _, _ = np.linalg.lstsq(selected, target, rcond=None)
+    predicted = selected @ parameters
+    return {
+        "columns": columns,
+        "parameters": parameters,
+        "predicted": predicted,
+        "rank": int(np.linalg.matrix_rank(selected)),
+        "condition_number": float(np.linalg.cond(selected)),
+        "rmse": rmse(target, predicted),
+        "r_squared": r_squared(target, predicted),
+    }
+
+
+def forward_block_test(traces: Sequence[dict], columns: List[int]) -> dict:
+    train_design = []
+    train_target = []
+    test_design = []
+    test_target = []
+    split_records = []
+    for trace in traces:
+        count = len(trace["projected_velocity"])
+        train_count = max(2, int(np.floor(0.6 * count)))
+        if count - train_count < 2:
+            raise ValueError("forward block test needs at least two held-out samples")
+        train_design.append(trace["design"][:train_count, columns])
+        train_target.append(trace["projected_velocity"][:train_count])
+        test_design.append(trace["design"][train_count:, columns])
+        test_target.append(trace["projected_velocity"][train_count:])
+        split_records.append({
+            "source": trace["spec"].filename,
+            "train_samples": train_count,
+            "test_samples": count - train_count,
+        })
+    train_design_array = np.vstack(train_design)
+    train_target_array = np.concatenate(train_target)
+    test_design_array = np.vstack(test_design)
+    test_target_array = np.concatenate(test_target)
+    parameters, _, _, _ = np.linalg.lstsq(
+        train_design_array, train_target_array, rcond=None)
+    train_prediction = train_design_array @ parameters
+    test_prediction = test_design_array @ parameters
+    return {
+        "split": "first 60% of each retained trace trains; final 40% tests",
+        "trace_splits": split_records,
+        "parameters": parameters.tolist(),
+        "train_rank": int(np.linalg.matrix_rank(train_design_array)),
+        "train_condition_number": float(np.linalg.cond(train_design_array)),
+        "train_rmse": rmse(train_target_array, train_prediction),
+        "test_rmse": rmse(test_target_array, test_prediction),
+        "test_r_squared": r_squared(test_target_array, test_prediction),
+    }
+
+
+def serializable_trace(trace: dict) -> dict:
+    spec = trace["spec"]
+    return {
+        "source_file": trace["source_file"],
+        "task": spec.task,
+        "solution_line": {
+            "gate_plus_slope": spec.solution_slope,
+            "gate_plus_intercept": spec.solution_intercept,
+        },
+        "sample_period_seconds": trace["sample_period_seconds"],
+        "retained_samples": int(len(trace["projected_velocity"])),
+        "retained_time_seconds": trace["selected_time"].tolist(),
+        "retained_gate_minus": trace["selected_gates"][:, 0].tolist(),
+        "retained_gate_plus": trace["selected_gates"][:, 1].tolist(),
+        "retained_line_distance": trace["selected_line_distance"].tolist(),
+        "path_coordinate": trace["path_coordinate"].tolist(),
+        "projected_drift_velocity_v_per_s": trace["projected_velocity"].tolist(),
+    }
+
+
+def analyze_pair(drift_dir: Path, spec: PairSpec, settings: dict) -> dict:
+    traces = [
+        load_trace(
+            drift_dir,
+            trace_spec,
+            spec.reference_gate,
+            **settings,
+        )
+        for trace_spec in spec.traces
+    ]
+    design = np.vstack([trace["design"] for trace in traces])
+    target = np.concatenate([trace["projected_velocity"] for trace in traces])
+    constant = fit_columns(design, target, [0, 1])
+    local_affine = fit_columns(design, target, [0, 1, 2, 3])
+    constant_forward = forward_block_test(traces, [0, 1])
+    affine_forward = forward_block_test(traces, [0, 1, 2, 3])
+    affine_parameters = local_affine["parameters"]
+    published = np.asarray(spec.published_bias_at_reference)
+    intercept = affine_parameters[:2]
+    return {
+        "reference_gate": list(spec.reference_gate),
+        "published_constant_bias_at_reference_v_per_s": published.tolist(),
+        "traces": [serializable_trace(trace) for trace in traces],
+        "constant_model": {
+            "definition": "B(G) = b",
+            "bias_v_per_s": constant["parameters"].tolist(),
+            "rank": constant["rank"],
+            "condition_number": constant["condition_number"],
+            "in_sample_rmse_v_per_s": constant["rmse"],
+            "in_sample_r_squared": constant["r_squared"],
+            "forward_block": constant_forward,
+        },
+        "local_affine_model": {
+            "definition": "B_i(G_i) = b_i + k_i (G_i - G_ref_i)",
+            "bias_at_reference_v_per_s": intercept.tolist(),
+            "local_slopes_per_s": affine_parameters[2:].tolist(),
+            "rank": local_affine["rank"],
+            "condition_number": local_affine["condition_number"],
+            "in_sample_rmse_v_per_s": local_affine["rmse"],
+            "in_sample_r_squared": local_affine["r_squared"],
+            "forward_block": affine_forward,
+            "reference_bias_error_v_per_s": (intercept - published).tolist(),
+        },
+        "diagnostic": {
+            "in_sample_rmse_ratio_affine_over_constant": (
+                local_affine["rmse"] / constant["rmse"]),
+            "forward_test_rmse_ratio_affine_over_constant": (
+                affine_forward["test_rmse"] / constant_forward["test_rmse"]),
+            "affine_improves_forward_test": bool(
+                affine_forward["test_rmse"] < constant_forward["test_rmse"]),
+        },
+        "_plot": {
+            "traces": traces,
+            "constant_parameters": constant["parameters"],
+            "affine_parameters": affine_parameters,
+        },
+    }
+
+
+def headroom_sensitivity(drift_dir: Path, base_settings: dict) -> List[dict]:
+    records = []
+    for headroom in (0.03, 0.05, 0.08, 0.10):
+        settings = dict(base_settings)
+        settings["terminal_headroom"] = headroom
+        pair_ratios = {}
+        for spec in PAIR_SPECS:
+            pair = analyze_pair(drift_dir, spec, settings)
+            pair_ratios[spec.name] = pair["diagnostic"][
+                "forward_test_rmse_ratio_affine_over_constant"]
+        records.append({
+            "terminal_headroom_v": headroom,
+            "forward_test_rmse_ratios_affine_over_constant": pair_ratios,
+            "affine_improves_both": bool(all(
+                ratio < 1.0 for ratio in pair_ratios.values())),
+        })
+    return records
+
+
+def plot_report(report: dict, output: Path) -> None:
+    colors = {"alpha": "#4477AA", "beta": "#EE6677"}
+    fig, axes = plt.subplots(2, 2, figsize=(9.2, 7.0))
+    for column, pair_name in enumerate(("experiment_1", "experiment_2")):
+        pair = report["pairs"][pair_name]
+        plot_data = pair["_plot"]
+        path_axis = axes[0, column]
+        velocity_axis = axes[1, column]
+        for trace in plot_data["traces"]:
+            spec = trace["spec"]
+            color = colors[spec.task]
+            path_axis.plot(
+                trace["raw_gate_minus"], trace["raw_gate_plus"],
+                color=color, alpha=0.24, linewidth=1.0)
+            gates = trace["selected_gates"]
+            path_axis.plot(
+                gates[:, 0], gates[:, 1], "o-", color=color,
+                markersize=3.0, linewidth=1.2, label=f"task {spec.task}")
+            x_line = np.linspace(gates[:, 0].min(), gates[:, 0].max(), 100)
+            path_axis.plot(
+                x_line, spec.solution_slope * x_line + spec.solution_intercept,
+                "--", color=color, linewidth=1.0)
+
+            order = np.argsort(trace["path_coordinate"])
+            coordinate = trace["path_coordinate"][order]
+            observed = trace["projected_velocity"][order]
+            design = trace["design"][order]
+            constant_prediction = design[:, :2] @ plot_data["constant_parameters"]
+            affine_prediction = design @ plot_data["affine_parameters"]
+            velocity_axis.plot(
+                coordinate, observed, "o", color=color, markersize=3.5,
+                label=f"task {spec.task} measured")
+            velocity_axis.plot(
+                coordinate, constant_prediction, ":", color=color,
+                linewidth=1.3)
+            velocity_axis.plot(
+                coordinate, affine_prediction, "-", color=color,
+                linewidth=1.5)
+
+        path_axis.set_title(
+            f"{chr(ord('A') + column)}  Published drift: {pair_name.replace('_', ' ')}")
+        path_axis.set_xlabel("gate voltage minus (V)")
+        path_axis.set_ylabel("gate voltage plus (V)")
+        path_axis.legend(frameon=False, fontsize=8)
+        path_axis.grid(alpha=0.18)
+
+        ratio = pair["diagnostic"]["forward_test_rmse_ratio_affine_over_constant"]
+        velocity_axis.set_title(
+            f"{chr(ord('C') + column)}  Held-out RMSE ratio affine/constant = {ratio:.2f}")
+        velocity_axis.set_xlabel("position along solution line (V)")
+        velocity_axis.set_ylabel("tangential drift (V/s)")
+        velocity_axis.legend(frameon=False, fontsize=7)
+        velocity_axis.grid(alpha=0.18)
+
+    fig.suptitle(
+        "Released physical drift varies with local gate state\n"
+        "dots: measured projection; dotted: constant bias; solid: per-edge affine bias",
+        fontsize=11)
+    fig.tight_layout()
+    output.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output, dpi=180)
+    plt.close(fig)
+
+
+def strip_plot_data(report: dict) -> dict:
+    clean = dict(report)
+    clean["pairs"] = {}
+    for name, pair in report["pairs"].items():
+        clean["pairs"][name] = {key: value for key, value in pair.items() if key != "_plot"}
+    return clean
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--artifact-root", type=Path, required=True,
+        help="Root of the extracted Zenodo source tree")
+    parser.add_argument(
+        "--json", type=Path,
+        default=Path("results/physical_bias/p0_state_dependence.json"))
+    parser.add_argument(
+        "--figure", type=Path,
+        default=Path("results/figs/physical_bias_state_dependence.png"))
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    root = args.artifact_root.resolve()
+    if root.name != SOURCE_TREE:
+        raise ValueError(f"expected source root {SOURCE_TREE}, received {root.name}")
+    drift_dir = root / "small network" / "DriftTests"
+    settings = {
+        "minimum_time": 1.05,
+        "terminal_headroom": 0.05,
+        "maximum_line_distance": 0.03,
+        "trim_selected_ends": 1,
+    }
+    report = {
+        "analysis": "published_physical_drift_state_dependence_diagnostic",
+        "confirmatory": False,
+        "supports_task_learning_improvement": False,
+        "provenance": {
+            "zenodo_record": ZENODO_RECORD,
+            "doi": ZENODO_DOI,
+            "release": RELEASE,
+            "source_tree": SOURCE_TREE,
+        },
+        "identification": {
+            "observable": "tangential gate velocity after reaching a one-task solution line",
+            "assumption": (
+                "the clean force has zero tangential projection on the published solution line"
+            ),
+            "model_restriction": "each edge bias is affine in its own local gate voltage",
+            "limitation": (
+                "the traces identify a local affine surrogate over their visited states, not the "
+                "complete physical vector field"
+            ),
+        },
+        "preprocessing": settings,
+        "pairs": {},
+    }
+    for spec in PAIR_SPECS:
+        report["pairs"][spec.name] = analyze_pair(drift_dir, spec, settings)
+    report["headroom_sensitivity"] = headroom_sensitivity(drift_dir, settings)
+    report["summary"] = {
+        "forward_test_rmse_ratios_affine_over_constant": [
+            report["pairs"][spec.name]["diagnostic"][
+                "forward_test_rmse_ratio_affine_over_constant"]
+            for spec in PAIR_SPECS
+        ],
+        "affine_improves_both_forward_tests": bool(all(
+            report["pairs"][spec.name]["diagnostic"][
+                "affine_improves_forward_test"]
+            for spec in PAIR_SPECS
+        )),
+        "affine_improves_both_for_all_headroom_thresholds": bool(all(
+            record["affine_improves_both"]
+            for record in report["headroom_sensitivity"]
+        )),
+    }
+    plot_report(report, args.figure)
+    serializable = strip_plot_data(report)
+    args.json.parent.mkdir(parents=True, exist_ok=True)
+    args.json.write_text(json.dumps(serializable, indent=2) + "\n")
+    print(json.dumps(serializable["summary"], indent=2))
+    print(f"wrote {args.json}")
+    print(f"wrote {args.figure}")
+
+
+if __name__ == "__main__":
+    main()