{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Unsupervised learning - Part 2 - Autoencoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Literature / Sources\n",
    "\n",
    "* [Reducing the Dimensionality of Data with Neural Networks, Hinton & Salakhutdinov, 2006](https://www.cs.toronto.edu/~hinton/science.pdf)\n",
    "<br>\n",
    "* [Variational Autoencoder](https://keras.io/examples/generative/vae/)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "\n",
    "from sklearn import preprocessing\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "\n",
    "import keras.backend as K\n",
    "\n",
    "DATA_DIR = os.path.join('..', 'data', 'rave')\n",
    "TRAINED_DIR = os.path.join('..', 'data', 'trained')\n",
    "LIVE = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.config.list_physical_devices('GPU')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spectra = np.load(os.path.join(DATA_DIR, 'ravedr6.spectra.npy'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rave = pd.read_csv(os.path.join(DATA_DIR, 'ravedr6.csv'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = pd.read_csv(os.path.join(DATA_DIR, 'ravedr6.apogee.csv'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#nspec_idx = rave.index\n",
    "#rand_sel = np.random.choice(range(len(nspec_idx)), 10000, replace=False)\n",
    "#spec_sample = spectra[nspec_idx][rand_sel]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_sample = spectra[rave[rave.RAVE_OBS_ID.isin(labels.RAVE_OBS_ID)].index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_sample.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wl_range = 8446.12176814 + np.cumsum(np.ones(1024) * 0.30025021)\n",
    "\n",
    "def plot_spectra(spectra, wl_range=wl_range, txt=None):\n",
    "    nspectra = min(10, len(spectra))\n",
    "    plt.figure(figsize=(14, 1.5 * nspectra))\n",
    "    ax = plt.subplot(111)\n",
    "    for i in range(nspectra):\n",
    "        ax.plot(wl_range, spectra[i] + i * 0.7, 'k-')\n",
    "    if txt is not None:\n",
    "        for i, t in enumerate(txt):\n",
    "            ax.text(wl_range[20], i * 0.7 + 1.1, t)\n",
    "    ax.set_ylim(0, 0.7 + 0.7 * nspectra)\n",
    "    ax.set_xlim(wl_range[0], wl_range[-1])\n",
    "    y0, y1 = ax.get_ylim()\n",
    "    for wl in (8498, 8542, 8662):\n",
    "        ax.plot([wl, wl], [y0, y1], 'k-', alpha=0.3)\n",
    "    ax.set_xlabel(r'Wavelength $\\mathrm{[\\AA]}$')\n",
    "    ax.set_ylabel(r'Flux')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://miro.medium.com/max/1000/0*uq2_ZipB9TqI9G_k)\n",
    "_Source: https://mc.ai/auto-encoder-in-biology/_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_dim = 2\n",
    "\n",
    "encoder_ipt = layers.Input(shape=(1024, 1))\n",
    "\n",
    "x = layers.Convolution1D(16, kernel_size=3, padding='same', name='convolution1')(encoder_ipt)\n",
    "x = layers.BatchNormalization()(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.MaxPooling1D(pool_size=2, padding='same')(x)\n",
    "                  \n",
    "x = layers.Convolution1D(32, kernel_size=3, padding='same', name='convolution2')(x)\n",
    "x = layers.BatchNormalization()(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.MaxPooling1D(pool_size=2, padding='same')(x)\n",
    "\n",
    "x = layers.Flatten()(x)\n",
    "\n",
    "x = layers.Dense(32, name='dense1')(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "x = layers.Dense(latent_dim, name='dense2')(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "encoder = keras.Model(encoder_ipt, x, name='encoder')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Conv1DTranspose(input_tensor, filters, kernel_size, strides=2, padding='same'):\n",
    "    # Tensorflow v2.3.0 has Conv1DTranspose, but 2.2.0 does not\n",
    "    x = layers.Lambda(lambda x: K.expand_dims(x, axis=2))(input_tensor)\n",
    "    x = layers.Conv2DTranspose(filters=filters, kernel_size=(kernel_size, 1), strides=(strides, 1), padding=padding)(x)\n",
    "    x = layers.Lambda(lambda x: K.squeeze(x, axis=2))(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_ipt = layers.Input(shape=(latent_dim,))\n",
    "\n",
    "x = layers.Dense(32, name='dense3')(decoder_ipt)\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "x = layers.Dense(8192, name='dense4')(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "x = layers.Reshape((256, 32))(x)\n",
    "\n",
    "x = Conv1DTranspose(x, 16, 3, padding='same')\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "x = Conv1DTranspose(x, 1, 3, padding='same')\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "decoder = keras.Model(decoder_ipt, x, name='decoder')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Writing a training loop from scratch](https://keras.io/guides/writing_a_training_loop_from_scratch/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AutoEncoder(keras.Model):\n",
    "    def __init__(self, encoder, decoder, **kwargs):\n",
    "        super(AutoEncoder, self).__init__(**kwargs)\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def train_step(self, data):\n",
    "        with tf.GradientTape() as tape:\n",
    "            reconstruction = self.decoder(self.encoder(data))\n",
    "            loss = tf.reduce_mean(keras.losses.mse(data, reconstruction))\n",
    "           \n",
    "        grads = tape.gradient(loss, self.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))\n",
    "        return {\"loss\": loss}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ae = AutoEncoder(encoder, decoder)\n",
    "ae.compile(optimizer='adam')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.expand_dims(spec_sample, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if LIVE:\n",
    "    history = ae.fit(data, epochs=1024, batch_size=32)\n",
    "else:\n",
    "    ae.load_weights(os.path.join(TRAINED_DIR, 'ae.h5'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ae.save_weights('ae.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = encoder.predict(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 6))\n",
    "extent = (-0.2, 1.5, -0.4, 1.8)\n",
    "plt.hexbin(res[:, 0], res[:, 1], gridsize=(60, 30), mincnt=1, C=labels.FeH, extent=extent)\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_spectra(spec_sample[np.argwhere((res[:, 1] > 0.45) & (res[:, 1] < 0.55))].squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[[i, 1.0] for i in np.arange(0.0, 1.5, 0.15)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#spec = decoder.predict([[i, 1.0] for i in np.arange(0.0, 1.5, 0.15)])\n",
    "spec = decoder.predict([[0.75, i] for i in np.arange(0.0, 1.5, 0.15)])\n",
    "plot_spectra(spec.reshape((10, 1024)));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 13\n",
    "\n",
    "spec = spec_sample[i].squeeze()\n",
    "gen_spec = decoder.predict(res[i:i + 1]).flatten()\n",
    "\n",
    "plt.figure(figsize=(14, 4))\n",
    "plt.plot(wl_range, spec, 'k')\n",
    "plt.plot(wl_range, gen_spec, 'r');\n",
    "\n",
    "plt.plot([8450, 8750], [0, 0], 'k--', lw=1)\n",
    "plt.plot(wl_range, spec - gen_spec)\n",
    "plt.text(8700, 0.15, r'$\\mu=%.3f\\quad\\sigma=%.3f$' % (np.mean(spec - gen_spec), np.std(spec - gen_spec)))\n",
    "plt.xlim(8450, 8750)\n",
    "plt.ylim(-0.2, 1.4);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(14, 8))\n",
    "ax = plt.axes([0.1, 0.1, 0.9, 0.9])\n",
    "plt.scatter(res[:, 0], res[:, 1], c=labels.FeH, s=20)\n",
    "\n",
    "plt.axis(extent)\n",
    "\n",
    "xx = np.linspace(*extent[:2], 5)\n",
    "yy = np.linspace(*extent[2:], 5)\n",
    "\n",
    "for j in range(5):\n",
    "    for i in range(5):\n",
    "        ax = plt.axes([0.1 + i * 0.19, 0.1 + j * 0.19, 0.14, 0.14])\n",
    "        s = decoder.predict([[xx[i], yy[j]]])\n",
    "        ax.plot(s[0].flatten());\n",
    "        ax.set_ylim(0, 1.4)\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "proj = encoder.predict(np.expand_dims(spectra, -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(14, 6))\n",
    "extent = (-1, 5, -2, 3.5)\n",
    "plt.subplot(121)\n",
    "plt.hexbin(res[:, 0], res[:, 1], gridsize=(60, 30), mincnt=1, extent=extent)\n",
    "plt.axis(extent)\n",
    "plt.subplot(122)\n",
    "plt.hexbin(proj[:, 0], proj[:, 1], gridsize=(60, 30), mincnt=1, extent=extent)\n",
    "plt.axis(extent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_spectra(spectra[np.argwhere((proj[:, 0] > -1) & (proj[:, 1] < -1))].squeeze());"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variational Autoencoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sampling(layers.Layer):\n",
    "    def call(self, inputs):\n",
    "        z_mean, z_log_var = inputs\n",
    "        batch = tf.shape(z_mean)[0]\n",
    "        dim = tf.shape(z_mean)[1]\n",
    "        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))\n",
    "        return z_mean + tf.exp(0.5 * z_log_var) * epsilon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_ipt = layers.Input(shape=(1024, 1))\n",
    "\n",
    "x = layers.Convolution1D(16, kernel_size=3, padding='same', name='convolution1')(encoder_ipt)\n",
    "x = layers.BatchNormalization()(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.MaxPooling1D(pool_size=2, padding='same')(x)\n",
    "                  \n",
    "x = layers.Convolution1D(32, kernel_size=3, padding='same', name='convolution2')(x)\n",
    "x = layers.BatchNormalization()(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.MaxPooling1D(pool_size=2, padding='same')(x)\n",
    "\n",
    "x = layers.Flatten()(x)\n",
    "\n",
    "x = layers.Dense(32, name='dense1')(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "z_mean = layers.Dense(latent_dim, name=\"z_mean\")(x)\n",
    "z_log_var = layers.Dense(latent_dim, name=\"z_log_var\")(x)\n",
    "z = Sampling()([z_mean, z_log_var])\n",
    "\n",
    "encoder = keras.Model(encoder_ipt, [z_mean, z_log_var, z], name=\"encoder\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE(keras.Model):\n",
    "    def __init__(self, encoder, decoder, **kwargs):\n",
    "        super(VAE, self).__init__(**kwargs)\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def train_step(self, data):\n",
    "        with tf.GradientTape() as tape:\n",
    "            z_mean, z_log_var, z = encoder(data)\n",
    "            reconstruction = decoder(z)\n",
    "            reconstruction_loss = tf.reduce_mean(\n",
    "                keras.losses.binary_crossentropy(data, reconstruction)\n",
    "            )\n",
    "            reconstruction_loss *= 1024\n",
    "            kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)\n",
    "            kl_loss = tf.reduce_mean(kl_loss)\n",
    "            kl_loss *= -0.5\n",
    "            total_loss = reconstruction_loss + kl_loss\n",
    "        grads = tape.gradient(total_loss, self.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))\n",
    "        return {\n",
    "            \"loss\": total_loss,\n",
    "            \"reconstruction_loss\": reconstruction_loss,\n",
    "            \"kl_loss\": kl_loss,\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = VAE(encoder, decoder)\n",
    "vae.compile(optimizer='adam')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if LIVE:\n",
    "    history = vae.fit(data, epochs=128, batch_size=32)\n",
    "else:\n",
    "    vae.load_weights(os.path.join(TRAINED_DIR, 'vae.h5'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = encoder.predict(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 6))\n",
    "plt.hexbin(res[0][:, 0], res[0][:, 1], gridsize=(60, 30), mincnt=1, C=labels.Teff)\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}