# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/ziatdinovmax/pyroVED/blob/main/examples/pyroVED_examples_C_cont_VAE.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="caW4JMQ2tcuY"
# # Conditioning VAE on a continuous "property" vector
#
# Prepared by <NAME>
#
#
# ---
#
# This notebook denonstrates how to condition VAE on a *known* property (orproperties) of the data to help it discover the unknown ones. We demnstrate it using synthetic 1D spectral data with random variations in peak position and width.
#
# ---
#
#
#
# E-mail: <EMAIL>

# + [markdown] id="-SZPuudhPDiu"
# Install pyroVED package:

# + id="DbarJgF9KPu0" colab={"base_uri": "https://localhost:8080/"} outputId="095f0ee3-bf52-4494-e108-938863cdb995"
# !pip install git+https://github.com/ziatdinovmax/pyroved

# + [markdown] id="yaevNk8GxmwM"
# Imports:

# + id="WK-h3pm5KZyd"
import pyroved as pv
import torch
import numpy as np
import matplotlib.pyplot as plt


# + [markdown] id="cwEybJEX0_7r"
# First, we prepare our synthetic data representing 1D Gaussian signals with different widths that is randomly shifted along the *x*-axis and is corrupted by noise.

# + id="KT3bq7NJxs1G"
def gaussian(x, mu, sig):
    return torch.exp(-torch.pow(x - mu, 2.) / (2 * torch.pow(sig, 2.)))

n_samples = 5000 # number of samples
l_signal = 100 # signal length

# Generate synthetic dataset with randomly shifted noisy 1D signal
torch.manual_seed(1)  # for reproducibility
x = torch.linspace(-12, 12, l_signal).expand(n_samples, l_signal)
noise = torch.randint(1, 100, (n_samples, 1)) / 1e3
mu = torch.randint(-30, 30, size=(n_samples, 1)) / 10
sig = torch.randint(50, 500, size=(n_samples, 1)) / 1e2
train_data = gaussian(x, mu, sig) + noise * torch.randn(size=(n_samples, l_signal))
# Normalize to (0, 1)
train_data = (train_data - train_data.min()) / (train_data.max() - train_data.min())

# + id="_ouqe1kjg9IJ"
# Initialize train loader
train_loader = pv.utils.init_dataloader(train_data.unsqueeze(1), mu, batch_size=64)

# + colab={"base_uri": "https://localhost:8080/", "height": 469} id="YE8vs9-VuKBA" outputId="3c7532ec-98b2-42b5-fe4e-63265dd4f561"
fig, axes = plt.subplots(8, 8, figsize=(8, 8),
                         subplot_kw={'xticks':[], 'yticks':[]},
                         gridspec_kw=dict(hspace=0.1, wspace=0.1))

for ax, (y, m) in zip(axes.flat, train_loader):
    ax.plot(x[0], y[0, 0])
    ax.set_ylim(-0.15, 1.0)
    ax.text(0.05, 0.05, str(np.around(m[0].item(), decimals=4)),
            transform=ax.transAxes, color='green')

# + [markdown] id="VR4LX4QQOL6i"
# Train VAE model where our prior is conditioned on a continuous vector $c$ as $p(z|c)$. Here we pass the information about shifts as our $c$ vector.

# + colab={"base_uri": "https://localhost:8080/"} id="DqC8GQPaK7RL" outputId="fefd3996-2b18-46e9-ba97-5011c892e5a4"
in_dim = (100,)

# Initialize VAE (we specify c_dim=1 to add extra-dimension corresponding to c-vector to the decoder)
vae = pv.models.iVAE(in_dim, latent_dim=2, c_dim=1, invariances=None, seed=0)

# Initialize SVI trainer
trainer = pv.trainers.SVItrainer(vae)

# Train for n epochs:
for e in range(250):
    trainer.step(train_loader)
    trainer.print_statistics()

# + [markdown] id="9cwX2eW1Ppx0"
# Encode the training data into the latent space:

# + colab={"base_uri": "https://localhost:8080/", "height": 396} id="B_jDOtumPm4C" outputId="f5e2a43e-89a4-4f0e-b655-4f4344d52328"
z_mean, z_sd = vae.encode(train_data)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
im1 = ax1.scatter(z_mean[:, -1], z_mean[:, -2], s=1, c=mu)
ax1.set_xlabel("$z_2$", fontsize=14)
ax1.set_ylabel("$z_1$", fontsize=14)
cbar1 = fig.colorbar(im1, ax=ax1, shrink=.8)
cbar1.set_label("$\mu$", fontsize=14)
ax1.set_xlim(-2, 2)
ax1.set_ylim(-2, 2)
im2 = ax2.scatter(z_mean[:, -1], z_mean[:, -2], s=1, c=sig)
ax2.set_xlabel("$z_2$", fontsize=14)
ax2.set_ylabel("$z_1$", fontsize=14)
cbar2 = fig.colorbar(im2, ax=ax2, shrink=.8)
cbar2.set_label("$\sigma$", fontsize=14)
ax2.set_xlim(-2, 2)
ax2.set_ylim(-2, 2);

