tl_vae_model.py

# -*- coding: utf-8 -*-
"""TL-VAE Model.ipynb

Automatically generated by Colaboratory.

Target-learning the Latent Space of a Variational Autoencoder model for the Inverse Design of Stable Perovskites.

This code was developed by E.T. Chenebuah
(05-April-2023)
"""

import numpy as np
from tensorflow.keras.optimizers import Adam
import keras
from keras import layers
from keras.layers import Input, Dense, Lambda, Conv2DTranspose, LeakyReLU
from keras.layers import Activation, Flatten, Reshape, BatchNormalization
from keras import backend as K


class Sampling(layers.Layer):
    """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim), 
                                                 mean=0, stddev=1)
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

# Shape of input samples i.e X_samples = (27587*60*88*3)

## Note that the concept behind the development of the invertible descriptor design for representing perovskite inputs, including relevant feature embedding 
# can be found at: https://github.com/chenebuah/high_dim_descriptor/blob/main/fourier_calc.ipynb

latent_dim = 256

#encoder
encoder_inputs = keras.Input(shape=(X_samples.shape[1], X_samples.shape[2], X_samples.shape[3]))
x = layers.Conv2D(32, (3, 3), activation=LeakyReLU(alpha=0.2), strides= 1, padding='same')(encoder_inputs)
x = BatchNormalization()(x)
x = layers.Conv2D(64, (3, 3), activation=LeakyReLU(alpha=0.2), strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = layers.Conv2D(64, (5, 5), activation=LeakyReLU(alpha=0.2), strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = layers.Flatten()(x)
x = layers.Dense(1024, activation="sigmoid")(x)
z_mean = layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)
z = Lambda(Sampling(), output_shape=(latent_dim,), name='z')([z_mean, z_log_var])

# classifier Ef
clf_latent_inputs = Activation("relu")(z)
x = Dense(128, activation="relu")(clf_latent_inputs)
x = Dense(64, activation="relu")(x)
x = Dense(32, activation="relu")(x)
clf_outputs = Dense(1, activation='sigmoid', name='class_output')(x)

# Regressor Ef
reg_latent_inputs = Activation("relu")(z)
x = Dense(256, activation="relu")(reg_latent_inputs)
x = Dense(128, activation="relu")(x)
x = Dense(64, activation="relu")(x)
x = Dense(32, activation="relu")(x)
reg_outputs = Dense(1, activation='linear', name='reg_output')(x)

# Decoder
latent_inputs = keras.Input(shape=(latent_dim,), name='z_sampling')
x = layers.Dense(1024, activation="sigmoid")(latent_inputs)
x = layers.Dense((21120), activation=LeakyReLU(alpha=0.2))(x)
x = BatchNormalization()(x)
x = layers.Reshape((15, 22, 64))(x)
x = layers.Conv2DTranspose(64, (5, 5), activation=LeakyReLU(alpha=0.2), strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = layers.Conv2DTranspose(64, (3, 3), activation=LeakyReLU(alpha=0.2), strides=2, padding='same')(x)
x = BatchNormalization()(x)
decoder_outputs = layers.Conv2DTranspose(3, (3, 3), activation='sigmoid', padding='same')(x)

# Models
encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
clf_supervised = keras.Model(clf_latent_inputs, clf_outputs, name='clf')
reg_supervised = keras.Model(reg_latent_inputs, reg_outputs, name='reg')
decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

outputs = [
    decoder(encoder(encoder_inputs)[2]),
    clf_supervised(encoder(encoder_inputs)[2]),
    reg_supervised(encoder(encoder_inputs)[2])
           ]

vae = keras.Model(encoder_inputs, outputs, name='vae_mlp')

reconstruction_loss = (
    (tf.reduce_mean(tf.reduce_sum(K.square(encoder_inputs - outputs[0]),
                                 axis=[1,2])
                                 ))*60*88*3
                      )

kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))

vae_loss = reconstruction_loss + kl_loss
vae.add_loss(vae_loss)

vae.compile(
    optimizer=Adam(learning_rate=1e-4, decay=1e-3 / 200), 
    loss={'clf': 'binary_crossentropy', 'reg': 'mean_squared_error'},
    metrics={'clf': 'accuracy', 'reg': 'mae'}, 
    loss_weights={'clf': 10, 'reg':10}
            )

# Train model
vae.fit(X_samples, {'clf': y_Ef_class, 'reg': y_Ef},  epochs=150)

# y_Ef_class are formation energy binary labels in training set.
# y_Ef are formation energy continuous values in training set.