{ "cells": [ { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "\n", "from tqdm import tqdm\n", "import random\n", "random.seed(1337)\n", "import matplotlib.pyplot as plt\n", "import argparse\n", "import numpy as np\n", "np.random.seed(1337)\n", "import pandas as pd\n", "import os\n", "import torch\n", "import sys\n", "\n", "from src.mrl_te_optimization.framepool import *\n", "from src.mrl_te_optimization.util import *\n", "import keras\n", "\n", "import random\n", "random.seed(1337)\n", "\n", "\n", "import scipy.stats as stats\n", "\n", "import tensorflow as tf\n", "from tensorflow.keras import backend as K\n", "from tensorflow.keras.models import load_model\n", "\n", "tf.compat.v1.enable_eager_execution()\n", "\n", "import pandas as pd\n", "import numpy as np\n", "import requests, sys\n", "\n", "DATA = './data/utrdb2.csv'\n", "motifs_path = './data/motifs.csv'\n", "BATCH_SIZE = 64\n", "TASK = \"mrl\"\n", "GPU = '-1'\n", "STEPS = 10\n", "\n", "if GPU == '-1':\n", " device = 'cpu'\n", "else:\n", " os.environ['CUDA_VISIBLE_DEVICES'] = GPU\n", " device = 'cuda'\n", " if ',' in GPU:\n", " device = 'cuda:1'\n", "\n", "def prepare_mttrans(seqs):\n", " seqs_init = torch.tensor(np.array(one_hot_all_motif(seqs),dtype=np.float32))\n", "\n", " seqs_init = torch.transpose(seqs_init, 1, 2)\n", " seqs_init = torch.tensor(seqs_init,dtype=torch.float32).to(device)\n", " return seqs_init\n", "\n", "def prepare_framepool(seqs):\n", " return tf.convert_to_tensor(np.array([encode_seq_framepool(seq) for seq in seqs]),dtype=tf.float32)\n", "\n", "\n", "\n", "DIM = 40\n", "SEQ_LEN = 128\n", "UTR_LEN = 128\n", "gpath = './models/checkpoint_3000.h5'\n", "\n", "\n", "if TASK == 'te':\n", " path = './src/mrl_te_optimization/script/checkpoint/RL_hard_share_MTL/3R/schedule_MTL-model_best_cv1.pth'\n", " OPT = 'TE'\n", "else:\n", " path = './models/utr_model_combined_residual_new.h5'\n", " OPT = 'FMRL'\n", "\n", "\n", "out_folder = './outputs/'\n", "os.makedirs(out_folder, exist_ok=True)" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "GPU is not available. Using CPU instead.\n" ] } ], "source": [ "# Check for GPU availability\n", "gpus = tf.config.list_physical_devices('GPU')\n", "\n", "if gpus:\n", " print(f\"GPU is available. Using GPU:{GPU} for computation.\")\n", " print(\"List of GPUs:\", gpus)\n", "else:\n", " print(\"GPU is not available. Using CPU instead.\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.\n" ] }, { "ename": "NameError", "evalue": "name 'LR' is not defined", "output_type": "error", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", "Cell \u001b[0;32mIn[7], line 49\u001b[0m\n\u001b[1;32m 47\u001b[0m DIM \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m40\u001b[39m\n\u001b[1;32m 48\u001b[0m MAX_LEN \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m128\u001b[39m\n\u001b[0;32m---> 49\u001b[0m LR \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mexp(\u001b[38;5;241m-\u001b[39m\u001b[43mLR\u001b[49m)\n\u001b[1;32m 51\u001b[0m tempnoise \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mnormal(shape\u001b[38;5;241m=\u001b[39m[BATCH_SIZE,DIM])\n\u001b[1;32m 52\u001b[0m selectednoise \u001b[38;5;241m=\u001b[39m tempnoise\n", "\u001b[0;31mNameError\u001b[0m: name 'LR' is not defined" ] } ], "source": [ "\n", "\n", "def select_best(scores, seqs):\n", " selected_scores = []\n", " selected_seqs = []\n", " for i in range(len(scores[0])):\n", " best = scores[0][i]\n", " best_seq = seqs[0][i]\n", " for j in range(len(scores)-1):\n", " if scores[j+1][i] > best:\n", " best = scores[j+1][i]\n", " best_seq = seqs[j+1][i]\n", " selected_scores.append(best)\n", " selected_seqs.append(best_seq)\n", "\n", " return selected_seqs, selected_scores\n", "\n", "if __name__ == '__main__':\n", " \n", " if OPT == 'FMRL':\n", " Optimize_FrameSlice = True\n", " else:\n", " Optimize_FrameSlice = False\n", "\n", "\n", "\n", " if Optimize_FrameSlice:\n", " model = load_framepool(path)\n", "\n", " else:\n", "\n", " model = torch.load(path,map_location=torch.device(device))['state_dict'] \n", " model.train() \n", " \n", "\n", " wgan = tf.keras.models.load_model(gpath)\n", "\n", " \"\"\"\n", " Data:\n", " \"\"\"\n", "\n", " tf.random.set_seed(33)\n", " np.random.seed(33)\n", "\n", " diffs = []\n", " init_exps = []\n", " opt_exps = []\n", " orig_vals = []\n", "\n", "\n", " LR = 0.001\n", " DIM = 40\n", " MAX_LEN = 128\n", " \n", "\n", " tempnoise = tf.random.normal(shape=[BATCH_SIZE,DIM])\n", " selectednoise = tempnoise\n", "\n", " best = 10\n", "\n", " LOW_START = False\n", "\n", "\n", " if LOW_START:\n", " \n", " for i in range(10000):\n", " tempnoise = tf.random.normal(shape=[BATCH_SIZE,DIM])\n", " sequences = wgan(tempnoise)\n", "\n", " seqs_gen = recover_seq(sequences, rev_rna_vocab)\n", " seqs_str = seqs_gen\n", "\n", " shape_ = tf.shape(np.array([encode_seq_framepool(seq) for seq in recover_seq(sequences, rev_rna_vocab)]))\n", "\n", " seqs = tf.convert_to_tensor(np.array([encode_seq_framepool(seq) for seq in recover_seq(sequences, rev_rna_vocab)]),dtype=tf.float32)\n", "\n", " \n", " pred = model(seqs)\n", "\n", " t = tf.reshape(pred,(-1))\n", " t = t.numpy().astype('float')\n", " score = np.mean(t)\n", "\n", " if score < best:\n", " best = score\n", " selectednoise = tempnoise\n", " noise = tf.Variable(selectednoise)\n", " else:\n", " noise = tf.Variable(tf.random.normal(shape=[BATCH_SIZE,DIM]))\n", " \n", "\n", " noise_small = tf.random.normal(shape=[BATCH_SIZE,DIM],stddev=1e-4)\n", "\n", " optimizer = tf.keras.optimizers.Adam(learning_rate=np.power(np.e,LR))\n", "\n", " '''\n", " Optimization takes place here.\n", " '''\n", "\n", " bind_scores_list = []\n", " bind_scores_means = []\n", " sequences_list = []\n", "\n", " means = []\n", " maxes = []\n", " iters_ = []\n", "\n", " OPTIMIZE = True\n", "\n", " DNA_SEL = False\n", "\n", "\n", " sequences_init = wgan(noise)\n", "\n", " gen_seqs_init = sequences_init.numpy().astype('float')\n", "\n", " seqs_gen_init = recover_seq(gen_seqs_init, rev_rna_vocab)\n", "\n", " init_pos, init_neg = motif_count(seqs_gen_init,motifs_path)\n", " \n", " if Optimize_FrameSlice:\n", " seqs = prepare_framepool(seqs_gen_init)\n", "\n", " seqs_init = prepare_mttrans(seqs_gen_init)\n", "\n", " pred_init = model(seqs)\n", " \n", " else:\n", "\n", "\n", " one_hots = one_hot_all_motif(np.array(seqs_gen_init))\n", " seqs = torch.tensor(one_hots,dtype=torch.double)\n", " seqs = torch.transpose(seqs, 1, 2)\n", " seqs = seqs.float().to(device)\n", "\n", "\n", " pred_init = model.forward(seqs)\n", " \n", " if Optimize_FrameSlice:\n", "\n", " t = tf.reshape(pred_init,(-1))\n", "\n", " init_t = t.numpy().astype('float')\n", " \n", " else:\n", " \n", " t = torch.flatten(pred_init)\n", " t.float()\n", " \n", " init_t = t.cpu().detach().numpy()\n", "\n", " init_exp = np.mean(init_t)\n", "\n", " max_init = np.max(init_t)\n", "\n", " min_init = np.min(init_t)\n", " \n", " predicted_mrls = []\n", "\n", " STEPS = STEPS\n", "\n", " seqs_collection = []\n", " scores_collection = []\n", " if OPTIMIZE:\n", " iter_ = 0\n", " for opt_iter in tqdm(range(int(STEPS))):\n", " \n", " with tf.GradientTape() as gtape:\n", " gtape.watch(noise)\n", " sequences = wgan(noise)\n", "\n", " seqs_gen = recover_seq(sequences, rev_rna_vocab)\n", " seqs_collection.append(seqs_gen)\n", " seqs_str = seqs_gen\n", " \n", " if Optimize_FrameSlice:\n", "\n", " seqs = tf.convert_to_tensor(np.array([encode_seq_framepool(seq) for seq in recover_seq(sequences, rev_rna_vocab)]),dtype=tf.float32)\n", " \n", " else:\n", " seqs = torch.tensor(np.array(one_hot_all_motif(seqs_gen),dtype=np.float32)) \n", "\n", " if Optimize_FrameSlice:\n", "\n", " with tf.GradientTape() as ptape:\n", " ptape.watch(seqs)\n", "\n", " pred = model(seqs)\n", " score = tf.reduce_mean(pred)\n", " t = tf.reshape(pred,(-1))\n", " mx = t.numpy().astype('float')\n", " scores_collection.append(mx)\n", " mx = np.max(mx)\n", " \n", " sum_ = tf.reduce_sum(t).numpy().astype('float')\n", " \n", " maxes.append(mx)\n", " predicted_mrls.append(sum_/BATCH_SIZE)\n", " means.append(sum_/BATCH_SIZE)\n", "\n", " g1 = ptape.gradient(score,seqs)\n", "\n", " OPTIMIZE_FULL = False\n", " if OPTIMIZE_FULL:\n", " tmp_g = g1.numpy().astype('float')\n", " tmp_seqs = seqs_gen\n", " tmp_lst = np.zeros(shape=(BATCH_SIZE,MAX_LEN,5))\n", " for i in range(len(tmp_seqs)):\n", " \n", " len_ = len(tmp_seqs[i])\n", " edited_g = tmp_g[i][:len_,:]\n", " edited_g = np.pad(edited_g,((0,MAX_LEN-len_),(0,1)),'constant') \n", " tmp_lst[i] = edited_g \n", " \n", " g1 = tf.convert_to_tensor(tmp_lst,dtype=tf.float32)\n", "\n", " else:\n", " \n", " g1 = tf.pad(g1,tf.constant([[0, 0], [0, 0], [0, 1]]),\"CONSTANT\")\n", "\n", " g1 = tf.math.scalar_mul(-1.0,g1)\n", "\n", " \n", " else:\n", " \n", " seqs = torch.transpose(seqs, 1, 2)\n", " seqs = seqs.float()\n", " seqs = torch.tensor(seqs.to(device), requires_grad=True)\n", " pred = model(seqs)\n", " pred = torch.flatten(pred)\n", " predicted_mrls.append(np.average(pred.cpu().detach().numpy()))\n", " scores_collection.append(pred.cpu().detach().numpy())\n", " score = torch.mean(pred)\n", " t = torch.flatten(pred)\n", " mx = t.cpu().detach().numpy()\n", " mx = np.max(mx)\n", " \n", " sum_ = torch.mean(t).cpu().detach().numpy()\n", " \n", " maxes.append(mx)\n", " means.append(sum_/BATCH_SIZE)\n", " pred.backward(torch.ones_like(pred))\n", " \n", " g1 = seqs.grad\n", " \n", " g1 = g1.cpu().detach().numpy()\n", " g1 = tf.convert_to_tensor(g1)\n", " g1 = tf.transpose(g1, perm=[0,2,1])\n", " g1 = tf.pad(g1,tf.constant([[0, 0], [0, 0], [0, 1]]),\"CONSTANT\")\n", " g1 = tf.math.scalar_mul(-1.0,g1)\n", " \n", " \n", " g2 = gtape.gradient(sequences,noise,output_gradients=g1)\n", "\n", " a1 = g2 + noise_small\n", " change = [(a1,noise)]\n", " optimizer.apply_gradients(change)\n", "\n", " iters_.append(iter_)\n", " iter_ += 1\n", "\n", " best_seqs, best_scores = select_best(scores_collection, seqs_collection)\n", "\n", " sequences_opt = wgan(noise)\n", " \n", " gen_seqs_opt = sequences_opt.numpy().astype('float')\n", "\n", " seqs_gen_opt = recover_seq(gen_seqs_opt, rev_rna_vocab)\n", "\n", " opt_pos, opt_neg = motif_count(seqs_gen_opt,motifs_path)\n", " \n", " if Optimize_FrameSlice:\n", " \n", " seqs_opt = prepare_framepool(seqs_gen_opt)\n", "\n", "\n", " \n", " else: \n", "\n", " one_hots = np.array(one_hot_all_motif(seqs_gen_opt))\n", " # print(np.shape(one_hots))\n", " seqs = torch.tensor(one_hots,dtype=torch.double)\n", " seqs = torch.transpose(seqs, 1, 2)\n", " seqs = seqs.float().to(device)\n", "\n", " pred_opt = model(seqs)\n", " \n", " if Optimize_FrameSlice:\n", "\n", " t = tf.reshape(pred_opt,(-1))\n", " \n", " opt_t = t.numpy().astype('float')\n", " \n", " else:\n", " \n", " t = torch.flatten(pred_opt)\n", " \n", " \n", " opt_t = t.cpu().detach().numpy()\n", "\n", " opt_exp = np.mean(opt_t)\n", "\n", " min_opt = np.min(opt_t)\n", " max_opt = np.max(opt_t)\n", "\n", " with open(f'./src/mrl_te_optimization/outputs/init_mrl_{OPT}.txt', 'w') as f:\n", " f.writelines([str(x)+'\\n' for x in init_t])\n", "\n", " with open(f'./src/mrl_te_optimization/outputs/opt_mrl_{OPT}.txt', 'w') as f:\n", " f.writelines([str(x)+'\\n' for x in best_scores])\n", "\n", " with open(f'./src/mrl_te_optimization/outputs/opt_seqs_{OPT}.txt', 'w') as f:\n", " f.writelines([str(x)+'\\n' for x in best_seqs])\n", "\n", " with open(f'./src/mrl_te_optimization/outputs/init_seqs_{OPT}.txt', 'w') as f:\n", " f.writelines([str(x)+'\\n' for x in seqs_gen_init])\n", " \n", "\n", " print(f\"Average Initial Pred: {np.average(init_t)}\")\n", " print(f\"Max Initial Pred: {np.max(init_t)}\")\n", " print(f\"Average Opt. Pred: {np.average(best_scores)}\")\n", " print(f\"Max Opt. Pred: {np.max(best_scores)}\")\n", "\n" ] } ], "metadata": { "kernelspec": { "display_name": "utrgan", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.8" } }, "nbformat": 4, "nbformat_minor": 2 }