VAEtrain.py

"""
Author: Navid Shervani-Tabar
"""
import torch
import pickle
import numpy as np
import torch.optim as optim

from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader

from VAEmodel import VAEmod
from utils import tools, MolecularGraphDataset as MGD

class VAEgraph(object):
    def __init__(self, args):

        # -- training parameters
        self.device = args.device
        self.epochs = args.epochs
        self.batch_size = args.batch_size
        self.dataset_name = args.data_dir
        self.res_dir = args.res_dir
        self.N     = args.N
        self.N_vis = args.N_vis
        self.log_interval = args.log_interval
        self.n_samples = args.n_samples
        self.vis = args.vis
        self.mu_reg_1 = args.mu_reg_1
        self.mu_reg_2 = args.mu_reg_2
        self.mu_reg_3 = args.mu_reg_3
        self.mu_reg_4 = args.mu_reg_4
        self.reg_flag = [bool(i != 0) for i in args.reg_vec]
        self.L = args.L

        self.train_hist = {}
        for file in ['Tl', 'KL', 'RC', 'R1', 'R2', 'R3', 'R4']:
            self.train_hist[file] = []

        self.y_id = args.y_id
        self.y_target = args.y_target

        # -- model loading parameters
        self.filemodel = args.loadtrainedmodel
        self.loadmodel = bool(self.filemodel)

        # -- graph parameters
        self.n_atom_features = args.n_atom_type
        self.n_bond_features = args.n_bond_type
        self.n_node = args.n_node

        # -- set seed
        if not args.seed == 0:
            torch.manual_seed(args.seed)
            if bool(args.gpu_mode):
                torch.cuda.manual_seed(args.seed)

        # -- network setting
        self.z_dim = args.z_dim
        self.TrainDataset = DataLoader(dataset=MGD(self.dataset_name, self.N, 0), batch_size=self.batch_size, shuffle=False)
        self.VisulDataset = DataLoader(dataset=MGD(self.dataset_name, self.N_vis, self.N), batch_size=self.N_vis, shuffle=False)

        props = self.TrainDataset.dataset[:]['properties']

        self.mean = self.TrainDataset.dataset.mean
        self.var = self.TrainDataset.dataset.var

        self.mu_prior = torch.mean(props, 0)
        self.cov_prior = torch.tensor(np.cov(props.T)).float()

        self.cond_dsgn = bool(self.y_target)

        self.model = VAEmod(args).to(self.device)
        self.optimizer = optim.Adam(self.model.parameters(), lr=1e-3)

        self.tools = tools(args)

    def constraints(self, reg_sig, reg_adj, batch_dim):
        """
            physical constraints.
        :param reg_sig: signal generated by sampling from latent space.
        :param reg_adj: adjacency generated by sampling from latent space.
        :param batch_dim: batch size.
        :return: regularization terms.
        """

        SM_f = nn.Softmax(dim=2)
        SM_W = nn.Softmax(dim=3)

        p_f = SM_f(reg_sig)
        p_W = SM_W(reg_adj)

        ReLU = nn.ReLU()

        # -- Constraint: Ghost Nodes and Valence
        if self.reg_flag[0]:
            h_vec = torch.arange(self.n_bond_features, device=self.device).float()
            inner_sum = torch.einsum('i,bjki->bjk', h_vec, p_W)
            V = (inner_sum - torch.diag_embed(torch.einsum('...ii->...i', inner_sum))).sum(2)

            valence_dict = torch.Tensor([4, 2, 3, 1, 0]).to(self.device)
            U = torch.einsum('k,bjk->bj', valence_dict, p_f)

            reg_1 = torch.sqrt(torch.mean(torch.sum(ReLU(V - U), 1).reshape(self.L, -1)**2, 0))
        else:
            reg_1 = torch.zeros(batch_dim//self.L, device=self.device)

        # -- Constraint: Connectivity
        if self.reg_flag[1]:
            Sig = nn.Sigmoid()

            q = 1 - p_f[:, :, 4]
            A = 1 - p_W[:, :, :, 0]

            A = A - torch.diag_embed(torch.einsum('...ii->...i', A))
            A_0 = torch.eye(self.n_node).unsqueeze(0)
            A_i = A.clone()

            B = A_0.repeat(reg_sig.size(0), 1, 1).to(self.device)

            for i in range(1, self.n_node):
                A_i = Sig(100 * (torch.bmm(A_i, A) - 0.5))
                B += A_i

            C = Sig(100 * (B - 0.5))

            reg_2 = torch.sqrt(torch.mean(torch.sum(torch.triu(torch.einsum('ij,ik,ijk->ijk', q, q, 1 - 2 * C) + C,
                                                               diagonal=1), (1, 2)).reshape(self.L, -1)**2, 0))
        else:
            reg_2 = torch.zeros(batch_dim//self.L, device=self.device)

        # -- Constraint: 3-member cycle
        if self.reg_flag[2]:
            A = 1 - p_W[:, :, :, 0]

            A = A - torch.diag_embed(torch.einsum('...ii->...i', A))
            A_i = A.clone()
            for i in range(2):
                A_i = torch.bmm(A, A_i)

            reg_3 = torch.sqrt(torch.mean((torch.einsum('bii->b', A_i) / 6.).reshape(self.L, -1)**2, 0))
        else:
            reg_3 = torch.zeros(batch_dim//self.L, device=self.device)

        # -- Constraint: Cycles with triple bonds
        if self.reg_flag[3]:
            A = 1 - p_W[:, :, :, 0]
            D = p_W[:, :, :, 3]

            C = torch.empty(batch_dim, self.n_node, self.n_node, device=self.device)
            nI = self.n_node * torch.eye(self.n_node, device=self.device).unsqueeze(0).repeat(batch_dim, 1, 1)

            for i in range(self.n_node):
                for j in range(i, self.n_node):
                    B = A.clone()
                    B[:, i, j] = B[:, j, i] = 0
                    C[:, i, j] = C[:, j, i] = torch.inverse(nI - B)[:, i, j]

            reg_4 = torch.sqrt(torch.mean(torch.sum(torch.triu(torch.einsum('bij,bij->bij', D, C), diagonal=1),
                                                    (1, 2)).reshape(self.L, -1)**2, 0))
        else:
            reg_4 = torch.zeros(batch_dim//self.L, device=self.device)

        return [self.L * reg_1, self.L * reg_2, self.L * reg_3, self.L * reg_4]

    def loss_function(self, recon_sig, recon_adj, weight_vec, signal, adj, reg_sig, reg_adj, mu, logvar):

        batch_dim = recon_sig.shape[0]

        target_sig = signal.reshape(-1, self.n_node, self.n_atom_features).view(-1, self.n_atom_features).argmax(1)
        target_adj = adj.view(-1)

        output_adj = recon_adj.reshape(-1, self.n_bond_features)
        output_sig = recon_sig.reshape(-1, self.n_atom_features)

        atm_class_weights = None
        loss_sig = torch.nn.CrossEntropyLoss(weight=atm_class_weights, reduction='none')

        bnd_class_weights = None
        loss_adj = torch.nn.CrossEntropyLoss(weight=bnd_class_weights, reduction='none')

        fcn_loss_1 = loss_sig(output_sig, target_sig.long()).view(-1, self.n_node)
        fcn_loss_2 = loss_adj(output_adj, target_adj.long()).view(-1, self.n_node, self.n_node)

        fcn_loss = torch.sum(fcn_loss_1, 1) + torch.sum(torch.triu(fcn_loss_2, diagonal=1), (1, 2))
        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), 1)

        [reg_1, reg_2, reg_3, reg_4] = self.constraints(reg_sig, reg_adj, batch_dim)

        loss = self.N / batch_dim * (torch.dot(fcn_loss, weight_vec) + torch.dot(KLD, weight_vec) +
                                     self.mu_reg_1 * torch.dot(reg_1, weight_vec[:batch_dim//self.L]) +
                                     self.mu_reg_2 * torch.dot(reg_2, weight_vec[:batch_dim//self.L]) +
                                     self.mu_reg_3 * torch.dot(reg_3, weight_vec[:batch_dim//self.L]) +
                                     self.mu_reg_4 * torch.dot(reg_4, weight_vec[:batch_dim//self.L]))

        self.train_hist['Tl'].append(loss)
        self.train_hist['RC'].append(self.N / batch_dim * torch.dot(fcn_loss, weight_vec))
        self.train_hist['KL'].append(self.N / batch_dim * torch.dot(KLD, weight_vec))
        self.train_hist['R1'].append(self.N / batch_dim * self.mu_reg_1 * torch.dot(reg_1, weight_vec[:batch_dim//self.L]))
        self.train_hist['R2'].append(self.N / batch_dim * self.mu_reg_2 * torch.dot(reg_2, weight_vec[:batch_dim//self.L]))
        self.train_hist['R3'].append(self.N / batch_dim * self.mu_reg_3 * torch.dot(reg_3, weight_vec[:batch_dim//self.L]))
        self.train_hist['R4'].append(self.N / batch_dim * self.mu_reg_4 * torch.dot(reg_4, weight_vec[:batch_dim//self.L]))

        return loss

    def trainepoch(self, epoch):

        self.model.train()
        train_loss = 0
        for batch_idx, (train_batch, weight_vec) in enumerate(zip(self.TrainDataset, self.weights_loader)):
            train_batch['signal'] = train_batch['signal'].to(self.device)
            train_batch['adjacency'] = train_batch['adjacency'].to(self.device)

            sig_aug = train_batch['signal'].clone()
            adj_aug = train_batch['adjacency'].clone()
            w_aug = weight_vec.clone()
            for l in range(self.L - 1):
                sig_aug = torch.cat((sig_aug, train_batch['signal']), 0)
                adj_aug = torch.cat((adj_aug, train_batch['adjacency']), 0)
                w_aug = torch.cat((w_aug, weight_vec), 0)

            train_batch['signal'] = Variable(sig_aug)
            train_batch['adjacency'] = Variable(adj_aug)
            weight_vec = Variable(w_aug)

            props = None
            if self.cond_dsgn:
                props = train_batch['properties'].to(self.device)

                prp_aug = props.clone()
                for l in range(self.L - 1):
                    prp_aug = torch.cat((prp_aug, props), 0)

                props = Variable(prp_aug)

            self.optimizer.zero_grad()

            [rec_sig, rec_adj], mu, logvar, [reg_sig, reg_adj] = self.model(train_batch['signal'], train_batch['adjacency'], props)
            loss = self.loss_function(rec_sig, rec_adj, weight_vec, train_batch['signal'], train_batch['adjacency'], reg_sig, reg_adj, mu, logvar)
            loss.backward()
            train_loss += loss.item()

            self.optimizer.step()

        print('Train Epoch: {}\tLoss: {:.6f}'.format(epoch, train_loss/len(self.TrainDataset.dataset)))

    def train(self, weights, model_name = '/model.pth'):

        if not self.loadmodel:

            self.weights_loader = DataLoader(weights.to(self.device), batch_size=self.batch_size, shuffle=False)
            for epoch in range(1, self.epochs + 1):
                self.trainepoch(epoch)

                if self.vis and epoch % self.log_interval == 0:
                    self.tools.visLatent(self.VisulDataset, self.model, epoch)

                if self.vis and epoch == self.epochs:
                    self.tools.pltLoss(self.train_hist, epoch)

            torch.save(self.model, self.res_dir + model_name)

        else:
            self.model = torch.load(self.filemodel + model_name)

        if self.vis:
            self.tools.visLatent(self.VisulDataset, self.model, self.epochs, TrainData=self.TrainDataset)
        self.model.eval()

    def get_samples(self, sample_name='/samples.data'):
        self.model.eval()

        if self.cond_dsgn:
            y_target = (self.y_target - self.mean[self.y_id]) / np.sqrt(self.var[self.y_id])

        with torch.no_grad():
            sample_z = torch.randn((self.n_samples, self.z_dim), device=self.device)

            if self.cond_dsgn:

                # -- conditional generation
                id2 = [self.y_id]
                id1 = np.setdiff1d([0, 1, 2], id2)

                mu1 = self.mu_prior[id1]
                mu2 = self.mu_prior[id2]

                cov11 = self.cov_prior[id1][:, id1]
                cov12 = self.cov_prior[id1][:, id2]
                cov22 = self.cov_prior[id2][:, id2]
                cov21 = self.cov_prior[id2][:, id1]

                cond_mu = np.transpose(mu1.T + np.matmul(cov12, np.linalg.inv(cov22)) * (y_target - mu2))[0]
                cond_cov = cov11 - np.matmul(np.matmul(cov12, np.linalg.inv(cov22)), cov21)

                sample_y = torch.empty(self.n_samples, 3, device=self.device)
                sample_y[:, id1] = torch.distributions.multivariate_normal.MultivariateNormal(cond_mu, cond_cov).sample(
                    (self.n_samples,)).to(self.device)
                sample_y[:, id2] = y_target

        if self.cond_dsgn:
            samplesTorch = self.model.decode(torch.cat((sample_z, sample_y), dim=1))
        else:
            samplesTorch = self.model.decode(sample_z)

        with torch.no_grad():
            samples_sig = torch.argmax(samplesTorch[0], dim=2)
            samples_adj = torch.argmax(samplesTorch[1], dim=3)
            samples_adj = samples_adj - torch.diag_embed(torch.einsum('...ii->...i', samples_adj))

        # -- store the samples
        with open(self.res_dir + sample_name, 'wb') as f:
            pickle.dump(samples_sig, f)
            pickle.dump(samples_adj, f)
            pickle.dump(sample_z, f)

        return samples_sig, samples_adj, sample_z