1import json
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 7
  6RNG = np.random.default_rng(SEED)
  7
  8
  9def paper_matrix_step(x1, x2, a, c, d, h, e, f, u):
 10    z = np.array([x1, x2, x1 * x1, 1.0])
 11    A = np.array([[a, 0, 0, 0], [0, c, d, h], [0, 0, a*a, 0], [0, 0, 0, 1.]])
 12    B = np.array([[0, 0, 0, 0], [e, 0, 0, f], [0, 0, 0, 0], [0, 0, 0, 0.]])
 13    matrix_next = A @ z + (B @ z) * u
 14    direct_next = np.array([a*x1, c*x2 + d*x1*x1 + h + (e*x1 + f)*u, (a*x1)**2, 1.0])
 15    return matrix_next, direct_next
 16
 17
 18def plant(x, u):
 19    return (0.72 + 0.58 * u) * x + 0.08 * u + 0.03
 20
 21
 22def make_data(n_sequences=160, horizon=24):
 23    rows, targets = [], []
 24    starts, controls = [], []
 25    for _ in range(n_sequences):
 26        x = RNG.uniform(-1.0, 1.0)
 27        starts.append(x)
 28        us = RNG.uniform(-1.0, 1.0, horizon)
 29        controls.append(us)
 30        for u in us:
 31            rows.append((x, u))
 32            x = plant(x, u)
 33            targets.append(x)
 34    return np.asarray(rows), np.asarray(targets), np.asarray(starts), np.asarray(controls)
 35
 36
 37def design(rows, kind):
 38    x, u = rows[:, 0], rows[:, 1]
 39    if kind == 'additive':
 40        return np.column_stack([x, u, np.ones_like(x)])
 41    if kind == 'bilinear':
 42        return np.column_stack([x, u, x*u, np.ones_like(x)])
 43    raise ValueError(kind)
 44
 45
 46def fit(rows, targets, kind):
 47    X = design(rows, kind)
 48    return np.linalg.lstsq(X, targets, rcond=None)[0]
 49
 50
 51def predict(x, u, coef, kind):
 52    return float((design(np.array([[x, u]]), kind) @ coef).item())
 53
 54
 55def evaluate(starts, controls, coef, kind):
 56    errors = []
 57    for x0, us in zip(starts, controls):
 58        x = x0
 59        for u in us:
 60            y = plant(x, u)
 61            x = predict(x, u, coef, kind)
 62            errors.append((x-y)**2)
 63    return float(np.mean(errors)), float(np.sqrt(np.mean(errors)))
 64
 65
 66def main():
 67    # Core math sanity check: matrix form equals the explicitly expanded update.
 68    diffs = []
 69    for _ in range(1000):
 70        vals = RNG.normal(size=9)
 71        diffs.append(np.max(np.abs(paper_matrix_step(*vals)[0] - paper_matrix_step(*vals)[1])))
 72    identity_error = max(diffs)
 73
 74    rows, targets, starts, controls = make_data()
 75    # Hold out complete trajectories, preserving a genuine rollout test.
 76    n = len(starts)
 77    split = int(0.75 * n)
 78    train_rows = []
 79    train_targets = []
 80    for i in range(split):
 81        x = starts[i]
 82        for u in controls[i]:
 83            train_rows.append((x, u))
 84            x = plant(x, u)
 85            train_targets.append(x)
 86    train_rows = np.asarray(train_rows)
 87    train_targets = np.asarray(train_targets)
 88    test_starts, test_controls = starts[split:], controls[split:]
 89
 90    results = {'identity_max_abs_error': identity_error, 'seed': SEED,
 91               'train_sequences': split, 'test_sequences': n-split, 'horizon': controls.shape[1]}
 92    for kind in ('additive', 'bilinear'):
 93        coef = fit(train_rows, train_targets, kind)
 94        one_step = float(np.mean((design(train_rows, kind) @ coef - train_targets)**2))
 95        mse, rmse = evaluate(test_starts, test_controls, coef, kind)
 96        results[kind] = {'one_step_mse': one_step, 'rollout_mse': mse,
 97                         'rollout_rmse': rmse, 'parameter_count': int(len(coef))}
 98        results[kind]['coefficients'] = coef.tolist()
 99
100    Path('results.json').write_text(json.dumps(results, indent=2))
101    print(json.dumps(results, indent=2))
102
103
104if __name__ == '__main__':
105    main()