1import json, math, random
 2import numpy as np
 3import torch
 4
 5SEED = 137
 6random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 7
 8def pce_components(teacher, student, sigma=0.35, eps=1e-8):
 9    # teacher/student: [n,2] birth-death coordinates
10    dist2 = ((teacher[:, None, :] - student[None, :, :]) ** 2).sum(-1)
11    sim = torch.exp(-dist2 / (2 * sigma * sigma))
12    a = sim.max(dim=1).values
13    persistence = (teacher[:, 1] - teacher[:, 0]).clamp_min(0)
14    p = persistence / (persistence.sum() + eps)
15    explained = p * a
16    unexplained = 1.0 - explained.sum()
17    loss = -(p * torch.log(p * a + eps)).sum()
18    entropy = -(p * torch.log(p + eps)).sum()
19    return loss, entropy, unexplained, a, p
20
21def chamfer_loss(teacher, student):
22    d = torch.cdist(teacher, student).pow(2)
23    return d.min(1).values.mean() + d.min(0).values.mean()
24
25def run(method, teacher, init, steps=500, lr=0.04):
26    student = torch.nn.Parameter(init.clone())
27    opt = torch.optim.Adam([student], lr=lr)
28    history=[]
29    for step in range(steps):
30        opt.zero_grad()
31        if method == 'pce':
32            loss, _, _, _, _ = pce_components(teacher, student)
33            # The entropy term is constant, but retaining it matches the proposed excess objective.
34            objective = loss
35        else:
36            objective = chamfer_loss(teacher, student)
37        objective.backward(); opt.step()
38        if step in (0, 99, 249, 499):
39            with torch.no_grad():
40                pc = pce_components(teacher, student)
41                history.append((step + 1, float(pc[2]), float(pc[0]), float(objective)))
42    with torch.no_grad():
43        loss, ent, unexplained, a, p = pce_components(teacher, student)
44    return student.detach(), history, (loss.item(), ent.item(), unexplained.item(), a.numpy(), p.numpy())
45
46def main():
47    # Four teacher events with unequal importance and only two student events.
48    teacher = torch.tensor([[0.10,0.90], [0.18,0.52], [0.62,0.96], [0.38,0.48]], dtype=torch.float32)
49    init = torch.tensor([[0.25,0.25], [0.75,0.75]], dtype=torch.float32)
50
51    # Core numerical verification: valid measure and excess identity for random diagrams.
52    torch.manual_seed(11)
53    T = torch.rand(7,2); T[:,1] = T[:,0] + torch.rand(7)
54    S = torch.rand(3,2); S[:,1] = S[:,0] + torch.rand(3)
55    l,e,u,a,p = pce_components(T,S,sigma=.27)
56    identity_error = abs((l-e).item()) # overwritten below; calculate stated excess exactly
57    excess_error = abs((l-e).item() - (-(p*torch.log(a+1e-8)).sum()).item())
58    validity = float((p >= 0).all() and (a >= 0).all() and (a <= 1+1e-6).all() and (u >= -1e-6).item() and (u <= 1+1e-6).item())
59
60    pce_student, pce_hist, pce_final = run('pce', teacher, init)
61    cham_student, cham_hist, cham_final = run('chamfer', teacher, init)
62    # Repeated starts test whether the directional signal survives initialization.
63    repeat = []
64    for seed in (3, 19, 71):
65        g = torch.Generator().manual_seed(seed)
66        start = torch.rand((2, 2), generator=g) * 0.8 + 0.1
67        ps = run('pce', teacher, start)[2]
68        cs = run('chamfer', teacher, start)[2]
69        repeat.append({'seed': seed, 'pce_unexplained': ps[2],
70                       'chamfer_unexplained': cs[2],
71                       'pce_excess': ps[0]-ps[1],
72                       'chamfer_excess': cs[0]-cs[1]})
73    out = {
74      'seed': SEED, 'teacher_points': teacher.tolist(), 'student_count': 2,
75      'repeated_starts': repeat,
76      'core_check': {'probability_valid': bool(validity), 'excess_identity_abs_error': excess_error,
77                     'unexplained_mass': u.item(), 'p_sum': p.sum().item(), 'explained_sum': (p*a).sum().item()},
78      'pce': {'student_points': pce_student.tolist(), 'trajectory': pce_hist,
79              'final_loss': pce_final[0], 'teacher_entropy': pce_final[1],
80              'final_excess': pce_final[0]-pce_final[1], 'unexplained_mass': pce_final[2],
81              'explanation_scores': pce_final[3].tolist()},
82      'chamfer_control': {'student_points': cham_student.tolist(), 'trajectory': cham_hist,
83              'pce_loss_evaluated': cham_final[0], 'teacher_entropy': cham_final[1],
84              'final_excess': cham_final[0]-cham_final[1], 'unexplained_mass': cham_final[2],
85              'explanation_scores': cham_final[3].tolist()},
86    }
87    with open('results.json','w') as f: json.dump(out,f,indent=2)
88    print(json.dumps(out,indent=2))
89
90if __name__ == '__main__': main()