path: root/experiments/verify_theory.py

"""Deterministic numerical checks for the claims in THEORY.md."""
import itertools
import math

import numpy as np


def simultaneous_mse(rng, depth, width, directions, trials=12000):
    gradients = rng.normal(size=(depth, width))
    gradients /= np.linalg.norm(gradients, axis=1, keepdims=True)
    estimate = np.zeros((trials, width), dtype=np.float64)
    for _ in range(directions):
        xi = rng.integers(0, 2, size=(trials, depth, width), dtype=np.int8)
        xi = 2.0 * xi - 1.0
        directional = np.einsum("tdw,dw->t", xi, gradients)
        estimate += xi[:, 0, :] * directional[:, None] / directions
    empirical = np.square(estimate - gradients[0]).sum(axis=1).mean()
    total_energy = np.square(gradients).sum()
    theoretical = (width * total_energy - np.square(gradients[0]).sum()) / directions
    return empirical, theoretical


def check_simultaneous_variance():
    print("SIMULTANEOUS RADEMACHER MSE")
    print("depth width K empirical theory ratio")
    cases = ((2, 8, 1), (4, 8, 1), (8, 8, 1),
             (4, 32, 1), (4, 32, 4), (4, 32, 16))
    for index, (depth, width, directions) in enumerate(cases):
        empirical, theoretical = simultaneous_mse(
            np.random.default_rng(100 + index), depth, width, directions)
        ratio = empirical / theoretical
        print(f"{depth:5d} {width:5d} {directions:2d} "
              f"{empirical:9.4f} {theoretical:9.4f} {ratio:6.3f}")
        assert abs(ratio - 1.0) < 0.06


def cubic_loss(z, linear, cubic):
    return linear @ z + cubic * np.power(z, 3).sum() / 6.0


def check_sigma_bias():
    # Exhaustive directions remove Monte Carlo error, exposing only the
    # centered-difference bias. For this separable cubic it is exactly O(sigma^2).
    dim = 8
    z = np.linspace(-0.4, 0.5, dim)
    linear = np.linspace(0.2, 0.9, dim)
    cubic = 1.7
    true_gradient = linear + 0.5 * cubic * np.square(z)
    directions = np.asarray(list(itertools.product((-1.0, 1.0), repeat=dim)))
    sigmas = np.asarray((0.005, 0.01, 0.02, 0.04))
    biases = []
    print("\nCENTERED FINITE-DIFFERENCE BIAS")
    print("sigma bias_norm")
    for sigma in sigmas:
        estimates = []
        for xi in directions:
            coefficient = (cubic_loss(z + sigma * xi, linear, cubic)
                           - cubic_loss(z - sigma * xi, linear, cubic)) / (2.0 * sigma)
            estimates.append(xi * coefficient)
        bias = np.linalg.norm(np.mean(estimates, axis=0) - true_gradient)
        biases.append(bias)
        print(f"{sigma:5.3f} {bias:.9e}")
    slope = np.polyfit(np.log(sigmas), np.log(biases), 1)[0]
    print(f"log-log slope={slope:.6f}")
    assert abs(slope - 2.0) < 1e-6


def check_descent_threshold():
    rng = np.random.default_rng(301)
    dim = 20
    gradient = rng.normal(size=dim)
    orthogonal = rng.normal(size=dim)
    orthogonal -= orthogonal.dot(gradient) * gradient / gradient.dot(gradient)
    orthogonal /= np.linalg.norm(orthogonal)
    neg_gradient = -gradient / np.linalg.norm(gradient)
    cosine = 0.35
    update = cosine * neg_gradient + math.sqrt(1.0 - cosine ** 2) * orthogonal
    beta = 2.3
    threshold = 2.0 * cosine * np.linalg.norm(gradient) / (beta * np.linalg.norm(update))

    def quadratic_change(step):
        return step * gradient.dot(update) + 0.5 * beta * step ** 2 * update.dot(update)

    below = quadratic_change(0.9 * threshold)
    above = quadratic_change(1.1 * threshold)
    print("\nSMOOTH DESCENT THRESHOLD")
    print(f"cos={cosine:.3f} eta_max={threshold:.6f} "
          f"delta_below={below:+.6f} delta_above={above:+.6f}")
    assert below < 0 < above


