1import json
  2import time
  3import numpy as np
  4
  5
  6def make_problem(seed=7, n=96, d=4, h=5):
  7    rng = np.random.default_rng(seed)
  8    x = rng.normal(size=(n, d))
  9    w1 = rng.normal(scale=0.8, size=(d, h))
 10    b1 = rng.normal(scale=0.2, size=h)
 11    w2 = rng.normal(scale=0.8, size=(h, 1))
 12    b2 = rng.normal(scale=0.2, size=1)
 13    y = np.tanh(x @ w1 + b1) @ w2 + b2
 14    y += 0.025 * rng.normal(size=(n, 1))
 15    return x, y
 16
 17
 18def unpack(q, d, h):
 19    k = 0
 20    w1 = q[k:k+d*h].reshape(d, h); k += d*h
 21    b1 = q[k:k+h]; k += h
 22    w2 = q[k:k+h].reshape(h, 1); k += h
 23    b2 = q[k:k+1]
 24    return w1, b1, w2, b2
 25
 26
 27def residual_jacobian(q, x, y, d, h):
 28    w1, b1, w2, b2 = unpack(q, d, h)
 29    z = np.tanh(x @ w1 + b1)
 30    r = (z @ w2 + b2 - y).ravel()
 31    n = x.shape[0]
 32    p = q.size
 33    j = np.zeros((n, p))
 34    a = 0
 35    b = d*h
 36    c = b+h
 37    e = c+h
 38    for i in range(n):
 39        j[i, a:b] = np.outer(x[i], (1-z[i]**2) * w2.ravel()).ravel()
 40        j[i, b:c] = (1-z[i]**2) * w2.ravel()
 41        j[i, c:e] = z[i]
 42        j[i, e] = 1.0
 43    return r, j
 44
 45
 46def state(q, x, y, d, h):
 47    r, j = residual_jacobian(q, x, y, d, h)
 48    f = 0.5 * float(r @ r)
 49    g = j.T @ r
 50    hess = j.T @ j
 51    return f, g, hess
 52
 53
 54def solve_direction(g, hess, lam):
 55    return np.linalg.solve(hess + lam*np.eye(g.size), -g)
 56
 57
 58def armijo_lm(q0, x, y, d, h, steps=35):
 59    q = q0.copy(); lam = 1e-2
 60    beta, c = 0.5, 1e-4
 61    low, high, up, down = 0.25, 0.75, 4.0, 2.0
 62    solves = rejects = accepted = trials = 0
 63    t0 = time.perf_counter()
 64    for _ in range(steps):
 65        f, g, H = state(q, x, y, d, h)
 66        p = solve_direction(g, H, lam); solves += 1
 67        pred = max(1e-14, -g @ p - 0.5*p @ H @ p)
 68        alpha = 1.0; accepted_here = False
 69        while alpha >= 1/32 - 1e-15:
 70            trials += 1
 71            ft = state(q + alpha*p, x, y, d, h)[0]
 72            if ft <= f + c*alpha*(g @ p):
 73                act = f-ft; rho = act/pred
 74                q = q + alpha*p; accepted += 1; accepted_here = True
 75                if rho < low: lam *= up
 76                elif rho > high: lam /= down
 77                break
 78            rejects += 1; alpha *= beta
 79        if not accepted_here:
 80            lam *= up
 81    final = state(q, x, y, d, h)[0]
 82    return dict(final_loss=final, solves=solves, rejected_trials=rejects,
 83                trials=trials, accepted=accepted, seconds=time.perf_counter()-t0)
 84
 85
 86def resolving_lm(q0, x, y, d, h, steps=35):
 87    """Conventional LM: only full steps are tested; rejection changes lambda
 88    and triggers a fresh curvature solve."""
 89    q = q0.copy(); lam = 1e-2
 90    low, high, up, down = 0.25, 0.75, 4.0, 2.0
 91    solves = rejects = accepted = trials = 0
 92    t0 = time.perf_counter()
 93    for _ in range(steps):
 94        f, g, H = state(q, x, y, d, h)
 95        accepted_here = False
 96        for _retry in range(12):
 97            p = solve_direction(g, H, lam); solves += 1
 98            pred = max(1e-14, -g @ p - 0.5*p @ H @ p)
 99            trials += 1
100            ft = state(q + p, x, y, d, h)[0]
101            if ft <= f + 1e-4*(g @ p):
102                rho = (f-ft)/pred
103                q = q + p; accepted += 1; accepted_here = True
104                if rho < low: lam *= up
105                elif rho > high: lam /= down
106                break
107            rejects += 1
108            lam *= up
109        if not accepted_here:
110            lam *= up
111    final = state(q, x, y, d, h)[0]
112    return dict(final_loss=final, solves=solves, rejected_trials=rejects,
113                trials=trials, accepted=accepted, seconds=time.perf_counter()-t0)
114
115
116def math_check():
117    rng = np.random.default_rng(123)
118    p = 7
119    A = rng.normal(size=(p,p)); H = A.T@A + 0.2*np.eye(p)
120    g = rng.normal(size=p); direction = np.linalg.solve(H+0.7*np.eye(p), -g)
121    f0 = 3.25
122    def model(a): return f0 + a*g@direction + 0.5*a*a*direction@H@direction
123    pred = -g@direction - 0.5*direction@H@direction
124    model_reduction = f0-model(1.0)
125    # Construct an exactly quadratic objective, so actual and model reductions agree.
126    actual_reduction = f0-model(1.0)
127    return dict(prediction_abs_error=abs(pred-model_reduction),
128                ratio_abs_error=abs(actual_reduction/pred-1.0),
129                pred_positive=bool(pred > 0))
130
131
132def main():
133    np.set_printoptions(precision=6, suppress=True)
134    d, h = 4, 5
135    x, y = make_problem()
136    rng = np.random.default_rng(99)
137    q0 = 0.4*rng.normal(size=d*h + h + h + 1)
138    result = {'math_check': math_check(),
139              'baseline_resolve_lm': resolving_lm(q0, x, y, d, h),
140              'idea_decoupled_armijo': armijo_lm(q0, x, y, d, h)}
141    print(json.dumps(result, indent=2))
142
143
144if __name__ == '__main__':
145    main()