scripts/optimize_te_mrl.py

import os
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
MODEL_ROOT = PROJECT_ROOT / "model"
MODULE_ROOT = MODEL_ROOT / "src" / "mrl_te_optimization"
for import_root in (MODEL_ROOT, MODULE_ROOT):
    if str(import_root) not in sys.path:
        sys.path.insert(0, str(import_root))

os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")

from tqdm import tqdm
import random
random.seed(1337)
import matplotlib.pyplot as plt
import argparse
import numpy as np
np.random.seed(1337)
import pandas as pd
import torch
from framepool import *
from util import *

import random
random.seed(1337)
import scipy.stats as stats

import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.models import load_model

tf.compat.v1.enable_eager_execution()

import pandas as pd
import numpy as np
import requests

parser = argparse.ArgumentParser()
parser.add_argument('-d', type=str, required=False,
                    default=str(PROJECT_ROOT / 'conf' / 'data' / 'utrdb2.csv'))
parser.add_argument('-bs', type=int, required=False ,default=64)
parser.add_argument('-lr', type=int, required=False ,default=1)
parser.add_argument('-task', type=str, required=False ,default="mrl")
parser.add_argument('-gpu', type=str, required=False ,default='-1')
parser.add_argument('-s', type=int, required=False ,default=10000)
parser.add_argument('--output-dir', type=str,
                    default=str(PROJECT_ROOT / 'outputs' / 'optimization'))
args = parser.parse_args()

if args.gpu == '-1':
    device = 'cpu'
else:
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    device = 'cuda'

def prepare_mttrans(seqs):
    seqs_init = torch.tensor(np.array(one_hot_all_motif(seqs),dtype=np.float32))

    seqs_init = torch.transpose(seqs_init, 1, 2)
    seqs_init = torch.tensor(seqs_init,dtype=torch.float32).to(device)
    return seqs_init

def prepare_framepool(seqs):
    return tf.convert_to_tensor(np.array([encode_seq_framepool(seq) for seq in seqs]),dtype=tf.float32)



BATCH_SIZE = args.bs
motifs_path = str(PROJECT_ROOT / 'conf' / 'data' / 'motifs.csv')
STEPS = args.s
LR = args.lr
DIM = 40
SEQ_LEN = 128
UTR_LEN = 128

TASK = args.task

gpath = str(PROJECT_ROOT / 'weight' / 'checkpoint_3000.h5')


if TASK == 'te':
    path = str(PROJECT_ROOT / 'weight' / 'mttrans' / 'RL_hard_share_MTL' /
               '3R' / 'schedule_MTL-model_best_cv1.pth')
    OPT = 'TE'
else:
    path = str(PROJECT_ROOT / 'weight' / 'utr_model_combined_residual_new.h5')
    OPT = 'FMRL'


# Check for GPU availability
gpus = tf.config.list_physical_devices('GPU')

if gpus:
  print(f"GPU is available. Using GPU:{args.gpu} for computation.")
  print("List of GPUs:", gpus)
else:
  print("GPU is not available. Using CPU instead.")

out_folder = str(Path(args.output_dir).expanduser().resolve())
os.makedirs(out_folder, exist_ok=True)



def select_best(scores, seqs):
    selected_scores = []
    selected_seqs = []
    for i in range(len(scores[0])):
        best = scores[0][i]
        best_seq = seqs[0][i]
        for j in range(len(scores)-1):
            if scores[j+1][i] > best:
                best = scores[j+1][i]
                best_seq = seqs[j+1][i]
        selected_scores.append(best)
        selected_seqs.append(best_seq)

    return selected_seqs, selected_scores

