# Diatomic Molecules and Spectroscopy
*Roberto Di Remigio*, *Luca Frediani*

Spectroscopy probes the electronic structure of atoms and molecules by measuring their interaction with light.
Different portions of the light spectrum can be used to explore different aspects of the electronic structure of a molecule. In this notebook, we will look closely at rotatonial and vibrational spectroscopy of diatomic molecules.
Some properties: http://hyperphysics.phy-astr.gsu.edu/hbase/Tables/diatomic.html

## Rotational Spectroscopy

Rotational spectroscopy uses _microwave radiation_ (300 MHz to 300 GHz) to investigate the rotational levels of molecule, _i.e._ the distribution of the nuclear masses in the molecule.
In a rotational spectroscopy experiment, transitions between the closely space rotational levels are observed and their energy differences can provide information on the molecular geometry and the masses of the nuclei.
For a diatomic molecule, the energy of the transition is:
\begin{equation}
F(J) = BJ(J+1)
\end{equation}
with intensity
\begin{equation}
|\mu_{J+1, J}|^2 = \frac{J+1}{2J+1}\mu_0^2
\end{equation}
$B$ and $\mu_0$ are the rotational constant and the dipole moment, respectively:
\begin{equation}
B = \frac{\hbar}{4\pi c I}; \quad I = \mu R_\mathrm{e}^2
\end{equation}
Thus a rotational transition will be deteceted only if the molecule has a nonzero dipole moment in the ground state. Moreover, the selection rule:
\begin{equation}
\Delta J = \pm 1
\end{equation}
applies.

## Vibrational Spectroscopy

Vibrational spectoscopy uses _infrared radiation_ (430 THz to 300 GHz) to investigate the relative motions of nuclei in molecules, _i.e._ molecular vibrations.
In the harmonic model of diatomic molecules, the wavenumber for the vibrational levels are obtained as:
\begin{equation}
G(v) = (v+\frac{1}{2})\tilde{\nu}, \quad \tilde{\nu} = \frac{1}{2\pi c}\sqrt{\frac{k}{\mu}}
\end{equation}
The intensity of a vibrational transition can be calculated as:
\begin{equation}
    I_{vw} = \left\langle \psi_v|\mu|\psi_w\right\rangle
\end{equation}
The first-order term in the Taylor expansion with respect to the nuclear displacements determines the selection rules for vibrational spectroscopy in the harmonic approximation.
Vibrational motion is never exactly harmonic. In the harmonic model we cannot, in fact, observe bond dissociation which is clearly in contrast with experiment.

## Computing the Rovibrational Spectrum

All vibrational levels have a _stack_ of rotational levels associated. Thus, when measuring vibrational spectra in _gas phase_ one can observe transitions between these stacks of rotational levels associated with different vibrational levels. The energy of such transitions is in general expressed as:
\begin{equation}
S(v, J) = G(v) + F_v(J)
\end{equation}
where $G(v)$ can contain anharmonic terms and $F_v(J)$ centrifugal distortion terms, _i.e._ depending on the vibrational quantum number $v$.

In the simplest case, these nonlinear terms can be neglected, obtaining the simplified formula:
\begin{equation}
S(v, J) = (v+\frac{1}{2})\tilde{\nu} + BJ(J+1)
\end{equation}
where the $\tilde{\nu}$ and $B$ are the vibrational and rotational constants, respectively.
Due to the small energy separation between rotational levels, the rovibrational spectrum will be temperature dependent. Each state, with energy $S(v, J)$, will be represented with a _weight_ according to the Boltzmann distribution:
\begin{equation}
w(v, J, T) = \frac{\mathrm{deg}(J)\exp(-\frac{S(v, J)}{k_\mathrm{B}T})}{\sum_{v,J}\mathrm{deg}(J)\exp(-\frac{S(v, J)}{k_\mathrm{B}T})}
\end{equation}
$\mathrm{deg}(J)$ is the degeneration of the rotational state with quantum number $J$: states with higher degeneration will be more represented in the spectrum.
The rovibrational spectrum captures the _rotational transitions_ that occur when the _vibrational transition_ $v+1\leftarrow v$ occurs. It consists of three branches:
- The **P branch** with all transitions where $\Delta J = -1$.
- The **Q branch** with the $\Delta J = 0$ transition. This can only occur with open-shell molecules.
- The **R branch** with all transitions where $\Delta J = +1$.

### Tasks

0. Derive a general formula for the energy difference (in wavenumbers) of transitions in the P, Q and R branches of the spectrum
1. Set up general script(s) to calculate and plot the rovibrational spectrum of a diatomic molecule. More in detail, the script(s) should:
    - Calculate the vibrational and rotational constants, given the masses, bond length and bond stifness constant.
    - Calculate the energies of the rotational levels for the P and R branches. **We assume the Q branch to be forbidden**.
    - Calculate the _ relative intensities_ of each transition in the P and R branches in terms of the Boltzmann distribution weight.
    - Plot the spectrum as wavenumber _vs_ relative intensity.
2. Compute the spectrum for the HCl molecule. An experimentally measured spectrum is reported below.

```python

```


```python
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import sympy

%matplotlib inline
```

```python
df = pd.read_csv('../data/raw/cities.csv', index_col=['CityId'])
primes = list(sympy.primerange(0, max(df.index)))
df['prime'] = df.index.isin(primes).astype(int)
```

```python
df_prime = df[(df.index == 0) | df['prime'] == 1]
```

```python
def read_link(filename):
    data = open(filename, 'r').read()
    data = data.replace('\n', ' ')
    data = np.fromstring(data, sep=' ', dtype=np.int32)
    if len(data) != data[0] + 1:
        raise Exception('Unrecognized format in %s' % filename)
    return np.concatenate((data[1:], [0]))
```

```python
tour_data = read_link('prime.tour')
tour_data_2 = np.array([df_prime.index[x] for x in tour_data])
```

```python
tour_data_2
```

    array([     0,   5333,  97649, ..., 148853, 153911,      0])

```python
def read_submission(filename):
    data = open(filename, 'r').read()
    data = data.replace('Path', '')
    data = data.replace('\n', ' ')
    data = np.fromstring(data, sep=' ', dtype=np.int32)
    return data
```

```python
best_tour = read_submission('submission.151557248.csv')
```

```python
best_tour.size
```

    197770

```python
si = tour_data_2[2]
ei = tour_data_2[3]
bsi = np.where(best_tour==si)[0][0]
bei = np.where(best_tour==ei)[0][0]
bei - bsi
```

    28

```python
c=0
for i in range(tour_data_2.size-1):
    si = tour_data_2[i]
    ei = tour_data_2[i+1]
    bsi = np.where(best_tour==si)[0][0]
    bei = np.where(best_tour==ei)[0][0]
    if abs(bei - bsi) == 5:
        c+=1

print(c)
```

    594


# Financial asset recommender - Association rule mining

## Introduction

This notebook is aimed to serve as an introduction to the creation of a recommender that operates on a proprietary dataset of user-financial asset transactions in order to recommend assets that frequently occur alongside each other in user transactions. It makes use of the Apriori algorithm to determine these frequent itemsets and then finds rules associating the assets together. A "rule" in this case is a relationship which states that if a specific asset X is bought, asset Y is likely to be bought alongside. This relationship is commonly used to provide recommendations to users about items they might like given their previous purchases.

## Dataset

We use a proprietary dataset of investment transactions provided by the National Bank of Greece. See the Infinitech Marketplace for more details.

If you are using your own dataset, it should be formatted as a comma-separated CSV file, where each line represents a transaction, and contains the following fields:
- **User id:** Unique identifier for the user.
- **Asset id:** Unique identifier for the asset.
- **Rating:** The value associated to the transaction.
- **Timestamp:** Timestamp of the transaction. Format: YYYY-MM-DD.

In order to load and preprocess the data, we use the Beta-Recsys library. For this we then create a dataset object from this raw data, which can be processed as interactions by Beta-RecSys.

```python
!pip install beta-recsys
```

```python
import sys, os, beta_rec
import pandas as pd
import numpy as np
from functools import partial

np_load_old = partial(np.load)
np.load = lambda *a,**k: np_load_old(*a, allow_pickle=True, **k)

from beta_rec.data import BaseData
from beta_rec.datasets.dataset_base import DatasetBase
from beta_rec.utils.constants import (
    DEFAULT_ITEM_COL,
    DEFAULT_RATING_COL,
    DEFAULT_TIMESTAMP_COL,
    DEFAULT_USER_COL,
)
```