def check_predictor_timescale():
    rng = np.random.default_rng(401)
    examples = 50000
    cells = 12
    h = rng.normal(size=(examples, cells))
    slopes = np.exp(rng.normal(scale=0.25, size=cells))
    traffic = h * slopes
    print("\nNEUTRAL PREDICTOR TIMESCALE")
    print("eta steps residual_power empirical theory")
    for eta in (0.002, 0.01, 0.05):
        for steps in (0, 20, 100, 500):
            predictor = slopes * (1.0 - (1.0 - eta) ** steps)
            residual = traffic - h * predictor
            empirical = np.square(residual).mean() / np.square(traffic).mean()
            theoretical = (1.0 - eta) ** (2 * steps)
            print(f"{eta:5.3f} {steps:5d} {empirical:14.8f} {theoretical:14.8f}")
            assert abs(empirical - theoretical) < 1e-12


def check_conditional_projection():
    """Finite-sample L2 projection realizes the innovation identities exactly."""
    rng = np.random.default_rng(351)
    examples = 4096
    h = rng.uniform(-1.5, 1.5, size=(examples, 2))
    basis = np.column_stack((
        np.ones(examples), h[:, 0], h[:, 1], np.square(h[:, 0]),
        np.square(h[:, 1]), h[:, 0] * h[:, 1], np.sin(h[:, 0]),
        np.cos(h[:, 1])))
    coefficients = rng.normal(size=(basis.shape[1], 3))
    conditional_mean = basis @ coefficients
    innovation = rng.normal(scale=0.4, size=conditional_mean.shape)
    # Project the finite-sample noise off every soma-measurable basis vector.
    innovation -= basis @ np.linalg.lstsq(basis, innovation, rcond=None)[0]
    nuisance = conditional_mean + innovation

    linear_basis = basis[:, :3]
    restricted = linear_basis @ np.linalg.lstsq(
        linear_basis, nuisance, rcond=None)[0]
    lhs = np.square(nuisance - restricted).mean()
    irreducible = np.square(innovation).mean()
    approximation = np.square(conditional_mean - restricted).mean()
    orthogonality = np.abs(basis.T @ innovation / examples).max()

    teaching = rng.normal(size=conditional_mean.shape)
    raw = teaching + nuisance
    residual = teaching + innovation
    alpha = (np.linalg.norm(residual, axis=1)
             / np.linalg.norm(raw, axis=1).clip(min=1e-12))
    matched = alpha[:, None] * raw

    def row_cosine(left, right):
        numerator = np.sum(left * right, axis=1)
        denominator = (np.linalg.norm(left, axis=1)
                       * np.linalg.norm(right, axis=1)).clip(min=1e-12)
        return numerator / denominator

    direction_difference = np.abs(
        row_cosine(raw, teaching) - row_cosine(matched, teaching)).max()
    norm_difference = np.abs(
        np.linalg.norm(matched, axis=1) - np.linalg.norm(residual, axis=1)).max()
    print("\nCONDITIONAL INNOVATION PROJECTION")
    print(f"orthogonality={orthogonality:.3e} pythagorean_error="
          f"{abs(lhs - irreducible - approximation):.3e}")
    print(f"norm_match_error={norm_difference:.3e} "
          f"direction_change={direction_difference:.3e}")
    assert orthogonality < 2e-14
    assert abs(lhs - irreducible - approximation) < 2e-14
    assert norm_difference < 2e-14
    assert direction_difference < 2e-14


def squared_cosine(x, y):
    return float((x.ravel() @ y.ravel()) ** 2
                 / ((x.ravel() @ x.ravel()) * (y.ravel() @ y.ravel())))


def check_innovation_identification():
    rng = np.random.default_rng(501)
    examples = 100000
    cells = 8
    h = rng.normal(size=(examples, cells))
    traffic_slopes = np.linspace(0.8, 1.5, cells)
    teaching_slopes = np.linspace(0.3, 0.7, cells)
    innovation = rng.normal(scale=0.6, size=(examples, cells))
    noise = rng.normal(scale=0.15, size=(examples, cells))
    teaching = h * teaching_slopes + innovation
    traffic = h * traffic_slopes
    apical = teaching + traffic + noise

    neutral_residual = apical - traffic
    task_coeff = (apical * h).mean(axis=0) / np.square(h).mean(axis=0)
    task_residual = apical - h * task_coeff
    raw_alignment = squared_cosine(apical, teaching)
    neutral_alignment = squared_cosine(neutral_residual, teaching)
    task_alignment = squared_cosine(task_residual, teaching)
    predictable_fraction = np.square(h * teaching_slopes).mean() / np.square(teaching).mean()
    retained_fraction = np.square(teaching - h * teaching_slopes).mean() / np.square(teaching).mean()
    print("\nINNOVATION IDENTIFICATION")
    print(f"squared cosine raw={raw_alignment:.4f} neutral={neutral_alignment:.4f} "
          f"task_fit={task_alignment:.4f}")
    print(f"teaching predictable={predictable_fraction:.4f} "
          f"retained_after_task_fit={retained_fraction:.4f}")
    assert neutral_alignment > raw_alignment + 0.15
    assert task_alignment < neutral_alignment - 0.15
    assert abs(retained_fraction - (1.0 - predictable_fraction)) < 0.01


