# Source: content/notes/math/optimization.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np

rng = np.random.default_rng(61)
X = rng.normal(size=(60, 8))
true = np.array([2., -1., 0., 0., 0., 0., 0., 0.])
y = X @ true
lam = 1.
L = np.linalg.norm(X, 2)**2
w = np.zeros(8)
objectives = []
for _ in range(600):
    v = w-X.T@(X@w-y)/L
    w = np.sign(v)*np.maximum(np.abs(v)-lam/L, 0.)
    objectives.append(.5*np.sum((X@w-y)**2)+lam*np.abs(w).sum())
assert np.max(np.diff(objectives)) < 1e-10
g = X.T@(X@w-y)
active = np.abs(w) > 1e-10
assert np.max(np.abs(g[active]+lam*np.sign(w[active]))) < 1e-7
assert np.max(np.abs(g[~active])) <= lam+1e-7
assert np.isclose(.1/np.sqrt(.001), np.sqrt(10))
assert np.isclose(np.sqrt(32/512), .25)

def f(x):
    return x**4-x**2
x = .1
gradient, hessian = 4*x**3-2*x, 12*x*x-2
assert f(x-gradient/hessian) > f(x)
p, alpha = -gradient, 1.
while f(x+alpha*p) > f(x)+1e-4*alpha*gradient*p:
    alpha *= .5
assert f(x+alpha*p) < f(x)
print("ISTA coefficients:", w, "Armijo step:", alpha)