#!/usr/bin/env python
# coding: utf-8

# # Descente de gradient: choix de pas, démonstration interactive
# 
# ***
# > __Auteur__: Alexandre Gramfort, Joseph Salmon <joseph.salmon@umontpellier.fr>

# On cherche à minimiser la fonctiont suivante:
# 
# $f(x_1, x_2) = (x_1^2 + x_2 - 11)^2 + (x_1 + x_2^2 - 7)^2$
# 
# **Question** : est-ce une fonction convexe?

# In[ ]:


import numpy as np


def f(x):
    x1, x2 = x
    return (x1**2 + x2 - 11)**2 + (x1 + x2**2 - 7)**2


# In[ ]:


get_ipython().run_line_magic('matplotlib', 'notebook')
import matplotlib.pyplot as plt
import matplotlib
cmap_reversed = matplotlib.cm.get_cmap('RdBu_r')
from mpl_toolkits import mplot3d


# In[ ]:


X1, X2 = np.meshgrid(np.linspace(-5.5, 5.5, 50),
                     np.linspace(-5.5, 5.5, 50))
Z = f([X1, X2])  # Altitude


# In[ ]:


fig = plt.figure(figsize=(6, 6))
ax = plt.axes(projection='3d')
ax.plot_surface(X1, X2, Z, rstride=1, cstride=1,
                cmap=cmap_reversed, edgecolor='none')
ax.set_xlim(-5, 5)
ax.set_ylim(-5, 5)
ax.set_zlim(0, 500)
plt.show()


# In[ ]:


def my_plot(xs=None, cmap=cmap_reversed):
    levels = list(1.7 ** np.linspace(0, 10, 30) - 1.) + [300]
    plt.figure(figsize=(5, 5))
    plt.contourf(X1, X2, np.sqrt(Z), levels=np.sqrt(
        levels), cmap=cmap)
    plt.colorbar(extend='both')
    if xs is not None:
        x1, x2 = np.array(xs).T
        plt.plot(x1, x2, 'k')
        plt.plot(x1, x2, 'o', color='purple')
    plt.show()


# In[ ]:


plot()


# In[ ]:


plot(cmap=matplotlib.cm.get_cmap('RdBu'))


# In[ ]:





# In[ ]:


def f_grad(x):
    x1, x2 = x
    df_x1 = 2 * (-7 + x1 + x2**2 + 2 * x1 * (-11 + x1**2 + x2))
    df_x2 = 2 * (-11 + x1**2 + x2 + 2 * x2 * (-7 + x1 + x2**2))
    return np.array([df_x1, df_x2])


x0 = [0.5, -4]


def grad_descent(x_init=x0, step_size=0.01, max_iter=20):
    """Descente de gradient avec un pas constant"""
    x = x_init
    xs = [x]
    for k in range(max_iter):
        x = ### XXXX
        xs.append(x)
    my_plot(xs)
    plt.show()


grad_descent(x_init=x0, step_size=0.01, max_iter=20)


# In[ ]:


from ipywidgets import interact, fixed
interact(grad_descent, x_init=fixed(x0), step_size=(0., .05, 0.005), max_iter=(0, 50, 1))