path: root/sdil/physical_grid.py

"""Circuit-faithful 4x4 physical learning network used by Dillavou et al.

The nonlinear conductance, periodic topology, and local voltage-square update
follow Eqs. (2)--(3) of arXiv:2505.22887v2.  Source and target node locations
match the released Figure-5 experiment objects.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


Array = np.ndarray


@dataclass(frozen=True)
class GridCircuit:
    rows: int = 4
    columns: int = 4
    conductance_scale: float = 8.0e-4
    threshold_voltage: float = 0.7
    measured_learning_rate: float = 2.5e3
    low_voltage: float = 0.0181
    high_voltage: float = 0.4351
    gate_minimum: float = 1.0
    gate_maximum: float = 5.2
    source_nodes: tuple[int, ...] = (15, 13, 7, 5)
    target_nodes: tuple[int, int] = (10, 8)

    @property
    def node_count(self) -> int:
        return self.rows * self.columns

    @property
    def edge_pairs(self) -> tuple[tuple[int, int], ...]:
        horizontal = []
        vertical = []
        for row in range(self.rows):
            for column in range(self.columns):
                node = row * self.columns + column
                horizontal.append((
                    node,
                    row * self.columns + (column + 1) % self.columns,
                ))
                vertical.append((
                    node,
                    ((row + 1) % self.rows) * self.columns + column,
                ))
        return tuple(horizontal + vertical)

    @property
    def edge_count(self) -> int:
        return len(self.edge_pairs)

    def source_values(self, input_one: float, input_two: float) -> Array:
        return np.asarray((
            input_one,
            input_two,
            self.low_voltage,
            self.high_voltage,
        ), dtype=float)


def edge_voltage_drops(circuit: GridCircuit, node_voltages: Array) -> Array:
    voltages = np.asarray(node_voltages, dtype=float)
    if voltages.shape != (circuit.node_count,):
        raise ValueError("node voltage vector has the wrong shape")
    return np.asarray([
        voltages[first] - voltages[second]
        for first, second in circuit.edge_pairs
    ])


def output_difference(circuit: GridCircuit, node_voltages: Array) -> float:
    positive, negative = circuit.target_nodes
    return float(node_voltages[positive] - node_voltages[negative])


def _residual_and_jacobian(
    circuit: GridCircuit, gates: Array, voltages: Array
) -> tuple[Array, Array]:
    residual = np.zeros(circuit.node_count, dtype=float)
    jacobian = np.zeros(
        (circuit.node_count, circuit.node_count), dtype=float)
    scale = circuit.conductance_scale
    threshold = circuit.threshold_voltage
    for gate, (first, second) in zip(gates, circuit.edge_pairs):
        voltage_first = voltages[first]
        voltage_second = voltages[second]
        conductance = scale * (
            gate - threshold - 0.5 * (voltage_first + voltage_second))
        current = conductance * (voltage_first - voltage_second)
        residual[first] += current
        residual[second] -= current
        derivative_first = scale * (gate - threshold - voltage_first)
        derivative_second = scale * (-gate + threshold + voltage_second)
        jacobian[first, first] += derivative_first
        jacobian[first, second] += derivative_second
        jacobian[second, first] -= derivative_first
        jacobian[second, second] -= derivative_second
    return residual, jacobian


def solve_grid_state(
    circuit: GridCircuit,
    gates: Array,
    source_values: Array,
    *,
    target_values: Array | None = None,
    initial_state: Array | None = None,
    tolerance: float = 1e-11,
    maximum_iterations: int = 20,
) -> Array:
    """Solve Kirchhoff's laws by Newton iteration with an analytic Jacobian."""
    gates = np.asarray(gates, dtype=float)
    sources = np.asarray(source_values, dtype=float)
    if gates.shape != (circuit.edge_count,):
        raise ValueError("gate vector has the wrong shape")
    if sources.shape != (len(circuit.source_nodes),):
        raise ValueError("source voltage vector has the wrong shape")
    fixed = dict(zip(circuit.source_nodes, sources))
    if target_values is not None:
        targets = np.asarray(target_values, dtype=float)
        if targets.shape != (2,):
            raise ValueError("target voltage vector must have shape (2,)")
        fixed.update(zip(circuit.target_nodes, targets))
    unknown = np.asarray([
        node for node in range(circuit.node_count) if node not in fixed
    ])
    voltages = np.full(
        circuit.node_count, float(np.mean(sources)), dtype=float)
    if initial_state is not None:
        initial = np.asarray(initial_state, dtype=float)
        if initial.shape != (circuit.node_count,):
            raise ValueError("initial state vector has the wrong shape")
        voltages[:] = initial
    for node, value in fixed.items():
        voltages[node] = value

    for _ in range(maximum_iterations):
        residual, jacobian = _residual_and_jacobian(
            circuit, gates, voltages)
        unknown_residual = residual[unknown]
        if np.linalg.norm(unknown_residual, ord=np.inf) <= tolerance:
            return voltages
        unknown_jacobian = jacobian[np.ix_(unknown, unknown)]
        step = np.linalg.solve(unknown_jacobian, unknown_residual)
        voltages[unknown] -= step
    residual, _ = _residual_and_jacobian(circuit, gates, voltages)
    raise RuntimeError(
        "grid state did not converge; residual="
        f"{np.linalg.norm(residual[unknown], ord=np.inf):.3e}")