# + [markdown] id="CEVlZG2VuL5Y"
# Our data had only two factors of variation, peak position and peak width. By conditioning the model on the "known" property (shift) it was able to clearly identify the unknown one (width). Notice that [without the conditioning](https://colab.research.google.com/github/ziatdinovmax/pyroVED/blob/master/examples/shiftVAE.ipynb?authuser=1#scrollTo=B_jDOtumPm4C&line=8&uniqifier=1), the two factors of variation wouldn't be completely disentangled.

# + [markdown] id="cIZlk3dfPw0g"
# Plot the learned latent manifolds conditioned on the shift:

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="w7qw_GUgPyuW" outputId="ee8de85c-3fb5-4b67-8b71-c8b8fc8a89c2"
y_all = [-2, 0, 2]
for y in y_all:
    y = torch.tensor(y).unsqueeze(0)
    vae.manifold2d(10, y)

# + [markdown] id="YqpMXN84BLXQ"
# Plot the latent variables vs. the ground truth:

# + colab={"base_uri": "https://localhost:8080/", "height": 610} id="abAd16hrB0_i" outputId="1b3d062f-f4dc-4687-d72d-90cf26742c32"
_, ax = plt.subplots(2, 2, figsize=(10, 10))
ax[0, 0].scatter(z_mean[:, 0], mu, c='blue')
ax[0, 0].set_xlabel("$\mu$", fontsize=14)
ax[0, 0].set_ylabel("Latent variable 1", fontsize=14)
ax[0, 1].scatter(z_mean[:, 1], mu, c='blue')
ax[0, 1].set_xlabel("$\mu$", fontsize=14)
ax[0, 1].set_ylabel("Latent variable 2", fontsize=14)
ax[1, 0].scatter(z_mean[:, 0], sig, c='blue')
ax[1, 0].set_xlabel("$\sigma$", fontsize=14)
ax[1, 0].set_ylabel("Latent variable 1", fontsize=14)
ax[1, 1].scatter(z_mean[:, 1], sig, c='blue')
ax[1, 1].set_xlabel("$\sigma$", fontsize=14)
ax[1, 1].set_ylabel("Latent variable 2", fontsize=14);

# + [markdown] id="pm8GgK5715Ek"
# Now let's condition on a 2-dimensional vector containing information on the peak shifts and widths.

# + id="sQjQyU7e0P76"
# Initialize train loader
train_loader = pv.utils.init_dataloader(
    train_data.unsqueeze(1), torch.cat([mu, sig], -1), batch_size=64)

# + id="OGONVp7Yx1mt" colab={"base_uri": "https://localhost:8080/"} outputId="9ed51f3d-c0a0-483b-82ad-3345fe60eed1"
in_dim = (100,)

# Initialize VAE (we specify c_dim=2 to add extra dimensions corresponding to c-vector to the decoder)
svae = pv.models.iVAE(in_dim, latent_dim=2, c_dim=2, seed=0)

# Initialize SVI trainer
trainer = pv.trainers.SVItrainer(svae)

# Train for n epochs:
for e in range(250):
    trainer.step(train_loader)
    trainer.print_statistics()

# + id="N-CEtHptyLki" colab={"base_uri": "https://localhost:8080/", "height": 396} outputId="7bfd62d9-a93d-4c87-b336-6de18b31affd"
z_mean, z_sd = svae.encode(train_data)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
im1 = ax1.scatter(z_mean[:, -1], z_mean[:, -2], s=1, c=mu)
ax1.set_xlabel("$z_2$", fontsize=14)
ax1.set_ylabel("$z_1$", fontsize=14)
cbar1 = fig.colorbar(im1, ax=ax1, shrink=.8)
cbar1.set_label("$\mu$", fontsize=14)
ax1.set_xlim(-2, 2)
ax1.set_ylim(-2, 2)
im2 = ax2.scatter(z_mean[:, -1], z_mean[:, -2], s=1, c=sig)
ax2.set_xlabel("$z_2$", fontsize=14)
ax2.set_ylabel("$z_1$", fontsize=14)
cbar2 = fig.colorbar(im2, ax=ax2, shrink=.8)
cbar2.set_label("$\sigma$", fontsize=14)
ax2.set_xlim(-2, 2)
ax2.set_ylim(-2, 2);

# + [markdown] id="iNZDLMzszrYN"
# As expected, the latent space is collapsed (there are no factors of variation left to discover).

# + id="jXKvLoHM2Mam"


StarcoderdataJupyterScriptsDedupFiltered