if __name__ == '__main__':
    
    if OPT == 'FMRL':
        Optimize_FrameSlice = True
    else:
        Optimize_FrameSlice = False



    if Optimize_FrameSlice:
        model = load_framepool(path)

    else:

        model = torch.load(path,map_location=torch.device(device))['state_dict']  
        model.train()   
   

    wgan = tf.keras.models.load_model(gpath)

    """
    Data:
    """

    tf.random.set_seed(33)
    np.random.seed(33)

    diffs = []
    init_exps = []
    opt_exps = []
    orig_vals = []

    DIM = 40
    MAX_LEN = 128
    LR = np.exp(-LR)

    tempnoise = tf.random.normal(shape=[BATCH_SIZE,DIM])
    selectednoise = tempnoise

    best = 10

    LOW_START = False


    if LOW_START:
    
        for i in range(10000):
            tempnoise = tf.random.normal(shape=[BATCH_SIZE,DIM])
            sequences = wgan(tempnoise)

            seqs_gen = recover_seq(sequences, rev_rna_vocab)
            seqs_str = seqs_gen

            shape_ = tf.shape(np.array([encode_seq_framepool(seq) for seq in recover_seq(sequences, rev_rna_vocab)]))

            seqs = tf.convert_to_tensor(np.array([encode_seq_framepool(seq) for seq in recover_seq(sequences, rev_rna_vocab)]),dtype=tf.float32)

            
            pred =  model(seqs)

            t = tf.reshape(pred,(-1))
            t = t.numpy().astype('float')
            score = np.mean(t)

            if score < best:
                best = score
                selectednoise = tempnoise
        noise = tf.Variable(selectednoise)
    else:
        noise = tf.Variable(tf.random.normal(shape=[BATCH_SIZE,DIM]))
    

    noise_small = tf.random.normal(shape=[BATCH_SIZE,DIM],stddev=1e-4)

    optimizer = tf.keras.optimizers.Adam(learning_rate=np.power(np.e,LR))

    '''
    Optimization takes place here.
    '''

    bind_scores_list = []
    bind_scores_means = []
    sequences_list = []

    means = []
    maxes = []
    iters_ = []

    OPTIMIZE = True

    DNA_SEL = False


    sequences_init = wgan(noise)

    gen_seqs_init = sequences_init.numpy().astype('float')

    seqs_gen_init = recover_seq(gen_seqs_init, rev_rna_vocab)

    init_pos, init_neg = motif_count(seqs_gen_init,motifs_path)
    
    if Optimize_FrameSlice:
        seqs = prepare_framepool(seqs_gen_init)

        seqs_init = prepare_mttrans(seqs_gen_init)

        pred_init = model(seqs)
        
    else:


        one_hots = one_hot_all_motif(np.array(seqs_gen_init))
        seqs = torch.tensor(one_hots,dtype=torch.double)
        seqs = torch.transpose(seqs, 1, 2)
        seqs = seqs.float().to(device)


        pred_init = model.forward(seqs)
    
    if Optimize_FrameSlice:

        t = tf.reshape(pred_init,(-1))

        init_t = t.numpy().astype('float')
        
    else:
        
        t = torch.flatten(pred_init)
        t.float()
        
        init_t = t.cpu().detach().numpy()

    init_exp = np.mean(init_t)

    max_init = np.max(init_t)

    min_init = np.min(init_t)
    
    predicted_mrls = []

    STEPS = STEPS

    seqs_collection = []
    scores_collection = []
    if OPTIMIZE:
        iter_ = 0
        for opt_iter in tqdm(range(int(STEPS))):
            
            with tf.GradientTape() as gtape:
                gtape.watch(noise)
                sequences = wgan(noise)

                seqs_gen = recover_seq(sequences, rev_rna_vocab)
                seqs_collection.append(seqs_gen)
                seqs_str = seqs_gen
                
                if Optimize_FrameSlice:

                    seqs = tf.convert_to_tensor(np.array([encode_seq_framepool(seq) for seq in recover_seq(sequences, rev_rna_vocab)]),dtype=tf.float32)
                
                else:
                    seqs = torch.tensor(np.array(one_hot_all_motif(seqs_gen),dtype=np.float32))    

                if Optimize_FrameSlice:

                    with tf.GradientTape() as ptape:
                        ptape.watch(seqs)

                        pred =  model(seqs)
                        score = tf.reduce_mean(pred)
                        t = tf.reshape(pred,(-1))
                        mx = t.numpy().astype('float')
                        scores_collection.append(mx)
                        mx = np.max(mx)
                        
                        sum_ = tf.reduce_sum(t).numpy().astype('float')
                        
                        maxes.append(mx)
                        predicted_mrls.append(sum_/BATCH_SIZE)
                        means.append(sum_/BATCH_SIZE)

                    g1 = ptape.gradient(score,seqs)

                    OPTIMIZE_FULL = False
                    if OPTIMIZE_FULL:
                        tmp_g = g1.numpy().astype('float')
                        tmp_seqs = seqs_gen
                        tmp_lst = np.zeros(shape=(BATCH_SIZE,MAX_LEN,5))
                        for i in range(len(tmp_seqs)):
                            
                            len_ = len(tmp_seqs[i])
                            edited_g = tmp_g[i][:len_,:]
                            edited_g = np.pad(edited_g,((0,MAX_LEN-len_),(0,1)),'constant')   
                            tmp_lst[i] = edited_g   
                        
                        g1 = tf.convert_to_tensor(tmp_lst,dtype=tf.float32)

                    else:
                        
                        g1 = tf.pad(g1,tf.constant([[0, 0], [0, 0], [0, 1]]),"CONSTANT")

                    g1 = tf.math.scalar_mul(-1.0,g1)

                
                else:
                    
                    seqs = torch.transpose(seqs, 1, 2)
                    seqs = seqs.float()
                    seqs = torch.tensor(seqs.to(device), requires_grad=True)
                    pred = model(seqs)
                    pred = torch.flatten(pred)
                    predicted_mrls.append(np.average(pred.cpu().detach().numpy()))
                    scores_collection.append(pred.cpu().detach().numpy())
                    score = torch.mean(pred)
                    t = torch.flatten(pred)
                    mx = t.cpu().detach().numpy()
                    mx = np.max(mx)
                    
                    sum_ = torch.mean(t).cpu().detach().numpy()
                    
                    maxes.append(mx)
                    means.append(sum_/BATCH_SIZE)
                    pred.backward(torch.ones_like(pred))
                    
                    g1 = seqs.grad
                    
                    g1 = g1.cpu().detach().numpy()
                    g1 = tf.convert_to_tensor(g1)
                    g1 = tf.transpose(g1, perm=[0,2,1])
                    g1 = tf.pad(g1,tf.constant([[0, 0], [0, 0], [0, 1]]),"CONSTANT")
                    g1 = tf.math.scalar_mul(-1.0,g1)
                
                
                g2 = gtape.gradient(sequences,noise,output_gradients=g1)

            a1 = g2 + noise_small
            change = [(a1,noise)]
            optimizer.apply_gradients(change)

            iters_.append(iter_)
            iter_ += 1

        best_seqs, best_scores = select_best(scores_collection, seqs_collection)

        sequences_opt = wgan(noise)
        
        gen_seqs_opt = sequences_opt.numpy().astype('float')

        seqs_gen_opt = recover_seq(gen_seqs_opt, rev_rna_vocab)

        opt_pos, opt_neg = motif_count(seqs_gen_opt,motifs_path)
        
        if Optimize_FrameSlice:
            
            seqs_opt = prepare_framepool(seqs_gen_opt)


        
        else: 

            one_hots = np.array(one_hot_all_motif(seqs_gen_opt))
            # print(np.shape(one_hots))
            seqs = torch.tensor(one_hots,dtype=torch.double)
            seqs = torch.transpose(seqs, 1, 2)
            seqs = seqs.float().to(device)

        pred_opt = model(seqs)
        
        if Optimize_FrameSlice:

            t = tf.reshape(pred_opt,(-1))
            
            opt_t = t.numpy().astype('float')
            
        else:
            
            t = torch.flatten(pred_opt)
        
        
            opt_t = t.cpu().detach().numpy()

        opt_exp = np.mean(opt_t)

        min_opt = np.min(opt_t)
        max_opt = np.max(opt_t)

        with open(os.path.join(out_folder, f'init_mrl_{OPT}.txt'), 'w') as f:
            f.writelines([str(x)+'\n' for x in init_t])

        with open(os.path.join(out_folder, f'opt_mrl_{OPT}.txt'), 'w') as f:
            f.writelines([str(x)+'\n' for x in best_scores])

        with open(os.path.join(out_folder, f'opt_seqs_{OPT}.txt'), 'w') as f:
            f.writelines([str(x)+'\n' for x in best_seqs])

        with open(os.path.join(out_folder, f'init_seqs_{OPT}.txt'), 'w') as f:
            f.writelines([str(x)+'\n' for x in seqs_gen_init])
    

        print(f"Average Initial Pred: {np.average(init_t)}")
        print(f"Max Initial Pred: {np.max(init_t)}")
        print(f"Average Opt. Pred: {np.average(best_scores)}")
        print(f"Max Opt. Pred: {np.max(best_scores)}")