tutorial_phuijse_evic2018.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "%%HTML\n",
    "<style>\n",
    ".container {\n",
    "    width:90% ! important;\n",
    "}\n",
    ".output_png {\n",
    "    display: table-cell;\n",
    "    text-align: center;\n",
    "    vertical-align: middle;\n",
    "}\n",
    ".rendered_html { \n",
    "    font-size:0.8em; \n",
    "}\n",
    ".rendered_html table{\n",
    "    width: 80%;\n",
    "    margin-left:auto; \n",
    "    margin-right:auto;\n",
    "    padding: 20px;\n",
    "    border: 0px solid black;    \n",
    "    background-color: #ff;\n",
    "}\n",
    ".rendered_html td, .rendered_html th \n",
    "{\n",
    "    vertical-align: top;\n",
    "    text-align: left;\n",
    "    font-size: 14px;\n",
    "    font-face: sans-serif;\n",
    "}\n",
    "</style>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "<center>\n",
    "<h1> Learning latent representations for<br> Astronomical data using Neural Networks</h1>\n",
    "<h2>Pablo Huijse H. (phuijse at inf dot uach dot cl)</h2>\n",
    "<h3>Universidad Austral de Chile & Millennium Institute of Astrophysics</h3>\n",
    "</center>\n",
    "<br>\n",
    "<center>\n",
    "LIVE at https://github.com/phuijse/tutorial_periodic_stars\n",
    "</center>\n",
    "\n",
    "Thanks to\n",
    "- The organizers\n",
    "- The Millennium Institute of Astrophysics\n",
    "- CONICYT FONDECYT 1170305 and PAI 79170017\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Variable stars\n",
    "***\n",
    "\n",
    "- Stars whose brightness change in time\n",
    "- Different reasons behind this\n",
    "\n",
    "### Pulsating variables\n",
    "- Some variable stars pulsate radially\n",
    "- They expand/heat and contract/cool regularly\n",
    "- Examples: Cepheid and RR Lyrae\n",
    "\n",
    "<center>\n",
    "<a href=\"https://www.youtube.com/watch?v=sXJBrRmHPj8\">\n",
    "    <img src=\"https://media.giphy.com/media/QP4taxvfVmVEI/giphy.gif\" width=\"400\">\n",
    "</a>\n",
    "</center>\n",
    "\n",
    "### Eclipsing Binaries\n",
    "\n",
    "- System of two stars\n",
    "- The rotational plane is aligned with us\n",
    "- From our point of view we see brightness decrease with the mutual eclipses\n",
    "<center>\n",
    "<table>\n",
    "    <tr><td>\n",
    "        <a href=\"http://www.physast.uga.edu/~rls/astro1020/ch16/ovhd.html\">\n",
    "            <img src=\"img/intro-eb.gif\" width=\"300\">\n",
    "        </a>\n",
    "    </td>\n",
    "    <td>\n",
    "        <a href=\"https://en.wikipedia.org/wiki/File:Algol_AB_movie_imaged_with_the_CHARA_interferometer_-_labeled.gif\">\n",
    "            <img src=\"https://media.giphy.com/media/aYb0Ob2GHJ280/giphy.gif\" width=\"300\">\n",
    "        </a>\n",
    "    </td></tr>\n",
    "</table>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Scientific motivation\n",
    "***\n",
    "\n",
    "- Variable stars as distance tracers: Milky-way maps\n",
    "<table>\n",
    "    <tr><td>   \n",
    "        <img src=\"img/period-luminosity-relation.gif\" width=\"400\">\n",
    "    </td><td>\n",
    "        <img src=\"img/intro-milky-way.jpg\" width=\"400\">\n",
    "    </td></tr>\n",
    "</table>\n",
    "\n",
    "\n",
    "- Variable star analysis and classification: **Astrophysics**\n",
    "<center>\n",
    "<a href=\"http://www.atnf.csiro.au/outreach/education/senior/astrophysics/variable_types.html\">\n",
    "    <img src=\"img/variable-star-classification.gif\" width=\"400\">\n",
    "</a>\n",
    "</center>\n",
    "- New methods to analyze astronomical data: **Signal processing** and **Data Science**\n",
    "    - Room for interdisciplinary research\n",
    "    - Astroinformatics and Astrostatistics\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Light curve\n",
    "***\n",
    "\n",
    "- Time series of a star's flux (brightness) on a given passband\n",
    "- The \"apparent\" brightness is estimated through **Photometry**\n",
    "- Variable stars are studied through their light curves\n",
    "\n",
    "<table><tr><td>\n",
    "    <img src=\"img/intro-vista.png\" width=\"250\">\n",
    "</td><td>\n",
    "    <img src=\"img/intro-sources.png\" width=\"300\">\n",
    "</td></tr></table>\n",
    "\n",
    "<center>\n",
    "    <img src=\"img/intro-sources-time.png\" width=\"600\">\n",
    "</center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "import gzip\n",
    "import pickle\n",
    "import numpy as np\n",
    "%matplotlib notebook\n",
    "import matplotlib.pylab as plt\n",
    "from matplotlib import rcParams, animation\n",
    "rcParams.update({'font.size': 12})\n",
    "rcParams.update({'axes.grid': True})\n",
    "\n",
    "# Get some light curves to play\n",
    "with gzip.open(\"data/lc_data.pgz\", mode=\"r\") as f:\n",
    "    lc_data = pickle.load(f)\n",
    "\n",
    "lc_periods = pickle.load(open(\"data/lc_periods.pkl\", \"rb\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "## Inspecting a light curve\n",
    "***\n",
    "\n",
    "In this case light curves are text files with three colums\n",
    "- **Modified Julian Data (MJD):** Corresponds to time \n",
    "- **Magnitude:** Corresponds to apparent brightness (log scale)\n",
    "- **Error:** Photometric error estimation of the magnitude"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(9, 4), tight_layout=True)\n",
    "mjd, mag, err = lc_data[6].T\n",
    "ax.errorbar(mjd, mag, err, fmt='o')\n",
    "ax.invert_yaxis(); \n",
    "ax.set_xlabel('Modified Julian Date (MJD)\\n ')\n",
    "ax.set_ylabel('Magnitude\\n(The smaller the brighter)');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "- Irregular sampling, data gaps\n",
    "- Heteroscedastic noise: Error variance change in time\n",
    "\n",
    "This light curve is actually from a periodic variable star..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "def fold(time, period):\n",
    "    \"\"\"\n",
    "    returns phase = time/period - floor(time/period)\n",
    "    \"\"\"\n",
    "    return np.mod(time, period)/period\n",
    "\n",
    "idx = 6\n",
    "mjd, mag, err = lc_data[idx].T\n",
    "fig, ax = plt.subplots(figsize=(9, 4), tight_layout=True)\n",
    "phi = fold(mjd, lc_periods[idx])\n",
    "ax.errorbar(np.hstack((phi, phi+1)), \n",
    "            np.hstack((mag, mag)), \n",
    "            np.hstack((err, err)), fmt='o')\n",
    "ax.invert_yaxis(); \n",
    "ax.set_ylabel('Magnitude\\n(The smaller the brighter)');\n",
    "ax.set_xlabel('Phase @ Period %0.6f' %(lc_periods[idx]));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "## Folding the light curve\n",
    "***\n",
    "- Technique used by astronomers to visually inspect periodic variables\n",
    "- You need a candidate period $P$ to perform the folding\n",
    "- The time axis is divided in chucks of size $P$ and plotted on top each other\n",
    "\n",
    "$$\n",
    "\\phi = \\text{modulo}(\\text{MJD}, P)/P\n",
    "$$\n",
    "- Then you plot the magnitude as a function of $\\phi$ \n",
    "    - If $P$ is close to the true period:  Nice periodic shape\n",
    "    - Otherwise: Noisy pattern\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(9, 4), tight_layout=True)\n",
    "period_grid = np.linspace(lc_periods[6]-0.001, lc_periods[6]+0.001, num=100)\n",
    "phi = fold(mjd, period_grid[0])\n",
    "line, caps, errorbars = ax.errorbar(np.hstack((phi, phi+1)), \n",
    "                                    np.hstack((mag, mag)), \n",
    "                                    np.hstack((err, err)), fmt='o')\n",
    "segs = errorbars[0].get_segments()\n",
    "ax.invert_yaxis(); \n",
    "ax.set_ylabel('Magnitude\\n(The smaller the brighter)');\n",
    "\n",
    "def update(n):\n",
    "    phi = fold(mjd, period_grid[n])\n",
    "    for i in range(len(segs)//2):\n",
    "        segs[i][:, 0] = phi[i]\n",
    "        segs[i+len(phi)][:, 0] = phi[i]+1\n",
    "\n",
    "    line.set_xdata(np.hstack((phi, phi+1))) \n",
    "    errorbars[0].set_segments(segs)\n",
    "    ax.set_xlabel('Phase @ Period %0.6f' %(period_grid[n]))\n",
    "\n",
    "anim = animation.FuncAnimation(fig, update, frames=100, interval=100, repeat=False, blit=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "## Periodograms\n",
    "***\n",
    "- We want to find the period (fundamental frequency) of the star\n",
    "- This is generally done using the **periodogram**\n",
    "- The periodogram is an estimator of the signal's **power as a function of frequency**\n",
    "- It can be computed as the FFT of the autocorrelation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(9, 4), tight_layout=True)\n",
    "time, dtime = np.linspace(0, 10, num=100, retstep=True)\n",
    "freq = np.fft.rfftfreq(d=dtime, n=100)\n",
    "noise = np.random.randn(len(time))\n",
    "def draw(step):\n",
    "    ax[0].cla(); ax[1].cla(); ax[0].set_ylim([-2, 2])\n",
    "    s = np.sin(2*np.pi*time - 2.0*np.pi*step/50)\n",
    "    ax[0].plot(time, s, lw=2); s+= noise*step/50; ax[0].plot(time, s, 'k.')\n",
    "    ax[1].plot(freq, np.abs(np.fft.rfft(s)/len(time)), 'k-', lw=2)\n",
    "anim = animation.FuncAnimation(fig, draw, frames=100, interval=100, repeat=False, blit=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "But the FFT and correlation assumme regular time sampling\n",
    "\n",
    "Estimating the period in irregularly sampled time series:\n",
    "1. Best least squares fit of sine wave: Lomb-Scargle periodogram\n",
    "1. Most \"ordered\" folded light curve: Conditional Entropy and **Mutual Information**\n",
    "1. ...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "!pip install P4J --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "import P4J\n",
    "my_per = P4J.periodogram(method='QMIEU') \n",
    "mjd, mag, err = lc_data[6].T\n",
    "my_per.set_data(mjd, mag, err, h_KDE_P=0.2)\n",
    "my_per.frequency_grid_evaluation(fmin=0.0, fmax=4.0, fresolution=1e-4)\n",
    "my_per.finetune_best_frequencies(fresolution=1e-5, n_local_optima=10)\n",
    "freq, per = my_per.get_periodogram()\n",
    "fbest, pbest  = my_per.get_best_frequencies()\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(9, 4), tight_layout=True)\n",
    "ax.plot(freq, per)\n",
    "ax.set_xlabel('Frequency [1/MJD]')\n",
    "ax.set_ylabel('Periodogram')\n",
    "print(\"Best period: %f days\" %(1.0/fbest[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(9, 4), tight_layout=True)\n",
    "phi = fold(mjd, 1.0/fbest[0])\n",
    "ax.errorbar(np.hstack((phi, phi+1)), \n",
    "            np.hstack((mag, mag)), \n",
    "            np.hstack((err, err)), fmt='o')\n",
    "ax.invert_yaxis(); \n",
    "ax.set_ylabel('Magnitude\\n(The smaller the brighter)');\n",
    "ax.set_xlabel('Phase @ Period %0.6f' %(1.0/fbest[0]));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Getting features from our periodic light curves\n",
    "\n",
    "- We want to train a autoencoder neural network for light curves\n",
    "- Given that we have the period we train on the folded light curve\n",
    "- We normalize and interpolate the folded light curve using kernel regression\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "phi_interp = np.linspace(0, 1, num=40)\n",
    "fig = plt.figure(figsize=(9, 5), tight_layout=True)\n",
    "ax1 = plt.subplot2grid((3,1), (0,0), rowspan=2, fig=fig)\n",
    "mjd, mag, err = lc_data[6].T\n",
    "phi = fold(mjd, 1.0/fbest[0]); w = 1.0/err**2;\n",
    "ax1.errorbar(phi, mag, err, fmt='.', c='k'); ax1.invert_yaxis();\n",
    "l1 = ax1.plot(phi_interp, [np.amax(mag)]*len(phi_interp), lw=4)\n",
    "ax2 = plt.subplot2grid((3,1), (2,0), rowspan=1, fig=fig)\n",
    "l2 = ax2.plot(phi_interp, np.exp(-0.5*phi_interp**2/0.1**2), lw=4)\n",
    "mag_interp = np.ones_like(phi_interp)*np.amax(mag)\n",
    "\n",
    "def update(idx):\n",
    "    gt = np.exp(-0.5*(phi_interp[idx]-phi)**2/0.05**2)\n",
    "    mag_interp[idx] = np.sum(w*gt*mag)/np.sum(w*gt)\n",
    "    l1[0].set_ydata(mag_interp)\n",
    "    l2[0].set_ydata(np.exp(-0.5*(phi_interp-phi_interp[idx])**2/0.05**2))\n",
    "anim = animation.FuncAnimation(fig, update, frames=40, interval=100, repeat=False, blit=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "from utils import featurize_lc, defeaturize_lc\n",
    "from IPython.display import display\n",
    "from ipywidgets import Button\n",
    "\n",
    "features = np.zeros(shape=(len(lc_data), len(phi_interp)))\n",
    "weights = np.zeros(shape=(len(lc_data), len(phi_interp)))\n",
    "norm = np.zeros(shape=(len(lc_data), 3))\n",
    "for i in range(len(lc_data)):\n",
    "    features[i, :], weights[i, :], norm[i, :] = featurize_lc(lc_data[i], lc_periods[i], phi_interp)\n",
    "    \n",
    "next_button = Button(description=\"Next\")\n",
    "idx = 4950\n",
    "fig, ax = plt.subplots(figsize=(6, 4), tight_layout=True)\n",
    "\n",
    "def plot_features(idx):\n",
    "    ax.cla(); \n",
    "    ax.set_title(\"Idx: %d\\nPeriod: %0.6f\" %(idx, lc_periods[idx]))\n",
    "    mag, err = defeaturize_lc(features[idx, :], weights[idx, :], norm[idx, :])\n",
    "    ax.plot(phi_interp, mag, lw=4)\n",
    "    ax.fill_between(phi_interp, (mag-err), (mag+err), alpha=0.5)\n",
    "    ax.set_xlabel('Phase'); ax.set_ylabel('Normalized magnitude');\n",
    "    mjd, mag, err = lc_data[idx][:, 0], lc_data[idx][:, 1], lc_data[idx][:, 2]\n",
    "    phi = fold(mjd, lc_periods[idx])\n",
    "    ax.errorbar(phi, mag, err, fmt='.', c='k', alpha=0.5, label='data'); \n",
    "    ax.invert_yaxis()\n",
    "    plt.legend();\n",
    "\n",
    "def on_nbutton_clicked(b):\n",
    "    global idx\n",
    "    idx += 1\n",
    "    plot_features(idx)\n",
    "                \n",
    "next_button.on_click(on_nbutton_clicked)\n",
    "plot_features(idx)\n",
    "next_button"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# [PyTorch](https://pytorch.org/) in a nutshell\n",
    "***\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "import torch \n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "class simple_MLP(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self, n_input=2, n_hidden=2, n_output=1):\n",
    "        super(simple_MLP, self).__init__()\n",
    "        # Encoder layers\n",
    "        self.hidden = torch.nn.Linear(n_input, n_hidden)\n",
    "        self.output = torch.nn.Linear(n_hidden, n_output)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        z = F.relu(self.hidden(x))\n",
    "        y = self.output(z)\n",
    "        return y\n",
    "\n",
    "hello_world_model = simple_MLP(n_hidden=6)\n",
    "print(hello_world_model)\n",
    "optimizer = torch.optim.Adam(hello_world_model.parameters(), lr=1e-1)\n",
    "\n",
    "data = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])\n",
    "label = np.array([0, 1, 1, 0])\n",
    "xor_dataset = TensorDataset(torch.from_numpy(data.astype('float32')),\n",
    "                            torch.from_numpy(label.astype('float32')))\n",
    "\n",
    "xor_loader = DataLoader(dataset=xor_dataset, batch_size=1, shuffle=True)\n",
    "\n",
    "for epoch in range(100):\n",
    "    for bdata, labels in xor_loader:\n",
    "        optimizer.zero_grad()\n",
    "        predictions = hello_world_model.forward(bdata)\n",
    "        loss = F.binary_cross_entropy_with_logits(predictions, labels.unsqueeze(1))        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "torch.sigmoid(hello_world_model(xor_dataset.tensors[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Training a variational autoencoder\n",
    "***\n",
    "\n",
    "1. In this part we will train an [autoencoder](https://docs.google.com/presentation/d/1IJ2n8X4w8pvzNLmpJB-ms6-GDHWthfsJTFuyUqHfXg8/edit?usp=sharing) to visualize the feature space \n",
    "- We will use PyTorch to create and train the model\n",
    "- We have light curves with unknown label and 50 light curves labeled as **RR Lyrae**\n",
    "- Can we find unlabeled light curves that belong to the RR Lyrae class?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def logsumexp(inputs, dim=None, keepdim=True):    \n",
    "    # From: https://github.com/YosefLab/scVI/issues/13\n",
    "    return (inputs - F.log_softmax(inputs, dim=dim)).sum(dim, keepdim=keepdim)\n",
    "\n",
    "class VAE(torch.nn.Module):\n",
    "    def __init__(self, n_input=40, n_hidden=[10, 10], n_latent=2, importance_sampling=False):\n",
    "        super(VAE, self).__init__()\n",
    "        self.importance = importance_sampling\n",
    "        # Encoder layers\n",
    "        self.enc_hidden1 = torch.nn.Linear(n_input, n_hidden[0])\n",
    "        self.enc_hidden2 = torch.nn.Linear(n_hidden[0], n_hidden[1])\n",
    "        #self.enc_hidden = torch.nn.Conv1d(in_channels=1, kernel_size=10, out_channels=n_hidden)\n",
    "        self.enc_mu = torch.nn.Linear(n_hidden[1], n_latent)\n",
    "        self.enc_logvar = torch.nn.Linear(n_hidden[1], n_latent)\n",
    "        # decoder layers\n",
    "        self.dec_hidden2 = torch.nn.Linear(n_latent, n_hidden[1]) \n",
    "        self.dec_hidden1 = torch.nn.Linear(n_hidden[1], n_hidden[0]) \n",
    "        self.dec_mu = torch.nn.Linear(n_hidden[0], n_input)\n",
    "        self.dec_logvar = torch.nn.Linear(n_hidden[0], 1)\n",
    "\n",
    "        \n",
    "    def encode(self, x):\n",
    "        h = F.relu(self.enc_hidden1(x))\n",
    "        h = F.relu(self.enc_hidden2(h))\n",
    "        #h = F.relu(self.enc_hidden(x.unsqueeze(1)))\n",
    "        #h = F.adaptive_avg_pool1d(h, output_size=1).view(-1, 20)\n",
    "        return self.enc_mu(h), self.enc_logvar(h)\n",
    "\n",
    "    def sample(self, mu, logvar, k=1):\n",
    "        batch_size, n_latent = logvar.shape\n",
    "        std = (0.5*logvar).exp()\n",
    "        eps = torch.randn(batch_size, k, n_latent, device=std.device, requires_grad=False)\n",
    "        return eps.mul(std.unsqueeze(1)).add(mu.unsqueeze(1))\n",
    "\n",
    "    def decode(self, z):\n",
    "        h = F.relu(self.dec_hidden2(z))\n",
    "        h = F.relu(self.dec_hidden1(h))\n",
    "        hatx, hatlogvar = self.dec_mu(h), self.dec_logvar(h)\n",
    "        return hatx, hatlogvar        \n",
    "\n",
    "    def forward(self, x, k=1):\n",
    "        enc_mu, enc_logvar = self.encode(x)\n",
    "        z = self.sample(enc_mu, enc_logvar, k)\n",
    "        dec_mu, dec_logvar = self.decode(z)\n",
    "        return dec_mu, enc_mu, enc_logvar, dec_logvar, z\n",
    "    \n",
    "    def ELBO(self, x, w, mc_samples=1):  \n",
    "        dec_mu, enc_mu, enc_logvar, dec_logvar, z = self.forward(x, mc_samples)\n",
    "        logpxz = -0.5*(2.*torch.log(w.unsqueeze(1) + (dec_logvar/2).exp()) \\\n",
    "                       + (x.unsqueeze(1) - dec_mu).pow(2)/(dec_logvar.exp() + w.pow(2).unsqueeze(1))).sum(dim=-1)    \n",
    "        #logpxz = -0.5*(dec_logvar + (x.unsqueeze(1) - dec_mu).pow(2)/dec_logvar.exp()).sum(dim=-1)    \n",
    "        \n",
    "        if self.importance: # Importance-Weighted autoencoder (IWAE)\n",
    "            logqzxpz = 0.5 * (z.pow(2) - z.sub(enc_mu.unsqueeze(1)).pow(2)/enc_logvar.unsqueeze(1).exp() - enc_logvar.unsqueeze(1)).sum(dim=-1)\n",
    "        else:  # Variational autoencoder\n",
    "            logqzxpz = -0.5 * (1.0 + enc_logvar - enc_mu.pow(2) - enc_logvar.exp()).sum(dim=-1).unsqueeze_(1)\n",
    "        ELBO = torch.sum(logsumexp(logqzxpz - logpxz, dim=1) + np.log(mc_samples))\n",
    "        return ELBO, logpxz.sum()/mc_samples, logqzxpz.sum()/logqzxpz.shape[1]\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "!pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "from utils import live_metric_plotter\n",
    "from tqdm import tqdm_notebook\n",
    "\n",
    "# Create light curve dataset from numpy arrays\n",
    "lc_dataset = TensorDataset(torch.from_numpy(features.astype('float32')), \n",
    "                           torch.from_numpy(weights.astype('float32')),\n",
    "                           torch.from_numpy(lc_periods.astype('float32')))\n",
    "\n",
    "batch_size_, nepochs, mc_samples = 32, 50, 32\n",
    "torch.manual_seed(0);  np.random.seed(0) # reproducible results, comment this for real world\n",
    "\n",
    "# Generate data loaders\n",
    "P = np.random.permutation(5000)\n",
    "train_loader = DataLoader(dataset=lc_dataset, batch_size=batch_size_, \n",
    "                          sampler=torch.utils.data.SubsetRandomSampler(P[:4000]))\n",
    "valid_loader = DataLoader(dataset=lc_dataset, batch_size=batch_size_, \n",
    "                          sampler=torch.utils.data.SubsetRandomSampler(P[4000:]))\n",
    "\n",
    "# instantiate model\n",
    "model = VAE(n_input=40, n_hidden=[20, 10], n_latent=2, importance_sampling=True)\n",
    "print(model)\n",
    "model_parameters = filter(lambda p: p.requires_grad, model.parameters())\n",
    "print(\"Number of trainable parameters: %d\" %(sum([np.prod(p.size()) for p in model_parameters])))\n",
    "\n",
    "# Set optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "plotter = live_metric_plotter(figsize=(9, 4))\n",
    "metrics = np.zeros(shape=(nepochs, 2, 2))\n",
    "\n",
    "# Train\n",
    "for epoch in tqdm_notebook(range(nepochs)):\n",
    "    for feature, weight, period in train_loader:\n",
    "        optimizer.zero_grad()        \n",
    "        loss, rec_loss, reg_loss = model.ELBO(feature, weight, mc_samples)        \n",
    "        loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)        \n",
    "        optimizer.step()\n",
    "        metrics[epoch, 0, 0] += rec_loss.detach().numpy()/len(train_loader.sampler)\n",
    "        metrics[epoch, 0, 1] += reg_loss.detach().numpy()/len(train_loader.sampler)\n",
    "    # Test\n",
    "    for feature, weight, period in valid_loader:\n",
    "        loss, rec_loss, reg_loss = model.ELBO(feature, weight, mc_samples)\n",
    "        metrics[epoch, 1, 0] += rec_loss.detach().numpy()/len(valid_loader.sampler)\n",
    "        metrics[epoch, 1, 1] += reg_loss.detach().numpy()/len(valid_loader.sampler)\n",
    "    \n",
    "    if epoch > 0:\n",
    "        plotter.update(epoch, metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "!pip install bokeh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "test_dec_mu, test_enc_mu, test_enc_logvar, test_dec_logvar, test_z = model.forward(lc_dataset.tensors[0], k=10)\n",
    "test_enc_mu, test_enc_sigma = test_enc_mu.detach().numpy(), (test_enc_logvar.detach()*0.5).exp().numpy()\n",
    "test_dec_mu = test_dec_mu.detach().numpy()\n",
    "\n",
    "from bokeh import __version__\n",
    "print(__version__)\n",
    "from bokeh.io import output_notebook, show, push_notebook\n",
    "from bokeh.layouts import row, column \n",
    "from bokeh.models import Whisker, Band, ColumnDataSource, HoverTool, CustomJS, Range1d\n",
    "from bokeh.plotting import figure\n",
    "output_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "TOOLS = \"pan, wheel_zoom, reset\"\n",
    "hover = HoverTool(tooltips=[(\"index\", \"$index\")])    \n",
    "main_plot = figure(tools=[hover, TOOLS], plot_width=600, plot_height=600, \n",
    "                   min_border=10, min_border_left=0,\n",
    "                   toolbar_location=\"above\", title=\"Latent space\")\n",
    "\n",
    "lc_plots, lc_sources = [], []\n",
    "p = figure(tools=TOOLS, plot_width=350, plot_height=250, toolbar_location=\"above\", title='Data')\n",
    "source = ColumnDataSource(data={'x': [0], 'y': [0], 'y+dy': [0.1], 'y-dy': [-0.1]}) \n",
    "p.circle('x', 'y', size=5, line_color=\"black\", fill_color=\"black\", source=source)\n",
    "p.add_layout(Whisker(source=source, base=\"x\", upper=\"y+dy\", lower=\"y-dy\"))\n",
    "lc_plots.append(p)\n",
    "lc_sources.append(source)\n",
    "\n",
    "p = figure(tools=TOOLS, plot_width=350, plot_height=250, toolbar_location=\"above\", title='Features')\n",
    "source = ColumnDataSource(data={'x': [0], 'y_rec': [0], \n",
    "                                'y_feat': [0], 'y_rec_lower': [-1], 'y_rec_upper': [1]})\n",
    "\n",
    "p.line('x', 'y_feat', line_width=4, color='blue', source=source)\n",
    "p.line('x', 'y_rec', line_width=4, color='red', source=source)\n",
    "band = Band(base='x', lower='y_rec_lower', upper='y_rec_upper', source=source, \n",
    "            level='underlay', fill_alpha=0.5, fill_color='red', line_width=None)\n",
    "p.add_layout(band)\n",
    "lc_plots.append(p)\n",
    "lc_sources.append(source)\n",
    "\n",
    "labels = np.zeros(shape=(5000,)); labels[4950:] = 1\n",
    "colors = ['red' if label else 'blue' for label in labels]\n",
    "sc = main_plot.scatter(test_enc_mu[:, 0], test_enc_mu[:, 1], \n",
    "                       size=5, alpha=0.5, fill_color=colors, line_color=None)\n",
    "\n",
    "\n",
    "callback = CustomJS(args={'segment': sc.data_source}, code=\"\"\" \n",
    "if (IPython.notebook.kernel !== undefined && cb_data.index.indices !== undefined) {\n",
    "    var kernel = IPython.notebook.kernel;\n",
    "    cmd = \"update_plots(\" + cb_data.index.indices[0] + \")\";\n",
    "    kernel.execute(cmd, {}, {})}; \n",
    "\"\"\")\n",
    "\n",
    "main_plot.hover.callback=callback\n",
    "\n",
    "t = show(row(main_plot, \n",
    "             column(lc_plots[0], lc_plots[1])), notebook_handle=True)\n",
    "\n",
    "def update_plots(idx):\n",
    "    #print(idx)\n",
    "    mjd, mag, err = lc_data[idx].T\n",
    "    phi = fold(mjd, lc_dataset.tensors[2][idx]).numpy()\n",
    "    lc_sources[0].data={'x': phi, 'y': -mag, 'y+dy': -mag+err, 'y-dy': -mag-err} \n",
    "    mu_dec = np.mean(test_dec_mu[idx], axis=0)\n",
    "    s_dec = 2*np.std(test_dec_mu[idx], axis=0)\n",
    "    lc_sources[1].data={'x': phi_interp, 'y_feat': -lc_dataset[idx][0].numpy(), \n",
    "                        'y_rec': -mu_dec, 'y_rec_lower':-mu_dec+s_dec, 'y_rec_upper':-mu_dec-s_dec} \n",
    "    \n",
    "    \n",
    "    push_notebook(handle=t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "z = model.sample(torch.tensor([[0., -1.]]), torch.tensor([[-2., -2.]]), k=2)\n",
    "hatmu, hats = model.decode(z)\n",
    "hatmu, hats = hatmu.detach().numpy()[0][0], (hats/2.).exp().detach().numpy()[0][0]\n",
    "fig, ax = plt.subplots(figsize=(9, 4))\n",
    "ax.plot(np.concatenate((phi_interp, phi_interp+1)), \n",
    "        np.concatenate((hatmu, hatmu)), lw=4)\n",
    "ax.invert_yaxis();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(20, 20, figsize=(8, 6), tight_layout=True)\n",
    "for i in range(20):\n",
    "    for j in range(20):\n",
    "        hatmu, hats = model.decode(torch.tensor([-3. + 6.*j/20, 3. - 6.*i/20]))\n",
    "        ax[i, j].plot(phi_interp, hatmu.detach().numpy(), lw=2)\n",
    "        ax[i, j].invert_yaxis(); ax[i, j].axis('off')\n",
    "                           "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "The figure shows the digits in latent space as a dot (mean of the variational posterior) with \n",
    "errorbars (standard deviation of the variational posterior). Each point is a distribution!\n",
    "\"\"\"\n",
    "fig = plt.figure(figsize=(10, 5), dpi=80)\n",
    "ax_main = plt.subplot2grid((2, 3), (0, 0), colspan=2, rowspan=2)\n",
    "ax_ori = plt.subplot2grid((2, 3), (0, 2))\n",
    "ax_rec = plt.subplot2grid((2, 3), (1, 2))\n",
    "a, b, c = ax_main.errorbar(x=test_enc_mu[:, 0], y=test_enc_mu[:, 1], \n",
    "                           xerr=test_enc_sigma[:, 0], yerr=test_enc_sigma[:, 1], \n",
    "                           fmt='none', alpha=0.2, zorder=-1)\n",
    "labels = np.zeros(shape=(5000,)); labels[4950:] = 1\n",
    "\n",
    "sc = ax_main.scatter(test_enc_mu[:, 0], test_enc_mu[:, 1], s=2, alpha=0.2, \n",
    "                     c=labels, cmap=plt.cm.RdBu_r)\n",
    "clb = plt.colorbar(sc, ax=ax_main)\n",
    "for i in range(2):\n",
    "    c[i].set_color(clb.to_rgba(labels))\n",
    "    \n",
    "    \n",
    "c_lim, r_lim = ax_main.get_xlim(), ax_main.get_ylim()\n",
    "plt.tight_layout()\n",
    "phi_interp = np.linspace(0, 1, num=40)\n",
    "def onclick(event):\n",
    "    z_closest = [event.xdata, event.ydata]\n",
    "    print(z_closest)\n",
    "    idx = np.argmin(np.sum((test_enc_mu[:, :2] - z_closest)**2, axis=1))\n",
    "    ax_ori.cla(); ax_ori.set_title(\"Idx:%d, Label:%d\" %(idx, labels[idx]))\n",
    "    mjd, mag, err = lc_data[idx].T\n",
    "    phi = fold(mjd, lc_dataset.tensors[2][idx])\n",
    "    ax_ori.errorbar(phi, mag, err, c='k', fmt='.')\n",
    "    mag, err = defeaturize_lc(lc_dataset.tensors[0][idx].numpy(), \n",
    "                              lc_dataset.tensors[1][idx].numpy(), norm[idx])\n",
    "    ax_ori.plot(phi_interp, mag, lw=2)\n",
    "    ax_ori.fill_between(phi_interp, mag - err, mag + err, alpha=0.5)\n",
    "    ax_ori.invert_yaxis(); \n",
    "    ax_rec.cla(); ax_rec.invert_yaxis(); #ax_rec.set_ylim([2.5, -1.5]); \n",
    "    mag, err = lc_dataset.tensors[0][idx].numpy(), lc_dataset.tensors[1][idx].numpy()\n",
    "    ax_rec.plot(phi_interp, mag, lw=2)\n",
    "    ax_rec.fill_between(phi_interp, mag - err, mag + err, alpha=0.5)\n",
    "    mu_dec = np.mean(test_dec_mu[idx], axis=0)\n",
    "    s_dec = np.std(test_dec_mu[idx], axis=0)\n",
    "    ax_rec.plot(phi_interp, mu_dec, c='r', lw=2)\n",
    "    ax_rec.fill_between(phi_interp, mu_dec-2*s_dec, mu_dec+2*s_dec, facecolor='r', alpha=0.5)\n",
    "    \n",
    "cid = fig.canvas.mpl_connect('button_press_event', onclick);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}