text: #!/usr/bin/env python
  # coding: utf-8
  
  # # Train a VAE on Cell Painting LINCS Data
  
  # In[1]:
  
  
  import sys
  import pathlib
  import numpy as np
  import pandas as pd
  sys.path.insert(0, "../../scripts")
  from utils import load_data
  
  
  from pycytominer.cyto_utils import infer_cp_features
  
  
  import matplotlib.pyplot as plt
  from matplotlib.pyplot import figure
  
  from sklearn.decomposition import PCA
  from tensorflow import keras
  
  from vae import VAE
  
  from tensorflow.keras.models import Model, Sequential
  import seaborn
  import random as python_random
  import tensorflow as tf
  import umap
  
  
  # In[2]:
  
  
  # %load_ext nb_black
  
  
  # In[3]:
  
  
  # np.random.seed(123)
  # python_random.seed(123)
  # tf.random.set_random_seed(1234)
  
  
  # In[4]:
  
  
  data_splits = ["train", "test", "valid", "complete"]
  data_dict = load_data(data_splits)
  
  
  # In[5]:
  
  
  # Prepare data for training
  meta_features = infer_cp_features(data_dict["train"], metadata=True)
  cp_features = infer_cp_features(data_dict["train"])
  
  train_features_df = data_dict["train"].reindex(cp_features, axis="columns")
  train_meta_df = data_dict["train"].reindex(meta_features, axis="columns")
  
  test_features_df = data_dict["test"].reindex(cp_features, axis="columns")
  test_meta_df = data_dict["test"].reindex(meta_features, axis="columns")
  
  valid_features_df = data_dict["valid"].reindex(cp_features, axis="columns")
  valid_meta_df = data_dict["valid"].reindex(meta_features, axis="columns")
  
  complete_features_df = data_dict["complete"].reindex(cp_features, axis="columns")
  complete_meta_df = data_dict["complete"].reindex(meta_features, axis="columns")
  
  
  # In[6]:
  
  
  print(train_features_df.shape)
  train_features_df.head(3)
  
  
  # In[7]:
  
  
  print(test_features_df.shape)
  test_features_df.head(3)
  
  
  # In[8]:
  
  
  print(complete_features_df.shape)
  complete_features_df.head(3)
  
  
  # In[171]:
  
  
  # VAE of one layer
  encoder_architecture = [250]
  decoder_architecture = [250]
  
  
  # In[366]:
  
  
  # model from optimal hyperparameters for onelayer obtained from 1.optimize
  cp_vae = VAE(
      input_dim=train_features_df.shape[1],
      latent_dim=10,
      batch_size=96,
      encoder_batch_norm=True,
      epochs=30,
      learning_rate=0.001,
      encoder_architecture=encoder_architecture,
      decoder_architecture=decoder_architecture,
      beta=0,
      lam=1000,
      verbose=True,
  )
  #100000
  #10000
  #200 ep 1000000 results in 1.20
  cp_vae.compile_vae()
  
  
  # In[367]:
  
  
  cp_vae.train(x_train=train_features_df, x_test=valid_features_df)
  
  
  # In[368]:
  
  
  cp_vae.vae
  
  
  # In[369]:
  
  
  # cp_vae.vae.evaluate(test_features_df)
  
  
  # In[370]:
  
  
  # Save training performance
  history_df = pd.DataFrame(cp_vae.vae.history.history)
  history_df
  
  
  # In[397]:
  
  
  history_df.to_csv('training_data/level5_training_mmd.csv')
  
  
  # In[373]:
  
  
  plt.figure(figsize=(10, 5))
  plt.plot(history_df["loss"], label="Training data")
  plt.plot(history_df["val_loss"], label="Validation data")
  plt.title("Loss for VAE training on Cell Painting Level 5 data")
  plt.ylabel("MSE + MMD")
  plt.xlabel("No. Epoch")
  plt.legend()
  plt.show()
  
  
  # In[374]:
  
  
  reconstruction = pd.DataFrame(cp_vae.vae.predict(test_features_df), columns=cp_features)
  (sum(sum((np.array(test_features_df) - np.array(reconstruction)) ** 2))) ** 0.5
  
  
  # In[375]:
  
  
  #latent space heatmap
  fig, ax = plt.subplots(figsize=(10, 10))
  encoder = cp_vae.encoder_block["encoder"]
  latent = np.array(encoder.predict(test_features_df)[2])
  seaborn.heatmap(latent, ax=ax)
  
  
  # In[377]:
  
  
  #original vs reconstructed heatmap
  reconstruction = pd.DataFrame(cp_vae.vae.predict(test_features_df), columns=cp_features)
  pca = PCA(n_components=2).fit(test_features_df)
  pca_reconstructed_latent_df = pd.DataFrame(pca.transform(reconstruction))
  pca_test_latent_df = pd.DataFrame(pca.transform(test_features_df))
  # fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 10))
  # seaborn.heatmap(test_features_df, ax=ax1)
  # seaborn.heatmap(reconstruction, ax=ax2)
  # plt.show()
  
  
  # In[378]:
  
  
  figure(figsize=(10, 10), dpi=80)
  plt.scatter(pca_test_latent_df[0],pca_test_latent_df[1], marker = ".", alpha = 0.5)
  plt.scatter(pca_reconstructed_latent_df[0],pca_reconstructed_latent_df[1], marker = ".", alpha = 0.5)
  
  
  # In[379]:
  
  
  pca = PCA(n_components=50).fit(test_features_df)
  pca_reconstructed_latent_df = pd.DataFrame(pca.transform(reconstruction))
  pca_test_latent_df = pd.DataFrame(pca.transform(test_features_df))
  
  
  # In[380]:
  
  
  figure(figsize=(10, 10), dpi=80)
  plt.scatter(pca_test_latent_df[2],pca_test_latent_df[3], marker = ".", alpha = 0.5)
  plt.scatter(pca_reconstructed_latent_df[2],pca_reconstructed_latent_df[3], marker = ".", alpha = 0.5)
  
  
  # In[389]:
  
  
  #difference between original and reconstructed heatmap
  difference = abs(reconstruction - test_features_df)
  fig, ax = plt.subplots(figsize=(10, 10))
  seaborn.heatmap(difference, ax=ax, cmap="Blues")
  
  
  # In[390]:
  
  
  #encoder heatmap
  weights = cp_vae.encoder_block["encoder"].get_weights()
  fig, ax = plt.subplots(figsize=(10, 10))
  seaborn.heatmap(weights[0], ax=ax)
  
  
  # In[381]:
  
  
  decoder = cp_vae.decoder_block["decoder"]
  
  
  # In[382]:
  
  
  
  pca_training = PCA(n_components=2).fit(train_features_df)
  
  simulated_df = pd.DataFrame(np.random.normal(size=(8294, 10)), columns=np.arange(0,10))
  reconstruction_of_simulated = decoder.predict(simulated_df)
  pca_reconstruction_of_simulated = pd.DataFrame(pca_training.transform(reconstruction_of_simulated))
  pca_train_latent_df = pd.DataFrame(pca_training.transform(train_features_df))
  
  fig, (ax1,ax2) = plt.subplots(1, 2, figsize=(16,8), sharey = True, sharex = True)
  ax1.scatter(pca_train_latent_df[0],pca_train_latent_df[1], marker = ".", alpha = 0.5)
  ax2.scatter(pca_reconstruction_of_simulated[0],pca_reconstruction_of_simulated[1], marker = ".", alpha = 0.5)
  
  
  # In[383]:
  
  
  from scipy.spatial.distance import directed_hausdorff
  max(directed_hausdorff(train_features_df, reconstruction_of_simulated)[0],directed_hausdorff(reconstruction_of_simulated,train_features_df)[0])
  
  
  # In[384]:
  
  
  #NOTE: IF YOU RUN THIS, YOU WILL NOT BE ABLE TO REPRODUCE THE EXACT RESULTS IN THE EXPERIMENT
  latent_complete = np.array(encoder.predict(complete_features_df)[2])
  latent_df = pd.DataFrame(latent_complete)
  latent_df.to_csv("../3.application/level5Latent_mmd.csv")
  
  
  # In[3]:
  
  
  reducer = umap.UMAP()
  latent_embedding = reducer.fit_transform(latent_df)
  
  
  # In[4]:
  
  
  figure(figsize=(10, 10), dpi=80)
  plt.scatter(
      latent_embedding[:, 0],
      latent_embedding[:, 1],
      marker = '.'
  )
  
  
  # In[385]:
  
  
  #NOTE: IF YOU RUN THIS, YOU WILL NOT BE ABLE TO REPRODUCE THE EXACT RESULTS IN THE EXPERIMENT
  decoder.save("models/level5Decoder_mmd")
  
  
  # In[386]:
  
  
  encoder.save("models/level5Encoder_mmd")
  
  