import numpy as np
n = 100
# Random
# A = np.random.randn(n, n)
# A = A.T.dot(A)
# Clustered eigenvalues
A = np.diagflat([np.ones(n//4), 10 * np.ones(n//4), 100*np.ones(n//4), 1000* np.ones(n//4)])
U = np.random.rand(n, n)
Q, _ = np.linalg.qr(U)
A = Q.dot(A).dot(Q.T)
A = (A + A.T) * 0.5
print("A is normal matrix: ||AA* - A*A|| =", np.linalg.norm(A.dot(A.T) - A.T.dot(A)))
b = np.random.randn(n)
# Hilbert matrix
# A = np.array([[1.0 / (i+j - 1) for i in range(1, n+1)] for j in range(1, n+1)])
# b = np.ones(n)

f = lambda x: 0.5 * x.dot(A.dot(x)) - b.dot(x)
grad_f = lambda x: A.dot(x) - b
x0 = np.zeros(n)

"""
Explanation: Beyond gradient descent
Метод сопряжённых градиентов
Система линейных уравнений vs. задача безусловной минимизации
Рассмотрим задачу
$$
\min_{x \in \mathbb{R}^n} \frac{1}{2}x^{\top}Ax - b^{\top}x,
$$
где $A \in \mathbb{S}^n_{++}$.
Из необходимого условия экстремума имеем
$$
Ax^* = b
$$
Также обозначим $f'(x_k) = Ax_k - b = r_k$
Как решить систему $Ax = b$?

Прямые методы основаны на матричных разложениях:
Плотная матрица $A$: для размерностей не больше нескольких тысяч
Разреженная (sparse) матрица $A$: для размерностей порядка $10^4 - 10^5$


Итерационные методы: хороши во многих случаях, единственный подход для задач с размерностью $ > 10^6$

Метод сопряжённых направлений
В градиентном спуске направления убывания - анти-градиенты, но для функций с плохо обусловленным гессианом сходимость медленная.
Идея: двигаться вдоль направлений, которые гарантируют сходимость за $n$ шагов.
Определение. Множество ненулевых векторов ${p_0, \ldots, p_l}$ называется сопряжённым относительно матрицы $A \in \mathbb{S}^n_{++}$, если 
$$
p^{\top}_iAp_j = 0, \qquad i \neq j
$$
Утверждение. Для любой $x_0 \in \mathbb{R}^n$ последовательность ${x_k}$, генерируемая методом сопряжённых направлений, сходится к решению системы $Ax = b$ максимум за $n$ шагов.
```python
def ConjugateDirections(x0, A, b, p):
x = x0

r = A.dot(x) - b

for i in range(len(p)):

    alpha = - (r.dot(p[i])) / (p[i].dot(A.dot(p[i])))

    x = x + alpha * p[i]

    r = A.dot(x) - b

return x

```
Примеры сопряжённых направлений

Собственные векторы матрицы $A$
Для любого набора из $n$ векторов можно провести аналог ортогонализации Грама-Шмидта и получить сопряжённые направления

Вопрос: что такое ортогонализация Грама-Шмидта? :)
Метод сопряжённых градиентов
Идея: новое направление $p_k$ ищется в виде $p_k = -r_k + \beta_k p_{k-1}$, где $\beta_k$ выбирается, исходя из требования сопряжённости $p_k$ и $p_{k-1}$:
$$
\beta_k = \dfrac{p^{\top}{k-1}Ar_k}{p^{\top}{k-1}Ap_{k-1}}
$$
Таким образом, для получения следующего сопряжённого направления $p_k$ необходимо хранить только сопряжённое направление $p_{k-1}$ и остаток $r_k$ с предыдущей итерации. 
Вопрос: как находить размер шага $\alpha_k$?
Сопряжённость сопряжённых градиентов
Теорема
Пусть после $k$ итераций $x_k \neq x^*$. Тогда 

$\langle r_k, r_i \rangle = 0, \; i = 1, \ldots k - 1$
$\mathtt{span}(r_0, \ldots, r_k) = \mathtt{span}(r_0, Ar_0, \ldots, A^kr_0)$
$\mathtt{span}(p_0, \ldots, p_k) = \mathtt{span}(r_0, Ar_0, \ldots, A^kr_0)$
$p_k^{\top}Ap_i = 0$, $i = 1,\ldots,k-1$

Теоремы сходимости
Теорема 1. Если матрица $A$ имеет только $r$ различных собственных значений, то метод сопряжённых градиентов cойдётся за $r$ итераций.
Теорема 2. Имеет место следующая оценка сходимости
$$
\| x_{k} - x^ \|_A \leq 2\left( \dfrac{\sqrt{\kappa(A)} - 1}{\sqrt{\kappa(A)} + 1} \right)^k \|x_0 - x^\|_A,
$$
где $\|x\|^2_A = x^{\top}Ax$ и $\kappa(A) = \frac{\lambda_1(A)}{\lambda_n(A)}$ - число обусловленности матрицы $A$, $\lambda_1(A) \geq ... \geq \lambda_n(A)$ - собственные значения матрицы $A$
Замечание: сравните коэффициент геометрической прогрессии с аналогом в градиентном спуске.
Интерпретации метода сопряжённых градиентов

Градиентный спуск в пространстве $y = Sx$, где $S = [p_0, \ldots, p_n]$, в котором матрица $A$ становится диагональной (или единичной в случае ортонормированности сопряжённых направлений)
Поиск оптимального решения в Крыловском подпространстве $\mathcal{K}_k(A) = {b, Ab, A^2b, \ldots A^{k-1}b}$

$$
x_k = \arg\min_{x \in \mathcal{K}_k} f(x)
$$

Однако естественный базис Крыловского пространства неортогональный и, более того, плохо обусловлен.

Упражнение Проверьте численно, насколько быстро растёт обусловленность матрицы из векторов ${b, Ab, ... }$

Поэтому его необходимо ортогонализовать, что и происходит в методе сопряжённых градиентов

Основное свойство
$$ 
A^{-1}b \in \mathcal{K}_n(A)
$$
Доказательство

Теорема Гамильтона-Кэли: $p(A) = 0$, где $p(\lambda) = \det(A - \lambda I)$
$p(A)b = A^nb + a_1A^{n-1}b + \ldots + a_{n-1}Ab + a_n b = 0$
$A^{-1}p(A)b = A^{n-1}b + a_1A^{n-2}b + \ldots + a_{n-1}b + a_nA^{-1}b = 0$
$A^{-1}b = -\frac{1}{a_n}(A^{n-1}b + a_1A^{n-2}b + \ldots + a_{n-1}b)$

Улучшенная версия метода сопряжённых градиентов
На практике используются следующие формулы для шага $\alpha_k$ и коэффициента $\beta_{k}$:
$$
\alpha_k = \dfrac{r^{\top}k r_k}{p^{\top}{k}Ap_{k}} \qquad \beta_k = \dfrac{r^{\top}k r_k}{r^{\top}{k-1} r_{k-1}}
$$
Вопрос: чем они лучше базовой версии?
Псевдокод метода сопряжённых градиентов
```python
def ConjugateGradientQuadratic(x0, A, b, eps):
r = A.dot(x0) - b

p = -r

while np.linalg.norm(r) &gt; eps:

    alpha = r.dot(r) / p.dot(A.dot(p))

    x = x + alpha * p

    r_next = r + alpha * A.dot(p)

    beta = r_next.dot(r_next) / r.dot(r)

    p = -r_next + beta * p

    r = r_next

return x

```
Метод сопряжённых градиентов для неквадратичной функции
Идея: использовать градиенты $f'(x_k)$ неквадратичной функции вместо остатков $r_k$ и линейный поиск шага $\alpha_k$ вместо аналитического вычисления. Получим метод Флетчера-Ривса.
```python
def ConjugateGradientFR(f, gradf, x0, eps):
x = x0

grad = gradf(x)

p = -grad

while np.linalg.norm(gradf(x)) &gt; eps:

    alpha = StepSearch(x, f, gradf, **kwargs)

    x = x + alpha * p

    grad_next = gradf(x)

    beta = grad_next.dot(grad_next) / grad.dot(grad)

    p = -grad_next + beta * p

    grad = grad_next

    if restart_condition:

        p = -gradf(x)

return x

```
Теорема сходимости
Теорема. Пусть 
- множество уровней $\mathcal{L}$ ограничено
- существует $\gamma > 0$: $\| f'(x) \|_2 \leq \gamma$ для $x \in \mathcal{L}$
Тогда
$$
\lim_{j \to \infty} \| f'(x_{k_j}) \|_2 = 0
$$
Перезапуск (restart)

Для ускорения метода сопряжённых градиентов используют технику перезапусков: удаление ранее накопленной истории и перезапуск метода с текущей точки, как будто это точка $x_0$
Существуют разные условия, сигнализирующие о том, что надо делать перезапуск, например
$k = n$
$\dfrac{|\langle f'(x_k), f'(x_{k-1}) \rangle |}{\| f'(x_k) \|_2^2} \geq \nu \approx 0.1$


Можно показать (см. Nocedal, Wright Numerical Optimization, Ch. 5, p. 125), что запуск метода Флетчера-Ривза без использования перезапусков на некоторых итерациях может приводить к крайне медленной сходимости! 
Метод Полака-Рибьера и его модификации лишены подобного недостатка.

Комментарии

Замечательная методичка "An Introduction to the Conjugate Gradient Method Without the Agonizing Pain" размещена тут
Помимо метода Флетчера-Ривса существуют другие способы вычисления $\beta_k$: 
метод Полака-Рибьера $\beta_k = \frac{f'(x_k)^{\top}(f'(x_k) - f'(x_{k-1}))}{\| f'(x_{k-1})\|^2_2}$ 
метод Хестенса-Штифеля $\beta_k = \frac{f'(x_k)^{\top}(f'(x_k) - f'(x_{k-1}))}{-p_{k-1}^{\top}(f'(x_k) - f'(x_{k-1}))}$


Для метода сопряжённых градиентов требуется 4 вектора: каких?
Самой дорогой операцией является умножение матрицы на вектор

Эксперименты
Квадратичная целевая функция
End of explanation
"""


%matplotlib inline
import matplotlib.pyplot as plt
plt.rc("text", usetex=True)
plt.rc("font", family='serif')
eigs = np.linalg.eigvalsh(A)
plt.semilogy(np.unique(eigs))
plt.ylabel("Eigenvalues", fontsize=20)
plt.xticks(fontsize=18)
_ = plt.yticks(fontsize=18)

"""
Explanation: Распределение собственных значений
End of explanation
"""


import scipy.optimize as scopt

def callback(x, array):
    array.append(x)

scopt_cg_array = []
scopt_cg_callback = lambda x: callback(x, scopt_cg_array)
x = scopt.minimize(f, x0, method="CG", jac=grad_f, callback=scopt_cg_callback)
x = x.x
print("||f'(x*)|| =", np.linalg.norm(A.dot(x) - b))
print("f* =", f(x))

"""
Explanation: Правильный ответ
End of explanation
"""


def ConjugateGradientQuadratic(x0, A, b, tol=1e-8, callback=None):
    x = x0
    r = A.dot(x0) - b
    p = -r
    while np.linalg.norm(r) > tol:
        alpha = r.dot(r) / p.dot(A.dot(p))
        x = x + alpha * p
        if callback is not None:
            callback(x)
        r_next = r + alpha * A.dot(p)
        beta = r_next.dot(r_next) / r.dot(r)
        p = -r_next + beta * p
        r = r_next
    return x

import liboptpy.unconstr_solvers as methods
import liboptpy.step_size as ss

max_iter = 70
print("\t CG quadratic")
cg_quad = methods.fo.ConjugateGradientQuad(A, b)
x_cg = cg_quad.solve(x0, tol=1e-7, max_iter=max_iter, disp=True)

print("\t Gradient Descent")
gd = methods.fo.GradientDescent(f, grad_f, ss.ExactLineSearch4Quad(A, b))
x_gd = gd.solve(x0, tol=1e-7, max_iter=max_iter, disp=True)

print("Condition number of A =", abs(max(eigs)) / abs(min(eigs)))

"""
Explanation: Реализация метода сопряжённых градиентов
End of explanation
"""


plt.figure(figsize=(8,6))
plt.semilogy([np.linalg.norm(grad_f(x)) for x in cg_quad.get_convergence()], label=r"$\|f'(x_k)\|^{CG}_2$", linewidth=2)
plt.semilogy([np.linalg.norm(grad_f(x)) for x in scopt_cg_array[:max_iter]], label=r"$\|f'(x_k)\|^{CG_{PR}}_2$", linewidth=2)
plt.semilogy([np.linalg.norm(grad_f(x)) for x in gd.get_convergence()], label=r"$\|f'(x_k)\|^{G}_2$", linewidth=2)
plt.legend(loc="best", fontsize=20)
plt.xlabel(r"Iteration number, $k$", fontsize=20)
plt.ylabel("Convergence rate", fontsize=20)
plt.xticks(fontsize=18)
_ = plt.yticks(fontsize=18)

print([np.linalg.norm(grad_f(x)) for x in cg_quad.get_convergence()])

plt.figure(figsize=(8,6))
plt.plot([f(x) for x in cg_quad.get_convergence()], label=r"$f(x^{CG}_k)$", linewidth=2)
plt.plot([f(x) for x in scopt_cg_array], label=r"$f(x^{CG_{PR}}_k)$", linewidth=2)
plt.plot([f(x) for x in gd.get_convergence()], label=r"$f(x^{G}_k)$", linewidth=2)
plt.legend(loc="best", fontsize=20)
plt.xlabel(r"Iteration number, $k$", fontsize=20)
plt.ylabel("Function value", fontsize=20)
plt.xticks(fontsize=18)
_ = plt.yticks(fontsize=18)

"""
Explanation: График сходимости
End of explanation
"""


import numpy as np
import sklearn.datasets as skldata
import scipy.special as scspec
import jax
import jax.numpy as jnp
from jax.config import config
config.update("jax_enable_x64", True)

n = 300
m = 1000

X, y = skldata.make_classification(n_classes=2, n_features=n, n_samples=m, n_informative=n//3, random_state=0)
X = jnp.array(X)
y = jnp.array(y)
C = 1
@jax.jit
def f(w):
    return jnp.linalg.norm(w)**2 / 2 +  C * jnp.mean(jnp.logaddexp(jnp.zeros(X.shape[0]), -y * (X @ w)))

autograd_f = jax.jit(jax.grad(f))
def grad_f(w):
    denom = scspec.expit(-y * X.dot(w))
    return w - C * X.T.dot(y * denom) / X.shape[0]

x0 = jax.random.normal(jax.random.PRNGKey(0), (n,))
print("Initial function value = {}".format(f(x0)))
print("Initial gradient norm = {}".format(jnp.linalg.norm(autograd_f(x0))))
print("Initial gradient norm = {}".format(jnp.linalg.norm(grad_f(x0))))

"""
Explanation: Неквадратичная функция
$$
f(w) = \frac12 \|w\|2^2 + C \frac1m \sum{i=1}^m \log (1 + \exp(- y_i \langle x_i, w \rangle)) \to \min_w
$$
End of explanation
"""


def ConjugateGradientFR(f, gradf, x0, num_iter=100, tol=1e-8, callback=None, restart=False):
    x = x0
    grad = gradf(x)
    p = -grad
    it = 0
    while np.linalg.norm(gradf(x)) > tol and it < num_iter:
        alpha = utils.backtracking(x, p, method="Wolfe", beta1=0.1, beta2=0.4, rho=0.5, f=f, grad_f=gradf)
        if alpha < 1e-18:
            break
        x = x + alpha * p
        if callback is not None:
            callback(x)
        grad_next = gradf(x)
        beta = grad_next.dot(grad_next) / grad.dot(grad)
        p = -grad_next + beta * p
        grad = grad_next.copy()
        it += 1
        if restart and it % restart == 0:
            grad = gradf(x)
            p = -grad
    return x

"""
Explanation: Реализация метода Флетчера-Ривса
End of explanation
"""


import scipy.optimize as scopt
import liboptpy.restarts as restarts

n_restart = 60
tol = 1e-5
max_iter = 600

scopt_cg_array = []
scopt_cg_callback = lambda x: callback(x, scopt_cg_array)
x = scopt.minimize(f, x0, tol=tol, method="CG", jac=autograd_f, callback=scopt_cg_callback, options={"maxiter": max_iter})
x = x.x
print("\t CG by Polak-Rebiere")
print("Norm of garient = {}".format(np.linalg.norm(autograd_f(x))))
print("Function value = {}".format(f(x)))

print("\t CG by Fletcher-Reeves")
# ss.Backtracking("Armijo", rho=0.5, beta=0.001, init_alpha=1.)
cg_fr = methods.fo.ConjugateGradientFR(f, autograd_f, ss.Backtracking("Wolfe", rho=0.9, beta1=0.3, beta2=0.8, init_alpha=1.))
x = cg_fr.solve(x0, tol=tol, max_iter=max_iter, disp=True)

print("\t CG by Fletcher-Reeves with restart n")
cg_fr_rest = methods.fo.ConjugateGradientFR(f, autograd_f, ss.Backtracking("Wolfe", rho=0.9, beta1=0.3, beta2=0.8, 
                                         init_alpha=1.), restarts.Restart(n // n_restart))
x = cg_fr_rest.solve(x0, tol=tol, max_iter=max_iter, disp=True)

print("\t Gradient Descent")
gd = methods.fo.GradientDescent(f, autograd_f, ss.Backtracking("Wolfe", rho=0.9, beta1=0.1, beta2=0.8, init_alpha=1.))
x = gd.solve(x0, max_iter=max_iter, tol=tol, disp=True)

plt.figure(figsize=(8, 6))
plt.semilogy([np.linalg.norm(grad_f(x)) for x in cg_fr.get_convergence()], label=r"$\|f'(x_k)\|^{CG_{FR}}_2$ no restart", linewidth=2)
plt.semilogy([np.linalg.norm(grad_f(x)) for x in cg_fr_rest.get_convergence()], label=r"$\|f'(x_k)\|^{CG_{FR}}_2$ restart", linewidth=2)
plt.semilogy([np.linalg.norm(grad_f(x)) for x in scopt_cg_array], label=r"$\|f'(x_k)\|^{CG_{PR}}_2$", linewidth=2)

plt.semilogy([np.linalg.norm(grad_f(x)) for x in gd.get_convergence()], label=r"$\|f'(x_k)\|^{G}_2$", linewidth=2)
plt.legend(loc="best", fontsize=16)
plt.xlabel(r"Iteration number, $k$", fontsize=20)
plt.ylabel("Convergence rate", fontsize=20)
plt.xticks(fontsize=18)
_ = plt.yticks(fontsize=18)

"""
Explanation: График сходимости
End of explanation
"""


%timeit scopt.minimize(f, x0, method="CG", tol=tol, jac=grad_f, options={"maxiter": max_iter})
%timeit cg_fr.solve(x0, tol=tol, max_iter=max_iter)
%timeit cg_fr_rest.solve(x0, tol=tol, max_iter=max_iter)
%timeit gd.solve(x0, tol=tol, max_iter=max_iter)

"""
Explanation: Время выполнения
End of explanation
"""