```python
class NBGSampleInstTrans(DatasetBase):
    """NBGSample Instruments Transactions Dataset."""

    def __init__(self, dataset_name="nbg_marketplace_cf_1", min_u_c=0, min_i_c=0, root_dir=None):
        """Init NBGSampleInstTrans Class."""
        super().__init__(
            dataset_name=dataset_name,
            min_u_c=min_u_c,
            min_i_c=min_i_c,
            root_dir=root_dir,
        )

    def preprocess(self):
        """Preprocess the raw file.

        Preprocess the file, convert it to a dataframe consisting of the user-item
        interactions and save it in the processed directory.
        """
        file_name = '/tmp/iPythonNotebooks/nbg_instruments_large.csv'

        data = pd.read_csv(
            file_name,
            skiprows=[0],
            engine="python",
            names=[
                DEFAULT_USER_COL,
                DEFAULT_ITEM_COL,
                DEFAULT_RATING_COL,
                DEFAULT_TIMESTAMP_COL,
            ],
        )
        data[DEFAULT_TIMESTAMP_COL] =  pd.to_datetime(data[DEFAULT_TIMESTAMP_COL], format='%Y-%m-%d')

        self.save_dataframe_as_npz(
            #
            data,
           os.path.join(self.processed_path, f"{self.dataset_name}_interaction.npz"),
        )

        return os.path.join(self.processed_path, f"{self.dataset_name}_interaction.npz"), data

nbgdata = NBGSampleInstTrans()
result, data = nbgdata.preprocess()
```

```python
data
```

      col_user
      col_item
      col_rating
      col_timestamp

      0
      -1000399803742468017
      12541
      1.0
      2018-05-09

      1
      -1000476326435445839
      11907
      1.0
      2018-02-01

      2
      -1000476326435445839
      13233
      1.0
      2018-10-12

      3
      -1000476326435445839
      11907
      1.0
      2018-10-12

      4
      -1000570356359468369
      12917
      1.0
      2018-07-18

      ...
      ...
      ...
      ...
      ...

      312999
      998868334061509693
      89
      1.0
      2021-05-07

      313000
      998868334061509693
      268
      1.0
      2021-05-07

      313001
      998987671268040482
      12885
      1.0
      2021-02-09

      313002
      998987671268040482
      12885
      1.0
      2021-02-15

      313003
      999456452282545623
      17605
      1.0
      2021-04-16

313004 rows × 4 columns

## Data splitting
We then perform a temporal split on the data, through the data object we just created. Beta-RecSys provides methods to split data in several different ways, such as leave-one-out, random, and several types of basket splits. A temporal data split splits the data into training, testing, and validation sets on the basis of the timestamps associated with each data point, where earlier data points are used for training and later ones for validation and testing, in order to simulate a more real-world setting where future trends are learned from past events. This is suitable for our current data, where investment transactions can change depending on temporal factors. However, if needed, one can also try out other splitting strategies.

```python
split_dataset = nbgdata.load_temporal_split(n_test=1, n_negative=0)
new_data =  BaseData(split_dataset)
```

    --------------------------------------------------------------------------------
    Loaded training set statistics
    +---------+------------+------------+--------------+-----------------+
    |         | col_user   | col_item   | col_rating   | col_timestamp   |
    |---------+------------+------------+--------------+-----------------|
    | count   | 214521     | 214521     | 214521       | 214521          |
    | nunique | 10950      | 896        | 2            | 733             |
    +---------+------------+------------+--------------+-----------------+
    valid_data_0 statistics
    +---------+------------+------------+--------------+-----------------+
    |         | col_user   | col_item   | col_rating   | col_timestamp   |
    |---------+------------+------------+--------------+-----------------|
    | count   | 11721      | 11721      | 11721        | 11721           |
    | nunique | 3793       | 401        | 2            | 1               |
    +---------+------------+------------+--------------+-----------------+
    test_data_0 statistics
    +---------+------------+------------+--------------+-----------------+
    |         | col_user   | col_item   | col_rating   | col_timestamp   |
    |---------+------------+------------+--------------+-----------------|
    | count   | 13868      | 13868      | 13868        | 13868           |
    | nunique | 4821       | 395        | 2            | 1               |
    +---------+------------+------------+--------------+-----------------+
    --------------------------------------------------------------------------------
    After intersection, testing set [0] statistics
    +---------+------------+------------+--------------+-----------------+
    |         | col_user   | col_item   | col_rating   | col_timestamp   |
    |---------+------------+------------+--------------+-----------------|
    | count   | 10485      | 10485      | 10485        | 10485           |
    | nunique | 3166       | 343        | 2            | 1               |
    +---------+------------+------------+--------------+-----------------+
    After intersection, validation set [0] statistics
    +---------+------------+------------+--------------+-----------------+
    |         | col_user   | col_item   | col_rating   | col_timestamp   |
    |---------+------------+------------+--------------+-----------------|
    | count   | 10665      | 10665      | 10665        | 10665           |
    | nunique | 3190       | 377        | 2            | 1               |
    +---------+------------+------------+--------------+-----------------+
    Filling alias table
    Filling alias table

```python
train = new_data.train
test = new_data.test[0]
```

```python
train
```

      col_user
      col_item
      col_rating
      col_timestamp

      0
      0
      0
      1.0
      1514851200000000000

      1
      1
      1
      1.0
      1514851200000000000

      2
      1
      2
      1.0
      1514851200000000000

      3
      2
      3
      1.0
      1514851200000000000

      4
      2
      4
      1.0
      1514851200000000000

      ...
      ...
      ...
      ...
      ...

      214516
      2567
      462
      1.0
      1606780800000000000

      214517
      2567
      16
      1.0
      1606780800000000000

      214518
      141
      20
      1.0
      1606780800000000000

      214519
      10197
      28
      1.0
      1606780800000000000

      214520
      168
      13
      0.0
      1606780800000000000

214521 rows × 4 columns

## Association rule mining

Association Rule Mining is a data mining technique that finds patterns that represent relationships between items. When this is used with sales data, it is referred to as Market Basket Analysis. It can be described as a two-step process: firstly, to find all frequent itemsets, and secondly, to generate strong association rules from the frequent itemsets. Association rules are created by searching data for if-then patterns and using the metric support and confidence to identify the most important relationships. Let a, b be the items purchased by a customer in their item-list, a ⇒ b an association rule, R be the overall dataset of purchases, customers, and item-lists, and m be the number of item-lists.

The support of a with respect to R is the proportion of item-lists which contains the item a. The confidence value of a rule a⇒b with respect to R is the proportion of item-lists that contains item a, which also contains item b - to the total number of item-lists containing item a. A third metric, called lift, can be used to compare the calculated confidence against expected confidence or how many times the rule is expected to be found true.

### Apriori algorithm

Apriori finds the frequent itemsets by a process called candidate itemset generation. This is an iterative approach, where k-itemsets are used to explore (k+1)-itemsets. First, the set of frequent 1-itemsets is found, then, frequent 2-itemsets, and so on, until no more frequent k-itemsets can be found. To improve the efficiency of the level-wise generation of frequent itemsets, an important property called the Apriori property, is used to reduce the search space. The Apriori Property states that “All non-empty subsets of a frequent itemset must also be frequent.” This means that if there is a frequent item then, its subsets will also be frequent.

In our code below we will import the apriori algorithm from the mlxtend library, which we first need to install using pip.

```python
!pip install mlxtend
```

```python
from mlxtend.frequent_patterns import association_rules, apriori

basket = (train
          .groupby([DEFAULT_USER_COL, DEFAULT_ITEM_COL])[DEFAULT_RATING_COL]
          .sum().unstack().reset_index().fillna(0)
          .set_index(DEFAULT_USER_COL))

def encode_units(x):
    if x <= 0:
        return 0
    if x > 0:
        return 1

basket_train = basket.applymap(encode_units)

frequent_itemsets_train = apriori(basket_train, min_support=0.007, use_colnames=True, max_len=2)
rules = association_rules(frequent_itemsets_train, metric='lift').sort_values(by='lift', ascending=False).reset_index()
```

We take a look at the association rules thus generated below. In this context, an antecedent is a frequent itemset formed of reoccurring items in user purchases. A consequent is another item that is likely to also be purchased given the antecedent itemset is present in the user purchases.

```python
rules
```

      index
      antecedents
      consequents
      antecedent support
      consequent support
      support
      confidence
      lift
      leverage
      conviction

      0
      1844
      (573)
      (574)
      0.008219
      0.008767
      0.007945
      0.966667
      110.260417
      0.007873
      29.736986

      1
      1845
      (574)
      (573)
      0.008767
      0.008219
      0.007945
      0.906250
      110.260417
      0.007873
      10.578995

      2
      1501
      (573)
      (555)
      0.008219
      0.008767
      0.007854
      0.955556
      108.993056
      0.007782
      22.302740

      3
      1500
      (555)
      (573)
      0.008767
      0.008219
      0.007854
      0.895833
      108.993056
      0.007782
      9.521096

      4
      1535
      (572)
      (556)
      0.009041
      0.008858
      0.008676
      0.959596
      108.325523
      0.008596
      24.530753

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      1865
      965
      (354)
      (40)
      0.075160
      0.105662
      0.007945
      0.105711
      1.000461
      0.000004
      1.000054

      1866
      328
      (8)
      (762)
      0.106575
      0.071233
      0.007306
      0.068552
      0.962362
      -0.000286
      0.997122

      1867
      329
      (762)
      (8)
      0.071233
      0.106575
      0.007306
      0.102564
      0.962362
      -0.000286
      0.995530

      1868
      318
      (8)
      (354)
      0.106575
      0.075160
      0.007489
      0.070266
      0.934883
      -0.000522
      0.994736

      1869
      319
      (354)
      (8)
      0.075160
      0.106575
      0.007489
      0.099635
      0.934883
      -0.000522
      0.992292

