[Sebastian Raschka](http://www.sebastianraschka.com)

```
%load_ext watermark
```

```
%watermark -v -d -u
```

    Last updated: 18/07/2014

    CPython 3.4.1
    IPython 2.1.0

[More information](http://nbviewer.ipython.org/github/rasbt/python_reference/blob/master/ipython_magic/watermark.ipynb) about the `watermark` magic command extension.

This IPython notebook can be found in the GitHub repository [rasbt/algorithms_in_ipython_notebooks](https://github.com/rasbt/algorithms_in_ipython_notebooks)

I would be happy to hear your comments and suggestions.
Please feel free to drop me a note via
[twitter](https://twitter.com/rasbt), [email](mailto:bluewoodtree@gmail.com), or [google+](https://plus.google.com/118404394130788869227).

# Dixon's Q test to identify outliers for small sample sizes

Dixon's Q test 1 is a convenient way to quickly identify outliers in datasets that only contains a small number of observations: typically $3 \leq n \leq 10$.

1 R. B. Dean and W. J. Dixon (1951) ["Simplified Statistics for Small Numbers of Observations"](http://pubs.acs.org/doi/abs/10.1021/ac60052a025). Anal. Chem., 1951, 23 (4), 636–638

# Sections

- [Application](#Application)

- [Application](#Application)

- [Criticism](#Criticism)

- [Method](#Method)

    - [1) Arange values for observations in ascending order](#step1)

    - [2) Calculate Q](#step2)

    - [3) Compare calculated $Q_{exp}$ to the tabulated critical Q-value $Q_{crit}$](#step3)

    - [4) Example](#step4)

- [Implementation in Python](#Implementation)

- [Plotting the Data](#Plotting-the-Data)

    - [Bar plot of the sample means with standard deviation](#Bar-plot-of-the-sample-means-with-standard-deviation)

    - [Boxplot](#Boxplot)

## Application

[[back to top]](#Sections)

Although (at least in my opinion), the removal of outliers is a very questionable practice, this test is quite popular in the field of chemistry to "objectively" detect and reject outliers that are due to systematic errors by the experimentalist.

## Criticism

[[back to top]](#Sections)

If we want to use this test to legitimately remove (potential) outliers from a dataset, we should keep in mind that

- our data has to be normal distributed,

- and that we are not supposed to use this test more than once the same data set.

In my opinion, the Dixon Q-test should only be used with great caution, since this simple statistic is based on the assumption that the data is normal distributed, which can be quite challenging to predict for small sample sizes (if no prior/additional information is provided).
Personally, I would use the Dixon Q-test to only **detect** outliers and not to **remove** those, which can help with the identification of uncertainties in the data set or problems in experimental procedures.
Intuitively, this is quite similar to an approach of identifying samples that have a large standard deviation.

For example, if I tested ~1000 chemical compounds in some sort of activity assay - each compound 5 times, I would mark compounds that contain Q-test outliers for re-testing, because there might have been some problem in the measurement procedure that could have caused this inconsistency.

## Method

[[back to top]](#Sections)

### 1) Arange values for observations in ascending order

[[back to top]](#Sections)

First, we arange the data for our sample in ascending order (from the lowest to the highest value):

$x_1 < x_2 < . . . < x_N$

### 2) Calculate Q

[[back to top]](#Sections)

Next, we calculate the experimental Q-value ($Q_{exp}$).

Note that in a later paper in 1953, Dixon and Dean 3 revisted the calculation of the Q-value and reported different equations for different scenarios:

(from: Rorabacher, David B., 1991 2)

- $r_{10} \;\; for \;\; 3 \geq n \leq 7$
- $r_{11} \;\; for \;\; 8 \geq n \leq 10$
- $r_{21} \;\; for \;\; 11 \geq n \leq 13$
- $r_{22} \;\; for \;\; n  \leq 14$

However, according to a statement/observation in a more recent paper (Rorabacher, David B., 1991 2): "The *rl0* ratio is commonly designated as 'Q' and is generally considered to be the most convenient, legitimate, statistical test available for the rejection of deviant values from a small sample conforming to a Gaussian distribution. (It is equally well suited to larger data sets if only one outlier is present.)"

Therefore, I will use *rl0* for the following implementation of the Dixon Q-test:

\begin{equation} Q_{exp} = \frac{x_{2} - x_{1}}{x_{n} - x_{1}} \end{equation}

where it is assumed that the data is arranged in ascending order: $x_1 < x_2 < . . . < x_N$

2 Rorabacher, David B. (1991) “[Statistical Treatment for Rejection of Deviant Values: Critical Values of Dixon’s‘ Q’ Parameter and Related Subrange Ratios at the 95% Confidence Level.](http://pubs.acs.org/doi/pdf/10.1021/ac00002a010)” Analytical Chemistry 63, no. 2 (1991): 139–46.

3 W.J. Dixon: ["Processing data for outliers Reference"](http://webspace.ship.edu/pgmarr/Geo441/Readings/Dixon%201953%20-%20Processing%20Data%20for%20Outliers.pdf): J. Biometrics 9 (1953) 74-89

### 3) Compare calculated $Q_{exp}$ to the tabulated critical Q-value $Q_{crit}$

[[back to top]](#Sections)

In the next step, we will compare the calculated $Q_{exp}$ value to the to the tabulated critical Q-value $Q_{crit}$ for a chosen confidence interval.
If the calculated Q-value for a particular observation is larger than the critical Q-value ($Q_{exp}$ > $Q_{crit}$), this observation is considered to be an outlier according to the Q-test.

**r10 Critical values for Dixon's two-tailored Q-Test for 3 different confidence levels**

| N  | Q90%  | Q95%  | Q99%  |
|----|-------|-------|-------|
| 3  | 0.941 | 0.97  | 0.994 |
| 4  | 0.765 | 0.829 | 0.926 |
| 5  | 0.642 | 0.71  | 0.821 |
| 6  | 0.56  | 0.625 | 0.74  |
| 7  | 0.507 | 0.568 | 0.68  |
| 8  | 0.468 | 0.526 | 0.634 |
| 9  | 0.437 | 0.493 | 0.598 |
| 10 | 0.412 | 0.466 | 0.568 |
| 11 | 0.392 | 0.444 | 0.542 |
| 12 | 0.376 | 0.426 | 0.522 |
| 13 | 0.361 | 0.41  | 0.503 |
| 14 | 0.349 | 0.396 | 0.488 |
| 15 | 0.338 | 0.384 | 0.475 |
| 16 | 0.329 | 0.374 | 0.463 |
| 17 | 0.32  | 0.365 | 0.452 |
| 18 | 0.313 | 0.356 | 0.442 |
| 19 | 0.306 | 0.349 | 0.433 |
| 20 | 0.3   | 0.342 | 0.425 |
| 21 | 0.295 | 0.337 | 0.418 |
| 22 | 0.29  | 0.331 | 0.411 |
| 23 | 0.285 | 0.326 | 0.404 |
| 24 | 0.281 | 0.321 | 0.399 |
| 25 | 0.277 | 0.317 | 0.393 |
| 26 | 0.273 | 0.312 | 0.388 |
| 27 | 0.269 | 0.308 | 0.384 |
| 28 | 0.266 | 0.305 | 0.38  |
| 29 | 0.263 | 0.301 | 0.376 |
| 30 | 0.26  | 0.29  | 0.372 |

### 4) Example

[[back to top]](#Sections)

Let's consider the following sample consisting of 5 observations:

0.142, 0.153, 0.135, 0.002, 0.175

- First, we sort it in ascending order:
0.002, 0.135, 0.142, 0.153, 0.175

- Next, we calculate the Q-value:

\begin{equation} Q = \frac{0.135 - 0.002}{0.175-0.002} \approx 0.7687 \end{equation}

- Now, we look up the critical value for n=5 for a confidence level 95% in the Q-table $\Rightarrow 0.71$
and we conclude that 0.002 (since 0.7687 > 0.71), that the observation 0.002 is an outlier at a confidence level of 95% according to Dixon's Q-test.

# Implementation

[[back to top]](#Sections)

### Building dictionaries for Q-value look-up

[[back to top]](#Sections)

We will build a simple set of dictionaries for different confidence intervals from the tabulated data in David B. Rorabacher's paper:
Rorabacher, David B. (1991) “[Statistical Treatment for Rejection of Deviant Values: Critical Values of Dixon’s‘ Q’ Parameter and Related Subrange Ratios at the 95% Confidence Level.](http://pubs.acs.org/doi/pdf/10.1021/ac00002a010)” Analytical Chemistry 63, no. 2 (1991): 139–46.

which we will use to look up the critical Q-values (dictionary values) for different sample sizes (dictionary keys).

```
q90 = [0.941, 0.765, 0.642, 0.56, 0.507, 0.468, 0.437,
       0.412, 0.392, 0.376, 0.361, 0.349, 0.338, 0.329,
       0.32, 0.313, 0.306, 0.3, 0.295, 0.29, 0.285, 0.281,
       0.277, 0.273, 0.269, 0.266, 0.263, 0.26
      ]

q95 = [0.97, 0.829, 0.71, 0.625, 0.568, 0.526, 0.493, 0.466,
       0.444, 0.426, 0.41, 0.396, 0.384, 0.374, 0.365, 0.356,
       0.349, 0.342, 0.337, 0.331, 0.326, 0.321, 0.317, 0.312,
       0.308, 0.305, 0.301, 0.29
      ]

q99 = [0.994, 0.926, 0.821, 0.74, 0.68, 0.634, 0.598, 0.568,
       0.542, 0.522, 0.503, 0.488, 0.475, 0.463, 0.452, 0.442,
       0.433, 0.425, 0.418, 0.411, 0.404, 0.399, 0.393, 0.388,
       0.384, 0.38, 0.376, 0.372
       ]

Q90 = {n:q for n,q in zip(range(3,len(q90)+1), q90)}
Q95 = {n:q for n,q in zip(range(3,len(q95)+1), q95)}
Q99 = {n:q for n,q in zip(range(3,len(q99)+1), q99)}
```

### Implementing a Dixon Q-test Function

[[back to top]](#Sections)

Below, I wrote some simple Python code to test one data row for Dixon Q-test outliers:

```
def dixon_test(data, left=True, right=True, q_dict=Q95):
    """
    Keyword arguments:
        data = A ordered or unordered list of data points (int or float).
        left = Q-test of minimum value in the ordered list if True.
        right = Q-test of maximum value in the ordered list if True.
        q_dict = A dictionary of Q-values for a given confidence level,
            where the dict. keys are sample sizes N, and the associated values
            are the corresponding critical Q values. E.g.,
            {3: 0.97, 4: 0.829, 5: 0.71, 6: 0.625, ...}

    Returns a list of 2 values for the outliers, or None.
    E.g.,
       for [1,1,1] -> [None, None]
       for [5,1,1] -> [None, 5]
       for [5,1,5] -> [1, None]

    """
    assert(left or right), 'At least one of the variables, `left` or `right`, must be True.'
    assert(len(data) >= 3), 'At least 3 data points are required'
    assert(len(data) <= max(q_dict.keys())), 'Sample size too large'

    sdata = sorted(data)
    Q_mindiff, Q_maxdiff = (0,0), (0,0)

    if left:
        Q_min = (sdata[1] - sdata[0])
        try:
            Q_min /= (sdata[-1] - sdata[0])
        except ZeroDivisionError:
            pass
        Q_mindiff = (Q_min - q_dict[len(data)], sdata[0])

    if right:
        Q_max = abs((sdata[-2] - sdata[-1]))
        try:
            Q_max /= abs((sdata[0] - sdata[-1]))
        except ZeroDivisionError:
            pass
        Q_maxdiff = (Q_max - q_dict[len(data)], sdata[-1])

    if not Q_mindiff[0] > 0 and not Q_maxdiff[0] > 0:
        outliers = [None, None]

    elif Q_mindiff[0] == Q_maxdiff[0]:
        outliers = [Q_mindiff[1], Q_maxdiff[1]]

    elif Q_mindiff[0] > Q_maxdiff[0]:
        outliers = [Q_mindiff[1], None]

    else:
        outliers = [None, Q_maxdiff[1]]

    return outliers
```

### Assertion Tests

[[back to top]](#Sections)

Some simple assertion tests to make sure that the Dixon Q-test function behaves as expected/desired.

```
test_data1 = [0.142, 0.153, 0.135, 0.002, 0.175]
test_data2 = [0.542, 0.153, 0.135, 0.002, 0.175]

assert(dixon_test(test_data1) == [0.002, None]), 'expect [0.002, None]'
assert(dixon_test(test_data1, right=False) == [0.002, None]), 'expect [0.002, None]'
assert(dixon_test(test_data2) == [None, None]), 'expect [None, None]'
assert(dixon_test(test_data2, q_dict=Q90) == [None, 0.542]), 'expect [None, 0.542]'

print('ok')
```

    ok

# Example application

[[back to top]](#Sections)

Below, I want to go through a naive example for our Dixon Q-test function using an example CSV file.
In "real" application I would prefer `NumPy` and/or `pandas`, however for this simple case the in-built Python `csv` library should suffice.

Below the example CSV file is shown that we are going to read in:

```
%%writefile ../../data/dixon_test_in.csv
,x1,x2,x3,x4,x5
id1,0.95,-0.65,0.6,0.82,NaN
id2,2.08,NaN,-1.43,0.38,NaN
id3,-0.46,NaN,-1.25,-2.62,0.22
id4,0.24,1.88,-0.49,-0.73,-0.49
id5,-1.65,2.1,-0.09,NaN,0.8
id6,-0.44,0.93,0.19,-4.36,-0.88
id7,0.36,-0.47,NaN,0.4,2.12
id8,1.29,-0.48,-0.6,-0.38,0.27
id9,-1.25,-1.35,1.13,1.7,-0.81
id10,0.04,1.98,NaN,NaN,NaN
```

    Overwriting ../../data/dixon_test_in.csv

```
import csv

def csv_to_list(csv_file, delimiter=','):
    """
    Reads in a CSV file and returns the contents as list,
    where every row is stored as a sublist, and each element
    in the sublist represents 1 cell in the table.

    """
    with open(csv_file, 'r') as csv_con:
        reader = csv.reader(csv_con, delimiter=delimiter)
        return list(reader)

def print_csv(csv_content):
    """ Prints CSV file to standard output."""
    print(50*'-')
    for row in csv_content:
        row = [str(e) for e in row]
        print('\t'.join(row))
    print(50*'-')

def convert_cells_to_floats(csv_cont):
    """
    Converts cells to floats if possible
    (modifies input CSV content list).

    """
    for row in range(len(csv_cont)):
        for cell in range(len(csv_cont[row])):
            try:
                csv_cont[row][cell] = float(csv_cont[row][cell])
            except ValueError:
                pass
```

```
csv_cont = csv_to_list('../../data/dixon_test_in.csv')
convert_cells_to_floats(csv_cont)
print_csv(csv_cont)
```

    --------------------------------------------------
    	x1	x2	x3	x4	x5
    id1	0.95	-0.65	0.6	0.82	nan
    id2	2.08	nan	-1.43	0.38	nan
    id3	-0.46	nan	-1.25	-2.62	0.22
    id4	0.24	1.88	-0.49	-0.73	-0.49
    id5	-1.65	2.1	-0.09	nan	0.8
    id6	-0.44	0.93	0.19	-4.36	-0.88
    id7	0.36	-0.47	nan	0.4	2.12
    id8	1.29	-0.48	-0.6	-0.38	0.27
    id9	-1.25	-1.35	1.13	1.7	-0.81
    id10	0.04	1.98	nan	nan	nan
    --------------------------------------------------

Now, let us add a new `outlier` column and apply the Dixon Q-test function to our data set.

```
import math

csv_cont[0].append('outlier')

for row in csv_cont[1:]: # skips header
    nan_removed = [i for i in row[1:] if not math.isnan(i)]
    if len(nan_removed) >= 3:
        row.append(dixon_test(nan_removed, left=True, right=True, q_dict=Q90))
    else:
        row.append('NaN')
```

```
print_csv(csv_cont)
```

    --------------------------------------------------
    	x1	x2	x3	x4	x5	outlier
    id1	0.95	-0.65	0.6	0.82	nan	[-0.65, None]
    id2	2.08	nan	-1.43	0.38	nan	[None, None]
    id3	-0.46	nan	-1.25	-2.62	0.22	[None, None]
    id4	0.24	1.88	-0.49	-0.73	-0.49	[None, None]
    id5	-1.65	2.1	-0.09	nan	0.8	[None, None]
    id6	-0.44	0.93	0.19	-4.36	-0.88	[-4.36, None]
    id7	0.36	-0.47	nan	0.4	2.12	[None, None]
    id8	1.29	-0.48	-0.6	-0.38	0.27	[None, None]
    id9	-1.25	-1.35	1.13	1.7	-0.81	[None, None]
    id10	0.04	1.98	nan	nan	nan	NaN
    --------------------------------------------------

As we can see in the table above, we have 2 potential outliers in our data set.

Finally, we let us write the results to a new CSV file for future reference:

```
def write_csv(dest, csv_cont):
    """ Writes a comma-delimited CSV file. """

    with open(dest, 'w') as out_file:
        writer = csv.writer(out_file, delimiter=',')
        for row in csv_cont:
            writer.writerow(row)

write_csv('../../data/dixon_test_out.csv', csv_cont)
```

# Plotting the Data

[[back to top]](#Sections)

To get a visual impression of how our data looks like, let us make some simple plots.

## Bar plot of the sample means with standard deviation

[[back to top]](#Sections)

First, let's create a bar plot with standard deviation, since such bar plots with standard deviation or standard error bars are probably the most common plots in any field. Although it is not always approapriate it is certainly the kind of data visualization we are most familiar with - due to the frequent exposure when reading scientific research articles.

```
%matplotlib inline
```

```
import numpy as np
from matplotlib import pyplot as plt

all_means = [np.nanmean(row[1:6]) for row in csv_cont[1:]]
all_stddevs = [np.nanstd(row[1:6]) for row in csv_cont[1:]]

fig = plt.figure(figsize=(8,6))

y_pos = np.arange(len(csv_cont[1:]))
y_pos = [x for x in y_pos]
plt.yticks(y_pos, [row[0] for row in csv_cont[1:]], fontsize=10)
plt.xlabel('measurement x')
t = plt.title('Bar plot with standard deviation')

plt.grid()
plt.barh(y_pos, all_means, xerr=all_stddevs, align='center', alpha=0.4, color='g')

plt.show()
```

As we would expect, basically every sample has a very (relatively) large standard deviation. However, we can't derive any detail about a possible cause for this large deviation: whether it is just due to 1 outlier, or if the whole data widely spread in general.

## Boxplot

[[back to top]](#Sections)

A more useful plot in my opinion is Tukey's boxplot 4. Boxplots are in facts one of my preferred approaches to quickly and visually indicate outliers in a Gaussian data set. However, also boxplots have to be used with real caution and might also not very informative for small sample sizes.

 4 Robert McGill, John W. Tukey and Wayne A. Larsen: "[The American Statistician](http://www.jstor.org/discover/10.2307/2683468?uid=3739256&uid=2&uid=4&sid=21104147331297)"
Vol. 32, No. 1 (Feb., 1978), pp. 12-16

```
csv_nonan = [[x for x in row[1:6] if not math.isnan(x)] for row in csv_cont[1:]]
fig = plt.figure(figsize=(8,6))
plt.boxplot(csv_nonan,0,'rs',0)
plt.yticks([y+1 for y in y_pos], [row[0] for row in csv_cont[1:]])
plt.xlabel('measurement x')
t = plt.title('Box plot')
plt.show()
```

The red squares indicate the outliers here. Quite interestingly, both outliers for the samples "id6" and "id1" where also picked up in our previous Dixon Q-test. However, the outliers in "id4" and "id7" were not indicated as outliers by Dixon's outlier test.

**I really don't want to draw any conclusion about which approach is right or wrong here, since in my opinion, drawing any conclusion from a data set that is based on such a small number of observations simply just doesn't make sense!**

So you may wonder why I wasted your time if you read this article up to this point? Since Dixon's Q-test is still quite popular in certain scientific fields (e.g., chemistry) that it is important to understand its principles in order to draw your own conclusion of the presented research data that you might stumble upon in research articles or scientific talks.


# MIT License
# Copyright (c) 2020 F. Anders (ICCUB)
# Reference: https://ui.adsabs.harvard.edu/abs/2020arXiv200601690A/abstract
# Completeness of the Gaia DR2 open cluster census

This notebook provides a complementary analysis for the Cluster Age Function analysis of the Gaia DR2 Open Custer census performed in the Appendix of [this A&A Letter](https://ui.adsabs.harvard.edu/abs/2020arXiv200601690A/abstract). In particular, we characterise the completeness of the [Cantat-Gaudin+2020](https://ui.adsabs.harvard.edu/abs/2020arXiv200407274C/abstract) catalogue of Galactic open clusters.

We make the following approximation:

\begin{align}
p({\it Gaia}\, {\rm DR2\ cluster\, has\, CG2020\, parameters}) &= p({\it Gaia}\, {\rm DR2\, detection}) \cdot p({\rm ANN\ converged} | {\it Gaia}\, {\rm DR2\, detection}) \\
&\approx p({\it Gaia}\, {\rm DR2\, detection} | l, b, d, \log t, A_V) \cdot p({\rm ANN\ converged} | d, \log t)\\
&\approx {\rm completeness\ of\ CastroGinard2020} \cdot {\rm completeness\ of\ CantatGaudin2020\ ANN}
\end{align}

In other words, the completeness of the age census of Gaia DR2 OCs depends on 1. our detection efficiency and 2. on our ability to derive OC parameters once we know that there is a cluster.

The implicit assumptions in this approximation are:
* The Gaia DR2 clusters contained in [Cantat-Gaudin+2020](https://ui.adsabs.harvard.edu/abs/2020arXiv200407274C/abstract) are all true physical OCs. Due to the exquisite astrometric quality of Gaia and the conservative requirements of CG+2020, this is probably a good assumption.
* The experiment of [Castro-Ginard+2020](https://ui.adsabs.harvard.edu/abs/2020A%26A...635A..45C/abstract) is representative for the full catalogue. This is hard to estimate, but the best we can do at present: On the one hand, this estimate is pessimistic, because we only estimate the completeness of 1 method (DBSCAN), while the full catalogue of Gaia-detected OCs was compiled from various different search methods (traditional OCs, serendipitous discoveries, FoF search, etc.). It is also pessimistic in the sense that the method has added more than 500 new OCs! On the other hand, one might say that the DBSCAN completeness test of [Castro-Ginard+2020](https://ui.adsabs.harvard.edu/abs/2020A%26A...635A..45C/abstract) was slightly optimistic, since it only used the catalogue of [Cantat-Gaudin et al. (2018)](https://ui.adsabs.harvard.edu/abs/2018A%26A...618A..93C/abstract) as an input, and not the most compilation of Gaia DR2-detected OCs. In summary, it is the best we can do at present.

## I. Analysing the completeness of Castro-Ginard et al. (2020) clusters

```python
import numpy as np
%matplotlib notebook
import matplotlib
import matplotlib.pylab as plt
import os.path
import matplotlib.gridspec as gridspec
from matplotlib import colors
from matplotlib.colors import LogNorm
from matplotlib import patches
from astropy.table import Table, Column, MaskedColumn
import astropy.table
import matplotlib
from matplotlib import rcParams
rcParams['text.usetex'] = False
rcParams["axes.grid"]   = False
from scipy.stats import gaussian_kde
from scipy import interpolate

import abj2016
```

Load list of Gaia DR2 clusters as of April 2020 (Cantat-Gaudin et al. 2020):

```python
nntable = Table.read("./data/Table1_CDS.fits")
nntable["oc"] = [np.char.strip(a, ' ') for a in nntable["oc"]]
```

[Castro-Ginard+2020](https://ui.adsabs.harvard.edu/abs/2020A%26A...635A..45C/abstract) investigated the recovery fraction of their catalogue by comparing their blind-search candidates obtained with DBSCAN to the OC list of [Cantat-Gaudin et al. (2018)](https://ui.adsabs.harvard.edu/abs/2018A%26A...618A..93C/abstract). They categorise these clusters into recovered/half-recovered/non-recovered - see [Fig. 1 of Castro-Ginard et al. (2020)](https://ui.adsabs.harvard.edu/abs/2020A%26A...635A..45C/graphics).

Here we read the output data of this test and join them into a table:

```python
fail = np.load("./data/clustering_no.npy")
half = np.load("./data/clustering_yes_ANN_no.npy")
good = np.load("./data/clustering_yes_ANN_yes.npy")
print(len(fail), "not found;", len(half), "found but rejected;", len(good), "found & okay.")
print(len(fail) + len(half) + len(good), "in total")

alfred = Table([np.hstack([fail, half, good]),
                np.hstack([np.zeros(len(fail)), .5*np.ones(len(half)), np.ones(len(good))])],
                names=("oc", "foundness"), )
```

    265 not found; 276 found but rejected; 688 found & okay.
    1229 in total

Join both tables:

```python
joint = astropy.table.join(alfred, nntable)
joint.filled()
```

Table length=1225

oc	foundness	ra	dec	l	b	r50	nbstars07	pmra	sigpmra	pmdec	sigpmdec	par	sigpar	flag	ageNN	avNN	dmNN	dist	X	Y	Z	Rgc

str17	float64	float32	float32	float32	float32	float32	int16	float32	float32	float32	float32	float32	float32	bytes14	float32	float32	float32	int16	int16	int16	int16	int16

ASCC_10	1.0	51.87	34.981	155.723	-17.77	0.558	65	-1.737	0.159	-1.368	0.143	1.459	0.1	NN	8.42	0.63	9.1	660	-573	258	-201	8917

ASCC_101	1.0	288.399	36.369	68.028	11.608	0.372	69	0.934	0.205	1.288	0.258	2.488	0.057	NN	8.69	0.19	8.08	412	151	374	83	8197

ASCC_105	0.0	295.548	27.366	62.825	2.063	0.648	113	1.464	0.162	-1.635	0.145	1.783	0.065	NN	7.87	0.41	8.77	567	259	504	20	8096

ASCC_107	1.0	297.164	21.987	58.904	-1.901	0.174	57	-0.155	0.101	-5.156	0.115	1.109	0.049	NN	7.23	1.69	9.7	870	449	744	-28	7925

ASCC_108	1.0	298.306	39.349	74.378	6.074	0.537	188	-0.519	0.099	-1.69	0.129	0.838	0.048	NN	8.03	0.34	10.32	1160	310	1110	122	8105

ASCC_11	1.0	53.056	44.856	150.546	-9.224	0.312	245	0.926	0.163	-3.03	0.147	1.141	0.061	NN	8.39	0.6	9.69	867	-745	421	-139	9095

ASCC_110	1.0	300.742	33.528	70.411	1.378	0.203	65	0.271	0.064	-3.132	0.05	0.497	0.036	NN	8.9	0.75	11.47	1971	660	1856	47	7900

ASCC_111	0.0	302.891	37.515	74.714	2.056	0.537	136	-1.15	0.151	-1.524	0.154	1.166	0.059	NN	8.44	0.39	9.65	851	224	820	30	8157

ASCC_113	1.0	317.933	38.638	82.877	-6.589	0.529	175	0.8	0.125	-3.679	0.163	1.762	0.041	NN	8.46	0.13	8.75	563	69	555	-64	8289

ASCC_114	1.0	324.99	53.997	97.082	1.028	0.216	134	-3.716	0.11	-3.421	0.115	1.066	0.037	NN	7.68	0.92	9.79	908	-111	901	16	8499

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

Turner_9	1.0	296.165	29.494	64.943	2.648	0.151	34	0.346	0.092	-2.399	0.075	0.563	0.034	NN	8.49	0.85	11.25	1774	750	1605	81	7757

Waterloo_1	1.0	64.682	52.863	151.279	1.799	0.065	47	0.492	0.133	-0.689	0.081	0.214	0.071	NN	7.22	2.6	13.35	4670	-4094	2243	146	12634

Waterloo_7	1.0	111.532	-15.09	230.275	0.623	0.024	35	-1.514	0.099	1.695	0.072	0.27	0.06	NN	7.48	1.48	12.76	3560	-2275	-2738	38	10963

Westerlund_1	0.0	251.76	-45.852	339.546	-0.401	0.023	308	-2.325	0.571	-3.654	0.427	0.252	0.192	tooRed	nan	nan	nan	-32768	-32768	-32768	-32768	-32768

Westerlund_2	0.5	156.01	-57.758	284.272	-0.328	0.02	168	-5.246	0.262	2.895	0.302	0.209	0.118	tooRed	nan	nan	nan	-32768	-32768	-32768	-32768	-32768

vdBergh_1	0.5	99.275	3.078	208.556	-1.781	0.03	68	0.396	0.18	-0.771	0.164	0.523	0.067	NN	7.61	1.89	11.45	1946	-1709	-930	-60	10092

vdBergh_80	0.0	97.739	-9.625	219.247	-8.891	0.143	75	-3.298	0.397	0.418	0.385	1.025	0.103	NN	6.81	1.53	9.8	910	-696	-569	-140	9054

vdBergh_83	0.0	100.025	-27.189	236.448	-14.329	0.151	48	-2.896	0.193	3.199	0.234	1.034	0.093	NN	7.49	0.51	9.76	897	-480	-724	-222	8850

vdBergh_85	0.5	101.718	1.32	211.237	-0.414	0.04	29	-1.0	0.155	0.328	0.223	0.552	0.045	NN	6.64	1.35	11.14	1690	-1445	-876	-12	9824

vdBergh_92	1.0	106.038	-11.475	224.587	-2.427	0.112	168	-4.54	0.198	1.618	0.226	0.842	0.081	NN	6.89	1.27	10.38	1193	-849	-837	-50	9227

Plot some projections highlighting the recovery rate of the Castro-Ginard et al. (2020) method (similar to their Fig. 1 that showed this distribution as a function of $l$ and $b$):

```python
fig = plt.figure(figsize=(15.0, 5.0))
ax  = plt.subplot(131)

mask0 = (joint["foundness"]==0.0)
mask1 = (joint["foundness"]==0.5)
mask2 = (joint["foundness"]==1.)
s=10

plt.scatter(joint[mask0]["X"]/1000., joint[mask0]["Y"]/1000., c="r", s=s)
plt.scatter(joint[mask1]["X"]/1000., joint[mask1]["Y"]/1000., c="y", s=s)
plt.scatter(joint[mask2]["X"]/1000., joint[mask2]["Y"]/1000., c="g", s=s)
plt.axis([-4,4,-4,4])
ax.set_xlabel(r"$X_{\rm sol}$  [kpc]", fontsize=16)
ax.set_ylabel(r"$Y_{\rm sol}$    [kpc]", fontsize=16)

ax  = plt.subplot(132)
joint["dxy"] = np.sqrt(joint["X"]**2.+joint["Y"]**2.)/1000.

plt.scatter(joint[mask0]["ageNN"], np.sqrt(joint[mask0]["X"]**2.+joint[mask0]["Y"]**2.)/1000., c="r", s=s)
plt.scatter(joint[mask1]["ageNN"], np.sqrt(joint[mask1]["X"]**2.+joint[mask1]["Y"]**2.)/1000., c="y", s=s)
plt.scatter(joint[mask2]["ageNN"], np.sqrt(joint[mask2]["X"]**2.+joint[mask2]["Y"]**2.)/1000., c="g", s=s)
plt.axis([6,10,0,4])
ax.set_xlabel(r"$\log t$", fontsize=16)
ax.set_ylabel(r"$d_{\rm xy}$    [kpc]", fontsize=16)

ax  = plt.subplot(133)
plt.scatter(joint[mask0]["ageNN"], joint[mask0]["avNN"], c="r", s=s)
plt.scatter(joint[mask1]["ageNN"], joint[mask1]["avNN"], c="y", s=s)
plt.scatter(joint[mask2]["ageNN"], joint[mask2]["avNN"], c="g", s=s)
plt.axis([6,10,0,5])
ax.set_xlabel(r"$\log t$", fontsize=16)
ax.set_ylabel(r"$A_V$    [mag]", fontsize=16)

plt.savefig("./im/combinedmap_clusters_recovered_4kpc.png")
```

    <IPython.core.display.Javascript object>

    /home/hpc/comas/fanders/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:17: RuntimeWarning: invalid value encountered in sqrt
    /home/hpc/comas/fanders/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:19: RuntimeWarning: invalid value encountered in sqrt
    /home/hpc/comas/fanders/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:20: RuntimeWarning: invalid value encountered in sqrt
    /home/hpc/comas/fanders/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:21: RuntimeWarning: invalid value encountered in sqrt

### Putting it together

We see that the recovery rate does not depend significantly on $A_V$ (beyond the intrinsic correlation of this parameter with distance). Can we get the selection fraction in bins of [$l, b, d_{xy}, \log t$] ?

Yes: we just count the recovered, half-recovered, and not recovered clusters in wide bins of these parameters:

```python
lgrid = np.linspace(0, 360, 5)
bgrid = np.linspace(-20,20, 3)
tgrid = np.linspace(6,10, 5)
dgrid = np.linspace(0.5,3,6)

generalmask = (joint["dxy"] > dgrid[0]) * (joint["dxy"] < dgrid[-1])

selfunc = np.zeros((len(lgrid) - 1, len(bgrid) - 1, len(tgrid) - 1, len(dgrid) - 1))
numstars= np.zeros((len(lgrid) - 1, len(bgrid) - 1, len(tgrid) - 1, len(dgrid) - 1))
sumstars= np.zeros((len(lgrid) - 1, len(bgrid) - 1, len(tgrid) - 1, len(dgrid) - 1))

for ii in np.arange(len(lgrid) - 1):
    for jj in np.arange(len(bgrid) - 1):
        for kk in np.arange(len(tgrid) - 1):
            for ll in np.arange(len(dgrid) - 1):
                mask = ( joint["l"] > lgrid[ii] ) & ( joint["l"] < lgrid[ii+1] ) & \
                       ( joint["b"] > bgrid[jj] ) & ( joint["b"] < bgrid[jj+1] ) & \
                       ( joint["ageNN"] > tgrid[kk] ) & ( joint["ageNN"] < tgrid[kk+1] ) & \
                       ( joint["dxy"]   > dgrid[ll] ) &  ( joint["dxy"] < dgrid[ll+1] ) & generalmask
                selfunc[ii,jj,kk,ll]  = np.sum(joint["foundness"][mask]) / np.sum(mask)

np.save("./data/castro-ginard-completeness-selfrac.npy", selfunc)
```

    /home/hpc/comas/fanders/anaconda3/lib/python3.7/site-packages/astropy/table/column.py:1020: RuntimeWarning: invalid value encountered in greater
      result = getattr(super(), op)(other)
    /home/hpc/comas/fanders/anaconda3/lib/python3.7/site-packages/astropy/table/column.py:1020: RuntimeWarning: invalid value encountered in less
      result = getattr(super(), op)(other)
    /home/hpc/comas/fanders/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:20: RuntimeWarning: invalid value encountered in double_scalars

```python
fig = plt.figure(figsize=(12,10))
gs = gridspec.GridSpec(3,3)
gs.update(hspace=0.22, wspace=0.35, top=0.95)
gridlist = [lgrid, bgrid, tgrid, dgrid]
gridname = ["$l$ [deg]", "$b$ [deg]", "$\log\ t$", "$d_{xy}$ [kpc]"]
coordname= ["l", "b", "ageNN", "dxy"]

mask0 = generalmask * (joint["foundness"]==0.0)
mask1 = generalmask * (joint["foundness"]==0.5)
mask2 = generalmask * (joint["foundness"]==1.)

for ii in [0,1,2]:
    for jj in [1,2,3]:
        xx = gridlist[ii]
        yy = gridlist[jj]
        if ii==jj or [ii,jj]==[2,1]:
            pass
        else:
            print(gridname[ii], gridname[jj])
            ax = plt.subplot(gs[2-ii,2-(jj-1)])
            # Find out over which axes we have to marginalise:
            nlist = [0,1,2,3]; nlist.remove(ii); nlist.remove(jj)
            # Integrate over these 2 dimensions
            red_selfunc = np.nansum(selfunc / ( (len(gridlist[nlist[0]])-1) * (len(gridlist[nlist[1]])-1) ),
                                    axis = tuple(nlist))
            plt.imshow(np.fliplr(red_selfunc).T, aspect="auto", cmap="inferno_r", vmax=1, vmin=0.2,
                       extent=[np.min(xx), np.max(xx), np.min(yy), np.max(yy)])
            plt.xlabel(gridname[ii], fontsize=14)
            plt.ylabel(gridname[jj], fontsize=14)
            # Overplot the data as dots, open dots, and crosses
            plt.scatter(joint[mask2][coordname[ii]], joint[mask2][coordname[jj]], c="k", marker="o", s=5,
                        label="Found by DBSCAN & ANN")
            plt.scatter(joint[mask1][coordname[ii]], joint[mask1][coordname[jj]], c="grey", marker="v", s=5,
                        label="Found by DBSCAN")
            plt.scatter(joint[mask0][coordname[ii]], joint[mask0][coordname[jj]], c="white", marker="+", s=5,
                        label="Not found by DBSCAN")
            if ii == 1 and jj == 2:
                plt.legend(bbox_to_anchor=(1.2, 1.2), loc='lower right', borderaxespad=0. , fontsize=16,
                           facecolor="moccasin")
            if ii == 0:
                ax.set_xticks([0,90,180,270,360])
            if ii == 1:
                ax.set_xticks([-20,0,20])
            if jj == 1:
                ax.set_yticks([-20,0,20])
            plt.axis([np.min(xx), np.max(xx), np.min(yy), np.max(yy)])

plt.suptitle("Recovery fraction of known clusters (Castro-Ginard+2020)", fontsize=20)
cax = fig.add_axes([0.82, 0.4, 0.02, 0.5])
norm = matplotlib.colors.Normalize(vmin=0.2, vmax=1)
cb = matplotlib.colorbar.ColorbarBase(cax, cmap=plt.cm.inferno_r, norm=norm) #, cmap="inferno", norm=norm, spacing='proportional'
cb.set_label('Recovery fraction', fontsize=16)
plt.savefig("im/cluster_recovery_experiment_alfred.png", dpi=150)
```

    <IPython.core.display.Javascript object>

    $l$ [deg] $b$ [deg]
    $l$ [deg] $\log\ t$
    $l$ [deg] $d_{xy}$ [kpc]
    $b$ [deg] $\log\ t$
    $b$ [deg] $d_{xy}$ [kpc]
    $\log\ t$ $d_{xy}$ [kpc]

Now we can define a completeness fraction as a function of $\log t$ out to a limiting distance $d_{xy}^{\rm max}$ by numerically integrating the selection fraction over the $d_{xy}, l$ and $b$ dimensions, and linear interpolation between the $\log t$ bins. Since the method of Castro-Ginard et al. (2020) was optimised for more distant clusters, we extrapolate the recovery fraction from 0.5 kpc inwards.

We also estimate the uncertainty associated to the completeness estimates.
In order to do this, we can [assume Poissonian distributions of the counts in each age bin](https://online.stat.psu.edu/stat504/node/48/) (for each group: non-detected, half-detected, and detected). This allows us to determine the completeness uncertainties numerically.

```python
dmask = (joint["dxy"] < 2) & (joint["dxy"] > 0.5)
tgrid = np.array([6.5, 7, 7.5, 8, 8.5, 9, 10])

tnumbers = np.zeros(len(tgrid)-1)
tnumbers0= np.zeros(len(tgrid)-1)
tnumbers1= np.zeros(len(tgrid)-1)
tnumbers2= np.zeros(len(tgrid)-1)

for ii in range(len(tnumbers)):
    tnumbers[ii]  = np.sum(dmask & (joint["ageNN"] > tgrid[ii]) & (joint["ageNN"] < tgrid[ii+1]))
    tnumbers0[ii] = np.sum(dmask & (joint["ageNN"] > tgrid[ii]) & (joint["ageNN"] < tgrid[ii+1]) & (joint["foundness"]==0.))
    tnumbers1[ii] = np.sum(dmask & (joint["ageNN"] > tgrid[ii]) & (joint["ageNN"] < tgrid[ii+1]) & (joint["foundness"]==1.))
    tnumbers2[ii] = np.sum(dmask & (joint["ageNN"] > tgrid[ii]) & (joint["ageNN"] < tgrid[ii+1]) & (joint["foundness"]==0.5))
```

```python
completeness_per_agebin = np.zeros((len(tgrid)-1, 3))

for ii in np.arange(len(tgrid)-1):
    poisson_0, poisson1, poisson2 = np.random.poisson(lam=tnumbers0[ii], size=1000), \
                                    np.random.poisson(lam=tnumbers1[ii], size=1000), \
                                    np.random.poisson(lam=tnumbers2[ii], size=1000)
    poisson_ii =( poisson1 + poisson2 * 0.5 ) / (poisson_0 + poisson1 + poisson2)
    completeness_per_agebin[ii,:] = np.around(np.percentile(poisson_ii, [16,50,84]),2)

def completeness_agebin_castroginard(logt, dmax=2, getquantiles=False):
    """
    Gets you an interpolated completeness value at for a given age value - inferred from the
    the data behind Castro-Ginard+2020, Fig. 1.
    """
    if dmax==2:
        comp = completeness_per_agebin
    f = interpolate.interp1d(0.5*(tgrid[1:]+tgrid[:-1]), comp[:,1], kind="linear",
                             fill_value=(comp[0,1],  comp[-1,1]), bounds_error=False)
    fm= interpolate.interp1d(0.5*(tgrid[1:]+tgrid[:-1]), comp[:,0], kind="linear",
                             fill_value=(comp[0,0],  comp[-1,0]), bounds_error=False)
    fp= interpolate.interp1d(0.5*(tgrid[1:]+tgrid[:-1]), comp[:,2], kind="linear",
                             fill_value=(comp[0,2],  comp[-1,2]), bounds_error=False)
    if getquantiles:
        return fm(logt), f(logt), fp(logt)
    else:
        return f(logt)

print("DBSCAN rceovery fraction for OCs out to 2kpc at log age [6.5, 7.5, 8.5, 9.5]:\n",
      completeness_agebin_castroginard((6.5,7.5,8.5,9.5)))
```

    DBSCAN rceovery fraction for OCs out to 2kpc at log age [6.5, 7.5, 8.5, 9.5]:
     [0.53 0.65 0.81 0.88]

```python
completeness_per_agebin
```

    array([[0.43, 0.53, 0.62],
           [0.49, 0.55, 0.61],
           [0.7 , 0.75, 0.8 ],
           [0.74, 0.78, 0.82],
           [0.8 , 0.84, 0.88],
           [0.82, 0.88, 0.93]])

## II. Completeness of the ANN results of Cantat-Gaudin+2020 (1,867 clusters) with respect to the full catalogue of 2,017 clusters

First read the catalogue and add some interesting columns:

```python
# Cantat-Gaudin+2020
nntable = Table.read("./data/Table1_CDS.fits")
for col in ["X", "Y", "Z", "Rgc"]:
    nntable[col] = nntable[col].astype(float)
mask = nntable["ageNN"]>0
nnfails = nntable[~mask ]
nntable = nntable[ mask ]
# 1st: estimate the position of the clusters not fitted by the ANN:
distpdf = abj2016.distpdf(nnfails["par"]+0.029, np.sqrt( nnfails["sigpar"]**2. + 0.05**2. ) )
nnfails["d"] = distpdf.modedist

nnfails["X"] = np.cos(nnfails["b"] * np.pi / 180.) * np.cos(nnfails["l"] * np.pi / 180.) * nnfails["d"]
nnfails["Y"] = np.cos(nnfails["b"] * np.pi / 180.) * np.sin(nnfails["l"] * np.pi / 180.) * nnfails["d"]
nnfails["Z"] = np.sin(nnfails["b"] * np.pi / 180.) * nnfails["d"]
nnfails["dxy"] = np.sqrt( nnfails["X"]**2. + nnfails["Y"]**2. )

# Also Kharchenko+2013
kharchenko = Table.read("./data/Kharchenko2013.fits")
```

    /home/hpc/comas/fanders/gaiadr2-caf/abj2016.py:64: RuntimeWarning: divide by zero encountered in true_divide
      return 1/np.sqrt(2*np.pi*sigma_pi**2.) * np.exp(-1/(2*sigma_pi**2.) * (pi[np.newaxis, :] - 1./d[:, np.newaxis])**2. )

Plot the Galactic distribution - focussing on the clusters with missing ANN parameters (right hand plot).
It may take a minute to compute the 2D KDE background distributions.

```python
kharchenko["X"] = np.cos(kharchenko["GLAT"] * np.pi / 180.) * np.cos(kharchenko["GLON"] * np.pi / 180.) \
                  * kharchenko["d"]
kharchenko["Y"] = np.cos(kharchenko["GLAT"] * np.pi / 180.) * np.sin(kharchenko["GLON"] * np.pi / 180.) \
                  * kharchenko["d"]
# estimate the position of the clusters not fitted by the ANN:
distpdf = abj2016.distpdf(nnfails["par"]+0.029, np.sqrt( nnfails["sigpar"]**2. + 0.05**2. ) )
nnfails["d"] = distpdf.modedist

nnfails["X"] = np.cos(nnfails["b"] * np.pi / 180.) * np.cos(nnfails["l"] * np.pi / 180.) * nnfails["d"]
nnfails["Y"] = np.cos(nnfails["b"] * np.pi / 180.) * np.sin(nnfails["l"] * np.pi / 180.) * nnfails["d"]
nnfails["Z"] = np.sin(nnfails["b"] * np.pi / 180.) * nnfails["d"]

fig = plt.figure(figsize=(10., 5.2))

# Left: Kharchenko+2013, right: Cantat-Gaudin+2020
subplots = [121,122]
data = [ np.vstack([kharchenko["X"]/1000., kharchenko["Y"]/1000.]),
         np.vstack([nntable["X"]/1000., nntable["Y"]/1000.]) ]
text = ["MWSC (Kharchenko et al. 2013)", "Gaia DR2 (Cantat-Gaudin et al. 2020)"]
circl= [1.8,2]

for ii in [0,1]:
    ax  = plt.subplot(subplots[ii])
    kde = gaussian_kde(data[ii], bw_method=0.025)

    xrange=(-3,3)
    yrange=(-3,3)
    grid=400
    # evaluate on a regular grid
    xgrid = np.linspace(xrange[0], xrange[1], grid)
    ygrid = np.linspace(yrange[0], yrange[1], grid)
    Xgrid, Ygrid = np.meshgrid(xgrid, ygrid)
    Z = kde.evaluate(np.vstack([Xgrid.ravel(), Ygrid.ravel()]))

    # Plot the result as an image
    plt.imshow(Z.reshape(Xgrid.shape),
           origin='lower', aspect='auto',
           extent=[xrange[0], xrange[1],yrange[0], yrange[1]],
           cmap='magma_r', vmin=0.01)#,norm=LogNorm())

    plt.axis([xrange[0], xrange[1], yrange[0], yrange[1]])
    ax.set_xlabel(r"$X_{\rm sol}$  [kpc]", fontsize=20)
    if ii == 0:
        ax.set_ylabel(r"$Y_{\rm sol}$    [kpc]", fontsize=20)
    else:
        ax.scatter(nnfails["X"], nnfails["Y"], marker="x", s=10, c="k")

    ax.text(-2.6,3.25, text[ii], fontsize=14)

    # Overplot the completeness limit considered
    circle1 = plt.Circle((0, 0), circl[ii], color='grey', ls="--", fill=False)
    ax.add_artist(circle1)
plt.tight_layout()
```

    <IPython.core.display.Javascript object>

```python
# Join the 2 tables to estimate the parameters of the missing objects
kharchenko["oc"] = kharchenko["Name"]
k13_cg20_match = astropy.table.join(nntable, kharchenko, join_type='inner', keys="oc", uniq_col_name='{col_name}_{table_name}',
                                    table_names=['1', '2'], metadata_conflicts='warn')
k13_cg20f_match = astropy.table.join(nnfails, kharchenko, join_type='inner', keys="oc", uniq_col_name='{col_name}_{table_name}',
                                    table_names=['1', '2'], metadata_conflicts='warn')
kharchenko["X"] = np.cos(kharchenko["GLAT"] * np.pi / 180.) * np.cos(kharchenko["GLON"] * np.pi / 180.) \
                  * kharchenko["d"]
kharchenko["Y"] = np.cos(kharchenko["GLAT"] * np.pi / 180.) * np.sin(kharchenko["GLON"] * np.pi / 180.) \
                  * kharchenko["d"]

```

    WARNING: MergeConflictWarning: Cannot merge meta key 'EXTNAME' types <class 'str'> and <class 'str'>, choosing EXTNAME='/home/friedel/Downloads/asu.fit#1' [astropy.utils.metadata]
    WARNING: MergeConflictWarning: Cannot merge meta key 'DATE-HDU' types <class 'str'> and <class 'str'>, choosing DATE-HDU='2020-05-01T18:57:47' [astropy.utils.metadata]

Now plot the direct comparison between the age scales of Kharchenko+2013 and CG2020 for the objects within a limiting distance `dmax`.

```python
dmax = 2.
distmask = ( np.sqrt(k13_cg20_match["X_1"]**2. + k13_cg20_match["Y_1"]**2.) < dmax*1000 )

plt.figure(figsize=(8, 8))
ax=plt.subplot(111)
plt.plot(np.array([6,6]), np.array([10,10]), ls="dashed", c="grey", lw=1)
plt.errorbar(k13_cg20_match["ageNN"][distmask], k13_cg20_match["logt"][distmask],
             xerr=0.15*np.ones(np.sum(distmask)), fmt="o", c="k",
             alpha=0.5)
ax.set_xlabel(r"log $t$ (Cantat-Gaudin+2020)")
ax.set_ylabel(r"log $t$ (Kharchenko+2013)")
for ii in np.arange(np.sum(distmask)):
    if abs(k13_cg20_match["ageNN"][distmask][ii] - k13_cg20_match["logt"][distmask][ii]) > 1.2:
        if k13_cg20_match["oc"][distmask][ii] in [b"Dolidze_8", b"NGC_7039", b"Trumpler28"]:
            ax.text(k13_cg20_match["ageNN"][distmask][ii]-.5, k13_cg20_match["logt"][distmask][ii]-.15,
                    k13_cg20_match["Name"][distmask][ii])
        else:
            ax.text(k13_cg20_match["ageNN"][distmask][ii], k13_cg20_match["logt"][distmask][ii],
                    k13_cg20_match["Name"][distmask][ii])

ax.axis([6,10,6,10])
plt.show()
```

    <IPython.core.display.Javascript object>

What does this tell us?

Well, that the differences in the CAF cannot be explained by a different age scale of the Kharchenko+2013 catalogue: in fact it is surprisingly consistent, except for a few outliers. Especially the number of old clusters in the overlap sample is very similar.

The problem, therefore, must lie in the objects that are not part of both samples:

1. The CG2020 catalogue has a lot of new clusters within 2 kpc - what ages do they typically sample?
2. The K13 objects that CG2020 do not recover - which ages do they (allegedly) have?

```python
#Do statistics for the dmax-limited sample
dmax = 2

print("OC age statistics in the extended solar vicinity (not completeness-corrected):")
##############################
print("-------------------------------")
distmask_k13 = ( np.sqrt(kharchenko["X"]**2. + kharchenko["Y"]**2.) < dmax*1000 )
print(len(kharchenko[ distmask_k13 ]), "K13 OCs within", dmax, "kpc")
print(len(kharchenko[ distmask_k13 & ( kharchenko["logt"]>9 )]), "out of those have log t_K13 > 9 (",
      int(100*len(kharchenko[ distmask_k13 & ( kharchenko["logt"]>9 )])/len(kharchenko[ distmask_k13 ])), "% )")
##############################
distmask_nn = ( np.sqrt(nntable["X"]**2. + nntable["Y"]**2.) < dmax*1000 )
print("-------------------------------")
print(len(nntable[ distmask_nn ]), "CG20 OCs within", dmax, "kpc")
print(len(nntable[ distmask_nn & ( nntable["ageNN"]>9 )]), "out of those have log t_CG20 > 9 (",
      int(1000*len(nntable[ distmask_nn & ( nntable["ageNN"]>9 )])/len(nntable[ distmask_nn ]))/10, "% )")
###############################
print("-------------------------------")
print(len(k13_cg20_match), "overlap between K13 & CG20 NN")

distmask = ( np.sqrt(k13_cg20_match["X_1"]**2. + k13_cg20_match["Y_1"]**2.) < dmax*1000 )
print("-------------------------------")
print(len(k13_cg20_match[ distmask ]), "overlap within", dmax, "kpc")
print(len(k13_cg20_match[ distmask & ( k13_cg20_match["logt"]>9 )]), "out of those have log t_K13 > 9")
print(len(k13_cg20_match[ distmask & ( k13_cg20_match["ageNN"]>9 )]), "out of those have log t_CG20 > 9")
###############################
distmask_nnf = ( nnfails["dxy"] < dmax )
print("-------------------------------")
print(len(nnfails[ distmask_nnf ]), "CG20 (no NN) OCs within", dmax, "kpc")
###############################
print("-------------------------------")
print(len(k13_cg20f_match), "overlap between K13 & CG20 (no NN)")

distmaskf = ( k13_cg20f_match["d_2"]*np.cos( 180*k13_cg20f_match["GLAT"]/np.pi ) < dmax*1000 )
distmaskf = ( k13_cg20f_match["dxy"] < dmax )
print("-------------------------------")
print(len(k13_cg20f_match[ distmaskf ]), "overlap within", dmax, "kpc")
print(len(k13_cg20f_match[ distmaskf & ( k13_cg20f_match["logt"]>9 )]), "out of those have log t_K13 > 9")
```

    OC age statistics in the extended solar vicinity (not completeness-corrected):
    -------------------------------
    1461 K13 OCs within 2 kpc
    241 out of those have log t_K13 > 9 ( 16 % )
    -------------------------------
    834 CG20 OCs within 2 kpc
    54 out of those have log t_CG20 > 9 ( 6.4 % )
    -------------------------------
    1007 overlap between K13 & CG20 NN
    -------------------------------
    391 overlap within 2 kpc
    29 out of those have log t_K13 > 9
    32 out of those have log t_CG20 > 9
    -------------------------------
    33 CG20 (no NN) OCs within 2 kpc
    -------------------------------
    85 overlap between K13 & CG20 (no NN)
    -------------------------------
    16 overlap within 2 kpc
    0 out of those have log t_K13 > 9

Now, from the non-converged ANN clusters in the overlap sample, determine the age-dependent completeness of the ANN sample with respect to the full CG2020 sample:

```python
def completeness_agebin_cantat2020_ann(logt, dmax=dmax):
    """
    Gets you an interpolated completeness value at for a given age value - inferred from the
    the Kharchenko+2013 ages for the clusters that lack ANN parameters
    """
    agebins = np.linspace(6,10,5)
    comp_nn = np.zeros(4)
    if dmax==2:
        comp_nn = [0.7692, 0.9932, 0.9955 , 1.]
    #else:
    for ii in np.arange(len(comp_nn)):
        maskok = distmask_nn  & ( nntable["ageNN"]>agebins[ii] )  & ( nntable["ageNN"]<agebins[ii+1] )
        maskf  = distmaskf    & ( k13_cg20f_match["logt"]>agebins[ii] ) & ( k13_cg20f_match["logt"]<agebins[ii+1] )
        comp_nn[ii] = np.sum(maskok) / ( np.sum(maskok) + np.sum(maskf) )
    f = interpolate.interp1d(0.5*(agebins[1:]+agebins[:-1]), comp_nn, kind="linear",
                             fill_value=(comp_nn[0],  comp_nn[-1]), bounds_error=False)
    return f(logt)
```

```python
print("ANN completeness out to 2kpc at log age [6.5, 7.5, 8.5, 9.5]:", completeness_agebin_cantat2020_ann((6.5,7.5,8.5,9.5)))
```

    ANN completeness out to 2kpc at log age [6.5, 7.5, 8.5, 9.5]: [0.76923077 0.99322034 0.9955157  1.        ]

```python

```


## Setup

```python
!git clone https://github.com/sayakpaul/Denoised-Smoothing-TF
!cd Denoised-Smoothing-TF && tar xf models.tar.gz
```

    Cloning into 'Denoised-Smoothing-TF'...
    remote: Enumerating objects: 93, done.
    remote: Counting objects: 100% (93/93), done.
    remote: Compressing objects: 100% (71/71), done.
    remote: Total 93 (delta 37), reused 62 (delta 21), pack-reused 0
    Unpacking objects: 100% (93/93), done.

```python
import sys
sys.path.append("Denoised-Smoothing-TF")

from models import dncnn

from tensorflow.keras import layers
import tensorflow as tf

import matplotlib.pyplot as plt
import numpy as np

tf.random.set_seed(42)
np.random.seed(42)
```

With a TPU, you would need to have a GCS Bucket to serialize and load the `SavedModel` files.

```python
try:
    tpu = None
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.TPUStrategy(tpu)
except ValueError:
    strategy = tf.distribute.MirroredStrategy()

print("Number of accelerators: ", strategy.num_replicas_in_sync)
```

    INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
    Number of accelerators:  4

## Load dataset

```python
(x_train, _), (x_test, _) = tf.keras.datasets.cifar10.load_data()
print(f"Total training examples: {len(x_train)}")
print(f"Total test examples: {len(x_test)}")
```

    Total training examples: 50000
    Total test examples: 10000

## Define constants

```python
BATCH_SIZE = 128 * strategy.num_replicas_in_sync
EPOCHS = 90
START_LR = 1e-3
AUTO = tf.data.AUTOTUNE
```

## Prepare data loaders

```python
# Augmentation pipeline
data_augmentation = tf.keras.Sequential(
    [
        layers.experimental.preprocessing.RandomCrop(32, 32),
        layers.experimental.preprocessing.RandomFlip("horizontal")
    ]
)

# Now, map the augmentation pipeline to our training dataset
train_ds = (
    tf.data.Dataset.from_tensor_slices(x_train)
    .shuffle(BATCH_SIZE * 100)
    .map(lambda x: (tf.image.convert_image_dtype(x, tf.float32)),
         num_parallel_calls=AUTO)
    .repeat()
    .batch(BATCH_SIZE)
    .map(data_augmentation, num_parallel_calls=AUTO)
    .prefetch(AUTO)
)

# Test dataset
test_ds = (
    tf.data.Dataset.from_tensor_slices(x_test)
    .map(lambda x: (tf.image.convert_image_dtype(x, tf.float32)),
         num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)

# Compute the mean and the variance of the training data for normalization
norm_layer = layers.experimental.preprocessing.Normalization()
norm_layer.adapt(x_train/255.) # Notice the scaling step
```

## Define training loop

```python
class Trainer(tf.keras.Model):
    def __init__(self, denoiser, classifier, sigma=0.25):
        super(Trainer, self).__init__()
        self.denoiser = denoiser
        self.classifier = classifier
        self.sigma = sigma

    def train_step(self, clean_inputs):
        batch_size = tf.shape(clean_inputs)[0]
        clean_inputs = tf.cast(clean_inputs, tf.float32)

        if self.classifier is not None:
            targets = self.classifier(clean_inputs, training=False)
        else:
            targets = clean_inputs

        with tf.GradientTape() as tape:
            noise = tf.experimental.numpy.random.randn(batch_size, 32, 32, 3) * self.sigma
            noisy_inputs = norm_layer(clean_inputs + noise) # Add noise BEFORE normalization
            clean_predictions = self.denoiser(noisy_inputs, training=True)

            if self.classifier is not None:
                clean_predictions = self.classifier(clean_predictions,
                                                    training=False)
                loss = self.compiled_loss(targets, clean_predictions)
                loss = tf.nn.compute_average_loss(loss,
                                                  global_batch_size=BATCH_SIZE)
            else:
                loss = self.compiled_loss(targets, clean_predictions)
                loss /= tf.cast(tf.reduce_prod(tf.shape(targets)[1:]), tf.float32)
                loss = tf.nn.compute_average_loss(loss,
                                                  global_batch_size=BATCH_SIZE)

        trainable_vars = self.denoiser.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        return {"loss": loss}

    def test_step(self, clean_inputs):
        batch_size = tf.shape(clean_inputs)[0]
        clean_inputs = tf.cast(clean_inputs, tf.float32)

        if self.classifier is not None:
            targets = self.classifier(clean_inputs, training=False)
        else:
            targets = clean_inputs

        noise = tf.experimental.numpy.random.randn(batch_size, 32, 32, 3) * self.sigma
        noisy_inputs = norm_layer(clean_inputs + noise)
        clean_predictions = self.denoiser(noisy_inputs, training=False)

        if self.classifier is not None:
            clean_predictions = self.classifier(clean_predictions,
                                                            training=False)
            loss = self.compiled_loss(targets, clean_predictions)
            loss = tf.nn.compute_average_loss(loss,
                                                  global_batch_size=BATCH_SIZE)
        else:
            loss = self.compiled_loss(targets, clean_predictions)
            loss /= tf.cast(tf.reduce_prod(tf.shape(targets)[1:]), tf.float32)
            loss = tf.nn.compute_average_loss(loss,
                                                global_batch_size=BATCH_SIZE)

        return {"loss": loss}
```

## LR schedule

```python
# https://github.com/cszn/DnCNN/blob/master/TrainingCodes/dncnn_keras/main_train.py
def lr_schedule(epoch):
    initial_lr = START_LR
    if epoch<=30:
        lr = initial_lr
    elif epoch<=60:
        lr = initial_lr/10
    elif epoch<=80:
        lr = initial_lr/20
    else:
        lr = initial_lr/20
    return lr

lr_scheduler = tf.keras.callbacks.LearningRateScheduler(lr_schedule)
```

## Training with the stability objective

From [this blog post](https://www.microsoft.com/en-us/research/blog/denoised-smoothing-provably-defending-pretrained-classifiers-against-adversarial-examples/):

> given a classifier $F$ attached to the denoiser $\mathcal{D}_{\theta}$, minimize the cross-entropy loss between the prediction of $F(\mathcal{D}_{theta})$ at the noisy input $x_i + \delta
$ and the prediction of the classifier at the clean data point $x_i$

$$\begin{align}L_{\text{Stab}}&=\mathbb{E}_{\mathcal{S}, \delta} \mathcal{\ell_\text{CE}}(F(\mathcal{D}_\theta(x_i + \delta)), f(x_i)) \quad \text{where} \; \delta \sim \mathcal{N}(0, \sigma^2 I) \; \end{align}$$

One can obtain the initial weights that I used from [here](https://github.com/sayakpaul/Denoised-Smoothing-TF/raw/main/models.tar.gz).

```python
with strategy.scope():
    pretrained_classifier = tf.keras.models.load_model("Denoised-Smoothing-TF/resnet20_classifier")
    denoiser = Trainer(dncnn.get_dncnn(), pretrained_classifier)
    denoiser.denoiser.save("initial_denoiser")
    denoiser.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=START_LR),
                     loss=tf.keras.losses.CategoricalCrossentropy(reduction="none")
    )

history = denoiser.fit(train_ds,
            steps_per_epoch=len(x_train)//BATCH_SIZE,
            validation_data=test_ds,
            epochs=EPOCHS,
            callbacks=[lr_scheduler])
```

    INFO:tensorflow:Assets written to: initial_denoiser/assets
    Epoch 1/90
    INFO:tensorflow:batch_all_reduce: 62 all-reduces with algorithm = nccl, num_packs = 1
    INFO:tensorflow:batch_all_reduce: 62 all-reduces with algorithm = nccl, num_packs = 1
    97/97 [==============================] - 42s 135ms/step - loss: 0.6532 - val_loss: 0.2795
    Epoch 2/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.5013 - val_loss: 0.2699
    Epoch 3/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.4726 - val_loss: 0.2519
    Epoch 4/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.4553 - val_loss: 0.2569
    Epoch 5/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.4462 - val_loss: 0.2513
    Epoch 6/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.4304 - val_loss: 0.2535
    Epoch 7/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.4230 - val_loss: 0.2321
    Epoch 8/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.4154 - val_loss: 0.2375
    Epoch 9/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.4066 - val_loss: 0.2319
    Epoch 10/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.4013 - val_loss: 0.2191
    Epoch 11/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3943 - val_loss: 0.2176
    Epoch 12/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3924 - val_loss: 0.2400
    Epoch 13/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3847 - val_loss: 0.2264
    Epoch 14/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3811 - val_loss: 0.2324
    Epoch 15/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3781 - val_loss: 0.2334
    Epoch 16/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3759 - val_loss: 0.2223
    Epoch 17/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3713 - val_loss: 0.2315
    Epoch 18/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3658 - val_loss: 0.2134
    Epoch 19/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3666 - val_loss: 0.2367
    Epoch 20/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3636 - val_loss: 0.2326
    Epoch 21/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3618 - val_loss: 0.2329
    Epoch 22/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3567 - val_loss: 0.2203
    Epoch 23/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3581 - val_loss: 0.2186
    Epoch 24/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3548 - val_loss: 0.2143
    Epoch 25/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3512 - val_loss: 0.2224
    Epoch 26/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3509 - val_loss: 0.2111
    Epoch 27/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3493 - val_loss: 0.2153
    Epoch 28/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3461 - val_loss: 0.2137
    Epoch 29/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3479 - val_loss: 0.2345
    Epoch 30/90
    97/97 [==============================] - 8s 82ms/step - loss: 0.3430 - val_loss: 0.2128
    Epoch 31/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3401 - val_loss: 0.2196
    Epoch 32/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3327 - val_loss: 0.2110
    Epoch 33/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3235 - val_loss: 0.2101
    Epoch 34/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3221 - val_loss: 0.2111
    Epoch 35/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3187 - val_loss: 0.2076
    Epoch 36/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3181 - val_loss: 0.2115
    Epoch 37/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3189 - val_loss: 0.1944
    Epoch 38/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3182 - val_loss: 0.2054
    Epoch 39/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3145 - val_loss: 0.2028
    Epoch 40/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3122 - val_loss: 0.2087
    Epoch 41/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3138 - val_loss: 0.2031
    Epoch 42/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3163 - val_loss: 0.1981
    Epoch 43/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3140 - val_loss: 0.2010
    Epoch 44/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3182 - val_loss: 0.2069
    Epoch 45/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3113 - val_loss: 0.2061
    Epoch 46/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3151 - val_loss: 0.2141
    Epoch 47/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3139 - val_loss: 0.2027
    Epoch 48/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3109 - val_loss: 0.1982
    Epoch 49/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3094 - val_loss: 0.2024
    Epoch 50/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3091 - val_loss: 0.2064
    Epoch 51/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3102 - val_loss: 0.1992
    Epoch 52/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3101 - val_loss: 0.2031
    Epoch 53/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3110 - val_loss: 0.2003
    Epoch 54/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3119 - val_loss: 0.2055
    Epoch 55/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3105 - val_loss: 0.2195
    Epoch 56/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3088 - val_loss: 0.1989
    Epoch 57/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3091 - val_loss: 0.1896
    Epoch 58/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3119 - val_loss: 0.1906
    Epoch 59/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3111 - val_loss: 0.2044
    Epoch 60/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3083 - val_loss: 0.2074
    Epoch 61/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3097 - val_loss: 0.1889
    Epoch 62/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3039 - val_loss: 0.1984
    Epoch 63/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3064 - val_loss: 0.1929
    Epoch 64/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3043 - val_loss: 0.2015
    Epoch 65/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3036 - val_loss: 0.2009
    Epoch 66/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3039 - val_loss: 0.1891
    Epoch 67/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3039 - val_loss: 0.1967
    Epoch 68/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3006 - val_loss: 0.2048
    Epoch 69/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3061 - val_loss: 0.1803
    Epoch 70/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3019 - val_loss: 0.2107
    Epoch 71/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3026 - val_loss: 0.1998
    Epoch 72/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3040 - val_loss: 0.2036
    Epoch 73/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3053 - val_loss: 0.2055
    Epoch 74/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3018 - val_loss: 0.1936
    Epoch 75/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3044 - val_loss: 0.2043
    Epoch 76/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3049 - val_loss: 0.2061
    Epoch 77/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3039 - val_loss: 0.2021
    Epoch 78/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3041 - val_loss: 0.1979
    Epoch 79/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3007 - val_loss: 0.2009
    Epoch 80/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.2997 - val_loss: 0.1949
    Epoch 81/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3025 - val_loss: 0.1965
    Epoch 82/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3053 - val_loss: 0.1914
    Epoch 83/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3007 - val_loss: 0.1911
    Epoch 84/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.2999 - val_loss: 0.1965
    Epoch 85/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3022 - val_loss: 0.1954
    Epoch 86/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3007 - val_loss: 0.1997
    Epoch 87/90
    97/97 [==============================] - 7s 75ms/step - loss: 0.3019 - val_loss: 0.2068
    Epoch 88/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3006 - val_loss: 0.1987
    Epoch 89/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.3018 - val_loss: 0.2050
    Epoch 90/90
    97/97 [==============================] - 7s 76ms/step - loss: 0.2983 - val_loss: 0.2073

```python
plt.plot(history.history["loss"], label="train loss")
plt.plot(history.history["val_loss"], label="test loss")
plt.grid()
plt.legend()
plt.show()

denoiser.denoiser.save("denoiser_stability")
```

## With MSE objective

```python
with strategy.scope():
    initial_denoiser = tf.keras.models.load_model("initial_denoiser")
    denoiser = Trainer(initial_denoiser, classifier=None)
    denoiser.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=START_LR),
                     loss=tf.keras.losses.MeanSquaredError(reduction="none")
    )

history = denoiser.fit(train_ds,
            steps_per_epoch=len(x_train)//BATCH_SIZE,
            validation_data=test_ds,
            epochs=EPOCHS,
            callbacks=[lr_scheduler])
```

    WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
    Epoch 1/90
    INFO:tensorflow:batch_all_reduce: 62 all-reduces with algorithm = nccl, num_packs = 1
    INFO:tensorflow:batch_all_reduce: 62 all-reduces with algorithm = nccl, num_packs = 1
    97/97 [==============================] - 24s 79ms/step - loss: 0.0604 - val_loss: 0.0183
    Epoch 2/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0144 - val_loss: 0.0076
    Epoch 3/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0068 - val_loss: 0.0055
    Epoch 4/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0050 - val_loss: 0.0041
    Epoch 5/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0042 - val_loss: 0.0029
    Epoch 6/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0037 - val_loss: 0.0039
    Epoch 7/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0037 - val_loss: 0.0020
    Epoch 8/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0031 - val_loss: 0.0024
    Epoch 9/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0030 - val_loss: 0.0060
    Epoch 10/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0028 - val_loss: 0.0033
    Epoch 11/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0027 - val_loss: 0.0013
    Epoch 12/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0026 - val_loss: 0.0051
    Epoch 13/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0025 - val_loss: 0.0016
    Epoch 14/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0023 - val_loss: 0.0011
    Epoch 15/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0023 - val_loss: 0.0042
    Epoch 16/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0022 - val_loss: 0.0025
    Epoch 17/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0022 - val_loss: 0.0039
    Epoch 18/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0021 - val_loss: 0.0017
    Epoch 19/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0020 - val_loss: 0.0016
    Epoch 20/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0019 - val_loss: 0.0018
    Epoch 21/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0019 - val_loss: 0.0016
    Epoch 22/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0018 - val_loss: 0.0011
    Epoch 23/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0018 - val_loss: 0.0019
    Epoch 24/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0017 - val_loss: 0.0022
    Epoch 25/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0017 - val_loss: 0.0014
    Epoch 26/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0017 - val_loss: 0.0014
    Epoch 27/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0016 - val_loss: 0.0015
    Epoch 28/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0016 - val_loss: 9.2881e-04
    Epoch 29/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0016 - val_loss: 0.0026
    Epoch 30/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0016 - val_loss: 0.0010
    Epoch 31/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0014 - val_loss: 0.0013
    Epoch 32/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0012 - val_loss: 5.1350e-04
    Epoch 33/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0012 - val_loss: 6.1331e-04
    Epoch 34/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0012 - val_loss: 5.8201e-04
    Epoch 35/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 5.5720e-04
    Epoch 36/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0012 - val_loss: 7.2087e-04
    Epoch 37/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 6.5761e-04
    Epoch 38/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 6.0990e-04
    Epoch 39/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0011 - val_loss: 5.5000e-04
    Epoch 40/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 5.5300e-04
    Epoch 41/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0012 - val_loss: 5.3825e-04
    Epoch 42/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 6.0790e-04
    Epoch 43/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0011 - val_loss: 5.1722e-04
    Epoch 44/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0012 - val_loss: 9.8958e-04
    Epoch 45/90
    97/97 [==============================] - 6s 62ms/step - loss: 0.0011 - val_loss: 5.0852e-04
    Epoch 46/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0011 - val_loss: 5.4033e-04
    Epoch 47/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 5.3346e-04
    Epoch 48/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0011 - val_loss: 5.2196e-04
    Epoch 49/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 6.4571e-04
    Epoch 50/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 5.7214e-04
    Epoch 51/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0011 - val_loss: 4.8099e-04
    Epoch 52/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 5.3950e-04
    Epoch 53/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 6.8400e-04
    Epoch 54/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 4.7000e-04
    Epoch 55/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 8.0983e-04
    Epoch 56/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0011 - val_loss: 5.7719e-04
    Epoch 57/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 5.8855e-04
    Epoch 58/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 6.2172e-04
    Epoch 59/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 5.6920e-04
    Epoch 60/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 6.0252e-04
    Epoch 61/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0011 - val_loss: 5.0472e-04
    Epoch 62/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0011 - val_loss: 8.1341e-04
    Epoch 63/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 5.5949e-04
    Epoch 64/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 4.6341e-04
    Epoch 65/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 4.4023e-04
    Epoch 66/90
    97/97 [==============================] - 5s 55ms/step - loss: 9.9303e-04 - val_loss: 4.9605e-04
    Epoch 67/90
    97/97 [==============================] - 5s 54ms/step - loss: 9.9123e-04 - val_loss: 6.1753e-04
    Epoch 68/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0010 - val_loss: 5.4595e-04
    Epoch 69/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0010 - val_loss: 4.4822e-04
    Epoch 70/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0011 - val_loss: 9.0038e-04
    Epoch 71/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 5.4294e-04
    Epoch 72/90
    97/97 [==============================] - 5s 54ms/step - loss: 9.7593e-04 - val_loss: 5.0562e-04
    Epoch 73/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 5.4990e-04
    Epoch 74/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 5.3130e-04
    Epoch 75/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 4.7490e-04
    Epoch 76/90
    97/97 [==============================] - 5s 54ms/step - loss: 9.8308e-04 - val_loss: 4.8572e-04
    Epoch 77/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0010 - val_loss: 4.4648e-04
    Epoch 78/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0010 - val_loss: 5.2068e-04
    Epoch 79/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0010 - val_loss: 5.0260e-04
    Epoch 80/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 5.6244e-04
    Epoch 81/90
    97/97 [==============================] - 5s 54ms/step - loss: 0.0010 - val_loss: 4.7123e-04
    Epoch 82/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0010 - val_loss: 5.6051e-04
    Epoch 83/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0010 - val_loss: 4.6816e-04
    Epoch 84/90
    97/97 [==============================] - 5s 54ms/step - loss: 9.8784e-04 - val_loss: 4.8585e-04
    Epoch 85/90
    97/97 [==============================] - 5s 54ms/step - loss: 9.8716e-04 - val_loss: 4.6904e-04
    Epoch 86/90
    97/97 [==============================] - 5s 54ms/step - loss: 9.5673e-04 - val_loss: 4.4165e-04
    Epoch 87/90
    97/97 [==============================] - 5s 55ms/step - loss: 0.0010 - val_loss: 5.8562e-04
    Epoch 88/90
    97/97 [==============================] - 5s 54ms/step - loss: 9.6637e-04 - val_loss: 4.5835e-04
    Epoch 89/90
    97/97 [==============================] - 5s 55ms/step - loss: 9.8395e-04 - val_loss: 5.6031e-04
    Epoch 90/90
    97/97 [==============================] - 5s 54ms/step - loss: 9.7118e-04 - val_loss: 5.6422e-04

```python
plt.plot(history.history["loss"], label="train loss")
plt.plot(history.history["val_loss"], label="test loss")
plt.grid()
plt.legend()
plt.show()

denoiser.denoiser.save("denoiser_mse")
```


```python
# Mount Google Drive
from google.colab import drive # import drive from google colab

ROOT = "/content/drive"     # default location for the drive
print(ROOT)                 # print content of ROOT (Optional)

drive.mount(ROOT,force_remount=True)
```

    /content/drive
    Mounted at /content/drive

# Neuromatch Academy: Week 1, Day 5, Tutorial 1
# Dimensionality Reduction: Geometric view of data

__Content creators:__ Alex Cayco Gajic, John Murray

__Content reviewers:__ Roozbeh Farhoudi, Matt Krause, Spiros Chavlis, Richard Gao, Michael Waskom

---
# Tutorial Objectives

In this notebook we'll explore how multivariate data can be represented in different orthonormal bases. This will help us build intuition that will be helpful in understanding PCA in the following tutorial.

Overview:
 - Generate correlated multivariate data.
 - Define an arbitrary orthonormal basis.
 - Project the data onto the new basis.

```python
# @title Video 1: Geometric view of data
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="THu9yHnpq9I", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=THu9yHnpq9I

---
# Setup

```python
# Import
import numpy as np
import matplotlib.pyplot as plt
```

```python
# @title Figure Settings
import ipywidgets as widgets  # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Helper functions

def get_data(cov_matrix):
  """
  Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian.

  Note that samples are sorted in ascending order for the first random variable

  Args:
    cov_matrix (numpy array of floats): desired covariance matrix

  Returns:
    (numpy array of floats) : samples from the bivariate Gaussian, with each
                              column corresponding to a different random
                              variable
  """

  mean = np.array([0, 0])
  X = np.random.multivariate_normal(mean, cov_matrix, size=1000)
  indices_for_sorting = np.argsort(X[:, 0])
  X = X[indices_for_sorting, :]

  return X

def plot_data(X):
  """
  Plots bivariate data. Includes a plot of each random variable, and a scatter
  plot of their joint activity. The title indicates the sample correlation
  calculated from the data.

  Args:
    X (numpy array of floats) :   Data matrix each column corresponds to a
                                  different random variable

  Returns:
    Nothing.
  """

  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(X[:, 0], color='k')
  plt.ylabel('Neuron 1')
  plt.title('Sample var 1: {:.1f}'.format(np.var(X[:, 0])))
  ax1.set_xticklabels([])
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(X[:, 1], color='k')
  plt.xlabel('Sample Number')
  plt.ylabel('Neuron 2')
  plt.title('Sample var 2: {:.1f}'.format(np.var(X[:, 1])))
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(X[:, 0], X[:, 1], '.', markerfacecolor=[.5, .5, .5],
           markeredgewidth=0)
  ax3.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(X[:, 0], X[:, 1])[0, 1]))
  plt.show()

def plot_basis_vectors(X, W):
  """
  Plots bivariate data as well as new basis vectors.

  Args:
    X (numpy array of floats) :   Data matrix each column corresponds to a
                                  different random variable
    W (numpy array of floats) :   Square matrix representing new orthonormal
                                  basis each column represents a basis vector

  Returns:
    Nothing.
  """

  plt.figure(figsize=[4, 4])
  plt.plot(X[:, 0], X[:, 1], '.', color=[.5, .5, .5], label='Data')
  plt.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.plot([0, W[0, 0]], [0, W[1, 0]], color='r', linewidth=3,
           label='Basis vector 1')
  plt.plot([0, W[0, 1]], [0, W[1, 1]], color='b', linewidth=3,
           label='Basis vector 2')
  plt.legend()
  plt.show()

def plot_data_new_basis(Y):
  """
  Plots bivariate data after transformation to new bases.
  Similar to plot_data but with colors corresponding to projections onto
  basis 1 (red) and basis 2 (blue). The title indicates the sample correlation
  calculated from the data.

  Note that samples are re-sorted in ascending order for the first
  random variable.

  Args:
    Y (numpy array of floats):   Data matrix in new basis each column
                                 corresponds to a different random variable

  Returns:
    Nothing.
  """
  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(Y[:, 0], 'r')
  plt.xlabel
  plt.ylabel('Projection \n basis vector 1')
  plt.title('Sample var 1: {:.1f}'.format(np.var(Y[:, 0])))
  ax1.set_xticklabels([])
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(Y[:, 1], 'b')
  plt.xlabel('Sample number')
  plt.ylabel('Projection \n basis vector 2')
  plt.title('Sample var 2: {:.1f}'.format(np.var(Y[:, 1])))
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(Y[:, 0], Y[:, 1], '.', color=[.5, .5, .5])
  ax3.axis('equal')
  plt.xlabel('Projection basis vector 1')
  plt.ylabel('Projection basis vector 2')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(Y[:, 0], Y[:, 1])[0, 1]))
  plt.show()
```

---
# Section 1: Generate correlated multivariate data

```python
# @title Video 2: Multivariate data
video = YouTubeVideo(id="jcTq2PgU5Vw", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=jcTq2PgU5Vw

To gain intuition, we will first use a simple model to generate multivariate data. Specifically, we will draw random samples from a *bivariate normal distribution*. This is an extension of the one-dimensional normal distribution to two dimensions, in which each $x_i$ is marginally normal with mean $\mu_i$ and variance $\sigma_i^2$:

\begin{align}
x_i \sim \mathcal{N}(\mu_i,\sigma_i^2).
\end{align}

Additionally, the joint distribution for $x_1$ and $x_2$ has a specified correlation coefficient $\rho$. Recall that the correlation coefficient is a normalized version of the covariance, and ranges between -1 and +1:

\begin{align}
\rho = \frac{\text{cov}(x_1,x_2)}{\sqrt{\sigma_1^2 \sigma_2^2}}.
\end{align}

For simplicity, we will assume that the mean of each variable has already been subtracted, so that $\mu_i=0$. The remaining parameters can be summarized in the covariance matrix, which for two dimensions has the following form:

\begin{equation*}
{\bf \Sigma} =
\begin{pmatrix}
 \text{var}(x_1) & \text{cov}(x_1,x_2) \\
 \text{cov}(x_1,x_2) &\text{var}(x_2)
\end{pmatrix}.
\end{equation*}

In general, $\bf \Sigma$ is a symmetric matrix with the variances $\text{var}(x_i) = \sigma_i^2$ on the diagonal, and the covariances on the off-diagonal. Later, we will see that the covariance matrix plays a key role in PCA.

## Exercise 1: Draw samples from a distribution

We have provided code to draw random samples from a zero-mean bivariate normal distribution. Throughout this tutorial, we'll imagine these samples represent the activity (firing rates) of two recorded neurons on different trials. Fill in the function below to calculate the covariance matrix given the desired variances and correlation coefficient. The covariance can be found by rearranging the equation above:

\begin{align}
\text{cov}(x_1,x_2) = \rho \sqrt{\sigma_1^2 \sigma_2^2}.
\end{align}

Use these functions to generate and plot data while varying the parameters. You should get a feel for how changing the correlation coefficient affects the geometry of the simulated data.

**Steps**
* Fill in the function `calculate_cov_matrix` to calculate the desired covariance.
* Generate and plot the data for $\sigma_1^2 =1$, $\sigma_1^2 =1$, and $\rho = .8$. Try plotting the data for different values of the correlation coefficent: $\rho = -1, -.5, 0, .5, 1$.

```python
help(plot_data)
help(get_data)
```

    Help on function plot_data in module __main__:

    plot_data(X)
        Plots bivariate data. Includes a plot of each random variable, and a scatter
        plot of their joint activity. The title indicates the sample correlation
        calculated from the data.

        Args:
          X (numpy array of floats) :   Data matrix each column corresponds to a
                                        different random variable

        Returns:
          Nothing.

    Help on function get_data in module __main__:

    get_data(cov_matrix)
        Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian.

        Note that samples are sorted in ascending order for the first random variable

        Args:
          cov_matrix (numpy array of floats): desired covariance matrix

        Returns:
          (numpy array of floats) : samples from the bivariate Gaussian, with each
                                    column corresponding to a different random
                                    variable

```python
def calculate_cov_matrix(var_1, var_2, corr_coef):
  """
  Calculates the covariance matrix based on the variances and correlation
  coefficient.

  Args:
    var_1 (scalar)          : variance of the first random variable
    var_2 (scalar)          : variance of the second random variable
    corr_coef (scalar)      : correlation coefficient

  Returns:
    (numpy array of floats) : covariance matrix
  """

  #################################################
  ## TODO for students: calculate the covariance matrix
  # Fill out function and remove
  #raise NotImplementedError("Student excercise: calculate the covariance matrix!")
  #################################################

  # Calculate the covariance from the variances and correlation
  cov = corr_coef*np.sqrt([var_1*var_2])

  cov_matrix = np.array([[var_1, cov], [cov, var_2]])

  return cov_matrix

###################################################################
## TO DO for students: generate and plot bivariate Gaussian data with variances of 1
## and a correlation coefficients of: 0.8
## repeat while varying the correlation coefficient from -1 to 1
###################################################################
np.random.seed(2020)  # set random seed
variance_1 = 1
variance_2 = 1
corr_coef = 0.8

# Uncomment to test your code and plot
cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
X = get_data(cov_matrix)
plot_data(X)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_57497711.py)

*Example output:*

---
# Section 2: Define a new orthonormal basis

```python
# @title Video 3: Orthonormal bases
video = YouTubeVideo(id="PC1RZELnrIg", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=PC1RZELnrIg

Next, we will define a new orthonormal basis of vectors ${\bf u} = [u_1,u_2]$ and ${\bf w} = [w_1,w_2]$. As we learned in the video, two vectors are orthonormal if:

1.   They are orthogonal (i.e., their dot product is zero):
\begin{equation}
{\bf u\cdot w} = u_1 w_1 + u_2 w_2 = 0
\end{equation}
2.   They have unit length:
\begin{equation}
||{\bf u} || = ||{\bf w} || = 1
\end{equation}

In two dimensions, it is easy to make an arbitrary orthonormal basis. All we need is a random vector ${\bf u}$, which we have normalized. If we now define the second basis vector to be ${\bf w} = [-u_2,u_1]$, we can check that both conditions are satisfied:
\begin{equation}
{\bf u\cdot w} = - u_1 u_2 + u_2 u_1 = 0
\end{equation}
and
\begin{equation}
{|| {\bf w} ||} = \sqrt{(-u_2)^2 + u_1^2} = \sqrt{u_1^2 + u_2^2} = 1,
\end{equation}
where we used the fact that ${\bf u}$ is normalized. So, with an arbitrary input vector, we can define an orthonormal basis, which we will write in matrix by stacking the basis vectors horizontally:

\begin{equation}
{{\bf W} } =
\begin{pmatrix}
 u_1 & w_1 \\
 u_2 & w_2
\end{pmatrix}.
\end{equation}

## Exercise 2: Find an orthonormal basis

In this exercise you will fill in the function below to define an orthonormal basis, given a single arbitrary 2-dimensional vector as an input.

**Steps**
* Modify the function `define_orthonormal_basis` to first normalize the first basis vector $\bf u$.
* Then complete the function by finding a basis vector $\bf w$ that is orthogonal to $\bf u$.
* Test the function using initial basis vector ${\bf u} = [3,1]$. Plot the resulting basis vectors on top of the data scatter plot using the function `plot_basis_vectors`. (For the data, use  $\sigma_1^2 =1$, $\sigma_2^2 =1$, and $\rho = .8$).

```python
help(plot_basis_vectors)
```

    Help on function plot_basis_vectors in module __main__:

    plot_basis_vectors(X, W)
        Plots bivariate data as well as new basis vectors.

        Args:
          X (numpy array of floats) :   Data matrix each column corresponds to a
                                        different random variable
          W (numpy array of floats) :   Square matrix representing new orthonormal
                                        basis each column represents a basis vector

        Returns:
          Nothing.

```python
a=[1,2]
b = np.flip(a)
b[0] = -b[0]
print(b)
```

    [-2  1]

```python
def define_orthonormal_basis(u):
  """
  Calculates an orthonormal basis given an arbitrary vector u.

  Args:
    u (numpy array of floats) : arbitrary 2-dimensional vector used for new
                                basis

  Returns:
    (numpy array of floats)   : new orthonormal basis
                                columns correspond to basis vectors
  """

  #################################################
  ## TODO for students: calculate the orthonormal basis
  # Fill out function and remove
  #raise NotImplementedError("Student excercise: implement the orthonormal basis function")
  #################################################

  # normalize vector u
  u = u/np.sqrt(np.transpose(u)@u)
  # calculate vector w that is orthogonal to w
  w = [-u[1],u[0]]

  W = np.column_stack([u, w])
  print(W)

  return W

np.random.seed(2020)  # set random seed
variance_1 = 1
variance_2 = 1
corr_coef = 0.8

cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
X = get_data(cov_matrix)
u = np.array([3, 1])

# Uncomment and run below to plot the basis vectors
W = define_orthonormal_basis(u)
plot_basis_vectors(X, W)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_7a9640ef.py)

*Example output:*

---
# Section 3: Project data onto new basis

```python
# @title Video 4: Change of basis
video = YouTubeVideo(id="Mj6BRQPKKUc", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=Mj6BRQPKKUc

Finally, we will express our data in the new basis that we have just found. Since $\bf W$ is orthonormal, we can project the data into our new basis using simple matrix multiplication :

\begin{equation}
{\bf Y = X W}.
\end{equation}

We will explore the geometry of the transformed data $\bf Y$ as we vary the choice of basis.

## Exercise 3: Define an orthonormal basis
In this exercise you will fill in the function below to define an orthonormal basis, given a single arbitrary vector as an input.

**Steps**
* Complete the function `change_of_basis` to project the data onto the new basis.
* Plot the projected data using the function `plot_data_new_basis`.
* What happens to the correlation coefficient in the new basis? Does it increase or decrease?
* What happens to variance?

```python
def change_of_basis(X, W):
  """
  Projects data onto new basis W.

  Args:
    X (numpy array of floats) : Data matrix each column corresponding to a
                                different random variable
    W (numpy array of floats) : new orthonormal basis columns correspond to
                                basis vectors

  Returns:
    (numpy array of floats)    : Data matrix expressed in new basis
  """

  #################################################
  ## TODO for students: project the data onto o new basis W
  # Fill out function and remove
  #raise NotImplementedError("Student excercise: implement change of basis")
  #################################################

  # project data onto new basis described by W
  Y = X@W

  return Y

# Unomment below to transform the data by projecting it into the new basis
Y = change_of_basis(X, W)
plot_data_new_basis(Y)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_a1124bbc.py)

*Example output:*

## Interactive Demo: Play with the basis vectors
To see what happens to the correlation as we change the basis vectors, run the cell below. The parameter $\theta$ controls the angle of $\bf u$ in degrees. Use the slider to rotate the basis vectors.

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

def refresh(theta=0):
  u = [1, np.tan(theta * np.pi / 180)]
  W = define_orthonormal_basis(u)
  Y = change_of_basis(X, W)
  plot_basis_vectors(X, W)
  plot_data_new_basis(Y)

_ = widgets.interact(refresh, theta=(0, 90, 5))
```

    interactive(children=(IntSlider(value=0, description='theta', max=90, step=5), Output()), _dom_classes=('widge…

## Questions

* What happens to the projected data as you rotate the basis?
* How does the correlation coefficient change? How does the variance of the projection onto each basis vector change?
* Are you able to find a basis in which the projected data is **uncorrelated**?

---
# Summary

- In this tutorial, we learned that multivariate data can be visualized as a cloud of points in a high-dimensional vector space. The geometry of this cloud is shaped by the covariance matrix.

- Multivariate data can be represented in a new orthonormal basis using the dot product. These new basis vectors correspond to specific mixtures of the original variables - for example, in neuroscience, they could represent different ratios of activation  across a population of neurons.

- The projected data (after transforming into the new basis) will generally have a different geometry from the original data. In particular, taking basis vectors that are aligned with the spread of cloud of points decorrelates the data.

* These concepts - covariance, projections, and orthonormal bases - are key for understanding PCA, which we be our focus in the next tutorial.