feat: reconstruct the Dillavou physical grid

1 files changed, 314 insertions, 0 deletions

diff --git a/sdil/physical_grid.py b/sdil/physical_grid.py
new file mode 100644
index 0000000..10d667b
--- /dev/null
+++ b/sdil/physical_grid.py
@@ -0,0 +1,314 @@
+"""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
+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))