1870 rows × 10 columns

## Recommendation
Now that we have a set of association rules generated between the assets, we have to decide how to use these to provide recommendations. As we have antecedents and consequents available to us from the rules, we design it as follows: if a user has invested in an asset or assets present in the antecedents, we treat the consequents for that rule as a recommendation. We make sure to rank the rules in descending order of lift, so that we have a way of ranking the recommendations - those consequents with a higher lift score will be ranked higher in the recommendation list.

```python
from itertools import chain, combinations
from operator import itemgetter
import math
import numpy as np

# We first filter the test set: we only keep those positive user / item pairs, and we only keep the user id, asset id and
# the amount of euros the customer spent on that asset.
test = test[[DEFAULT_USER_COL, DEFAULT_ITEM_COL, DEFAULT_RATING_COL]]
test = test[test[DEFAULT_RATING_COL] >= 1]
```

```python
# We get the set of users
custs = test[DEFAULT_USER_COL].unique()
# We obtain the relations between the elements in rules.
idx = rules['index'].values.tolist()
ants = []
cons = dict()
lifts = dict()
for i in idx:
    ants.append((i, list(rules['antecedents'][rules['index'] == i].values.tolist()[0])))
    cons[i] = list(rules['consequents'][rules['index'] == i].values.tolist()[0])
    lifts[i] = list(rules['lift'][rules['index'] == i])[0]
```

```python
# First, we find the rules to use for generating recommendations for each different user
cust_recs = {}

for cu in custs:
    recs = []
    # cats contains the set of items that user cu has acquired during test? I am not so sure about this...
    cats = train[DEFAULT_ITEM_COL][train[DEFAULT_USER_COL] == cu].values.tolist()
    for a in ants: # if the antecendent is in the rated items by the user, add the rule.
        if set(a[1]).issubset(set(cats)):
            recs.append(a[0])

    if len(recs) > 0:
        cust_recs[cu] = list(set(recs))
```

```python
# Then, for each user, we find the weights of those rules.
cust_recs_with_lifts = {}

# Here, there is something weird: why don't use maps instead of lists? They are much more efficient...
for cr in cust_recs:
    newlifts = []

    for item_id in cust_recs[cr]:
        newlifts.append(lifts[item_id])
    crl = zip(cust_recs[cr], newlifts) # then zips the corresponding list of the customer.
    cust_recs_with_lifts[cr] = sorted(list(crl), key=itemgetter(1), reverse=True) # and sorts it by lift value
    cust_recs_with_lifts[cr] = cust_recs_with_lifts[cr] # is this sentence even necessary? It does not do anything...
```

```python
# Next, we obtain a list containing [customer id, instrument id, rec. score, rank]
records = []

for c in cust_recs_with_lifts:
    for cl in cust_recs_with_lifts[c]:
        con_names = cons[cl[0]]
        for con_name in con_names: # for each consequent
            r = [c, con_name, cl[1]] # we add a list containing: customer id, item, value, rank
            records.append(r)

records_df = pd.DataFrame(records, columns=[DEFAULT_USER_COL, DEFAULT_ITEM_COL, 'col_lift'])
```

```python
# Now, there is a problem: some items appear several times in the same ranking, and we have to prevent this from happening.
# Because of this, we just group by CUST_ID_MASKED, INSTRMNT_
records_df = records_df.groupby([DEFAULT_USER_COL,DEFAULT_ITEM_COL], as_index=False, sort=False)['col_lift'].sum()
records_df
records_df["rank"] = (records_df.groupby(DEFAULT_USER_COL).cumcount() + 1)

rank = 1
current_cust = -1
rank_list = []
for index, row in records_df.iterrows():
    cust = row[DEFAULT_USER_COL]
    if cust != current_cust:
        current_cust = cust
        rank = 1
    rank_list.append(rank)
    rank += 1
records_df['rank']=rank_list

real_id_records_df = records_df.copy()
real_id_records_df = real_id_records_df.replace({DEFAULT_ITEM_COL : new_data.id2item})
real_id_records_df = real_id_records_df.replace({DEFAULT_USER_COL : new_data.id2user})
real_id_records_df
```

      col_user
      col_item
      col_lift
      rank

      0
      -9068932257782722186
      89
      25.934380
      1

      1
      -9068932257782722186
      207
      24.576910
      2

      2
      -9068932257782722186
      227
      20.620898
      3

      3
      -9068932257782722186
      8581
      22.079652
      4

      4
      -9068932257782722186
      267
      19.094932
      5

      ...
      ...
      ...
      ...
      ...

      127066
      1435227703229020590
      6699
      10.374733
      47

      127067
      1435227703229020590
      196
      2.021172
      48

      127068
      1435227703229020590
      108
      1.926762
      49

      127069
      1435227703229020590
      14753
      0.962362
      50

      127070
      1435227703229020590
      12363
      0.934883
      51

127071 rows × 4 columns

## Evaluation
Once we produce such ranked lists for every user, we proceed to calculate the following metrics for evaluating their performance.

1. Normalized discounted cumulative gain (NDCG): This is a measure of ranking quality normalized across predictions, and calculates the performance of the recommender on the basis that highly relevant items should be recommended at higher ranks, as they are more useful as recommendations than other, less relevant items. It is the ratio of the discounted cumulative gain (DCG) and the idealized DCG. DCG is represented below

\begin{equation}
\text{DCG} = \sum_{i=1}^N \frac{rel_{i}}{\log_2 (i+1)}
\end{equation}

\begin{equation}
\text{NDCG @ position p} = \frac{DCG_{p}}{IDCG_{p}}
\end{equation}

2. Precision: This commonly used metric refers to the number of true positive, or relevant instances among retrieved items.

\begin{equation}
\text{Precision} = \frac{True \ Positives}{True \ Positives \ + \ False \ Positives}
\end{equation}

3. Recall: Also referred to as sensitivity, this refers to the fraction of all relevant instances that were retrieved.

\begin{equation}
\text{Recall} = \frac{True \ Positives}{True \ Positives \ + \ False \ Negatives}
\end{equation}

4. Mean average precision (MAP): The mean average precision of a recommender is a metric that summarizes the precision-recall curve into a metric that represents the area under this curve and therefore is representative of the average of all precisions. MAP is the average of the AP value described below.

\begin{equation}
\text{Average Precision} = \sum_n (R_n - R_{n-1}) P_n
\end{equation}

These metrics are calculated across three ranges: 5, 10 and 20 recommendations per user respectively.

```python
ks = [5,10,20]
results = []
for k in ks:
    # This function just takes the records_df list and limits the rankings to a certain length.
    cb_pred_df = (
            records_df.groupby(DEFAULT_USER_COL, as_index=False)
            .apply(lambda x: x.nsmallest(k, 'rank'))
            .reset_index(drop=True)
        )

    cb_pred_df = cb_pred_df[[DEFAULT_USER_COL, DEFAULT_ITEM_COL, 'rank']]
    cb_pred_df

    common_users = set(test[DEFAULT_USER_COL]).intersection(set(train[DEFAULT_USER_COL]))
    n_users = len(common_users)

    # We obtain the test ratings for those users that appear both in the test set and the recommendation.
    test_common = test[test[DEFAULT_USER_COL].isin(common_users)]
    pred_common = cb_pred_df[cb_pred_df[DEFAULT_USER_COL].isin(common_users)]

    # We find the correctly predicted (user, item) pairs, and how many relevant items they have in the test set:
    df_hit = pd.merge(test_common, pred_common, on=[DEFAULT_USER_COL,DEFAULT_ITEM_COL])[[DEFAULT_USER_COL, DEFAULT_ITEM_COL,
                                                                                'rank']]
    df_hit.drop_duplicates(inplace=True)
    df_hit = df_hit.sort_values([DEFAULT_USER_COL, "rank"])
    df_hit["dcg"] = 1.0 / np.log1p(df_hit["rank"])
    df_hit["rr"] = (df_hit.groupby(DEFAULT_USER_COL).cumcount() + 1) / df_hit["rank"]

    # Aggregate by user
    df_hit_count = pd.merge(
            df_hit.groupby(DEFAULT_USER_COL, as_index=False)[DEFAULT_USER_COL].agg({"hit": "count"}),
            test_common.groupby(DEFAULT_USER_COL, as_index=False)[DEFAULT_USER_COL].agg({"total": "count"}),
            on=DEFAULT_USER_COL,
        )

    df_dcg = df_hit.groupby(DEFAULT_USER_COL, as_index=False, sort=False).agg({"dcg": "sum"})

    df_hit_count = pd.merge(
            df_hit_count,
            df_dcg,
            on=DEFAULT_USER_COL,
        )

    df_hit_count["idcg"] = df_hit_count["total"].apply(
        lambda x: sum(1 / np.log1p(range(1, min(x, k) + 1)))
    )

    df_ap = df_hit.groupby(DEFAULT_USER_COL, as_index=False, sort=False).agg({"rr": "sum"})
    df_hit_count = pd.merge(
            df_hit_count,
            df_ap,
            on=DEFAULT_USER_COL,
        )

    precision = (df_hit_count["hit"] / k).sum() / n_users
    recall = (df_hit_count["hit"]/df_hit_count["total"]).sum() / n_users
    ndcg = (df_hit_count["dcg"]/df_hit_count["idcg"]).sum() / n_users
    mapre = (df_hit_count["rr"]/df_hit_count["total"]).sum() / n_users

    results.append(("precision", k, precision))
    results.append(("recall", k, recall))
    results.append(("ndcg", k, ndcg))
    results.append(("map",k,mapre))
```

