1import json
  2from pathlib import Path
  3import numpy as np
  4from scipy.optimize import least_squares
  5
  6SEED = 1411
  7TRUE_A = 0.70
  8TRUE_B = 1.20
  9TRUE_TAU = 0.80
 10NOISE = 0.01
 11
 12
 13def disturbance(u, alpha, tau):
 14    return alpha * np.tanh(u / tau)
 15
 16
 17def make_data(alpha, n, u_limit, seed, tau=TRUE_TAU):
 18    rng = np.random.default_rng(seed)
 19    x = rng.uniform(-1.0, 1.0, n)
 20    u = rng.uniform(-u_limit, u_limit, n)
 21    y = TRUE_A * x + TRUE_B * (u + disturbance(u, alpha, tau))
 22    y += rng.normal(0.0, NOISE, n)
 23    return x, u, y
 24
 25
 26def fit_baseline(x, u, y):
 27    X = np.column_stack([x, u, np.ones_like(x)])
 28    p, *_ = np.linalg.lstsq(X, y, rcond=None)
 29    return p
 30
 31
 32def fit_joint(x, u, y, tau_hint=TRUE_TAU):
 33    def residual(p):
 34        a, b, alpha, tau, c = p
 35        pred = a * x + b * (u + disturbance(u, alpha, tau)) + c
 36        return (pred - y) / NOISE
 37    p0 = np.array([0.6, 1.1, 0.15, tau_hint, 0.0])
 38    lo = np.array([-3., -3., -2., 0.15, -1.])
 39    hi = np.array([3., 3., 2., 3.0, 1.])
 40    result = least_squares(residual, p0, bounds=(lo, hi), max_nfev=1000)
 41    return result.x
 42
 43
 44def mse(y, pred):
 45    return float(np.mean((y - pred) ** 2))
 46
 47
 48def main():
 49    # Prediction 1: max_u |d'(u)| = |alpha|/tau, attained at u=0.
 50    grid = np.linspace(-8, 8, 20001)
 51    alpha, tau = 0.37, 0.63
 52    analytic = alpha / tau / np.cosh(grid / tau) ** 2
 53    numeric = np.gradient(disturbance(grid, alpha, tau), grid)
 54    derivative_check = {
 55        "predicted_max": float(abs(alpha) / tau),
 56        "observed_max": float(np.max(np.abs(numeric))),
 57        "relative_error": float(abs(np.max(np.abs(numeric)) - abs(alpha) / tau) / (abs(alpha) / tau)),
 58        "bound_location": float(grid[np.argmax(np.abs(numeric))]),
 59        "analytic_numeric_max_gap": float(np.max(np.abs(analytic - numeric))),
 60    }
 61
 62    # Prediction 2: for small alpha, omitted-variable bias in the linear action
 63    # coefficient is approximately B*alpha*E[u*tanh(u/tau)]/E[u^2], hence linear in alpha.
 64    alphas = np.array([0.0, 0.05, 0.10, 0.20, 0.40, 0.80])
 65    rows = []
 66    for i, a in enumerate(alphas):
 67        x, u, y = make_data(a, 5000, 1.5, SEED + i)
 68        pb = fit_baseline(x, u, y)
 69        pjoint = fit_joint(x, u, y)
 70        # Test on a fresh broad command range to expose saturation extrapolation.
 71        xt, ut, yt = make_data(a, 3000, 2.2, SEED + 100 + i)
 72        pred_b = pb[0] * xt + pb[1] * ut + pb[2]
 73        pred_j = pjoint[0] * xt + pjoint[1] * (ut + disturbance(ut, pjoint[2], pjoint[3])) + pjoint[4]
 74        ratio = np.mean(u * np.tanh(u / TRUE_TAU)) / np.mean(u * u)
 75        predicted_bias = TRUE_B * a * ratio
 76        rows.append({
 77            "alpha": float(a),
 78            "predicted_linear_bias": float(predicted_bias),
 79            "observed_baseline_b_bias": float(pb[1] - TRUE_B),
 80            "joint_train_mse": mse(y, pjoint[0] * x + pjoint[1] * (u + disturbance(u, pjoint[2], pjoint[3])) + pjoint[4]),
 81            "baseline_cross_regime_mse": mse(yt, pred_b),
 82            "joint_cross_regime_mse": mse(yt, pred_j),
 83            "fitted_alpha": float(pjoint[2]),
 84            "fitted_tau": float(pjoint[3]),
 85        })
 86
 87    # Prediction 3: no actuator disturbance (alpha=0) removes the mechanism;
 88    # the two models should have comparable cross-regime error.
 89    zero_rows = []
 90    for i in range(8):
 91        x, u, y = make_data(0.0, 2000, 1.5, SEED + 300 + i)
 92        pb = fit_baseline(x, u, y)
 93        pj = fit_joint(x, u, y)
 94        xt, ut, yt = make_data(0.0, 2000, 2.2, SEED + 400 + i)
 95        zero_rows.append({
 96            "baseline_mse": mse(yt, pb[0] * xt + pb[1] * ut + pb[2]),
 97            "joint_mse": mse(yt, pj[0] * xt + pj[1] * (ut + disturbance(ut, pj[2], pj[3])) + pj[4]),
 98            "fitted_alpha": float(pj[2]),
 99        })
100
101    # Fit slopes through the origin for the small-alpha predicted/observed bias.
102    small = rows[:4]
103    alpha_small = np.array([r["alpha"] for r in small])
104    observed_bias_small = np.array([r["observed_baseline_b_bias"] for r in small])
105    predicted_bias_small = np.array([r["predicted_linear_bias"] for r in small])
106    observed_slope = float(np.dot(alpha_small, observed_bias_small) / np.dot(alpha_small, alpha_small))
107    predicted_slope = float(np.dot(alpha_small, predicted_bias_small) / np.dot(alpha_small, alpha_small))
108    zero_base = float(np.mean([r["baseline_mse"] for r in zero_rows]))
109    zero_joint = float(np.mean([r["joint_mse"] for r in zero_rows]))
110    output = {
111        "seed": SEED,
112        "derivative_bound_check": derivative_check,
113        "bias_scaling_prediction": {
114            "predicted_slope": predicted_slope,
115            "observed_slope": observed_slope,
116            "slope_ratio_observed_over_predicted": observed_slope / predicted_slope,
117            "rows": rows,
118        },
119        "zero_disturbance_prediction": {
120            "mean_baseline_mse": zero_base,
121            "mean_joint_mse": zero_joint,
122            "joint_over_baseline": zero_joint / zero_base,
123            "rows": zero_rows,
124        },
125        "interpretation": "Joint model should recover saturation and improve held-out command-regime error when alpha is nonzero; at alpha=0 no systematic benefit is expected.",
126    }
127    Path("results.json").write_text(json.dumps(output, indent=2))
128    print(json.dumps(output, indent=2))
129
130
131if __name__ == "__main__":
132    main()