# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + id="DZjcrHYLxwj7" colab_type="code" colab={}
import numpy as onp
import jax
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

# plot functions
import matplotlib.pyplot as plt
plt.style.use(['ggplot'])

# + [markdown] id="klhuTjJlSeoX" colab_type="text"
# ## Resources
#
# **Theory**
#
# * Kernel Methods in Machine Learning - Lecture 2: Regularized learning and kernels - [Slides](https://mycourses.aalto.fi/pluginfile.php/359553/mod_resource/content/1/CS-4830_Lecture2.pdf)
# * Intro to ML Course - [Colab](https://www.papernot.fr/teaching/w20-ml)
#
#

# + [markdown] id="nzw_F7ucC0OJ" colab_type="text"
# ## Data

# + id="HDy9nrB03ZHU" colab_type="code" cellView="form" colab={"base_uri": "https://localhost:8080/", "height": 355} outputId="bf8047a4-ff48-4b17-b016-047f6b02d497"
#@title 1D Data

rng_key = random.PRNGKey(10)

n_samples, n_features = 100, 1

# generate X
X = random.normal(rng_key, shape=(n_samples, n_features))

# generate true w
true_w = random.normal(rng_key+1, shape=(n_features,))
print("True Weights: ", true_w)
# generate y = Xw + eps
eps = 0.5
y_true = jnp.dot(X, true_w) + eps * random.normal(rng_key+2, shape=(n_samples,))

X_plot = jnp.linspace(X.min(), X.max(), 1_000)[:, None]
print(X.shape, X_plot.shape, y_true.shape)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(X, y_true)
plt.show()


# + [markdown] id="iKzxp1EcyTi5" colab_type="text"
# ## Theory
#
# ### Linear Regression
# Regression is a form of curve fitting or function approximation. Regression models typically try to solve the following problem: We want to find some latent function $f(\cdot)$ s.t.
#
# $$\hat{y}=f(\mathbf{x,w}) + \epsilon_y$$
#
# where:
# * $\mathbf{w}$ - weights, unknown parameters (`n_features`)
# * $\mathbf{x}$ - are independent variables (`n_features`)
#     * $\mathbf{X}$ - are independent variables stacked with samples/realizations (`n_samples, n_features`)
# * $\mathbf{y}$ - are dependent variables (`n_targets`)
#     * $\mathbf{Y}$ - are dependent variables stacked on top of each other (`n_samples, n_targets`)
#
# We hope to approximate $m$- dimensional continuous functions via said $f(\cdot)$. So if $\mathbf{x} \in \mathbb{R}^{D}$, we want the function to map from to the target space $\mathbf{y} \in \mathbb{R}^{P}$. 
#
#
# **How do we choose the weights?**
#
# We can use **cross-validation** or we can **minimize a cost function** either by gradient descent or exact solutions.
#
#

# + id="4_Fy3M9z3rT3" colab_type="code" colab={}
#@title Linear Model

# Linear Model
def linear_model(params, x):
    """"""
    return jnp.dot(params['w'], x) #+ params['b']



# + [markdown] id="qN0aKRVu5LID" colab_type="text"
# ### Exact Solution
#
# Firstly, we can solve this explicitly in closed-form:
#
# $$\hat{\mathbf{w}} = (X^\top X)^{-1}X^\top\mathbf{y}$$
#
# Fortunately, the jax library has a `np.linalg.pinv` implementation that allows us to calculate the gradients.

# + id="I4C5xI715MaI" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 337} outputId="2412b138-0ce3-463c-dc60-65ea70bc4c16"
# closed-form solution
w_linalg = jnp.dot(jnp.linalg.pinv(X), y_true)
print('Weights: ', w_linalg)

# add parameters to the weights
params = {'w': w_linalg}

# create linear model
model = linear_model

# vectorize the model, w.r.t. inputs
model = vmap(model, in_axes=(None, 0))

# evaluate linear model
y_linalg = model(params, X)