```python

```

```python
eval_df = pd.DataFrame(results, columns = ['metric','cutoff','value'])
eval_df
```

      metric
      cutoff
      value

      0
      precision
      5
      0.053166

      1
      recall
      5
      0.089077

      2
      ndcg
      5
      0.077163

      3
      map
      5
      0.044160

      4
      precision
      10
      0.051710

      5
      recall
      10
      0.160595

      6
      ndcg
      10
      0.104348

      7
      map
      10
      0.056826

      8
      precision
      20
      0.048273

      9
      recall
      20
      0.310296

      10
      ndcg
      20
      0.151732

      11
      map
      20
      0.072189


# Ejemplo Básico del Algoritmo de Maximización de la Esperanza

> En el notebook anterior derivamos todas las ecuaciones del algoritmo de maximización de la esperanza. Sin duda, este proceso requirió de bastantes matemáticas y concentración, pero lo logramos. Ya que tenemos las ecuaciones, comencemos a utilizarlas con un ejemplo básico.

> **Objetivos:**
> - Aplicar el algoritmo de maximización de la esperanza a un ejemplo sencillo.

> **Referencias:**
> - Bayesian Methods for Machine Learning course, HSE University, Coursera.

## 1. Enunciado: Mezcla discreta

Supongamos que tenemos datos correspondientes al siguiente histograma:

```python
# Importamos matplotlib.pyplot
from matplotlib import pyplot as plt
# Importamos scipy.stats
from scipy import stats
# Importamos numpy
import numpy as np
```

```python
# Graficamos
plt.figure(figsize=(6, 4))
plt.bar([0, 1, 2], [0.35, 0.25, 0.4], color='k')
plt.xlabel('$x$')
plt.ylabel('Freq. rel.')
```

```python
# Variable aleatoria "real"
X = stats.rv_discrete(values=([0, 1, 2], [0.35, 0.25, 0.4]))
data = X.rvs(size=500)
data
```

    array([2, 0, 1, 2, 2, 0, 0, 0, 2, 0, 0, 2, 1, 1, 0, 2, 1, 1, 2, 2, 0, 1,
           2, 2, 1, 2, 0, 0, 1, 2, 2, 0, 0, 0, 2, 2, 2, 2, 0, 0, 0, 0, 2, 0,
           2, 2, 2, 2, 0, 2, 1, 0, 2, 0, 0, 2, 1, 1, 0, 2, 0, 0, 2, 2, 0, 0,
           1, 2, 2, 1, 2, 0, 2, 1, 2, 2, 0, 2, 2, 0, 1, 0, 1, 2, 0, 2, 2, 2,
           1, 1, 0, 2, 2, 2, 0, 0, 1, 1, 0, 0, 0, 2, 2, 1, 2, 1, 0, 0, 1, 2,
           1, 0, 1, 0, 0, 1, 2, 0, 2, 0, 1, 2, 0, 0, 0, 0, 0, 2, 1, 1, 0, 2,
           0, 2, 2, 0, 0, 2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 1, 1, 2, 1, 0, 2, 0,
           0, 1, 0, 0, 0, 2, 0, 0, 2, 0, 0, 0, 2, 2, 1, 0, 1, 2, 0, 0, 0, 0,
           0, 2, 1, 2, 1, 0, 0, 2, 1, 0, 1, 2, 1, 1, 0, 0, 1, 2, 2, 2, 0, 2,
           0, 2, 1, 0, 1, 2, 0, 0, 0, 0, 2, 0, 2, 0, 2, 1, 0, 1, 0, 2, 2, 0,
           2, 2, 0, 2, 0, 1, 2, 0, 0, 0, 1, 2, 0, 2, 2, 1, 2, 1, 2, 0, 1, 0,
           0, 0, 0, 2, 2, 0, 1, 2, 2, 1, 0, 0, 2, 2, 2, 2, 0, 1, 1, 1, 2, 0,
           1, 2, 0, 0, 2, 2, 2, 0, 1, 0, 2, 1, 0, 0, 2, 2, 2, 2, 1, 2, 2, 1,
           1, 1, 0, 2, 2, 1, 0, 1, 2, 0, 1, 0, 0, 0, 0, 2, 0, 0, 0, 1, 0, 0,
           0, 1, 0, 2, 2, 1, 0, 2, 2, 0, 2, 2, 0, 1, 0, 1, 2, 1, 1, 0, 0, 1,
           0, 1, 1, 1, 0, 0, 0, 0, 2, 0, 0, 0, 1, 1, 0, 1, 0, 2, 0, 2, 1, 1,
           2, 2, 2, 0, 0, 2, 2, 2, 2, 1, 0, 2, 0, 2, 0, 0, 0, 2, 1, 1, 0, 1,
           1, 2, 1, 0, 1, 2, 1, 1, 1, 0, 1, 2, 1, 0, 2, 2, 2, 0, 2, 1, 2, 1,
           2, 1, 0, 2, 0, 2, 2, 1, 2, 0, 0, 2, 2, 1, 2, 0, 0, 0, 2, 2, 0, 0,
           2, 0, 2, 0, 2, 2, 2, 2, 0, 2, 0, 1, 0, 0, 2, 0, 0, 2, 0, 1, 2, 2,
           2, 1, 1, 0, 2, 2, 0, 0, 0, 0, 2, 0, 0, 2, 1, 1, 2, 2, 1, 0, 2, 1,
           0, 1, 0, 1, 2, 1, 2, 1, 0, 1, 1, 2, 2, 2, 2, 2, 0, 1, 0, 2, 2, 1,
           2, 2, 2, 0, 0, 1, 0, 2, 2, 1, 1, 2, 2, 2, 2, 2])

Queremos ajustar una distribución de probabilidad a estos datos, pero **suponiendo que estos datos vienen de una mezcla de distribuciones:**

$$
p(x_i) = \alpha p_1(x_i) + (1 - \alpha) p_2(x_i),
$$

donde

| Distribution | $x=0$    | $x=1$       | $x=2$        |
| ------------ | -------- | ----------- | -----------  |
| $p_1(x)$     | $\beta$  | $1 - \beta$ | $0$          |
| $p_2(x)$     | $0$      | $\delta$    | $1 - \delta$ |

Entonces, una suposición plausible es que tenemos un modelo con variable latente de la siguiente manera:

donde $t_i$ es una variable aleatoria que puede tomar dos valores $1, 2$, y básicamente nos dice a qué distribución pertenece el dato $x_i$.

De esta manera, la distribución previa sobre $t_i$ es:

$$
p(t_i) = \left\lbrace \begin{array}{ccc}
                      \alpha     & \text{si} & t_i = 1 \\
                      1 - \alpha & \text{si} & t_i = 2
                      \end{array}\right.
$$

## 2. Algoritmo de maximización de la esperanza

Para inicializar nuestro algoritmo de maximización de la esperanza definimos:

$$
\alpha^0 = \beta^0 = \delta^0 = 0.5
$$

### 1. E-step

Recordamos que en este paso, queremos encontrar la distribución posterior:

$$
q(t_i = c) = p(t_i=c | x_i)
$$

¿A qué es igual $p(t_i | x_i)$? Usando Bayes (recordar que: $p(x_i | t_i=c) = p_c(x_i)$):

\begin{align}
p(t_i | x_i) & = \frac{p(x_i | t_i) p(t_i)}{p(x_i)} \\
             & = \frac{p(x_i | t_i) p(t_i)}{p(x_i | t_i=1)p(t_i=1) + p(x_i | t_i=2)p(t_i=2)} \\
             & = \frac{p(x_i | t_i) p(t_i)}{\alpha p_1(x_i)) + (1 - \alpha)p_2(x_i)}.