@dataclass(frozen=True)
class GridSquareLawImperfection:
    free_gain: Array
    clamped_gain: Array
    free_input_offset_v: Array
    clamped_input_offset_v: Array
    multiplier_output_offset_v_per_s: Array

    def __post_init__(self) -> None:
        shapes = {
            np.asarray(value).shape
            for value in (
                self.free_gain,
                self.clamped_gain,
                self.free_input_offset_v,
                self.clamped_input_offset_v,
                self.multiplier_output_offset_v_per_s,
            )
        }
        if len(shapes) != 1:
            raise ValueError("grid imperfection arrays disagree")
        shape = next(iter(shapes))
        if len(shape) != 1 or shape[0] < 1:
            raise ValueError("grid imperfection arrays must be nonempty vectors")

    @classmethod
    def ideal(cls, edge_count: int) -> "GridSquareLawImperfection":
        return cls(
            free_gain=np.ones(edge_count),
            clamped_gain=np.ones(edge_count),
            free_input_offset_v=np.zeros(edge_count),
            clamped_input_offset_v=np.zeros(edge_count),
            multiplier_output_offset_v_per_s=np.zeros(edge_count),
        )

    @classmethod
    def sample_appendix_c(
        cls,
        edge_count: int,
        seed: int,
        *,
        gain_standard_deviation: float = 0.01,
        twin_mismatch_standard_deviation_v: float = 0.001,
        multiplier_offset_standard_deviation_v_per_s: float = 2.3,
    ) -> "GridSquareLawImperfection":
        rng = np.random.default_rng(seed)
        common_gain = 1.0 + rng.normal(
            0.0, gain_standard_deviation, edge_count)
        differential_gain = rng.normal(
            0.0, gain_standard_deviation, edge_count)
        common_offset = rng.normal(
            0.0, twin_mismatch_standard_deviation_v, edge_count)
        differential_offset = rng.normal(
            0.0, twin_mismatch_standard_deviation_v, edge_count)
        return cls(
            free_gain=common_gain + 0.5 * differential_gain,
            clamped_gain=common_gain - 0.5 * differential_gain,
            free_input_offset_v=common_offset + 0.5 * differential_offset,
            clamped_input_offset_v=common_offset - 0.5 * differential_offset,
            multiplier_output_offset_v_per_s=rng.normal(
                0.0, multiplier_offset_standard_deviation_v_per_s,
                edge_count),
        )

    def observed_rate(
        self,
        learning_rate: float,
        free_drops: Array,
        clamped_drops: Array,
    ) -> Array:
        measured_free = (
            self.free_gain * free_drops + self.free_input_offset_v)
        measured_clamped = (
            self.clamped_gain * clamped_drops + self.clamped_input_offset_v)
        return (
            learning_rate
            * (np.square(measured_free) - np.square(measured_clamped))
            + self.multiplier_output_offset_v_per_s
        )

    @staticmethod
    def ideal_rate(
        learning_rate: float, free_drops: Array, clamped_drops: Array
    ) -> Array:
        return learning_rate * (
            np.square(free_drops) - np.square(clamped_drops))

    def neutral_bias(self, learning_rate: float, free_drops: Array) -> Array:
        return self.observed_rate(learning_rate, free_drops, free_drops)