# Plot Results
fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(X, y_true, label='Data', color='red')
ax.plot(X, y_linalg, label='LinAlg Soln', color='black')
ax.legend(fontsize=12)
plt.show()

# + id="MCwhiMPigTu_" colab_type="code" colab={}
from dataclasses import dataclass 

@dataclass
class Params:
    # kernel params
    sigmas = [0.01, 0.1, 1.0]
    var = [10, 100, 1_000]
    # experimental params
    


# + id="hPw0T7mhgisi" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="e6414223-14f7-4b4f-fc27-f16091319176"
Params.sigmas, Params.var

# + [markdown] id="bLwLFMYcetal" colab_type="text"
# ```python
# def f(w, b, x):
#     return jnp.dot(x, w) + b
#
# df_dw = jax.grad(f, argnums=(0))
# df_db = jax.grad(f, argnums=(1))
# df_dx = jax.grad(f, argnums=(2))
# df_dw, df_db = jax.grad(f, argnums=(0, 1))
# ```
#
# ```python
# def linear_model(params, x):
#     return jnp.dot(x, params['w']) + params['b']
# df_dw = jax.grad(f, argnums=(0))
# ```

# + id="b1-PFwTgAhoC" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 337} outputId="6269a25a-781b-4f98-b5d7-aac3ae587e38"
# get gradient of model w.r.t. inputs (NOT params)
model_gradx = jax.grad(linear_model, argnums=(1))

# vectorize grad function
model_gradx = vmap(model_gradx, in_axes=(None, 0))

# get gradients w.r.t X
dydx = model_gradx(params, X)
print(jnp.mean(dydx))

# Plot
fig, ax = plt.subplots(figsize=(7, 5))
ax.plot(X, dydx, label='Gradient', color='red')
ax.plot(X, y_linalg, label='LinAlg Soln', color='black')
ax.legend(fontsize=12)
plt.show()


# + [markdown] id="UhEOVs0W5nf7" colab_type="text"
# ### Gradient Descent
#
#
# We can also use gradient descent to tackle this problem. If we initialize $\mathbf{w}_0$ to be something random, then we can update:
#
# $$\mathbf{w}_j = \mathbf{w}_{j-1} + \alpha \nabla \mathcal{L}(\mathbf{w}_{j-1})$$
#
# where $\nabla$ is the gradient operator and $\alpha$ is the learning rate.

# + id="UotJHY1t7cM6" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="9d9f554f-8701-4099-f3b4-9d26af80f894"
# initialize params
def init_params(x):
    return {'w': jnp.zeros(len(x))}

params = init_params(X[0, :])
print(params)

# + [markdown] id="C5XK9VdOLfpE" colab_type="text"
# ### Loss Function
#
# Recall the model for simple:
#
# $$y = X\mathbf{w}$$
#
# We can define a loss function to describe the penalization:
#
# $$\mathcal{L}(\mathbf{w}) 
# = \left( f(\mathbf{w,x}) - y \right)^2$$
#
# where $\mathbf{x,w} \in \mathbb{R}^D$ and $\mathbf{y} \in \mathbb{R}^P$

# + id="RaNNm5f57b0f" colab_type="code" colab={}
from typing import Callable, Dict

# Linear Function
@jax.jit
def linear_model(params, x):
    """"""
    return jnp.dot(params['w'], x) #+ params['b']

# MSE Loss Function
@jax.jit
def squared_loss(
    params,
    x,
    y):
    """calculates the squared loss"""
    # output of model
    out = linear_model(params, x)

    # calculate the squared loss
    se = jnp.square(out - y)

    return  se


# + id="_eOmtaUMHtXd" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="93aaed53-9b78-417d-9d1c-bc9d7437bb7b"
# initialize loss function, squared loss
loss = squared_loss

# test loss function, single sample x
params = {'w': jnp.zeros(n_features)}
iloss = loss(params, X[0, :], y_true[0])

print(iloss, type(iloss))