\end{align}

Resumamos lo anterior en una tabla:

| $x_i$ \ $t_i$ | $1$                                                          | $2$                                                                  |
| ------------- | ------------------------------------------------------------ | -------------------------------------------------------------------- |
| $0$           | $\frac{\beta\alpha}{\alpha\beta + (1-\alpha)0} = 1$          | $\frac{0(1-\alpha)}{\alpha\beta + (1-\alpha)0}=0$                    |
| $1$           | $\frac{(1-\beta)\alpha}{\alpha(1-\beta) + (1-\alpha)\delta}$ | $\frac{\delta(1-\alpha)}{\alpha(1-\beta) + (1-\alpha)\delta}$        |
| $2$           | $\frac{0 \alpha}{\alpha 0 + (1-\alpha)(1 - \delta)} = 0$     | $\frac{(1-\delta)(1-\alpha)}{\alpha 0 + (1-\alpha)(1 - \delta)} = 1$ |

En particular, para los parámetros iniciales $\alpha^0 = \beta^0 = \delta^0 = 0.5$, tenemos que:

$$
\frac{(1-\beta^0)\alpha^0}{\alpha^0(1-\beta^0) + (1-\alpha^0)\delta^0} = \frac{0.5 \times 0.5}{0.5 \times 0.5 + 0.5 \times 0.5} = 0.5,
$$

y

$$
\frac{\delta^0(1-\alpha^0)}{\alpha^0(1-\beta^0) + (1-\alpha^0)\delta^0} = \frac{0.5 \times 0.5}{0.5 \times 0.5 + 0.5 \times 0.5} = 0.5.
$$

En resumen:

$$
q^1(t_i = 1) = p(t_i=1 | x_i) = \left \lbrace \begin{array}{ccc}
                                              1   & \text{si} & x_i=0 \\
                                              0.5 & \text{si} & x_i=1 \\
                                              0   & \text{si} & x_i=2
                                              \end{array}\right .
$$

y $q^1(t_i = 2) = 1 - q^1(t_i = 1)$.

### 2. M-step

Ahora, queremos maximizar:

\begin{align}
\max_{\alpha, \beta, \delta} \sum_{i=1}^{N} \mathbb{E}_{q^{1}(t_i)} \left[\log p(x_i, t_i)\right]
& = \max_{\alpha, \beta, \delta} \sum_{i=1}^{N} \mathbb{E}_{q^{1}(t_i)} \left[\log p(x_i | t_i) p(t_i)\right] \\
& = \max_{\alpha, \beta, \delta} \sum_{i=1}^{N} \left[q^1(t_i = 1) \log p(x_i | t_i=1) p(t_i = 1) + q^1(t_i = 2) \log p(x_i | t_i=2) p(t_i = 2)\right] \\
& = \max_{\alpha, \beta, \delta} \sum_{i=1}^{N} \left[q^1(t_i = 1) \log p_1(x_i) \alpha + q^1(t_i = 2) \log p_2(x_i) (1 - \alpha)\right]
\end{align}

Notemos que para este punto, la expresión de arriba dependerá de:

- $N_0$: cantidad de datos que tenemos $x=0$,
- $N_1$: cantidad de datos que tenemos $x=1$,
- $N_2$: cantidad de datos que tenemos $x=2$.

De esta manera:

\begin{align}
\sum_{i=1}^{N}\mathbb{E}_{q^{1}(t_i)} \left[\log p(x_i, t_i)\right]
& = \sum_{i=1}^{N} \left[q^1(t_i = 1) \log p_1(x_i) \alpha + q^1(t_i = 2) \log p_2(x_i) (1 - \alpha)\right] \\
& = N_0 \left[\log \beta \alpha\right] + N_1 \left[0.5 \log (1 - \beta) \alpha + 0.5 \log \delta (1 - \alpha)\right] + N_2 \left[\log (1 - \delta) (1 - \alpha)\right]
\end{align}

Para maximizar lo anterior respecto a los parámetros $\alpha, \beta, \delta$, derivamos respecto a cada parámetro e igualamos a cero:

$$
\frac{\partial}{\partial \alpha}\sum_{i=1}^{N}\mathbb{E}_{q^{1}(t_i)} \left[\log p(x_i, t_i)\right] = N_0 \frac{1}{\alpha} + 0.5 N_1 \frac{1}{\alpha} - 0.5 N_1 \frac{1}{1 - \alpha} - N_2 \frac{1}{1 - \alpha} = 0 \Rightarrow \alpha^1 = \frac{N_0 + 0.5 N_1}{N_0 + N_1 + N_2}.
$$

## Tarea

Demostrar que:

$$
\beta^1 = \frac{2 N_0}{2 N_0 + N_1}, \qquad \text{ y } \qquad \delta^1 = \frac{N_1}{N_1 + 2 N_2}
$$

## 3. Visualización

Primero veamos en nuestra inicialización, cómo luce la aproximación:

```python
# Parámetros en la inicialización
alpha0 = 0.5
beta0 = 0.5
delta0 = 0.5
alpha0, beta0, delta0
```

    (0.5, 0.5, 0.5)

```python
# Aproximación de la distribución en la inicialización
p1 = np.array([beta0, 1 - beta0, 0])
p2 = np.array([0, delta0, 1 - delta0])
p_aprox = alpha0 * p1 + (1 - alpha0) * p2
```

```python
# Graficamos
plt.figure(figsize=(6, 4))
plt.bar([0, 1, 2], [0.35, 0.25, 0.4], color="k", label="Real")
plt.bar([0, 1, 2], p_aprox, alpha=0.2, color="r", width=0.5, label="Aproximada")
plt.xlabel('$x$')
plt.ylabel('Freq. rel.')
plt.legend(loc="upper left", bbox_to_anchor=(1.05, 1))
```

Ahora, en nuestra **primera iteración**:

```python
# Hacemos las cuentas
N0 = (data == 0).sum()
N1 = (data == 1).sum()
N2 = (data == 2).sum()
N0, N1, N2
```

    (187, 120, 193)

```python
# Parámetros en la primera iteración
alpha1 = (N0 + 0.5 * N1) / (N0 + N1 + N2)
beta1 = (2 * N0) / (2 * N0 + N1)
delta1 = (N1) / (N1 + 2 * N2)
alpha1, beta1, delta1
```

    (0.494, 0.757085020242915, 0.23715415019762845)

```python
# Aproximación de la distribución en la primera iteración
p1 = np.array([beta1, 1 - beta1, 0])
p2 = np.array([0, delta1, 1 - delta1])
p_aprox = alpha1 * p1 + (1 - alpha1) * p2
```

```python
# Graficamos
plt.figure(figsize=(6, 4))
plt.bar([0, 1, 2], [0.35, 0.25, 0.4], color="k", label="Real")
plt.bar([0, 1, 2], p_aprox, alpha=0.2, color="r", width=0.5, label="Aproximada")
plt.xlabel('$x$')
plt.ylabel('Freq. rel.')
plt.legend(loc="upper left", bbox_to_anchor=(1.05, 1))
```

¡Solo en la primera iteración ya tenemos una aproximación bastante aceptable!

<footer id="attribution" style="float:right; color:#808080; background:#fff;">
Created with Jupyter by Esteban Jiménez Rodríguez.
</footer>


```python
import os
from src.constants import ATMOS_DATA_PATH, DATA_PATH
import xarray as xr
import matplotlib.pyplot as plt
```

```python
xr.open_dataarray(os.path.join(ATMOS_DATA_PATH, "rh-CMIP6-clim60.nc")).plot()
```

```python
os.listdir(ATMOS_DATA_PATH)
```

    ['ps-ECMWF-clim.nc',
     'ts-ECMWF-clim60.nc',
     'sfcWind-ECMWF-clim60.nc',
     'pr-ECMWF-trend.nc',
     'sst-ECMWF-clim.nc',
     'mask-360x181.nc',
     'ts-ECMWF-clim.nc',
     'mask-360x180.nc',
     'clt-ECMWF-clim60.nc',
     'ts-ECMWF-trend.nc',
     'sst-ECMWF-trend.nc',
     'rh-fixed-clim60.nc',
     'rh-ECMWF-clim60.nc',
     'sfcWind-ECMWF-clim.nc',
     'pr-ECMWF-clim.nc',
     'rh-CMIP5-clim60.nc',
     'sfcWind-CMIP5-clim.nc',
     'sfcWind-CMIP5-clim60.nc',
     'rh-CMIP6-clim60.nc',
     'sfcWind-CMIP6-clim.nc',
     'sfcWind-CMIP6-clim60.nc']

```python
xr.open_dataarray(os.path.join(DATA_PATH, "nc_mean", "hur.nc")).plot()
```