@dataclass(frozen=True)
class AutozeroSampleHold:
    """Nonideal local sample-and-hold used for hardware auto-zeroing.

    The sampled input is the learning-circuit output while the free and
    clamped edge voltages are equal.  No device parameters are exposed to the
    sampler or to the learning rule.
    """

    sample_gain: float = 1.0
    pedestal_offset_v_per_s: float | Array = 0.0
    sample_noise_standard_deviation_v_per_s: float = 0.0
    refresh_interval_updates: int = 1

    def __post_init__(self) -> None:
        if not np.isfinite(self.sample_gain) or self.sample_gain < 0.0:
            raise ValueError("sample gain must be finite and nonnegative")
        if (
            not np.isfinite(self.sample_noise_standard_deviation_v_per_s)
            or self.sample_noise_standard_deviation_v_per_s < 0.0
        ):
            raise ValueError("sample noise must be finite and nonnegative")
        if self.refresh_interval_updates < 1:
            raise ValueError("refresh interval must be positive")

    def sample(self, neutral_output: Array, rng: np.random.Generator) -> Array:
        neutral = np.asarray(neutral_output, dtype=float)
        offset = np.asarray(self.pedestal_offset_v_per_s, dtype=float)
        try:
            offset = np.broadcast_to(offset, neutral.shape)
        except ValueError as error:
            raise ValueError(
                "sample-and-hold pedestal cannot broadcast to the edge vector"
            ) from error
        noise = rng.normal(
            0.0,
            self.sample_noise_standard_deviation_v_per_s,
            neutral.shape,
        )
        return self.sample_gain * neutral + offset + noise


@dataclass(frozen=True)
class CorrelatedDoubleSampleHold:
    """Matched local sampling of neutral and active learning outputs.

    A common pedestal is added to both samples and therefore cancels in their
    difference.  The mismatch fields model the remaining difference between
    the two sampling phases rather than exposing correction parameters to the
    learning rule.
    """

    common_sample_gain: float = 1.0
    sample_gain_mismatch: float | Array = 0.0
    common_pedestal_offset_v_per_s: float | Array = 0.0
    pedestal_mismatch_v_per_s: float | Array = 0.0
    sample_noise_standard_deviation_v_per_s: float = 0.0
    refresh_interval_updates: int = 1

    def __post_init__(self) -> None:
        if (
            not np.isfinite(self.common_sample_gain)
            or self.common_sample_gain < 0.0
        ):
            raise ValueError("common sample gain must be finite and nonnegative")
        if (
            not np.isfinite(self.sample_noise_standard_deviation_v_per_s)
            or self.sample_noise_standard_deviation_v_per_s < 0.0
        ):
            raise ValueError("sample noise must be finite and nonnegative")
        if self.refresh_interval_updates < 1:
            raise ValueError("refresh interval must be positive")

    @staticmethod
    def _edge_vector(value: float | Array, shape: tuple[int, ...]) -> Array:
        array = np.asarray(value, dtype=float)
        try:
            return np.broadcast_to(array, shape)
        except ValueError as error:
            raise ValueError(
                "correlated-sampling parameter cannot broadcast to edges"
            ) from error

    def sample_neutral(
        self, neutral_output: Array, rng: np.random.Generator
    ) -> Array:
        neutral = np.asarray(neutral_output, dtype=float)
        gain_mismatch = self._edge_vector(
            self.sample_gain_mismatch, neutral.shape)
        common_pedestal = self._edge_vector(
            self.common_pedestal_offset_v_per_s, neutral.shape)
        pedestal_mismatch = self._edge_vector(
            self.pedestal_mismatch_v_per_s, neutral.shape)
        noise = rng.normal(
            0.0,
            self.sample_noise_standard_deviation_v_per_s,
            neutral.shape,
        )
        return (
            (self.common_sample_gain - 0.5 * gain_mismatch) * neutral
            + common_pedestal - 0.5 * pedestal_mismatch + noise
        )

    def sample_active(
        self, active_output: Array, rng: np.random.Generator
    ) -> Array:
        active = np.asarray(active_output, dtype=float)
        gain_mismatch = self._edge_vector(
            self.sample_gain_mismatch, active.shape)
        common_pedestal = self._edge_vector(
            self.common_pedestal_offset_v_per_s, active.shape)
        pedestal_mismatch = self._edge_vector(
            self.pedestal_mismatch_v_per_s, active.shape)
        noise = rng.normal(
            0.0,
            self.sample_noise_standard_deviation_v_per_s,
            active.shape,
        )
        return (
            (self.common_sample_gain + 0.5 * gain_mismatch) * active
            + common_pedestal + 0.5 * pedestal_mismatch + noise
        )


