import math
import torch
import pickle
import numpy as np
import pandas as pd
import torch.nn as nn
from rdkit import Chem
from rdkit import rdBase
from GeneVAE import GeneVAE
from utils import Tokenizer, get_device, mean_similarity
from SmilesVAE import Smiles_DataLoader, EncoderRNN, DecoderRNN, SmilesVAE
rdBase.DisableLog('rdApp.error')
# ============================================================================
# Load data
def load_smiles_data(tokenizer, args):
# Load smiles and gene values
smiles_loader = Smiles_DataLoader(
args.gene_expression_file,
args.cell_name,
tokenizer,
args.gene_num,
batch_size=args.gene_batch_size,
train_rate=args.train_rate,
variant=args.variant
)
train_dataloader, valid_dataloader = smiles_loader.get_dataloader()
return train_dataloader, valid_dataloader
# ============================================================================
# Generate Smiles using learned gene representations
def train_smiles_vae(
trained_gene_vae,
train_dataloader,
valid_dataloader,
tokenizer,
args
):
"""
trained_gene_vae: the pretrained GeneVAE model for gene feature extraction
train_dataloader: splited training data (encoded smiles, genes)
tokenizer: SMILES vocabulary
"""
# Define EncoderRNN
encoder = EncoderRNN(
args.emb_size,
args.hidden_size,
args.num_layers,
args.smiles_latent_size,
args.bidirectional,
tokenizer
).to(get_device())
# Define DecoderRNN
decoder = DecoderRNN(
args.emb_size,
args.hidden_size,
args.num_layers,
args.smiles_latent_size,
args.gene_latent_size, # condition_size = gene_latent_size
tokenizer
).to(get_device())
# Define SmilesVAE
smiles_vae = SmilesVAE(encoder, decoder).to(get_device())
# Optimizer
optimizer = torch.optim.Adam(
smiles_vae.parameters(),
lr=args.smiles_lr
)
# Prepare file to save results
with open(args.smiles_vae_train_results, 'a+') as wf:
wf.truncate(0)
wf.write('{},{},{},{},{},{},{},{},{},{},{},{}\n'.format(\
'Epoch',
'Joint_loss',
'Rec_loss',
'KLD_loss',
'Total',
'Valid',
'Valid_rate',
'Unique',
'Unique_rate',
'Novelty',
'Novel_rate',
'Diversity'
))
print('\n')
print('Training Information:')
for epoch in range(args.smiles_epochs):
total_joint_loss = 0
total_rec_loss = 0
total_kld_loss = 0
smiles_vae.train()
# Operate on a batch of data
for _, (smiles, genes) in enumerate(train_dataloader):
smiles, genes = smiles.to(get_device()), genes.to(get_device())
# Extract gene expression features
trained_gene_vae.eval()
gene_latent_vectors, _ = trained_gene_vae(genes) # [batch_size, gene_latent_size]
# Apply SmilesVAE (MolVAE)
z, decoded = smiles_vae(smiles, gene_latent_vectors, args.temperature)
alphas = torch.cat([
torch.linspace(0.99, 0.5, int(args.smiles_epochs/2)),
0.5 * torch.ones(args.smiles_epochs - int(args.smiles_epochs/2))
]).double().to(get_device())
joint_loss, rec_loss, kld_loss = smiles_vae.joint_loss(
decoded,
targets=smiles,
alpha=alphas[epoch],
beta=1.
)
optimizer.zero_grad()
joint_loss.backward()
optimizer.step()
total_joint_loss += joint_loss.item()
total_rec_loss += rec_loss.item()
total_kld_loss += kld_loss.item()
mean_joint_loss = total_joint_loss / smiles.size(0)
mean_rec_loss = total_rec_loss / (smiles.size(0))
mean_kld_loss = total_kld_loss / (smiles.size(0))
# Evaluate valid and unique SMILES
smiles_vae.eval()
valid_smiles = []
label_smiles = []
total_num_data = len(valid_dataloader.dataset)
for _, (smiles, genes) in enumerate(valid_dataloader):
smiles, genes = smiles.to(get_device()), genes.to(get_device())
trained_gene_vae.eval()
# Extracted Gx as condition
gene_latent_vectors, _ = trained_gene_vae(genes)
# Random values as smiles_z
rand_z = torch.randn(genes.size(0), args.smiles_latent_size).to(get_device())
dec_sampled_char = smiles_vae.generation(rand_z, gene_latent_vectors, args.max_len, tokenizer)
output_smiles = ["".join(tokenizer.decode(\
dec_sampled_char[i].squeeze().detach().cpu().numpy()
)).strip("^$ ") for i in range(dec_sampled_char.size(0))]
#valid_smiles.extend([smi for smi in output_smiles if Chem.MolFromSmiles(smi) and Chem.MolFromSmiles(smi)!=])
for i in range(len(output_smiles)):
mol = Chem.MolFromSmiles(output_smiles[i])
if mol != None and mol.GetNumAtoms() > 1 and Chem.MolToSmiles(mol) != ' ':
valid_smiles.extend([output_smiles[i]])
label_smiles.extend(["".join(tokenizer.decode(smiles[i].squeeze().detach().cpu().numpy())).strip("^$ ")])
unique_smiles = list(set(valid_smiles))
# Novel Smiles
novel_smiles = [smi for smi in unique_smiles if smi not in label_smiles]
# Save valid Smiles to file
valid_csv = pd.DataFrame(valid_smiles).to_csv(args.valid_smiles_file, index=False)
valid_num = len(valid_smiles)
valid_rate = 100*len(valid_smiles)/total_num_data
unique_num = len(unique_smiles)
novel_num = len(novel_smiles)
if valid_num != 0:
unique_rate = 100*unique_num/valid_num
diversity = 1 - mean_similarity(valid_smiles, label_smiles)
else:
unique_rate = 100*unique_num/(valid_num+1)
diversity = 1
if unique_num != 0:
novel_rate = 100*novel_num/unique_num
else:
novel_rate = 100*novel_num/(unique_num+1)
print('Epoch: {:d} / {:d}, joint_loss: {:.3f}, rec_loss: {:.3f}, kld_loss: {:.3f}, Total: {:d}, valid: {:d} ({:.2f}), unique: {:d} ({:.2f}), novel: {:d} ({:.2f}), diversity: {:.3f}'.format(\
epoch+1,
args.smiles_epochs,
mean_joint_loss,
mean_rec_loss,
mean_kld_loss,
total_num_data,
valid_num,
valid_rate,
unique_num,
unique_rate,
novel_num,
novel_rate,
diversity
))
# Save trained results to file
with open(args.smiles_vae_train_results, 'a+') as wf:
wf.write('{},{:.3f},{:.3f},{:.3f},{},{},{:.2f},{},{:.2f},{},{:.2f},{:.2f}\n'.format(\
epoch+1,
mean_joint_loss,
mean_rec_loss,
mean_kld_loss,
total_num_data,
valid_num,
valid_rate,
unique_num,
unique_rate,
novel_num,
novel_rate,
diversity
))
# Save predicted and label SMILES into file
final_smiles = {'predict': valid_smiles, 'label': label_smiles}
final_smiles = pd.DataFrame(final_smiles)
# Save to file
final_smiles.to_csv(args.valid_smiles_file, index=False)
print('='*50)
# Save the trained SmilesVAE
smiles_vae.save_model(args.saved_smiles_vae)
print('Trained SmilesVAE is saved in {}'.format(args.saved_smiles_vae))
return smiles_vae