```python
xr.open_dataarray(os.path.join(DATA_PATH, "nc_mean", "wsp.nc")).plot()
```

    <matplotlib.collections.QuadMesh at 0x7fd4e105bf40>

```python
import matplotlib.pyplot as plt

with plt.xkcd():
    # This figure will be in XKCD-style
    fig1 = plt.figure()
    plt.plot([0, 1], [1, 0])
    plt.xlim([0, 1])
    plt.ylim([0, 1])
    plt.xlabel("Data/Information")
    plt.ylabel("Assumptions")
    plt.xticks([])
    plt.yticks([])
```

```python
from jupyterthemes import jtplot
jtplot.style(theme="grade3", context="notebook", ticks=True, grid=False)
from src.plot_utils import ps_defaults
ps_defaults(use_tex=False, dpi=150)
from IPython.display import Latex
```

```python
! pip install sympy
```

    Collecting sympy
      Downloading sympy-1.8-py3-none-any.whl (6.1 MB)
         |████████████████████████████████| 6.1 MB 7.1 MB/s eta 0:00:01     |████████████▎                   | 2.3 MB 2.8 MB/s eta 0:00:02
    Collecting mpmath>=0.19
      Downloading mpmath-1.2.1-py3-none-any.whl (532 kB)
         |████████████████████████████████| 532 kB 7.5 MB/s eta 0:00:01
    Installing collected packages: mpmath, sympy
    Successfully installed mpmath-1.2.1 sympy-1.8

```python
from sympy import solve, dsolve, Poly, Eq, Function, exp, Symbol
```

```python
V =  Function("V")
r = Symbol("r")
A = Symbol("A")
r0 = Symbol("r0")
f =  Symbol("f")
```

```python
f =  Function("f")
x = Symbol("x")

solve(f(x).diff(x) - f(x) - x, f(x).diff(x))
```

    [x + f(x)]

```python
dsolve(
    V(r) + r * (V(r)).diff(r) + f * r - A * r ** 2 * V(r) ** 2 / (r0 ** 2 - r ** 2), r
)
```

```python
dsolve(
    V(r) + r * (V(r)).diff(r) + f * r - A * r ** 2 * V(r) ** 2 / (r0 ** 2 - r ** 2), r
)
```

```python

```


# Optimizer tweaks

```python
%load_ext autoreload
%autoreload 2

%matplotlib inline
```

```python
#export
from exp.nb_08 import *
```

## Imagenette data

We grab the data from the previous notebook.

```python
path = datasets.untar_data(datasets.URLs.IMAGENETTE_160)
```

```python
tfms = [make_rgb, ResizeFixed(128), to_byte_tensor, to_float_tensor]
bs=128

il = ImageList.from_files(path, tfms=tfms)
sd = SplitData.split_by_func(il, partial(grandparent_splitter, valid_name='val'))
ll = label_by_func(sd, parent_labeler, proc_y=CategoryProcessor())
data = ll.to_databunch(bs, c_in=3, c_out=10, num_workers=4)
```

Then a model:

```python
nfs = [32,64,128,256]
```

```python
cbfs = [partial(AvgStatsCallback,accuracy), CudaCallback,
        partial(BatchTransformXCallback, norm_imagenette)]
```

This is the baseline of training with vanilla SGD.

```python
learn,run = get_learn_run(nfs, data, 0.4, conv_layer, cbs=cbfs)
```

```python
run.fit(1, learn)
```

    train: [1.7737157604796805, tensor(0.3847, device='cuda:0')]
    valid: [1.4919990234375, tensor(0.4820, device='cuda:0')]

## Refining the optimizer

In PyTorch, the base optimizer in `torch.optim` is just a dictionary that stores the hyper-parameters and references to the parameters of the model we want to train in parameter groups (different groups can have different learning rates/momentum/weight decay... which is what lets us do discriminative learning rates).

It contains a method `step` that will update our parameters with the gradients and a method `zero_grad` to detach and zero the gradients of all our parameters.

We build the equivalent from scratch, only ours will be more flexible. In our implementation, the step function loops over all the parameters to execute the step using stepper functions that we have to provide when initializing the optimizer.

```python
class Optimizer():
    def __init__(self, params, steppers, **defaults):
        # might be a generator
        self.param_groups = list(params)
        # ensure params is a list of lists
        if not isinstance(self.param_groups[0], list): self.param_groups = [self.param_groups]
#             this makes it so that there are hyper parameters for each paramater group.
#             start out here all being the same but different objects for reference and can change
#             (for learning rate annealing, etc.)
#                 the Scheduler is what goes through and actually changes these
        self.hypers = [{**defaults} for p in self.param_groups]
        print("self.hypers!!", self.hypers)
        self.steppers = listify(steppers)
        print("self.steppers!!", self.steppers)

    def grad_params(self):
#         convenience function for getting all the paramaters in all groups -- used to zero them all out in zero_grad
        res = [(p,hyper) for pg,hyper in zip(self.param_groups,self.hypers)
            for p in pg if p.grad is not None]
#         print("grad_params!!", res)
        return res

    def zero_grad(self):
        for p,hyper in self.grad_params():
            p.grad.detach_()
            p.grad.zero_()

    def step(self):
        for p,hyper in self.grad_params(): compose(p, self.steppers, **hyper)
```

To do basic SGD, this what a step looks like:

```python
#export
def sgd_step(p, lr, **kwargs):
#     add in pytorch will add the p.data to p.grad.data AFTER multiplying p.grad.data to the scalar of -lr
    p.data.add_(-lr, p.grad.data)
    return p
```

```python
opt_func = partial(Optimizer, steppers=[sgd_step])
```

```python
opt_func()
```

Now that we have changed the optimizer, we will need to adjust the callbacks that were using properties from the PyTorch optimizer: in particular the hyper-parameters are in the list of dictionaries `opt.hypers` (PyTorch has everything in the the list of param groups).

```python
#export
class Recorder(Callback):
    def begin_fit(self): self.lrs,self.losses = [],[]

    def after_batch(self):
        if not self.in_train: return
        self.lrs.append(self.opt.hypers[-1]['lr'])
        self.losses.append(self.loss.detach().cpu())

    def plot_lr  (self): plt.plot(self.lrs)
    def plot_loss(self): plt.plot(self.losses)

    def plot(self, skip_last=0):
        losses = [o.item() for o in self.losses]
        n = len(losses)-skip_last
        plt.xscale('log')
        plt.plot(self.lrs[:n], losses[:n])

class ParamScheduler(Callback):
    _order=1
    def __init__(self, pname, sched_funcs):
        self.pname,self.sched_funcs = pname,sched_funcs

    def begin_fit(self):
        if not isinstance(self.sched_funcs, (list,tuple)):
            self.sched_funcs = [self.sched_funcs] * len(self.opt.param_groups)

    def set_param(self):
        for f,h in zip(self.sched_funcs,self.opt.hypers):
            h[self.pname] = f(self.n_epochs/self.epochs)

    def begin_batch(self):
        if self.in_train: self.set_param()

class LR_Find(Callback):
    _order=1
    def __init__(self, max_iter=100, min_lr=1e-6, max_lr=10):
        self.max_iter,self.min_lr,self.max_lr = max_iter,min_lr,max_lr
        self.best_loss = 1e9

    def begin_batch(self):
        if not self.in_train: return
        pos = self.n_iter/self.max_iter
        lr = self.min_lr * (self.max_lr/self.min_lr) ** pos
        for pg in self.opt.hypers: pg['lr'] = lr

    def after_step(self):
        if self.n_iter>=self.max_iter or self.loss>self.best_loss*10:
            raise CancelTrainException()
        if self.loss < self.best_loss: self.best_loss = self.loss
```

So let's check we didn't break anything and that recorder and param scheduler work properly.

```python
sched = combine_scheds([0.3, 0.7], [sched_cos(0.3, 0.6), sched_cos(0.6, 0.2)])
```

```python
cbfs = [partial(AvgStatsCallback,accuracy),
        CudaCallback, Recorder,
        partial(ParamScheduler, 'lr', sched)]
```

```python
learn,run = get_learn_run(nfs, data, 0.4, conv_layer, cbs=cbfs, opt_func=opt_func)
```

    self.hypers!! [{'lr': 0.4}]
    self.steppers!! [<function sgd_step at 0x7f51e02af488>]

```python
%time run.fit(1, learn)
```

    train: [1.7521567080618892, tensor(0.3982, device='cuda:0')]
    valid: [1.335036376953125, tensor(0.5580, device='cuda:0')]
    CPU times: user 5.03 s, sys: 2.28 s, total: 7.31 s
    Wall time: 7.97 s

```python
run.recorder.plot_loss()
```

```python
run.recorder.plot_lr()
```

## Weight decay