@dataclass
class EdgePolynomialPredictor:
    feature_center: Array
    feature_scale: Array
    coefficients: Array

    @classmethod
    def zeros(
        cls, feature_center: Array, feature_scale: Array, *, degree: int
    ) -> "EdgePolynomialPredictor":
        center = np.asarray(feature_center, dtype=float)
        scale = np.asarray(feature_scale, dtype=float)
        if center.ndim != 1 or scale.shape != center.shape:
            raise ValueError("edge feature metadata disagree")
        if degree < 0 or np.any(scale <= 0.0):
            raise ValueError("invalid polynomial degree or feature scale")
        return cls(
            feature_center=center.copy(),
            feature_scale=scale.copy(),
            coefficients=np.zeros((len(center), degree + 1), dtype=float),
        )

    @property
    def degree(self) -> int:
        return int(self.coefficients.shape[1] - 1)

    def features(self, local_state: Array) -> Array:
        state = np.asarray(local_state, dtype=float)
        if state.shape != self.feature_center.shape:
            raise ValueError("edge local state has the wrong shape")
        normalized = (state - self.feature_center) / self.feature_scale
        return np.stack([
            normalized ** power for power in range(self.degree + 1)
        ], axis=1)

    def predict(self, local_state: Array) -> Array:
        return np.sum(self.coefficients * self.features(local_state), axis=1)

    def copy(self) -> "EdgePolynomialPredictor":
        return EdgePolynomialPredictor(
            feature_center=self.feature_center.copy(),
            feature_scale=self.feature_scale.copy(),
            coefficients=self.coefficients.copy(),
        )


def fit_edge_predictor(
    predictor: EdgePolynomialPredictor,
    local_states: Array,
    neutral_measurements: Array,
    *,
    ridge: float = 1e-12,
) -> int:
    states = np.asarray(local_states, dtype=float)
    measurements = np.asarray(neutral_measurements, dtype=float)
    if states.ndim != 2 or measurements.shape != states.shape:
        raise ValueError("edge calibration matrices disagree")
    if states.shape[1] != len(predictor.feature_center):
        raise ValueError("edge calibration width changed")
    features = np.asarray([predictor.features(state) for state in states])
    for edge in range(states.shape[1]):
        design = features[:, edge, :]
        gram = design.T @ design
        rhs = design.T @ measurements[:, edge]
        predictor.coefficients[edge] = np.linalg.solve(
            gram + ridge * np.eye(gram.shape[0]), rhs)
    return int(len(states))


@dataclass(frozen=True)
class RingClassificationDataset:
    inputs_v: Array
    labels_v: Array

    def __post_init__(self) -> None:
        inputs = np.asarray(self.inputs_v)
        labels = np.asarray(self.labels_v)
        if inputs.ndim != 2 or inputs.shape[1] != 2:
            raise ValueError("ring inputs must have shape (samples, 2)")
        if labels.shape != (len(inputs),):
            raise ValueError("ring labels must match the sample count")
        if np.any(labels == 0.0):
            raise ValueError("classification labels must be signed")


def evaluate_grid_classifier(
    circuit: GridCircuit,
    gates: Array,
    dataset: RingClassificationDataset,
    *,
    initial_states: list[Array | None] | None = None,
) -> tuple[dict, list[Array]]:
    if initial_states is None:
        initial_states = [None] * len(dataset.labels_v)
    states = []
    outputs = []
    for index, inputs in enumerate(dataset.inputs_v):
        state = solve_grid_state(
            circuit,
            gates,
            circuit.source_values(*inputs),
            initial_state=initial_states[index],
        )
        states.append(state)
        outputs.append(output_difference(circuit, state))
    outputs_array = np.asarray(outputs)
    labels = np.asarray(dataset.labels_v)
    errors = labels - outputs_array
    active = labels * errors > 0.0
    return {
        "classification_error": float(np.mean(
            np.sign(outputs_array) != np.sign(labels))),
        "hinge_loss_v2": float(np.mean(np.where(
            active, np.square(errors), 0.0))),
        "margin_success_fraction": float(np.mean(~active)),
        "outputs_v": outputs_array.tolist(),
    }, states