# + [markdown] id="BCxXEHSYMBmw" colab_type="text"
# #### Vectorizing
#
# In the above equation, we assumed that we only had one sample $\mathbf{x}$ of dimension D. What you'll typically see in your textbooks is the following: 
#
#
# $$\mathcal{L}(\mathbf{w}) 
# = \frac{1}{N}\sum_{i=1}^N\left( f(\mathbf{w,x_i}) - y_i \right)^2$$
#
# where $N$ is the number of samples, $\mathbf{X} \in \mathbb{R}^{N \times D}$, $\mathbf{w} \in \mathbb{R}^D$ and $\mathbf{y} \in \mathbb{R}^P$
#
# In JAX, we can vectorize this operation via the `vmap` function again so that we can deal with multiple samples.

# + id="_dlCSimzMC4r" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 382} outputId="108d73c4-feb5-4233-c112-e200bc67066e"
# initialize loss function, squared loss
loss = squared_loss

# vectorize loss function, w.r.t. inputs
loss = jax.jit(jax.vmap(loss, in_axes=(None, 0, 0), out_axes=0))

# test loss function, w.r.t. multiple inputs
nloss = loss(params, X, y_true)
print(nloss)

# mean squared loss
print(jnp.mean(nloss))

# + [markdown] id="X0bpA7kwOlUz" colab_type="text"
# ### Loss Gradient

# + id="aqm2i0SMKAff" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="bc48103a-0ae4-4872-862e-2f3defe849c8"
# create grad function from loss
grad_loss = jax.jit(jax.grad(squared_loss, argnums=(0)))

# vectorize grad_loss
grad_loss = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))

# gradients
iloss = grad_loss(params, X, y_true)

print(jnp.mean(iloss['w']))

# + [markdown] id="CpOkf46DuQ0T" colab_type="text"
# ### Training

# + id="M78q6YJGS3CL" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="62cd7705-98fb-48be-a750-ee0625f0d5b8"
iloss

# + id="esLJZYZg5owO" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 52} outputId="8bf052dc-900c-48af-cf81-90c7ce05d46d"
# do gradient descent
params = init_params(X[0, :])
print(params['w'])
losses = []
epsilon = 0.1
iterations = 100

for _ in range(iterations):
    
    # calculate the loss
    iloss = loss(params, X, y_true)
    iloss = jnp.mean(iloss)
    
    losses.append(iloss)

    # calculate the gradients of loss
    grad_losses = grad_loss(params, X, y_true)

    # average gradients over batch, update params
    params['w'] = params['w'] - epsilon * jnp.mean(grad_losses['w'], axis=0)
    
    # print(iloss, jnp.mean(grad_losses['w'], axis=0))
print(params)

# + id="34iqAnDZ6ujF" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 320} outputId="e5393bfd-44dc-4417-a360-f7915b83bb5a"


fig, ax = plt.subplots(ncols=2, figsize=(15, 5))

# plot the losses
ax[0].plot(losses, label='Losses')
ax[0].legend(fontsize=20)

# plot the resulting line
y_pred = vmap(linear_model, in_axes=(None, 0))(params, X_plot)

ax[1].scatter(X, y_true, label='Data', color='red')
ax[1].plot(X_plot, y_pred, label='Grad Soln', color='black')
ax[1].legend(fontsize=20)

plt.show()



# + [markdown] id="hV4rmmxn48lT" colab_type="text"
# ### Penalized Linear Regression: Ridge Regression
#
# For kernel methods, we add a penalty to the loss function.
#
# $$
# \mathcal{L}(\mathbf{w, \lambda}) 
# = \underset{\mathbf{w}}{\text{min}}
# \left\{ 
#     ||\mathbf{Y}-\mathbf{Xw} ||^2 + \lambda ||\mathbf{w}||^2 
# \right\}
# $$
#

# + id="zmzlQGvrbWdz" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="a44a04a8-43cd-44f1-8445-5a1954d83874"
# initialize params
def init_params(x):
    return {'w': jnp.ones(len(x)), 'lambda': 0.01 * jnp.ones((1,))}