By letting our model learn high parameters, it might fit all the data points in the training set with an over-complex function that has very sharp changes, which will lead to overfitting.

Weight decay comes from the idea of L2 regularization, which consists in adding to your loss function the sum of all the weights squared. Why do that? Because when we compute the gradients, it will add a contribution to them that will encourage the weights to be as small as possible.

Limiting our weights from growing too much is going to hinder the training of the model, but it will yield to a state where it generalizes better. Going back to the theory a little bit, weight decay (or just `wd`) is a parameter that controls that sum of squares we add to our loss:
``` python
loss_with_wd = loss + (wd/2) * (weights**2).sum()
```

In practice though, it would be very inefficient (and maybe numerically unstable) to compute that big sum and add it to the loss. If you remember a little bit of high schoool math, you should now that the derivative of `p**2` with respect to `p` is simple `2*p`, so adding that big sum to our loss is exactly the same as doing
``` python
weight.grad += wd * weight
```

for every weight in our model, which is equivalent to (in the case of vanilla SGD) updating the parameters
with
``` python
weight = weight - lr*(weight.grad + wd*weight)
```

This last formula explains why the name of this technique is weight decay, as each weight is decayed by a factor `lr * wd`.

This only works for standard SGD, as we have seen that with momentum, RMSProp or in Adam, the update has some additional formulas around the gradient. In those cases, the formula that comes from L2 regularization:
``` python
weight.grad += wd * weight
```
is different than weight decay
``` python
new_weight = weight - lr * weight.grad - lr * wd * weight
```

