Branch-Length Polynomial Fingerprint / experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
 1import numpy as np
 2from sklearn.linear_model import LogisticRegression
 3from sklearn.pipeline import make_pipeline
 4from sklearn.preprocessing import StandardScaler
 5from sklearn.model_selection import train_test_split
 6from sklearn.metrics import accuracy_score
 7
 8# A rooted tree is a tuple of (edge_length, child); a leaf is ().
 9def poly_eval(tree, ys, x=1.0):
10    ys = np.asarray(ys, dtype=float)
11    if len(tree) == 0:
12        return np.full(len(ys), x, dtype=float)
13    out = np.ones(len(ys), dtype=float)
14    for w, child in tree:
15        out *= w * ys + poly_eval(child, ys, x)
16    return out
17
18def ordinary_sum(tree):
19    # Standard scalar sum message: edge lengths plus terminal messages.
20    if len(tree) == 0:
21        return 1.0
22    return sum(float(w) + ordinary_sum(c) for w, c in tree)
23
24def count_edges(tree):
25    return sum(1 + count_edges(c) for _, c in tree)
26
27def fingerprint(tree, ys):
28    z = poly_eval(tree, ys)
29    return np.log(np.abs(z) + 1e-12)
30
31def make_a(vals):
32    # root -> two leaves and one internal node -> two leaves
33    a,b,c,d,e = vals
34    return ((a, ()), (b, ()), (c, ((d, ()), (e, ()))))
35
36def make_b(vals):
37    # root -> one leaf and one internal node -> three leaves
38    a,b,c,d,e = vals
39    return ((a, ()), (b, ((c, ()), (d, ()), (e, ()))))
40
41def main():
42    rng = np.random.default_rng(485)
43    ys = np.array([-3., -2., -1., -.5, .5, 1., 2., 3.])
44    heldout_y = np.array([-.8, -.3, .3, .8, 1.7, 2.4])
45
46    # Core claim sanity check: same ordinary sum, distinct polynomial values.
47    vals = np.array([.5, .8, 1.1, 1.4, 1.7])
48    ta, tb = make_a(vals), make_b(vals)
49    za, zb = poly_eval(ta, ys), poly_eval(tb, ys)
50    collision = np.isclose(ordinary_sum(ta), ordinary_sum(tb))
51    separation = float(np.max(np.abs(za - zb)))
52    assert count_edges(ta) == count_edges(tb) == 5
53    assert collision and separation > 1e-6
54
55    Xbase, Xfp, Xfp_hold, y = [], [], [], []
56    for label in (0, 1):
57        for _ in range(600):
58            # Both classes have the same edge count and leaf count, and use
59            # the same edge-length multiset: scalar sum is an exact collision.
60            v = rng.permutation(vals)
61            t = make_a(v) if label == 0 else make_b(v)
62            Xbase.append([ordinary_sum(t)])
63            Xfp.append(fingerprint(t, ys))
64            Xfp_hold.append(fingerprint(t, heldout_y))
65            y.append(label)
66    Xbase, Xfp, Xfp_hold, y = map(np.asarray, (Xbase, Xfp, Xfp_hold, y))
67    tr, te = train_test_split(np.arange(len(y)), test_size=.3, random_state=17, stratify=y)
68    base = make_pipeline(StandardScaler(), LogisticRegression(random_state=0, max_iter=1000))
69    fp = make_pipeline(StandardScaler(), LogisticRegression(random_state=0, max_iter=1000))
70    fp_hold = make_pipeline(StandardScaler(), LogisticRegression(random_state=0, max_iter=1000))
71    base.fit(Xbase[tr], y[tr]); fp.fit(Xfp[tr], y[tr]); fp_hold.fit(Xfp_hold[tr], y[tr])
72    acc_base = accuracy_score(y[te], base.predict(Xbase[te]))
73    acc_fp = accuracy_score(y[te], fp.predict(Xfp[te]))
74    acc_hold = accuracy_score(y[te], fp_hold.predict(Xfp_hold[te]))
75    print({
76        'ordinary_sum_collision': bool(collision),
77        'example_max_polynomial_difference': separation,
78        'baseline_accuracy': round(float(acc_base), 4),
79        'fingerprint_accuracy': round(float(acc_fp), 4),
80        'heldout_evaluation_grid_accuracy': round(float(acc_hold), 4),
81        'n_test': len(te), 'M': len(ys),
82    })
83
84if __name__ == '__main__':
85    main()