def check_residual_coupling_instability():
    rng = np.random.default_rng(551)
    rows, columns = 4, 3
    residual_coupling = rng.normal(size=(rows, rows))
    weight = rng.normal(size=(rows, columns))
    input_covariance = rng.normal(size=(columns, columns))
    input_covariance = input_covariance @ input_covariance.T / columns
    direct = residual_coupling @ weight @ input_covariance
    operator = np.kron(input_covariance.T, residual_coupling)
    vectorized = (operator @ weight.ravel(order="F")).reshape(
        weight.shape, order="F")
    identity_error = float(np.abs(direct - vectorized).max())

    eta = 0.1
    momentum = 0.9

    def radius(k):
        transition = np.asarray((
            (1.0 + eta * k, eta * momentum),
            (k, momentum)))
        return float(np.abs(np.linalg.eigvals(transition)).max())

    positive_k = 0.02
    negative_k = -0.20
    too_negative_k = -40.0
    positive_radius = radius(positive_k)
    negative_radius = radius(negative_k)
    too_negative_radius = radius(too_negative_k)
    polynomial_at_one = -eta * positive_k
    print("\nMULTIPLICATIVE RESIDUAL COUPLING")
    print(f"vectorization_error={identity_error:.3e} "
          f"rho(k={positive_k:+.3f})={positive_radius:.9f} "
          f"rho(k={negative_k:+.3f})={negative_radius:.9f} "
          f"rho(k={too_negative_k:+.1f})={too_negative_radius:.9f} "
          f"p_positive(1)={polynomial_at_one:+.3e}")
    assert identity_error < 2e-15
    assert polynomial_at_one < 0.0
    assert positive_radius > 1.0
    assert negative_radius < 1.0
    assert too_negative_k < -2.0 * (1.0 + momentum) / eta
    assert too_negative_radius > 1.0


def check_dynamic_neutral_projection():
    """A paired neutral fit nulls affine coupling across covariance scales."""
    rng = np.random.default_rng(571)
    examples = 128
    cells = 64
    soma = rng.normal(size=(examples, cells))
    coupling = rng.normal(scale=0.04, size=cells)
    offset = rng.normal(scale=0.01, size=cells)
    neutral_residual = soma * coupling + offset
    centered_soma = soma - soma.mean(axis=0)
    centered_residual = neutral_residual - neutral_residual.mean(axis=0)
    variance = np.square(centered_soma).mean(axis=0)
    correction = (centered_soma * centered_residual).mean(axis=0) / variance
    remainder = centered_residual - correction * centered_soma
    remainder_ratio = (np.linalg.norm(remainder)
                       / np.linalg.norm(neutral_residual))
    post_slope = ((centered_soma * remainder).mean(axis=0) / variance)
    maximum_post_slope = float(np.abs(post_slope).max())

    eta = 0.1
    momentum = 0.9
    decay = 1e-4

    def radius(k):
        transition = np.asarray((
            (1.0 + eta * k, eta * momentum),
            (k, momentum)))
        return float(np.abs(np.linalg.eigvals(transition)).max())

    # With zero residual coupling, covariance drops out and only decay remains.
    null_radii = [radius(-decay) for _ in (1e-3, 1.0, 1e3, 1e6)]
    # A fixed negative coefficient eventually crosses the lower Jury boundary
    # as the nonnegative covariance eigenvalue grows.
    fixed_coefficient = -0.03
    high_covariance_k = fixed_coefficient * 1e6 - decay
    high_covariance_radius = radius(high_covariance_k)
    print("\nDYNAMIC NEUTRAL PROJECTION")
    print(f"remainder_ratio={remainder_ratio:.3e} "
          f"post_slope={maximum_post_slope:.3e} "
          f"rho_null={max(null_radii):.9f} "
          f"rho_fixed_high_cov={high_covariance_radius:.3f}")
    assert remainder_ratio < 2e-15
    assert maximum_post_slope < 2e-16
    assert max(null_radii) < 1.0
    assert high_covariance_k < -2.0 * (1.0 + momentum) / eta
    assert high_covariance_radius > 1.0


