{ "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", "
\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 }