def train_grid_classifier(
    circuit: GridCircuit,
    initial_gates: Array,
    dataset: RingClassificationDataset,
    imperfection: GridSquareLawImperfection,
    *,
    method: str,
    epochs: int,
    predictor: EdgePolynomialPredictor | None = None,
    autozero_sample_hold: AutozeroSampleHold | None = None,
    correlated_sample_hold: CorrelatedDoubleSampleHold | None = None,
    autozero_seed: int = 0,
    standard_nudging: float = 128.0 / 129.0,
    standard_learning_time_seconds: float = 1.0e-3,
    overclamp_nudging: float = 32.0 / 129.0,
    overclamp_target_magnitude_v: float | None = None,
    overclamp_time_seconds_per_v: float = 0.05,
    record_every: int = 50,
    early_stop_perfect_checkpoints: int | None = None,
) -> dict:
    """Train the physical grid with explicit local voltage-square updates."""
    allowed = {
        "clean",
        "raw",
        "constant",
        "sdil",
        "oracle_neutral",
        "autozero_sdil",
        "cds_autozero_sdil",
        "overclamp_clean",
        "overclamp",
        "overclamp_constant",
        "overclamp_sdil",
        "overclamp_oracle_neutral",
        "overclamp_autozero_sdil",
        "overclamp_cds_autozero_sdil",
    }
    if method not in allowed:
        raise ValueError(f"unrecognized method {method}")
    if epochs < 1:
        raise ValueError("epochs must be positive")
    if method in {
        "constant", "sdil", "overclamp_constant", "overclamp_sdil"
    } and predictor is None:
        raise ValueError(f"{method} requires a predictor")
    autozero_methods = {
        "autozero_sdil", "overclamp_autozero_sdil"
    }
    correlated_methods = {
        "cds_autozero_sdil", "overclamp_cds_autozero_sdil"
    }
    if method in autozero_methods and autozero_sample_hold is None:
        autozero_sample_hold = AutozeroSampleHold()
    if method in correlated_methods and correlated_sample_hold is None:
        correlated_sample_hold = CorrelatedDoubleSampleHold()
    gates = np.asarray(initial_gates, dtype=float).copy()
    if gates.shape != (circuit.edge_count,):
        raise ValueError("initial gate vector has the wrong shape")
    active_predictor = predictor.copy() if predictor is not None else None
    target_magnitude = (
        circuit.high_voltage
        if overclamp_target_magnitude_v is None
        else overclamp_target_magnitude_v)
    is_overclamp = method.startswith("overclamp")
    free_cache: list[Array | None] = [None] * len(dataset.labels_v)
    clamped_cache: list[Array | None] = [None] * len(dataset.labels_v)
    trace = []
    cumulative_learning_time = 0.0
    clamp_l2_time = 0.0
    max_clamp_displacement = 0.0
    local_updates = 0
    clipped_updates = 0
    autozero_samples = 0
    autozero_active_samples = 0
    autozero_updates_since_sample = 0
    held_neutral_output: Array | None = None
    autozero_error_sum_squared = 0.0
    autozero_error_entries = 0
    autozero_applied_error_sum_squared = 0.0
    autozero_applied_error_entries = 0
    autozero_rng = np.random.default_rng(autozero_seed)
    perfect_checkpoints = 0
    completed_epochs = 0

    for epoch in range(epochs):
        for sample, (inputs, label) in enumerate(zip(
            dataset.inputs_v, dataset.labels_v
        )):
            sources = circuit.source_values(*inputs)
            free_state = solve_grid_state(
                circuit,
                gates,
                sources,
                initial_state=free_cache[sample],
            )
            free_cache[sample] = free_state
            output_free = output_difference(circuit, free_state)
            error = label - output_free
            if label * error <= 0.0:
                continue
            free_drops = edge_voltage_drops(circuit, free_state)
            if method in autozero_methods | correlated_methods:
                current_neutral_output = imperfection.observed_rate(
                    circuit.measured_learning_rate, free_drops, free_drops)
                refresh_interval = (
                    autozero_sample_hold.refresh_interval_updates
                    if method in autozero_methods
                    else correlated_sample_hold.refresh_interval_updates
                )
                if (
                    held_neutral_output is None
                    or autozero_updates_since_sample
                    >= refresh_interval
                ):
                    if method in autozero_methods:
                        held_neutral_output = autozero_sample_hold.sample(
                            current_neutral_output, autozero_rng)
                    else:
                        held_neutral_output = (
                            correlated_sample_hold.sample_neutral(
                                current_neutral_output, autozero_rng))
                    autozero_samples += 1
                    autozero_updates_since_sample = 0
                autozero_error_sum_squared += float(np.sum(np.square(
                    held_neutral_output - current_neutral_output)))
                autozero_error_entries += circuit.edge_count
                autozero_updates_since_sample += 1
            if is_overclamp:
                output_clamped = output_free + overclamp_nudging * (
                    target_magnitude * np.sign(error) - output_free)
                duration = overclamp_time_seconds_per_v * abs(error)
            else:
                output_clamped = output_free + standard_nudging * error
                duration = standard_learning_time_seconds
            target_mean = float(np.mean(
                free_state[list(circuit.target_nodes)]))
            target_values = np.asarray((
                target_mean + 0.5 * output_clamped,
                target_mean - 0.5 * output_clamped,
            ))
            clamped_state = solve_grid_state(
                circuit,
                gates,
                sources,
                target_values=target_values,
                initial_state=(
                    free_state if clamped_cache[sample] is None
                    else clamped_cache[sample]),
            )
            clamped_cache[sample] = clamped_state
            clamped_drops = edge_voltage_drops(circuit, clamped_state)
            ideal_rate = imperfection.ideal_rate(
                circuit.measured_learning_rate, free_drops, clamped_drops)
            observed_rate = imperfection.observed_rate(
                circuit.measured_learning_rate, free_drops, clamped_drops)
            if method in {"clean", "overclamp_clean"}:
                applied_rate = ideal_rate
            elif method in {
                "oracle_neutral", "overclamp_oracle_neutral"
            }:
                neutral_output = imperfection.observed_rate(
                    circuit.measured_learning_rate, free_drops, free_drops)
                applied_rate = observed_rate - neutral_output
            elif method in autozero_methods:
                applied_rate = observed_rate - held_neutral_output
            elif method in correlated_methods:
                sampled_active_rate = correlated_sample_hold.sample_active(
                    observed_rate, autozero_rng)
                autozero_active_samples += 1
                applied_rate = sampled_active_rate - held_neutral_output
                reference_rate = (
                    correlated_sample_hold.common_sample_gain
                    * (observed_rate - current_neutral_output))
                autozero_applied_error_sum_squared += float(np.sum(np.square(
                    applied_rate - reference_rate)))
                autozero_applied_error_entries += circuit.edge_count
            elif method in {
                "constant", "sdil", "overclamp_constant", "overclamp_sdil"
            }:
                applied_rate = observed_rate - active_predictor.predict(
                    free_drops)
            else:
                applied_rate = observed_rate
            proposed = gates + duration * applied_rate
            clipped = np.clip(
                proposed, circuit.gate_minimum, circuit.gate_maximum)
            clipped_updates += int(np.any(clipped != proposed))
            gates = clipped
            local_updates += 1
            cumulative_learning_time += duration
            displacement = abs(output_clamped - output_free)
            clamp_l2_time += duration * displacement * displacement
            max_clamp_displacement = max(
                max_clamp_displacement, displacement)
        if epoch % record_every == 0 or epoch == epochs - 1:
            metrics, free_cache = evaluate_grid_classifier(
                circuit, gates, dataset, initial_states=free_cache)
            trace.append({"epoch": epoch, **metrics})
            if metrics["classification_error"] == 0.0:
                perfect_checkpoints += 1
            else:
                perfect_checkpoints = 0
            if (
                early_stop_perfect_checkpoints is not None
                and perfect_checkpoints >= early_stop_perfect_checkpoints
            ):
                completed_epochs = epoch + 1
                break
        completed_epochs = epoch + 1

    final = trace[-1]
    return {
        "method": method,
        "requested_epochs": epochs,
        "completed_epochs": completed_epochs,
        "initial_gates_v": np.asarray(initial_gates).tolist(),
        "final_gates_v": gates.tolist(),
        "classification_error": final["classification_error"],
        "hinge_loss_v2": final["hinge_loss_v2"],
        "margin_success_fraction": final["margin_success_fraction"],
        "outputs_v": final["outputs_v"],
        "local_updates": local_updates,
        "cumulative_learning_time_seconds": float(cumulative_learning_time),
        "clamp_displacement_l2_time_v2_s": float(clamp_l2_time),
        "max_abs_clamp_displacement_v": float(max_clamp_displacement),
        "clipped_updates": clipped_updates,
        "autozero_samples": autozero_samples,
        "autozero_active_samples": autozero_active_samples,
        "autozero_sample_fraction_per_update": float(
            autozero_samples / local_updates if local_updates else 0.0),
        "autozero_baseline_rmse_v_per_s": float(np.sqrt(
            autozero_error_sum_squared / autozero_error_entries
        )) if autozero_error_entries else None,
        "autozero_applied_rate_rmse_v_per_s": float(np.sqrt(
            autozero_applied_error_sum_squared
            / autozero_applied_error_entries
        )) if autozero_applied_error_entries else None,
        "trace": trace,
    }