def check_intermittent_feedback_tracking():
    eta_m = 0.1
    cadence = 16
    per_step_change = 0.002
    block_change = cadence * per_step_change

    # Constant collinear forward-weight motion attains the norm bound.
    error = 0.0
    for _ in range(1000):
        error = (1.0 - eta_m) * error - block_change
    steady_state = -block_change / eta_m
    assert abs(error - steady_state) < 1e-14

    # Once forward motion stops, endpoint tracking can look excellent despite
    # the large error accumulated during the task-active trajectory.
    error_before_quiet_tail = abs(error)
    for _ in range(80):
        error *= 1.0 - eta_m
    expected_endpoint = error_before_quiet_tail * (1.0 - eta_m) ** 80
    print("\nINTERMITTENT FEEDBACK TRACKING")
    print(f"steady_error={abs(steady_state):.6f} bound={cadence * per_step_change / eta_m:.6f} "
          f"endpoint_after_quiet_tail={abs(error):.9e}")
    assert abs(abs(error) - expected_endpoint) < 1e-15


def check_kolen_pollack_difference_dynamics():
    rng = np.random.default_rng(601)
    eta = 0.07
    momentum = 0.9
    decay = 1e-3
    forward = rng.normal(size=(7, 5))
    reciprocal = rng.normal(size=(7, 5))
    forward_momentum = rng.normal(size=(7, 5))
    reciprocal_momentum = rng.normal(size=(7, 5))
    maximum_error = 0.0
    for _ in range(40):
        # Both paths independently obtain the same local activity product.
        correlation = rng.normal(size=(7, 5))
        difference = reciprocal - forward
        momentum_difference = reciprocal_momentum - forward_momentum
        predicted_momentum_difference = (
            momentum * momentum_difference - decay * difference)
        predicted_difference = difference + eta * predicted_momentum_difference

        forward_momentum = (momentum * forward_momentum
                            + correlation - decay * forward)
        reciprocal_momentum = (momentum * reciprocal_momentum
                               + correlation - decay * reciprocal)
        forward = forward + eta * forward_momentum
        reciprocal = reciprocal + eta * reciprocal_momentum
        maximum_error = max(
            maximum_error,
            float(np.abs((reciprocal_momentum - forward_momentum)
                         - predicted_momentum_difference).max()),
            float(np.abs((reciprocal - forward)
                         - predicted_difference).max()))
    print("\nKOLEN-POLLACK DIFFERENCE DYNAMICS")
    print(f"task-correlation cancellation error={maximum_error:.3e}")
    assert maximum_error < 3e-15


def check_bci_error_derivative_structure():
    """Separate causal role from temporal performance innovation."""
    rng = np.random.default_rng(641)
    role = np.concatenate((np.full(5, 0.2), np.full(5, -0.2),
                           np.zeros(30)))
    examples = 200_000
    xi = 2.0 * rng.integers(0, 2, size=(examples, role.size)) - 1.0
    cursor_direction = xi @ role
    role_estimate = (cursor_direction[:, None] * xi).mean(axis=0)
    role_error = float(np.abs(role_estimate - role).max())
    role_cosine = float(role_estimate @ role / (
        np.linalg.norm(role_estimate) * np.linalg.norm(role)))

    # Positive target error throughout the BCI trial. Improving events have
    # smaller current error; worsening events have larger current error.
    previous = np.asarray((0.8, 0.7, 0.6, 0.5, 0.4, 0.3))
    current = np.asarray((0.6, 0.5, 0.4, 0.7, 0.6, 0.5))
    improvement = previous - current
    improving = improvement > 0
    worsening = improvement < 0

    def sign_index(modulator):
        residual = modulator[:, None] * role[None, :]
        difference = residual[:, :5].mean(1) - residual[:, 5:10].mean(1)
        return 0.5 * (difference[improving].mean()
                      - difference[worsening].mean())

    instantaneous_index = float(sign_index(current))
    temporal_difference_index = float(sign_index(improvement))
    print("\nBCI ERROR-DERIVATIVE STRUCTURE")
    print(f"role_estimator_cosine={role_cosine:.9f} "
          f"max_error={role_error:.3e} "
          f"instantaneous_index={instantaneous_index:+.6f} "
          f"td_index={temporal_difference_index:+.6f}")
    assert role_cosine > 0.999
    assert role_error < 0.005
    assert instantaneous_index < 0.0
    assert temporal_difference_index > 0.0


def main():
    check_simultaneous_variance()
    check_sigma_bias()
    check_descent_threshold()
    check_conditional_projection()
    check_predictor_timescale()
    check_innovation_identification()
    check_residual_coupling_instability()
    check_dynamic_neutral_projection()
    check_intermittent_feedback_tracking()
    check_kolen_pollack_difference_dynamics()
    check_bci_error_derivative_structure()
    print("\nALL THEORY CHECKS PASSED")


if __name__ == "__main__":
    main()