1import json
 2import numpy as np
 3from sklearn.datasets import make_moons
 4from sklearn.neural_network import MLPClassifier
 5from sklearn.metrics import accuracy_score, log_loss
 6from convex_transport import QuadraticConvexTransport, change_of_variables_kl, standard_normal_logpdf
 7
 8SEED = 73
 9rng = np.random.default_rng(SEED)
10
11
12def exact_kl(d, s):
13    return 0.5 * d * (s*s - 1.0 - 2.0*np.log(s))
14
15
16def verify_math():
17    # Prediction 1: identity is the unique zero-KL member of this isotropic family.
18    scales = np.array([0.75, 0.9, 1.0, 1.1, 1.3])
19    identity_kls = np.array([exact_kl(3, s) for s in scales])
20    # Prediction 2: KL ~= d*(s-1)^2 around identity (Taylor expansion).
21    eps = np.array([0.005, 0.01, 0.02, 0.04])
22    small_kl = np.array([exact_kl(3, 1+e) for e in eps])
23    quad_ratio = small_kl / (3 * eps**2)
24    # Prediction 3: at fixed scale, KL is exactly linear in dimension.
25    dims = np.array([1, 2, 4, 8, 16])
26    dim_kls = np.array([exact_kl(int(d), 1.2) for d in dims])
27    dim_slope = np.polyfit(dims, dim_kls, 1)[0]
28    expected_slope = exact_kl(1, 1.2)
29    # Direct numerical change-of-variables check in 2D.
30    x = rng.normal(size=(300000, 2))
31    A = np.diag([1.15, 0.85])
32    t = QuadraticConvexTransport(A)
33    mc = change_of_variables_kl(x, t, standard_normal_logpdf)
34    analytic = 0.5 * np.sum(np.diag(A)**2 - 1 - 2*np.log(np.diag(A)))
35    return {
36        'identity_scales': scales.tolist(),
37        'identity_family_kl': identity_kls.tolist(),
38        'identity_min_at_scale_1': bool(np.argmin(identity_kls) == 2 and identity_kls[2] < 1e-12),
39        'small_eps': eps.tolist(),
40        'small_kl_over_d_eps2': quad_ratio.tolist(),
41        'quadratic_prediction_mean_ratio': 1.0,
42        'quadratic_observed_mean_ratio': float(np.mean(quad_ratio)),
43        'dimension_values': dims.tolist(),
44        'dimension_kls': dim_kls.tolist(),
45        'dimension_predicted_slope': float(expected_slope),
46        'dimension_observed_slope': float(dim_slope),
47        'cov_kl_analytic': float(analytic),
48        'cov_kl_monte_carlo': float(mc),
49        'cov_relative_error': float(abs(mc-analytic)/analytic),
50    }
51
52
53def train_moons():
54    X, y = make_moons(n_samples=1200, noise=0.20, random_state=SEED)
55    # fixed held-out shift: radial convex transport plus modest noise
56    tr, te = np.arange(900), np.arange(900, 1200)
57    # make_moons is ordered randomly by generator; use deterministic split
58    Xtr, ytr = X[tr], y[tr]
59    Xte, yte = X[te], y[te]
60    shift = QuadraticConvexTransport(np.eye(2)*1.18)
61    Xshift = shift.map(Xte) + rng.normal(0, .08, Xte.shape)
62    rho = exact_kl(2, 1.18)
63    # Use identical small MLP and data budget. Convex augmentation is a class-agnostic
64    # strongly convex gradient map, selected on the KL boundary.
65    configs = {
66        'ERM': Xtr,
67        'Gaussian': Xtr + rng.normal(0, .12, Xtr.shape),
68        'ConvexTransport': shift.map(Xtr),
69    }
70    out = {'rho': float(rho), 'n_train': len(tr)}
71    for name, Xa in configs.items():
72        model = MLPClassifier(hidden_layer_sizes=(32, 32), activation='tanh', solver='adam',
73                              alpha=1e-3, batch_size=64, learning_rate_init=2e-3,
74                              max_iter=350, random_state=SEED, early_stopping=False)
75        model.fit(Xa, ytr)
76        out[name] = {
77            'clean_accuracy': float(accuracy_score(yte, model.predict(Xte))),
78            'shift_accuracy': float(accuracy_score(yte, model.predict(Xshift))),
79            'clean_logloss': float(log_loss(yte, model.predict_proba(Xte))),
80        }
81    return out
82
83
84if __name__ == '__main__':
85    result = {'math_verification': verify_math(), 'mini_experiment': train_moons()}
86    with open('results.json', 'w') as f:
87        json.dump(result, f, indent=2)
88    print(json.dumps(result, indent=2))