Most libraries use the first one, but as it was pointed out in [Decoupled Weight Regularization](https://arxiv.org/pdf/1711.05101.pdf) by Ilya Loshchilov and Frank Hutter, it is better to use the second one with the Adam optimizer, which is why fastai made it its default.

Weight decay is subtracting `lr*wd*weight` from the weights. We need this function to have an attribute `_defaults` so that we are sure there is an hyper-parameter of the same name in our `Optimizer`.

```python
#export
def weight_decay(p, lr, wd, **kwargs):
    p.data.mul_(1 - lr*wd)
    return p
weight_decay._defaults = dict(wd=0.)
```

L2 regularization is adding `wd*weight` to the gradients.

```python
#export
def l2_reg(p, lr, wd, **kwargs):
    p.grad.data.add_(wd, p.data)
    return p
l2_reg._defaults = dict(wd=0.)
```

Let's allow steppers to add to our `defaults` (which are the default values of all the hyper-parameters). This helper function adds in `dest` the key/values it finds while going through `os` and applying `f` when they was no `key` of the same name.

```python
#export
def maybe_update(os, dest, f):
    for o in os:
        for k,v in f(o).items():
#             only updating if the key is not there at all
            if k not in dest: dest[k] = v

def get_defaults(d): return getattr(d,'_defaults',{})
```

This is the same as before, we just take the default values of the steppers when none are provided in the kwargs.

```python
#export
class Optimizer():
    def __init__(self, params, steppers, **defaults):
        self.steppers = listify(steppers)
#         will update defaults that have not been assigned with the defaults from the stepper defaults
        maybe_update(self.steppers, defaults, get_defaults)
        # might be a generator
        self.param_groups = list(params)
        # ensure params is a list of lists
        if not isinstance(self.param_groups[0], list): self.param_groups = [self.param_groups]
        self.hypers = [{**defaults} for p in self.param_groups]

    def grad_params(self):
        return [(p,hyper) for pg,hyper in zip(self.param_groups,self.hypers)
            for p in pg if p.grad is not None]

    def zero_grad(self):
        for p,hyper in self.grad_params():
            p.grad.detach_()
            p.grad.zero_()

    def step(self):
        for p,hyper in self.grad_params(): compose(p, self.steppers, **hyper)
```

```python
#export
sgd_opt = partial(Optimizer, steppers=[weight_decay, sgd_step])
```

```python
learn,run = get_learn_run(nfs, data, 0.4, conv_layer, cbs=cbfs, opt_func=sgd_opt)
```

Before trying to train, let's check the behavior works as intended: when we don't provide a value for `wd`, we pull the corresponding default from `weight_decay`.

```python
model = learn.model
```

```python
opt = sgd_opt(model.parameters(), lr=0.1)
test_eq(opt.hypers[0]['wd'], 0.)
test_eq(opt.hypers[0]['lr'], 0.1)
```

But if we provide a value, it overrides the default.

```python
opt = sgd_opt(model.parameters(), lr=0.1, wd=1e-4)
test_eq(opt.hypers[0]['wd'], 1e-4)
test_eq(opt.hypers[0]['lr'], 0.1)
```

Now let's fit.

```python
cbfs = [partial(AvgStatsCallback,accuracy), CudaCallback]
```

```python
learn,run = get_learn_run(nfs, data, 0.3, conv_layer, cbs=cbfs, opt_func=partial(sgd_opt, wd=0.01))
```

```python
run.fit(1, learn)
```

    train: [1.8099562720548317, tensor(0.3771, device='cuda:0')]
    valid: [1.860242919921875, tensor(0.3520, device='cuda:0')]

This is already better than the baseline!

## With momentum

Momentum requires to add some state. We need to save the moving average of the gradients to be able to do the step and store this inside the optimizer state. To do this, we introduce statistics. Statistics are object with two methods:
- `init_state`, that returns the initial state (a tensor of 0. for the moving average of gradients)
- `update`, that updates the state with the new gradient value

We also read the `_defaults` values of those objects, to allow them to provide default values to hyper-parameters.

```python
#export
class StatefulOptimizer(Optimizer):
    def __init__(self, params, steppers, stats=None, **defaults):
        self.stats = listify(stats)
        maybe_update(self.stats, defaults, get_defaults)
        super().__init__(params, steppers, **defaults)
        self.state = {}

    def step(self):
        for p,hyper in self.grad_params():
            if p not in self.state:
                #Create a state for p and call all the statistics to initialize it.
#                 we track state for EACH PARAMETER
                self.state[p] = {}
                maybe_update(self.stats, self.state[p], lambda o: o.init_state(p))
            state = self.state[p]
            for stat in self.stats: state = stat.update(p, state, **hyper)
            compose(p, self.steppers, **state, **hyper)
            self.state[p] = state
```

```python
#export
class Stat():
    _defaults = {}
    def init_state(self, p): raise NotImplementedError
    def update(self, p, state, **kwargs): raise NotImplementedError
```

Here is an example of `Stat`:

```python
class AverageGrad(Stat):
    _defaults = dict(mom=0.9)

    def init_state(self, p): return {'grad_avg': torch.zeros_like(p.grad.data)}
    def update(self, p, state, mom, **kwargs):
        state['grad_avg'].mul_(mom).add_(p.grad.data)
        return state
```

Then we add the momentum step (instead of using the gradients to perform the step, we use the average).

```python
#export
def momentum_step(p, lr, grad_avg, **kwargs):
    p.data.add_(-lr, grad_avg)
    return p
```

```python
sgd_mom_opt = partial(StatefulOptimizer, steppers=[momentum_step,weight_decay],
                  stats=AverageGrad(), wd=0.01)
```

```python
learn,run = get_learn_run(nfs, data, 0.3, conv_layer, cbs=cbfs, opt_func=sgd_mom_opt)
```

```python
run.fit(1, learn)
```

    train: [1.7859231704184118, tensor(0.3843, device='cuda:0')]
    valid: [1.96252978515625, tensor(0.3720, device='cuda:0')]

### Momentum experiments

What does momentum do to the gradients exactly? Let's do some plots to find out!

```python
x = torch.linspace(-4, 4, 200)
y = torch.randn(200) + 0.3
betas = [0.5, 0.7, 0.9, 0.99]
```

```python
def plot_mom(f):
    _,axs = plt.subplots(2,2, figsize=(12,8))
    for beta,ax in zip(betas, axs.flatten()):
        ax.plot(y, linestyle='None', marker='.')
        avg,res = None,[]
        for i,yi in enumerate(y):
            avg,p = f(avg, beta, yi, i)
            res.append(p)
        ax.plot(res, color='red')
        ax.set_title(f'beta={beta}')
```

This is the regular momentum.

```python
def mom1(avg, beta, yi, i):
    if avg is None: avg=yi
    res = beta*avg + yi
    return res,res
plot_mom(mom1)
```

As we can see, with a too high value, it may go way too high with no way to change its course.

Another way to smooth noisy data is to do an exponentially weighted moving average. In this case, there is a dampening of (1-beta) in front of the new value, which is less trusted than the current average.

### This is lerp in pytorch speak. It dappens the new item. In regular momentum we have beta*avg + yi

```python
def ewma(v1, v2, beta): return beta*v1 + (1-beta)*v2
```

```python
def mom2(avg, beta, yi, i):
    if avg is None: avg=yi
    avg = ewma(avg, yi, beta)
    return avg, avg
plot_mom(mom2)
```

We can see it gets to a zero-constant when the data is purely random. If the data has a certain shape, it will get that shape (with some delay for high beta).

```python
y = 1 - (x/3) ** 2 + torch.randn(200) * 0.1
```

```python
y[0]=0.5
```

```python
plot_mom(mom2)
```

### Debiasing is here to correct the wrong information we may have in the very first batch. The debias term corresponds to the sum of the coefficient in our moving average. At the time step i, our average is:

$\begin{align*}
avg_{i} &= \beta\ avg_{i-1} + (1-\beta)\ v_{i} = \beta\ (\beta\ avg_{i-2} + (1-\beta)\ v_{i-1}) + (1-\beta)\ v_{i} \\
&= \beta^{2}\ avg_{i-2} + (1-\beta)\ \beta\ v_{i-1} + (1-\beta)\ v_{i} \\
&= \beta^{3}\ avg_{i-3} + (1-\beta)\ \beta^{2}\ v_{i-2} + (1-\beta)\ \beta\ v_{i-1} + (1-\beta)\ v_{i} \\
&\vdots \\
&= (1-\beta)\ \beta^{i}\ v_{0} + (1-\beta)\ \beta^{i-1}\ v_{1} + \cdots + (1-\beta)\ \beta^{2}\ v_{i-2} + (1-\beta)\ \beta\  v_{i-1} + (1-\beta)\ v_{i}
\end{align*}$

and so the sum of the coefficients is

$\begin{align*}
S &=(1-\beta)\ \beta^{i} + (1-\beta)\ \beta^{i-1} + \cdots + (1-\beta)\ \beta^{2} + (1-\beta)\ \beta + (1-\beta) \\
&= (\beta^{i} - \beta^{i+1}) + (\beta^{i-1} - \beta^{i}) + \cdots + (\beta^{2} - \beta^{3}) + (\beta - \beta^{2}) + (1-\beta) \\
&= 1 - \beta^{i+1}
\end{align*}$

since all the other terms cancel out each other.

By dividing by this term, we make our moving average a true average (in the sense that all the coefficients we used for the average sum up to 1).

```python
def mom3(avg, beta, yi, i):
    if avg is None: avg=0
    avg = ewma(avg, yi, beta)
#     key to debiasing is this divide step
    return avg, avg/(1-beta**(i+1))
plot_mom(mom3)
```

## Adam and friends

In Adam, we use the gradient averages but with dampening (not like in SGD with momentum), so let's add this to the `AverageGrad` class.

```python
dict(somekey=0.9)
```

    {'somekey': 0.9}

```python
#export
class AverageGrad(Stat):
    _defaults = dict(mom=0.9)

    def __init__(self, dampening:bool=False): self.dampening=dampening
    def init_state(self, p): return {'grad_avg': torch.zeros_like(p.grad.data)}
    def update(self, p, state, mom, **kwargs):
        state['mom_damp'] = 1-mom if self.dampening else 1.
        state['grad_avg'].mul_(mom).add_(state['mom_damp'], p.grad.data)
        return state
```

We also need to track the moving average of the gradients squared.

```python
#export
class AverageSqrGrad(Stat):
    _defaults = dict(sqr_mom=0.99)

    def __init__(self, dampening:bool=True): self.dampening=dampening
    def init_state(self, p): return {'sqr_avg': torch.zeros_like(p.grad.data)}
    def update(self, p, state, sqr_mom, **kwargs):
        state['sqr_damp'] = 1-sqr_mom if self.dampening else 1.
        state['sqr_avg'].mul_(sqr_mom).addcmul_(state['sqr_damp'], p.grad.data, p.grad.data)
        return state
```

We will also need the number of steps done during training for the debiasing.

```python
#export
class StepCount(Stat):
    def init_state(self, p): return {'step': 0}
    def update(self, p, state, **kwargs):
        state['step'] += 1
        return state
```

This helper function computes the debias term. If we dampening, `damp = 1 - mom` and we get the same result as before. If we don't use dampening, (`damp = 1`) we will need to divide by `1 - mom` because that term is missing everywhere.

```python
#export
def debias(mom, damp, step): return damp * (1 - mom**step) / (1-mom)
```

Then the Adam step is just the following:

```python
#export
def adam_step(p, lr, mom, mom_damp, step, sqr_mom, sqr_damp, grad_avg, sqr_avg, eps, **kwargs):
    debias1 = debias(mom,     mom_damp, step)
    debias2 = debias(sqr_mom, sqr_damp, step)
    p.data.addcdiv_(-lr / debias1, grad_avg, (sqr_avg/debias2).sqrt() + eps)
    return p
adam_step._defaults = dict(eps=1e-5)
```

```python
#export
def adam_opt(xtra_step=None, **kwargs):
    return partial(StatefulOptimizer, steppers=[adam_step,weight_decay]+listify(xtra_step),
                   stats=[AverageGrad(dampening=True), AverageSqrGrad(), StepCount()], **kwargs)
```

```python
learn,run = get_learn_run(nfs, data, 0.001, conv_layer, cbs=cbfs, opt_func=adam_opt())
```

```python
run.fit(3, learn)
```

    train: [1.717051505302854, tensor(0.4114, device='cuda:0')]
    valid: [1.3037923583984374, tensor(0.5480, device='cuda:0')]
    train: [1.2101377093997208, tensor(0.6031, device='cuda:0')]
    valid: [1.061533203125, tensor(0.6660, device='cuda:0')]
    train: [0.9421420422047077, tensor(0.6971, device='cuda:0')]
    valid: [1.033464599609375, tensor(0.6620, device='cuda:0')]

## LAMB

It's then super easy to implement a new optimizer. This is LAMB from a [very recent paper](https://arxiv.org/pdf/1904.00962.pdf):

$\begin{align}
g_{t}^{l} &= \nabla L(w_{t-1}^{l}, x_{t}) \\
m_{t}^{l} &= \beta_{1} m_{t-1}^{l} + (1-\beta_{1}) g_{t}^{l} \\
v_{t}^{l} &= \beta_{2} v_{t-1}^{l} + (1-\beta_{2}) g_{t}^{l} \odot g_{t}^{l} \\
m_{t}^{l} &= m_{t}^{l} / (1 - \beta_{1}^{t}) \\
v_{t}^{l} &= v_{t}^{l} / (1 - \beta_{2}^{t}) \\
r_{1} &= \|w_{t-1}^{l}\|_{2} \\
s_{t}^{l} &= \frac{m_{t}^{l}}{\sqrt{v_{t}^{l} + \epsilon}} + \lambda w_{t-1}^{l} \\
r_{2} &= \| s_{t}^{l} \|_{2} \\
\eta^{l} &= \eta * r_{1}/r_{2} \\
w_{t}^{l} &= w_{t}^{l-1} - \eta_{l} * s_{t}^{l} \\
\end{align}$

```python
def lamb_step(p, lr, mom, mom_damp, step, sqr_mom, sqr_damp, grad_avg, sqr_avg, eps, wd, **kwargs):
    debias1 = debias(mom,     mom_damp, step)
    debias2 = debias(sqr_mom, sqr_damp, step)
    r1 = p.data.pow(2).mean().sqrt()
    step = (grad_avg/debias1) / ((sqr_avg/debias2).sqrt()+eps) + wd*p.data
    r2 = step.pow(2).mean().sqrt()
    p.data.add_(-lr * min(r1/r2,10), step)
    return p
lamb_step._defaults = dict(eps=1e-6, wd=0.)
```

```python
lamb = partial(StatefulOptimizer, steppers=lamb_step, stats=[AverageGrad(dampening=True), AverageSqrGrad(), StepCount()])
```

```python
learn,run = get_learn_run(nfs, data, 0.003, conv_layer, cbs=cbfs, opt_func=lamb)
```

```python
run.fit(3, learn)
```

    train: [1.846915534841787, tensor(0.3682, device='cuda:0')]
    valid: [1.482430908203125, tensor(0.5140, device='cuda:0')]
    train: [1.3096460502462386, tensor(0.5643, device='cuda:0')]
    valid: [1.1855213623046874, tensor(0.6220, device='cuda:0')]
    train: [1.0266159859963937, tensor(0.6672, device='cuda:0')]
    valid: [1.011562744140625, tensor(0.6740, device='cuda:0')]

Other recent variants of optimizers:
- [Large Batch Training of Convolutional Networks](https://arxiv.org/abs/1708.03888) (LARS also uses weight statistics, not just gradient statistics. Can you add that to this class?)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://arxiv.org/abs/1804.04235) (Adafactor combines stats over multiple sets of axes)
- [Adaptive Gradient Methods with Dynamic Bound of Learning Rate](https://arxiv.org/abs/1902.09843)

## Export

```python
!python notebook2script.py 09_optimizers.ipynb
```

    Converted 09_optimizers.ipynb to exp/nb_09.py

```python

```