params = init_params(X[0, :])
print(params)

# + id="BTd0c68T4xkA" colab_type="code" colab={}
from typing import Callable, Dict

# Linear Function
@jax.jit
def linear_model(params, x):
    """"""
    return jnp.dot(params['w'], x) #+ params['b']

# MSE Loss Function
@jax.jit
def l2_squared_loss(
    params,
    x,
    y,
    l2=0.01):
    """calculates the squared loss"""
    # output of model
    out = linear_model(params, x)

    # calculate the squared loss
    se = jnp.power(out - y, 2)

    # add l2 regularization
    # print(jnp.square(jnp.linalg.norm(params['w'], 2)))
    reg = l2 * jnp.square(jnp.linalg.norm(params['w'], 2))
    # print(reg)
    return se + reg


# + id="8s7dNgBRew-G" colab_type="code" colab={}
def make_loss_n_grad(loss_func, params, X, y_true, test=False, jitted=True):

    # =========================
    # Loss Function
    # =========================

    # vectorize loss function, w.r.t. inputs
    if jitted == True:
        loss = jax.jit(jax.vmap(loss_func, in_axes=(None, 0, 0), out_axes=0))
    else:
        loss = jax.vmap(loss_func, in_axes=(None, 0, 0), out_axes=0)

    if test == True:
        # test loss function, w.r.t. multiple inputs
        nloss = loss(params, X, y_true)

        # mean squared loss
        print(jnp.mean(nloss))
    
    # =========================
    # Grad Loss Function
    # =========================
    
    # grad_loss = jax.grad(l2_squared_loss, argnums=(0))
    if jitted:
        # create grad function from loss
        grad_loss = jax.jit(jax.grad(loss_func, argnums=(0)))

        # vectorize grad_loss
        grad_loss = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))
    else:
        # create grad function from loss
        grad_loss = jax.grad(loss_func, argnums=(0))

        # vectorize grad_loss
        grad_loss = jax.vmap(grad_loss, in_axes=(None, 0, 0))
    

    if test == True:
        # gradients
        iloss = grad_loss(params, X, y_true)

        print(jnp.mean(iloss['w']))
    return loss, grad_loss

# + id="ueXm-MpUisWD" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 69} outputId="99e4bdef-fbdd-4705-d153-c360320bf6f0"
# init params
params = init_params(X[0, :])
print("Params:", params)
# create loss functions
loss, grad_loss = make_loss_n_grad(l2_squared_loss, params, X, y_true, test=True)

# + id="Tx71CxiibOE9" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="567da6ee-5b4e-4f7a-b5d9-25ca0954b2a1"
losses = []
alpha = 0.1
iterations = 100

for _ in range(iterations):
    
    # calculate the loss
    iloss = loss(params, X, y_true)
    iloss = jnp.mean(iloss)
    
    losses.append(iloss)

    # calculate the gradients of loss
    grad_losses = grad_loss(params, X, y_true)

    # average gradients over batch, update params
    params['w'] = params['w'] - alpha * jnp.mean(grad_losses['w'], axis=0)
    # print(iloss, jnp.mean(grad_losses['w'], axis=0))
print(params)

# + id="b6PSGu1XcbC7" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 322} outputId="e07278fe-5acb-499e-a1d6-1b291afa5f5c"
fig, ax = plt.subplots(ncols=2, figsize=(15, 5))

# plot the losses
ax[0].plot(losses, label='Losses')
ax[0].legend(fontsize=20)

# plot the resulting line
y_pred = vmap(linear_model, in_axes=(None, 0))(params, X_plot)

ax[1].scatter(X, y_true, label='Data', color='red')
ax[1].plot(X_plot, y_pred, label='Grad Soln', color='black')
ax[1].legend(fontsize=20)

plt.show()

StarcoderdataJupyterScriptsDedupFiltered