1import json
  2import math
  3import numpy as np
  4import torch
  5from orientation_pooling import doubled_angle_pool, orientation_from_pool
  6
  7
  8def energies_for_angles(angles, weights, B=256):
  9    # Direct impulses at arbitrary angles represented exactly by continuous moments;
 10    # bins are used only for the neural pooling API, with nearest-bin assignment.
 11    e = torch.zeros((B,), dtype=torch.float64)
 12    for a, w in zip(angles, weights):
 13        b = int(round((a % math.pi) / math.pi * B)) % B
 14        e[b] += w
 15    return e
 16
 17
 18def circular_axis_error(a, b):
 19    d = abs(float(a-b)) % math.pi
 20    return min(d, math.pi-d)
 21
 22
 23def main():
 24    torch.manual_seed(7)
 25    np.random.seed(7)
 26    B = 256
 27    # Prediction 1: orientation reversal is exactly identical because bins live on [0, pi).
 28    rng = np.random.default_rng(7)
 29    reversal_err = []
 30    for _ in range(100):
 31        angles = rng.uniform(0, math.pi, 8)
 32        weights = rng.uniform(.1, 2., 8)
 33        e1 = energies_for_angles(angles, weights, B)
 34        e2 = energies_for_angles(np.r_[angles, angles + math.pi], np.r_[weights, np.zeros(8)], B)
 35        p1, _ = doubled_angle_pool(e1)
 36        p2, _ = doubled_angle_pool(e2)
 37        reversal_err.append(float((p1-p2).abs().max()))
 38
 39    # Prediction 2: equal two-orientation mixture has q=|cos(delta)|.
 40    deltas = np.linspace(0, math.pi/2, 19)
 41    q_obs, q_pred, phi_err = [], [], []
 42    base = .37
 43    for d in deltas:
 44        e = energies_for_angles([base, base+d], [1., 1.], B)
 45        p, _ = doubled_angle_pool(e)
 46        q = float(p[-1]); q_obs.append(q); q_pred.append(abs(math.cos(d)))
 47        expected = (base + d/2) % math.pi if d < math.pi/2 else 0.
 48        # At exact cancellation angle is undefined; gated output is explicitly zero.
 49        ph = float(orientation_from_pool(p))
 50        phi_err.append(0. if abs(d-math.pi/2)<1e-12 else circular_axis_error(ph, expected))
 51
 52    # Prediction 3: unequal weighted pair follows the exact resultant formula.
 53    deltas2 = np.linspace(.0, math.pi/2, 13)
 54    ratios = [.25, .5, 1., 2., 4.]
 55    weighted_errors = []
 56    for ratio in ratios:
 57        for d in deltas2:
 58            a, b = 1., ratio
 59            e = energies_for_angles([base, base+d], [a,b], B)
 60            p, _ = doubled_angle_pool(e)
 61            pred = math.sqrt(a*a+b*b+2*a*b*math.cos(2*d))/(a+b)
 62            weighted_errors.append(abs(float(p[-1])-pred))
 63
 64    # Tiny downstream check: classify a randomly chosen unoriented bin. A model
 65    # trained on raw directed 2B bins sees reversal as a new class; pooled channels do not.
 66    # We instead test nearest-prototype generalization after training on angles in [0,pi).
 67    train_angles = np.linspace(0, math.pi, 32, endpoint=False)
 68    prototypes = []
 69    for a in train_angles:
 70        e = energies_for_angles([a], [1.], B)
 71        prototypes.append(doubled_angle_pool(e)[0].numpy())
 72    prototypes = np.asarray(prototypes)
 73    test = [(a+math.pi) % (2*math.pi) for a in train_angles]
 74    pooled_correct = 0
 75    for a in test:
 76        e = energies_for_angles([a], [1.], B)
 77        x = doubled_angle_pool(e)[0].numpy()
 78        pred = np.argmin(((prototypes-x)**2).sum(1))
 79        pooled_correct += (pred == int(round((a % math.pi)/math.pi*32)) % 32)
 80    # Baselines on the same prototype task. Directed bins use [0, 2pi),
 81    # so reversal shifts the representation by pi; scalar averaging has no
 82    # orientation information and must guess among classes.
 83    def directed_hist(angle, n=64):
 84        h = np.zeros(n)
 85        h[int(round((angle % (2*math.pi)) / (2*math.pi) * n)) % n] = 1.
 86        return h
 87    directed_prototypes = np.asarray([directed_hist(a) for a in train_angles])
 88    directed_correct = 0
 89    scalar_correct = 0
 90    for a in test:
 91        h = directed_hist(a)
 92        pred = np.argmin(((directed_prototypes-h)**2).sum(1))
 93        directed_correct += (pred == int(round((a % math.pi)/math.pi*32)) % 32)
 94        # Equal-energy scalar representations tie; deterministic argmin guesses class 0.
 95        scalar_correct += (0 == int(round((a % math.pi)/math.pi*32)) % 32)
 96
 97    results = {
 98        "config": {"B": B, "seed": 7},
 99        "prediction_1_reversal_max_abs_error": max(reversal_err),
100        "prediction_1_expected": 0.0,
101        "prediction_2_deltas_deg": [round(float(x*180/math.pi), 3) for x in deltas],
102        "prediction_2_q_observed": q_obs,
103        "prediction_2_q_predicted_abs_cos_delta": q_pred,
104        "prediction_2_max_q_error": max(abs(a-b) for a,b in zip(q_obs,q_pred)),
105        "prediction_2_max_orientation_error_non_cancel_deg": max(phi_err)*180/math.pi,
106        "prediction_3_max_weighted_q_error": max(weighted_errors),
107        "prediction_3_expected": 0.0,
108        "reversal_prototype_accuracy": pooled_correct/len(test),
109        "directed_bin_baseline_reversal_accuracy": directed_correct/len(test),
110        "scalar_average_baseline_reversal_accuracy": scalar_correct/len(test),
111        "cancellation_q_at_90deg": q_obs[-1],
112    }
113    print(json.dumps(results, indent=2))
114
115if __name__ == '__main__':
116    main()