Upload folder using huggingface_hub (part 2)

model/src/mrl_te_optimization/models/Modules/_operator.py

@@ -0,0 +1,246 @@
+import torch 
+import math
+import numpy  as np
+from torch import nn
+from scipy import stats
+import torch.nn.functional as F
+from einops import rearrange
+from sklearn.metrics import roc_auc_score, r2_score
+from torch.nn.modules import activation
+from torch.nn.modules.dropout import Dropout
+
+class Conv1d_block(nn.Module):
+    """
+    the Convolution backbone define by a list of convolution block
+    """
+    def __init__(self,channel_ls,kernel_size,stride, padding_ls=None,diliation_ls=None,pad_to=None, activation='ReLU'):
+        """
+        Argument
+            channel_ls : list, [int] , channel for each conv layer
+            kernel_size : int
+            stride :  list , [int]
+            padding_ls :   list , [int]
+            diliation_ls : list , [int]
+        """
+        super(Conv1d_block,self).__init__()
+        ### property
+        self.activation = activation
+        self.channel_ls = channel_ls
+        self.kernel_size = kernel_size
+        self.stride = stride
+        if padding_ls is None:
+            self.padding_ls = [0] * (len(channel_ls) - 1)
+        else:
+            assert len(padding_ls) == len(channel_ls) - 1
+            self.padding_ls = padding_ls
+        if diliation_ls is None:
+            self.diliation_ls = [1] * (len(channel_ls) - 1)
+        else:
+            assert len(diliation_ls) == len(channel_ls) - 1
+            self.diliation_ls = diliation_ls
+        
+        self.encoder = nn.ModuleList(
+            #                   in_C         out_C           padding            diliation
+            [self.Conv_block(channel_ls[i],channel_ls[i+1],self.padding_ls[i],self.diliation_ls[i],self.stride[i]) for i in range(len(self.padding_ls))]
+        )
+        
+    def Conv_block(self,in_Chan,out_Chan,padding,dilation,stride): 
+        
+        activation_layer = eval(f"nn.{self.activation}")
+        
+        block = nn.Sequential(
+                nn.Conv1d(in_Chan,out_Chan,self.kernel_size,stride,padding,dilation),
+                nn.BatchNorm1d(out_Chan),
+                activation_layer())
+        
+        return block
+    
+    def forward(self,x):
+        if x.shape[2] == 4:
+            out = x.transpose(1,2)
+        else:
+            out = x
+        for block in self.encoder:
+            out = block(out)
+        return out
+    
+    def forward_stage(self,x,stage):
+        """
+        return the activation of each stage for exchanging information
+        """
+        assert stage < len(self.encoder)
+        
+        out = self.encoder[stage](x)
+        return out
+
+    def cal_out_shape(self,L_in=100,padding=0,diliation=1,stride=2):
+        """
+        For convolution 1D encoding , compute the final length 
+        """
+        L_out = 1+ (L_in + 2*padding -diliation*(self.kernel_size-1) -1)/stride
+        return L_out
+    
+    def last_out_len(self,L_in=100):
+        for i in range(len(self.padding_ls)):
+            padding = self.padding_ls[i]
+            diliation = self.diliation_ls[i]
+            stride = self.stride[i]
+            L_in = self.cal_out_shape(L_in,padding,diliation,stride)
+        # assert int(L_in) == L_in , "convolution out shape is not int"
+        
+        return int(L_in) if L_in >=0  else 1
+    
+class ConvTranspose1d_block(Conv1d_block):
+    """
+    the Convolution transpose backbone define by a list of convolution block
+    """
+    def __init__(self,channel_ls,kernel_size,stride,padding_ls=None,diliation_ls=None,pad_to=None):
+        channel_ls = channel_ls[::-1]
+        stride = stride[::-1]
+        padding_ls =  padding_ls[::-1] if padding_ls  is not None else  [0] * (len(channel_ls) - 1)
+        diliation_ls =  diliation_ls[::-1] if diliation_ls  is not None else  [1] * (len(channel_ls) - 1)
+        super(ConvTranspose1d_block,self).__init__(channel_ls,kernel_size,stride,padding_ls,diliation_ls,pad_to)
+        
+    def Conv_block(self,in_Chan,out_Chan,padding,dilation,stride): 
+        """
+        replace `Conv1d` with `ConvTranspose1d`
+        """
+        block = nn.Sequential(
+                nn.ConvTranspose1d(in_Chan,out_Chan,self.kernel_size,stride,padding,dilation=dilation),
+                nn.BatchNorm1d(out_Chan),
+                nn.ReLU())
+        
+        return block
+    
+    def cal_out_shape(self,L_in,padding=0,diliation=1,stride=1,out_padding=0):
+        #                  L_in=100,padding=0,diliation=1,stride=2
+        """
+        For convolution Transpose 1D decoding , compute the final length
+        """
+        L_out = (L_in -1 )*stride + diliation*(self.kernel_size -1 )+1-2*padding + out_padding 
+        return L_out
+
+
+class linear_block(nn.Module):
+    def __init__(self,in_Chan,out_Chan,dropout_rate=0.2):
+        """
+        building block func to define dose network
+        """
+        super(linear_block,self).__init__()
+        self.block = nn.Sequential(
+            nn.Linear(in_Chan,out_Chan),
+            nn.Dropout(dropout_rate),
+            nn.BatchNorm1d(out_Chan),
+            nn.ReLU()
+        )
+    def forward(self,x):
+        return self.block(x)
+
+
+class Self_Attention(nn.Module):
+    """
+    self attention operator for Conv1d sequences output
+    """
+    def __init__(self, in_dim:int, out_dim:int, qk_dim:int, n_head:int):
+        super().__init__()
+
+        self.n_head = n_head
+        self.total_qk_dim = qk_dim * n_head
+        self.transform = nn.ModuleDict({
+            k : nn.Linear(in_dim, self.total_qk_dim) for k in ['k', 'q', 'v']
+        })
+
+        self.fc_out = nn.Linear(self.total_qk_dim, out_dim)
+
+    def dim_rerrange(self, x):
+
+        # first break down total qk dimension
+        # then transpose length with heads
+        x1 = rearrange(x, "b l (n c) -> b n c l", n=self.n_head) 
+        return x1
+    
+    def _get_attention_map(self,X):
+        """
+        break the forward function to access attention mat
+        """
+        # assume we have a 3 dimension input X (b, len, in_dim)
+        # each out in qkv is also 3 dimension  (b, len , qk_dim)
+        qkv = [self.transform[key](X) for key in ['k', 'q', 'v']]
+        q, k, v = map(self.dim_rerrange, qkv)
+
+        # here i and j is the channel
+        sim = torch.einsum("b n c i, b n c j -> b n i j", q, k)
+        sim = sim - sim.amax(dim=-1, keepdim=True).detach() 
+        attn = sim.softmax(dim=-1)
+        return attn, v
+
+    def forward(self, X):
+    
+        attn, v = self._get_attention_map(X)
+
+        out = torch.einsum("b n i j, b n c j -> b n i c", attn, v)
+        out = rearrange(out, "b n i c -> b i (n c)")
+        return self.fc_out(out)
+
+
+class Self_Attention_for_GP(Self_Attention):
+    """
+    self attention operator for Conv1d sequences Global Pooling output
+    The input has 2 dimension (no length dim), 
+    """
+    def __init__(self, in_dim:int, out_dim:int, qk_dim:int, n_head:int):
+        super().__init__(in_dim, out_dim, qk_dim, n_head)
+
+    def _get_attention_map(self,X):
+        # assume we have a 3 dimension input X (b, len, in_dim)
+        # each out in qkv is also 3 dimension  (b, len , qk_dim)
+        qkv = [self.transform[key](X) for key in ['k', 'q', 'v']]
+        q, k, v = map(
+            lambda x : rearrange(x, "b (n c)-> b n c"), qkv
+        )
+
+        # here i and j is the channel
+        sim = torch.einsum("b n i, b n j -> b n i j", q, k).softmax(dim=-2, keepdim=True)
+        sim = sim - attn.amax(dim=-1, keepdim=True).detach() 
+        attn = sim.softmax(dim=-1)
+        return attn, v
+
+    def forward(self, X):
+        # 
+        attn, v = self._get_attention_map(X)
+
+        out = torch.einsum("b n i j, b n j -> b n i", attn, v)
+        out = rearrange(out, "b n i -> b (n i)")
+        return self.fc_out(out)
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+    
+    def forward(self, x, *args, **kwargs):
+        return self.fn(x, *args, **kwargs) + x
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.GroupNorm(1, dim)
+        
+    def forward(self, x):
+        x = self.norm(x.transpose(1,2))
+        return self.fn(x.transpose(1,2))
+
+class SinusoidalPositionEmbeddings(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+    
+    def forward(self, time):
+        device = time.device
+        half_dim = self.dim // 2 
+        embeddings = math.log(10000) / (half_dim -1) # why do we minus 1 ?
+        embeddings = torch.exp(torch.arange(half_dim, device=device)* -embeddings)
+        embeddings = time[:, None] * embeddings[None, :] # expand to 2 dimension
+        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
+        return embeddings

model/src/mrl_te_optimization/models/Modules/nonDeep.py

@@ -0,0 +1,69 @@
+import os
+import sys
+sys.path.append(os.path.abspath("../"))
+import numpy as np
+import pandas as pd
+import torch
+from torch import nn
+import reader
+
+
+class Kmer_LinReg(nn.Module):
+    """
+    simple kmer model detect motifs
+    """
+    def __init__(self, kmer_size, pad_to):
+        super().__init__()
+        self.k = kmer_size
+        self.input_length = pad_to
+
+        # define Conv then replace the parameters
+        channel_size = 4**kmer_size 
+        self.kmer_conv = nn.Conv1d(4, 1, kmer_size)
+        self.custom_conv()
+
+        self.kmer_binarizer = nn.ReLU()
+
+        out_length = self.compute_kmer_outshape()
+        self.fc_out = nn.Linear(channel_size * out_length, 1)
+
+    def compute_kmer_outshape(self):
+        """
+        default stride = 1, pad = 0
+        """
+        dilation = 1
+        padding = 0
+        stride = 1
+        L_in = self.input_length 
+        L_out = 1 + L_in + 2 * padding - dilation * (self.k - 1) - 1 
+        return L_out
+
+    def create_kmer(self):
+        all_kmer = {0:['']}
+        k = 1
+        while k <= self.k:
+            k_mer = []  # 1 ; 4 ; 2: 4**2 ...
+            for source in all_kmer[k-1]:
+                k_mer += [source + base for base in ['A','C','G','T']]
+
+            assert len(k_mer) == 4**k, f"new kmers {len(k_mer)}, not equal to {4**k}"
+            all_kmer[k] = k_mer
+            k += 1
+
+        all_kmer.pop(0)
+        return all_kmer
+    
+    def custom_conv(self):
+
+        # only detect 5-mer is kmersize is 5
+        # Zhang et al includes shorter kmers in their features
+        kmers = self.create_kmer()[self.k]
+        
+        matrix = [reader.one_hot(kmer).T for kmer in kmers]
+        kernels = np.stack(matrix)
+
+        kernels = matrix[0].reshape(1, 4, 3)
+
+        self.kmer_conv.weight = nn.Parameter(torch.from_numpy(kernels).long(), requires_grad=False)
+
+

model/src/mrl_te_optimization/models/ScheduleOptimizer.py

@@ -0,0 +1,105 @@
+import torch
+import os
+import sys
+import numpy as np 
+import torch
+from torch import nn
+from matplotlib import pyplot as plt
+from matplotlib import cm
+
+global scheduleoptim_text
+scheduleoptim_text="ScheduledOptim(optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), betas=(0.9, 0.98), eps=1e-09, weight_decay=1e-4, amsgrad=True),n_warmup_steps=20)"
+
+scheduleoptim_dict_str="""ScheduledOptim(optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),
+                                                    betas=(0.9, 0.98),
+                                                    eps=1e-09, 
+                                                    weight_decay={weight_decay}, 
+                                                    amsgrad={amsgrad}),
+                                         n_warmup_steps={n_warmup_steps})"""
+
+class ScheduledOptim(object):
+    """A simple wrapper class for learning rate scheduling"""
+
+    def __init__(self, optimizer, n_warmup_steps):
+        self.optimizer = optimizer
+        self.d_model = 128
+        self.n_warmup_steps = n_warmup_steps
+        self.n_current_steps = 0 
+        self.delta = 1
+
+    def state_dict(self):
+        return self.optimizer.state_dict()
+        
+    def load_state_dict(self,state):
+        self.optimizer.load_state_dict(state)
+
+    def step(self):
+        """Step by the inner optimizer"""
+        self.optimizer.step()
+
+    def zero_grad(self):
+        """Zero out the gradients by the inner optimizer"""
+        self.optimizer.zero_grad()
+
+    def increase_delta(self):
+        self.delta *= 2
+        self.delta = min(1024,self.delta)
+
+    def update_learning_rate(self):
+        """Learning rate scheduling per step"""
+
+        self.n_current_steps += self.delta
+        new_lr = np.power(self.d_model, -0.5) * np.min([
+            np.power(self.n_current_steps, -0.5),
+            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])
+            # -1.5
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = new_lr
+        return new_lr
+
+
+def find_lr(net,train_data,Variable):
+    criterion = torch.nn.CrossEntropyLoss()
+    
+    net.fc = nn.Linear(2048, 120)
+
+    with torch.cuda.device(0):
+        net = net.cuda()
+
+    basic_optim = torch.optim.SGD(net.parameters(), lr=1e-5)
+    optimizer = ScheduledOptim(basic_optim)
+
+
+    lr_mult = (1 / 1e-5) ** (1 / 100)
+    lr = []
+    losses = []
+    best_loss = 1e9
+    for data, label in train_data:
+        with torch.cuda.device(0):
+            data = Variable(data.cuda())
+            label = Variable(label.cuda())
+        # forward
+        out = net(data)
+        loss = criterion(out, label)
+        # backward
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        lr.append(optimizer.learning_rate)
+        losses.append(loss.data[0])
+        optimizer.set_learning_rate(optimizer.learning_rate * lr_mult)
+        if loss.data[0] < best_loss:
+            best_loss = loss.data[0]
+        if loss.data[0] > 4 * best_loss or optimizer.learning_rate > 1.:
+            break
+
+    plt.figure()
+    plt.xticks(np.log([1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1]), (1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1))
+    plt.xlabel('learning rate')
+    plt.ylabel('loss')
+    plt.plot(np.log(lr), losses)
+    plt.show()
+    plt.figure()
+    plt.xlabel('num iterations')
+    plt.ylabel('learning rate')
+    plt.plot(lr)

model/src/mrl_te_optimization/models/Self_attention.py

@@ -0,0 +1,102 @@
+import os
+import sys
+import torch
+from torch import nn
+import numpy as np
+
+class self_attention(nn.Module):
+    
+    def __init__(self,in_channel,n_head,d_k,d_v): 
+        super(self_attention,self).__init__()
+        self.n_head=n_head
+        self.d_k = d_k
+        self.d_v = d_v
+        self.W_dict = nn.ModuleDict({"Wq" : nn.Linear(in_channel,n_head*d_k),
+                                     "Wk" : nn.Linear(in_channel,n_head*d_k),
+                                     "Wv" : nn.Linear(in_channel,n_head*d_v)})
+
+    def forward(self,X):
+        
+        # some dimension 
+        dk_sqrt = int(np.sqrt(self.d_k))
+        
+        querys = self.W_dict['Wq'](X)     # B*X_dim*dk
+        keys = self.W_dict['Wk'](X)       # B* hs_out -> B*64
+        values = self.W_dict['Wv'](X)    # B* hs_out -> B*128
+        
+        sim_M = torch.bmm(querys,keys.transpose(1,2))/8  # B* X_dim * X_dim
+        attention = torch.softmax(sim_M,dim=-1)
+        # result
+        result = torch.bmm(attention,values).squeeze(2)    # B*X_dim*1 -> B*X_dim
+        
+        return result
+    
+class task_attention(nn.Module):
+    """
+    Task Attention Layer, it contains a task specific mask `Wq` as the Query. Key is computed from the input 
+    
+    Arguments:
+
+        d_v : int , dimension of the network width input , i.e : (Batch_size, d_v)
+        d_k : int , dimension of the Query and Key.
+    """
+    def __init__(self,d_v:int,d_k=64): 
+        super(task_attention,self).__init__()
+    
+        self.d_k = d_k
+        self.d_v = d_v
+        self.W_dict = nn.ModuleDict({"Wq" : nn.Linear(d_k,d_v),
+                                     "Wk" : nn.Linear(d_v,d_k)})
+
+    def forward(self,X):
+        # i.e X : (B ,128)
+        # if len(X.shape) == 2:
+        #     X = X
+        
+        # some dimension 
+        dk_sqrt = int(np.sqrt(self.d_k))
+        
+        keys = self.W_dict['Wk'](X).squeeze(2)       # (B, 128) -> (B ,64)
+        
+        # dot product similarity is used here , which is implemented by `nn.Linear`
+        sim_M =self.W_dict['Wq'](keys) /dk_sqrt  # (B,128,64) * (B,64,128)  ->  (B, 128, 128)
+        attention = torch.softmax(sim_M,dim=-1)
+        # result
+        result = torch.mul(attention,X)   # (B, 128)  * (B, 128) -> B*128
+        
+        return result
+
+class task_attention(nn.Module):
+    """
+    Task Attention Layer, it contains a task specific mask `Wq` as the Query. Key is computed from the input 
+    
+    Arguments:
+
+        d_v : int , dimension of the network width input , i.e : (Batch_size, d_v)
+        d_k : int , dimension of the Query and Key.
+    """
+    def __init__(self,d_v:int,d_k=64): 
+        super(task_attention,self).__init__()
+    
+        self.d_k = d_k
+        self.d_v = d_v
+        self.W_dict = nn.ModuleDict({"Wq" : nn.Linear(d_v,d_k),
+                                     "Wk" : nn.Linear(d_v,d_k)})
+
+    def forward(self,X):
+        # i.e X : (B ,128)
+        # if len(X.shape) == 2:
+        #     X = X
+        
+        # some dimension 
+        dk_sqrt = int(np.sqrt(self.d_k))
+        query = self.W_dict['Wq'](X).squeeze(2)
+        keys = self.W_dict['Wk'](X).squeeze(2)       # (B, 128) -> (B ,64)
+        
+        # dot product similarity is used here , which is implemented by `nn.Linear`
+        sim_M =self.W_dict['Wq'](keys) /dk_sqrt  # (B,128,64) * (B,64,128)  ->  (B, 128, 128)
+        attention = torch.softmax(sim_M,dim=-1)
+        # result
+        result = torch.mul(attention,X)   # (B, 128)  * (B, 128) -> B*128
+        
+        return result

model/src/mrl_te_optimization/models/bucket_sampler.py

@@ -0,0 +1,99 @@
+import torch
+import numpy as np
+import pandas as pd
+from torch.utils.data.sampler import Sampler,BatchSampler,SubsetRandomSampler
+
+class Basic_sampler(Sampler):
+    def __init__(self, data):
+        super().__init__(data)
+        self.data= data
+        
+    def __len__(self):
+        return self.data.shape[0]
+    
+    def __iter__(self):
+        return (i for i in range(self.data.shape[0]))
+    
+class sort_sampler(Sampler):
+    def __init__(self, data, sort_key="utr_len"):
+        super().__init__(data)
+        self.data = data
+        self.sort_key = sort_key
+        zip_ = [(i, seq_len) for i, seq_len in enumerate(data[sort_key].values)]
+        zip_ = sorted(zip_, key=lambda r: r[1])
+        self.sorted_indexes = [item[0] for item in zip_]
+
+    def __iter__(self):
+        return iter(self.sorted_indexes)
+
+    def __len__(self):
+        return len(self.data)
+    
+class Bucket_Sampler(BatchSampler):
+    """ `BucketBatchSampler` toggles between `sampler` batches and sorted batches.
+
+    Typically, the `sampler` will be a `RandomSampler` allowing the user to toggle between
+    random batches and sorted batches. A larger `bucket_size_multiplier` is more sorted and vice
+    versa.
+
+    Background:
+        ``BucketBatchSampler`` is similar to a ``BucketIterator`` found in popular libraries like
+        ``AllenNLP`` and ``torchtext``. A ``BucketIterator`` pools together examples with a similar
+        size length to reduce the padding required for each batch while maintaining some noise
+        through bucketing.
+
+        **AllenNLP Implementation:**
+        https://github.com/allenai/allennlp/blob/master/allennlp/data/iterators/bucket_iterator.py
+
+        **torchtext Implementation:**
+        https://github.com/pytorch/text/blob/master/torchtext/data/iterator.py#L225
+
+    Args:
+        sampler (torch.data.utils.sampler.Sampler):
+        batch_size (int): Size of mini-batch.
+        drop_last (bool): If `True` the sampler will drop the last batch if its size would be less
+            than `batch_size`.
+        sort_key (callable, optional): Callable to specify a comparison key for sorting.
+        bucket_size_multiplier (int, optional): Buckets are of size
+            `batch_size * bucket_size_multiplier`.
+
+    Example:
+        >>> from torchnlp.random import set_seed
+        >>> set_seed(123)
+        >>>
+        >>> from torch.utils.data.sampler import SequentialSampler
+        >>> sampler = SequentialSampler(list(range(10)))
+        >>> list(BucketBatchSampler(sampler, batch_size=3, drop_last=False))
+        [[6, 7, 8], [0, 1, 2], [3, 4, 5], [9]]
+        >>> list(BucketBatchSampler(sampler, batch_size=3, drop_last=True))
+        [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
+    """
+    def __init__(self,
+                 data,
+                 batch_size,
+                 drop_last=False,
+                 sort_key='utr_len',
+                 bucket_size_multiplier=100):
+        self.data = data
+        self.sampler = Basic_sampler(data)
+        super().__init__(self.sampler, batch_size, drop_last)
+        self.sort_key = sort_key
+        _bucket_size = batch_size * bucket_size_multiplier
+        if hasattr(self.sampler, "__len__"):
+            _bucket_size = min(_bucket_size, len(self.sampler))
+        self.bucket_sampler = BatchSampler(self.sampler, _bucket_size, False)
+
+    def __iter__(self):
+        for bucket in self.bucket_sampler:
+            sorted_sampler = sort_sampler(self.data.iloc[bucket], self.sort_key)
+            for batch in SubsetRandomSampler(
+                    list(BatchSampler(sorted_sampler, self.batch_size, self.drop_last))):
+                yield [bucket[i] for i in batch]
+
+    def __len__(self):
+        if self.drop_last:
+            return len(self.sampler) // self.batch_size
+        else:
+            return np.ceil(len(self.sampler) / self.batch_size)    
+    
+

model/src/mrl_te_optimization/models/log_and_save.py

@@ -0,0 +1,271 @@
+import os
+import torch
+import logging
+import re
+import pandas as pd
+import numpy as np
+from matplotlib import pyplot as plt
+import copy
+
+def snapshot(dir_path, run_name, state,logger):
+    snapshot_file = os.path.join(dir_path,
+                                 run_name + '-model_best.pth')
+    # torch.save can save any object
+    # dict type object in our cases
+    torch.save(state, snapshot_file)
+    logger.info("Snapshot saved to {}\n".format(snapshot_file))
+
+class Log_parser(object):
+    def __init__(self,log_path,val_split_line=False,use_line_as_valtest=1):
+        #           --------   read   --------
+        self.val_split_line = val_split_line
+        self.use_line_as_valtest = use_line_as_valtest
+        if os.path.exists(log_path):
+            with open(log_path,'r') as f:
+                log_file = f.readlines()
+                f.close()
+            # stripping 
+            log_file = np.array([line.strip() for line in log_file])
+        else:
+            print('log path error !')
+        self.log_file = log_file
+        
+#         self.possible_metric = ['LOSS','lr','Avg_ACC','teaching_rate','TOTAL','KLD','MSE','M_N','CrossEntropy','chimerla_weight','Total','TE','Loop','Match','MAE','RMSE','RL_loss','Recons_loss','Motif_loss','RL_Acc','Recons_Acc','Motif_Acc','Acc','Mean_Total', 'DTP_wt_RL','DTP_wt_Recons','DTP_wt_Motif']
+        
+        #          --------  basic  matcher   --------
+        self.epoch_line_matcher = r"\s.* epoch (\d{1,4}).*"
+        self.start_val_line_matcher = r"\s*.* start validation .*\s*"
+        self.match_logging_time = r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2},\d{3} -"
+        self.match_percentage = r"\s*\d{1,6} /\s*\d{1,6}\s*\((\d|\.){,6}%\):"
+        
+        self.match_sub_verbose = lambda x : r"\s*%s:\s*(?P<%s>(-|\d|\.|e|){,40})"%(x,x)
+        
+        #          -------- high level matcher --------
+        self.train_verbose_finder = self.match_logging_time + self.match_percentage
+        
+        #          --------- get output DF ---------
+        self.extract_training_verbose_data()
+        self.extract_val_verbose_data()
+        
+    def lines_to_json(self, line, sett):
+        remove_time = line.split('%):')[1].split() if sett=='train' else line.split(' - \t ')[1].split()
+        line_json = {metrics.split(':')[0]:metrics.split(':')[1] for metrics in remove_time}
+        return line_json
+        
+    def lines_matching(self,matcher):
+        """
+        return lines that can match certain syntax
+        """
+        return [line for line in self.log_file if re.match(matcher,line) is not None]
+    
+    def position_matching(self,matcher):
+        """
+        return position of the line that can match certain syntax
+        """
+        return [i for i,line in enumerate(self.log_file) if re.match(matcher,line) is not None]
+    
+#     def get_metrics_order(self):
+#         """
+#         get train verbose line and define train_verbose_matcher automatically
+#         """
+#          # find all the train verbose lines
+        
+#         test_t_v = self.train_verbose_lines[0]                     # a testing train verbose
+        
+#         # using the esting trainverbose to determine metric order
+# #         train_metric = np.array([metric for metric in self.possible_metric if metric in test_t_v])
+# #         train_metric = self.check_dup_metric(train_metric,test_t_v)
+        
+# #         train_metric_posi = np.array([test_t_v.index(metric) for metric in train_metric])
+        
+# #         order = train_metric_posi.argsort()
+# #         self.train_metric = train_metric[order]
+        
+# #         #   ----|| automatically determine train verbose matcher ||----
+# #         self.train_verbose_matcher = self.train_verbose_finder
+# #         for metric in self.train_metric:
+# #             self.train_verbose_matcher += self.match_sub_verbose(metric)
+    
+# #     def check_dup_metric(self,train_metric,test_t_v):
+# #         """
+# #         to deal with the problem of `MSE` and `RMSE` 
+# #         """
+# #         train_metric = list(train_metric)
+# #         if ("MSE" in train_metric) & ("RMSE" in train_metric):
+# #             if test_t_v.index('MSE') == test_t_v.index('RMSE')+1:
+# #                 train_metric.remove('MSE')
+# #         return np.array(train_metric)
+    
+    def extract_training_verbose_data(self):
+        """
+        regular expression to match the printed metric during training and save to pd.DataFrame
+        """
+        self.train_verbose_lines = self.lines_matching(self.train_verbose_finder)
+        
+        self.train_verbose_dict = [self.lines_to_json(line,'train') for line in self.train_verbose_lines]
+        
+        self.train_metric = list(self.train_verbose_dict[0].keys())
+        
+        self.train_verbose_DF = pd.json_normalize(self.train_verbose_dict).astype(float)
+        
+        # return self.train_verbose_DF 
+    
+    def extract_val_verbose_data(self):
+        """
+        regular expression to match the printed metric during training and save to pd.DataFrame
+        """
+        self.start_val_posi = self.position_matching(self.start_val_line_matcher)
+        val_verbose_posi = np.array(self.start_val_posi) +1 # observe from log
+        self.val_verbose_posi = val_verbose_posi[val_verbose_posi < len(self.log_file)]
+        self.val_verbose_lines = self.log_file[self.val_verbose_posi]
+        if self.val_split_line:
+            self.val_verbose_lines = ["\t".join(self.log_file[[posi,posi+1,posi+2,posi+3,posi+4]]) for posi in self.val_verbose_posi]
+
+#         test_v_v = self.val_verbose_lines[self.use_line_as_valtest]
+        
+         # using the esting trainverbose to determine metric order
+#         val_metric = np.array([metric for metric in self.possible_metric if metric in test_v_v])
+#         val_metric = self.check_dup_metric(val_metric,test_v_v)
+#         val_metric_posi = np.array([test_v_v.index(metric) for metric in val_metric])
+#         order = val_metric_posi.argsort()    # sort 
+#         self.val_metric = val_metric[order]
+        
+        
+#         #   ----|| automatically determine val verbose matcher ||----
+#         self.val_verbose_matcher = self.match_logging_time 
+        
+#         if re.match(self.match_logging_time + self.match_percentage,test_v_v) is not None:
+#             self.val_verbose_matcher += self.match_percentage        # detect whether validation set also get percentage info
+        
+#         for metric in self.val_metric:
+#             self.val_verbose_matcher += self.match_sub_verbose(metric)
+        
+        self.val_verbose_dict = [self.lines_to_json(line, 'val') for line in self.val_verbose_lines]
+        self.val_metric = list(self.val_verbose_dict[0].keys())
+#         np.array(
+#             [list(
+#                 re.match(self.val_verbose_matcher,line).groupdict().values()
+#                     ) for line in self.val_verbose_lines]
+#             ).astype(np.float64)
+        
+        self.val_verbose_DF = pd.json_normalize(self.val_verbose_dict).astype(float)
+        
+        # return self.val_verbose_DF 
+    
+    def plot_val_metric(self,fig=None,dataset='val'):
+        DF = self.val_verbose_DF if dataset == 'val' else self.train_verbose_DF
+        metrics = self.val_metric if dataset == 'val' else self.train_metric
+        n = len(metrics)
+        
+        if fig is None:
+            fig = plt.figure(figsize=(18,5*np.ceil(n/3)))
+        if n <=3:
+            axs = fig.subplots(1,n)
+            for i in range(n):
+                axs[i].plot(DF[metrics[i]].values)
+                axs[i].set_title(dataset.capitalize()+" "+metrics[i])  # TRAIN or VAL
+        else:
+            axs = fig.add_subplot(n//3+1,n,1+i)
+        
+
+def plot_a_exp_set(log_list,log_name_ls,dataset='val',fig=None,layout=None,check_time=10,start_from=0,mean_of_train=None,define_order=None,esubset=None, cycle_train=False,**kwargs):
+    all_metric = [logg.__getattribute__(dataset+"_metric") for logg in log_list]
+    share_metric = [all_metric[0]]
+    for logg_metric in all_metric[1:]:
+        share_metric = np.intersect1d(share_metric,logg_metric)
+    if define_order is not None:
+        assert set(define_order) == set(share_metric)
+    
+    n = len(share_metric) + 1 # val or train
+    fig = plt.figure(figsize=(20,5)) if fig is None else fig
+
+    
+    if layout is None:
+        axs = fig.subplots(1,n);
+    else:
+        row,column = layout 
+        axs = fig.subplots(row,column).flatten()
+    
+    
+    for i,metric in enumerate(share_metric):
+        # layout 
+        
+        ax = axs[i]
+        for st,log in enumerate(log_list):
+            DF = log.__getattribute__(dataset+"_verbose_DF")
+                
+            if (dataset == 'train') & (type(mean_of_train)==int):
+                DF = mean_of(mean_of_train,DF)
+            elif (dataset == 'train') & (type(esubset)==slice):
+                DF = subset_of(esubset,DF)
+            X = np.arange(DF.shape[0])*check_time if dataset == 'val' else np.arange(DF.shape[0])
+            ax.plot(X[start_from:],DF[metric].values[start_from:],**kwargs)
+            ax.set_title(" ".join([dataset.capitalize(),metric]))
+    for st,log in enumerate(log_list):
+        axs[-1].plot(0,0,label=log_name_ls[st])
+        axs[-1].axis('off')
+        axs[-1].legend()
+        
+def plot_cycle_exp_set(log_ls,log_name,dataset='val',**kwargs):
+    interval = 2 if dataset=='val' else 6
+    new_log_ls = []
+    new_log_name = []
+    for i in range(len(log_ls)):
+        log = log_ls[i]
+        DF  = log.__getattribute__(dataset+"_verbose_DF")
+        ds1_index = [i  for i in range(DF.shape[0]) if i//interval%2 ==0]
+        ds2_index = [i  for i in range(DF.shape[0]) if i//interval%2 ==1]
+        DF1 = DF.iloc[ds1_index]
+        DF2 = DF.iloc[ds2_index]
+        log1 = copy.deepcopy(log)
+        log2 = copy.deepcopy(log)
+        
+        log1.__setattr__(dataset+'_verbose_DF', DF1)
+        log2.__setattr__(dataset+'_verbose_DF', DF2) 
+        
+        new_log_ls.append(log1)
+        new_log_ls.append(log2)
+        
+        new_log_name.append(log_name[i]+"_ds1")
+        new_log_name.append(log_name[i]+"_ds2")
+    
+    plot_a_exp_set(new_log_ls, new_log_name, **kwargs)
+                
+        
+def subset_of(x,DF):
+    
+    values = DF.values
+    mean_ls = []
+    
+    for i in range(0,values.shape[0],x.stop): 
+        # x : slice , x.stop , the slice window of the 
+        mean_ls.append(values[i:i+x.stop][x])
+    
+    mean_ls = np.concatenate(mean_ls,axis=0)
+    mean_DF = pd.DataFrame(mean_ls,columns=DF.columns)
+    
+    return mean_DF
+
+def mean_of(x,DF):
+    
+    values = DF.values
+    mean_ls = []
+    
+    for i in range(0,values.shape[0],x):
+        
+        mean_ls.append(np.mean(values[i:i+x,:],axis=0))
+    
+    mean_ls = np.stack(mean_ls)
+    mean_DF = pd.DataFrame(mean_ls,columns=DF.columns)
+    
+    return mean_DF
+
+def read_log_of_a_dir(log_dir):
+    """
+    ...log_dir : abs path of log dir
+    """
+    file_ls = [file for file in os.listdir(log_dir) if ".log" in file]
+    log_path = [os.path.join(log_dir,file) for file in file_ls]
+    log_name = [file.replace('.log','') for file in file_ls]
+    log_ls = [Log_parser(file) for file in log_path]
+    return log_ls,log_name

model/src/mrl_te_optimization/models/loss.py

@@ -0,0 +1,86 @@
+import os
+import torch
+from torch import nn
+import numpy as np
+
+
+def softmax(x):
+    """
+    one dimensional softmax designed for numpy array
+    """
+    e_x = np.exp(x)
+    out = e_x / e_x.sum()
+    return out
+
+class Dynamic_Weight_Averaging():
+    def __init__(self,tasks,tau,init_weight):
+        """
+        Dynamic Weight Averaging implementation
+        """
+        # number of tasks
+        self.N = len(tasks)
+        # self.loss_list = np.array([loss_dict[t+'_loss'].detach().cpu().item() for t in tasks])
+        self.step = 0
+        self.omega = np.array([init_weight]*self.N ) # the weight , \omega_i(t)
+        self.tau = tau # † : a value to adjust the soft-max
+    
+    def _magnitude_adjust(self):
+        """
+        automatic adjustment of loss magnitude
+        """
+        self.relative_magnitude = self.loss_list.min() / self.loss_list
+    
+    def  _update(self,loss_dict):
+        # update r with L_i(t-1) 
+        self.step += 1
+        if self.step < 2:
+            self.loss_list = np.array([loss_dict[t+'_loss'].detach().cpu().item() for t in self.tasks])
+            return self.init_weight
+        else:
+            self._magnitude_adjust()
+            last_loss = self.loss_list
+            self.loss_list = np.array([loss_dict[t+'_loss'].detach().cpu().item() for t in self.tasks])
+            # computing DWA
+            r_t = np.divide(self.loss_list,last_loss) / self.tau
+            self.omega = self.N * softmax(r_t)
+            weight_t = np.multiply(self.relative_magnitude,self.omega)
+            return {self.tasks[i]:weight_t[i] for i in range(len(self.tasks))}
+    
+class Dynamic_Task_Priority(object):
+    def __init__(self,tasks,gamma,init_weight):
+        """
+        Dynamic Task Priority (DTP) wish to weight more on tasks with lower KPI
+        In our cases, we take accuracy as the KPI
+        """
+        self.gamma = np.array([gamma[t] for t in tasks]) if type(gamma) == dict else gamma # an adjustment param
+        self.tasks = tasks
+        # self.kappa = [loss_dict[t+'_Acc'] for t in tasks]
+        self.omega = init_weight 
+        self.init_weight = init_weight
+        self.step = 0
+            
+    
+    def _magnitude_adjust(self):
+        """
+        automatic adjustment of loss magnitude
+        """
+        # TODO : test this effectiveness
+        self.relative_magnitude = self.kappa.min() / (self.kappa+1e-8)
+        if np.any(self.kappa==1):
+            self.kappa[np.where(self.kappa==1)[0]] -= 1e-8
+    
+    def _update(self,loss_dict):
+        """
+        weight is updated like belowing
+        `math` : w_i(t) = -(1-\kappa_i(t))^{\gamma_i} log \kappa_i(t)
+        """
+        self.step += 1
+        self.kappa = np.array([loss_dict[t+'_Acc'] for t in self.tasks])
+        self._magnitude_adjust()
+        self.omega = -1 * np.multiply(np.power(1 - self.kappa,self.gamma),np.log(self.kappa+1e-8))
+        
+        if self.step < 2:
+            return self.init_weight
+        else:
+            weight_t = np.multiply(self.relative_magnitude,self.omega)
+            return {self.tasks[i]:weight_t[i] for i in range(len(self.tasks))}

model/src/mrl_te_optimization/models/max_activation_patch.py

@@ -0,0 +1,795 @@
+import os
+import sys
+import re
+sys.path.append(os.path.dirname(__file__))
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import utils
+import copy
+import torch
+import train_val
+import reader
+import logomaker
+from torch import nn
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from tqdm.auto import tqdm
+from importlib import reload
+from scipy import stats
+from matplotlib import pyplot as plt
+from sklearn import cluster
+from sklearn.manifold import TSNE
+import ruptures as rpt
+from popen import Auto_popen
+
+tensor_2_numpy = lambda x: x.detach().cpu().numpy()
+
+class Maximum_activation_patch(object):
+    def __init__(self, popen, which_layer, n_patch=9, kfold_index=None, device_string='cpu'):
+        self.popen = popen
+        self.layer = which_layer
+        self.popen.cuda_id = torch.device("cuda:%s"%device_string) if device_string.isdigit() else torch.device('cpu')
+        self.n_patch = n_patch
+        self.kfold_index = kfold_index
+        self.total_stride = np.product(popen.stride[:self.layer])
+        self.compute_reception_field()
+        self.compute_virtual_pad()
+    
+    def load_indexed_dataloader(self,task):
+        #        True or 'train_val'
+        assert (self.popen.kfold_cv!=False) == (self.kfold_index != None), \
+                "kfold CV should match with kfold index"
+        
+        self.popen.kfold_index = self.kfold_index
+        self.popen.shuffle = False
+        self.popen.pad_to = 57 if self.popen.cycle_set==None else 105
+        
+        tmp_popen = copy.copy(self.popen)
+        if self.popen.cycle_set != None:
+            base_path = copy.copy(self.popen.split_like)
+            base_csv = copy.copy(self.popen.csv_path)
+            # base_csv = "/mnt/sina/run/ml/gan/motif/MTtrans/test.csv"
+            
+            if base_path is not None:
+                assert task != None , "task is not defined !"
+                tmp_popen.split_like = [path.replace('cycle', task) for path in base_path]
+            else:
+                tmp_popen.csv_path = base_csv.replace('cycle', task)
+                # tmp_popen.csv_path = "/mnt/sina/run/ml/gan/motif/MTtrans/test.csv"
+                
+        # tmp_popen.csv_path = "/mnt/sina/run/ml/gan/motif/MTtrans/test.csv"   
+        # base_csv = "/mnt/sina/run/ml/gan/motif/MTtrans/test.csv"     
+                
+        return  reader.get_dataloader(tmp_popen) 
+    
+    def load_model(self):
+        if self.kfold_index is not None:
+            base_pth = self.popen.vae_log_path
+            self.popen.vae_log_path = base_pth.replace(".pth","_cv%s.pth"%self.kfold_index)
+        
+        model = utils.load_model(self.popen, None)
+        if self.kfold_index is not None:
+            self.popen.vae_log_path = base_pth
+        return model
+    
+    def get_filter_param(self, model):
+        Conv_layer = model.soft_share.encoder[self.layer-1]
+        
+        return tensor_2_numpy( next(Conv_layer[0][0].parameters()) )
+    
+    def loading(self,task, which_set):
+        model = self.load_model().to(self.popen.cuda_id)
+        dataloader = self.load_indexed_dataloader(task)[which_set]
+        self.df = dataloader.dataset.df
+        return model, dataloader
+    
+    @torch.no_grad()
+    def cumulative_rl_decision(self, task=None,which_set=0, extra_loader=None):
+        """
+        take out the memory h_i of each position and pass to output layer
+        Arg:
+            task : str
+            which_set : int, 0 : training set, 1 : val set, 2 : test set
+        """
+        model, dataloader= self.loading(task, which_set)
+        model.eval()
+        
+        Y_ls = []
+        for Data in tqdm(dataloader):
+            # iter each batch
+            x,y = train_val.put_data_to_cuda(Data,self.popen,False)
+            x = torch.transpose(x, 1, 2)
+            y_pred = model.predict_each_position(x)
+            Y_ls.append( tensor_2_numpy(y_pred) )
+        Y_ay = np.concatenate(Y_ls, axis=0)
+        
+        # release some cache
+        torch.cuda.empty_cache()
+        del model
+        return Y_ay.reshape(Y_ay.shape[0],-1)
+        
+    def extract_feature_map(self, task=None, which_set=0, extra_loader=None):
+        """
+        load trained model and unshuffled dataloader, make model forwarded
+        Arg:
+            task : str
+            which_set : int, 0 : training set, 1 : val set, 2 : test set
+        """
+        model, dataloader= self.loading(task, which_set)
+        if extra_loader is not None:
+            dataloader = extra_loader
+            self.df = extra_loader.dataset.df
+            
+        feature_map = []
+        X_ls = []
+        Y_ls = []
+        # print(type(dataloader))
+        model.eval()
+        with torch.no_grad():
+            print(len(dataloader))
+            for Data in tqdm(dataloader):
+                # print(Data)
+                # print(Data[0].shape)
+                # print(Data[1].shape)
+                # iter each batch
+                x,y = train_val.put_data_to_cuda(Data,self.popen,False)
+                # print(x)
+                # print(np.shape(Data[0].numpy()))
+                shape = np.shape(Data[0].numpy()[0])
+                xt = torch.tensor(np.reshape(Data[0].numpy()[0],(1,shape[0],shape[1])),dtype=torch.float64)
+                # print(type(xt))
+                # print(type(x))
+                xt = xt.float()
+                # xt = torch.transpose(xt, 1, 2)
+                # print(model.forward(xt))
+                # print(model.predict_each_position(x))
+                
+                # print(type(y))
+                # print(x.shape)
+                x = torch.transpose(x, 1, 2)
+#                 X_ls.append(x.numpy())
+                Y_ls.append( tensor_2_numpy(y))
+                
+                for layer in model.soft_share.encoder[:self.layer]:
+                    out = layer(x)
+                    x = out
+                feature_map.append( tensor_2_numpy(out))
+                
+                torch.cuda.empty_cache()
+            
+
+        feature_map_l = np.concatenate( feature_map, axis=0)
+        
+#         self.X_ls = np.concatenate(X_ls, axis=0)
+        self.Y_ls = np.concatenate(Y_ls, axis=0)
+
+        print("activation map of layer |%d|"%self.layer,feature_map_l.shape)
+#         print(self.X_ls.shape)
+        print("Y : ",self.Y_ls.shape)
+        self.feature_map = feature_map_l
+        
+        self.filters = self.get_filter_param(model)
+        
+        del model
+        
+        return feature_map_l
+    
+    def compute_reception_field(self):
+        r = 1
+        strides = self.popen.stride[:self.layer][::-1]
+
+        for i,s in enumerate(strides):
+            r = s*(r-1) + self.popen.kernel_size
+        
+        self.r = r
+        self.strides = strides
+        print('the reception filed is ', r)
+    
+    def compute_virtual_pad(self):
+        v = []
+        pd = self.popen.padding_ls[:self.layer]
+
+        for i,p in enumerate(pd):
+            v.append(p*np.product(self.popen.stride[:len(pd)-i-1]))
+        
+        self.virtual_pad = np.sum(v)
+        print('the virtual pad is', self.virtual_pad)
+    
+    def retrieve_input_site(self, high_layer_site):
+        """pad ?"""
+        
+        virtual_start = high_layer_site*self.total_stride - self.virtual_pad
+        start = max(0, virtual_start)
+        
+        end = virtual_start + self.r
+        return max(0,int(start)), int(end)
+    
+    def detect_changepoint(self,time_series):
+        bkpt = rpt.KernelCPD(kernel="linear", min_size=1).fit_predict(time_series, n_bkps=1)[0]
+        return bkpt - 1        
+    
+    def retrieve_featmap_at_changepoint(self, featmap, rl_chain, threshold=1, direction='less', detect_region=None):
+        """
+        Using the change point of rl series, to retrieve feature map at the same position
+        e.g. : to find negative change point threshold=-1, direction='less'
+        e.g. : to find positive change point threshold=0.5, direction='greater'
+        
+        featmap : np.ndarray, (n_sample, n_channel, n_position)
+        rl_chain : np.ndarray, (n_sample,  n_position)
+
+        threshold : the threshold define the rl after change point minus that ahead the point that are considered
+        direction : the direction of changes
+        detect_region : list of slice [], which region of the rl chain is used to detect change point, default None (full sequence is considered)
+        """
+        # take out the negative break point
+        
+        if direction == 'less':
+            condition = lambda x1, x2 : x1 - x2 < threshold
+        elif direction == 'greater':
+            condition = lambda x1, x2 : x1 - x2 > threshold
+
+        if detect_region is None:
+            detect_region = [slice(0, None)] * rl_chain.shape[0]
+
+        change_point_act = []
+        for i, trend in enumerate(rl_chain):
+            region = detect_region[i]
+            bkpt= self.detect_changepoint(trend[region]) + region.start
+
+            if condition(trend[bkpt+1] , trend[bkpt]): 
+                activation_vec = featmap[i, :, bkpt+1]
+                change_point_act.append(activation_vec)
+
+        chagnepoint_map = np.asarray(change_point_act)
+        return chagnepoint_map
+    
+    def locate_MA_seq(self, channel, feature_map=None):
+        """
+        
+        """
+        if feature_map is None:
+            feature_map = self.feature_map
+            
+        channel_feature = feature_map[:, channel,:]
+        F0_ay = channel_feature.max(axis=-1)
+
+        max_n_index = np.argpartition(F0_ay, -1*self.n_patch, axis=0)[-1*self.n_patch:]
+
+
+        # a patch of sequences
+        max_patch = self.df[self.popen.seq_col].values[max_n_index]
+
+        # find the location of maximally acitvated
+        Conv4_sites = np.argmax(channel_feature[max_n_index],axis=1)
+        
+        mapped_input_sites = [self.retrieve_input_site(site) for site in Conv4_sites]
+        
+        # the mapped region of the sequences
+        max_act_region = []
+        for utr, (start, end) in zip(max_patch, mapped_input_sites):
+            
+
+            pad_gap = self.popen.pad_to - len(utr)
+            field = utr[max(0, start-pad_gap): end-pad_gap]
+            
+            max_act_region.append(field)
+            
+#         print(mapped_input_sites)
+        
+        full_field = [len(field)==self.r for field in max_act_region]
+        
+        return np.array(max_act_region)[full_field], max_n_index[full_field], np.array(mapped_input_sites)[full_field]
+    
+    def sequence_to_matrix(self, max_act_region, weight=None, transformation='counts'):
+        
+        assert np.all([seq != "" for seq in max_act_region])
+        
+        max_len = max([len(seq) for seq in max_act_region])
+        
+        M = np.zeros((max_len,4))
+        if weight is None:
+            weight = np.ones((self.n_patch,))
+        for seq, w in zip(max_act_region, weight):
+            oh_M = reader.one_hot(seq)*w
+            M += np.concatenate([np.zeros((max_len - len(seq),4)), oh_M],axis=0)
+            
+
+        seq_logo_df = pd.DataFrame(M, columns=['A', 'C', 'G', 'U'])
+        if transformation!='counts':
+            seq_logo_df = logomaker.transform_matrix(seq_logo_df, from_type='counts', to_type=transformation)
+            
+        return seq_logo_df
+        
+    def plot_sequence_logo(self, seq_logo_df,  save_fig_path=None, ax=None):
+        """
+        input : max_act_region ; list of str
+        """
+        # plot
+        if ax is None:
+            fig = plt.figure(dpi=300)
+            ax = fig.gca()
+        MA_C = logomaker.Logo(seq_logo_df,ax=ax)
+            
+        ax.spines['right'].set_visible(False)
+        ax.spines['top'].set_visible(False)
+        ax.spines['bottom'].set_visible(False)
+        
+        # save
+        if save_fig_path is not None:
+            MA_C.fig.savefig(save_fig_path,transparent=True,dpi=600)
+            save_dir = os.path.dirname(save_fig_path)
+            
+            try:
+                self.save_dir
+            except:
+                # which is the first time we save
+                self.save_dir = save_dir
+                print('fig saved to',self.save_dir)
+            plt.close(MA_C.fig)
+    
+
+        
+
+        
+    def fast_logo(self, channel, feature_map=None, n_patch=None, transformation='information', save_fig_path=None, title=None,ax=None):
+        
+        if n_patch is not None:
+            self.n_patch = n_patch
+        max_act_region, _, _ = self.locate_MA_seq(channel, feature_map)
+        M = self.sequence_to_matrix(max_act_region, transformation=transformation)
+        
+        # print(max_act_region)
+        
+        F0_ay, (spr,pr) = self.activation_density(channel, feature_map=feature_map, to_print=False, to_plot=False)
+        # print(np.shape(F0_ay))
+        # print(type(F0_ay))
+        
+        self.plot_sequence_logo(M,  save_fig_path=save_fig_path, ax=ax)
+        if ax is None:
+            ax=plt.gca()
+        if title is None:
+            title = "filter {} : $r =$ {}".format(channel, spr)
+        
+        ax.set_title(title, fontsize=35)
+        return M, spr
+         
+    
+    def activation_density(self, channel, to_print=True, to_plot=True, feature_map=None, **kwargs):
+        if feature_map is None:
+            feature_map = self.feature_map
+            
+        channel_feature = feature_map[:, channel,:]
+        F0_ay = channel_feature.max(axis=-1)
+        if to_plot:
+            sns.kdeplot(F0_ay, **kwargs);
+        if to_print:
+            print("num of max acti: %s"%np.sum(F0_ay==F0_ay.max()))
+            
+            n_gtr = [np.sum(F0_ay > q) for q in np.quantile(F0_ay,[0.5,0.8,0.95])]
+            print("quantiles: 50% {} 80% {} 95% {}".format(*n_gtr))
+        
+        spr = stats.spearmanr(F0_ay, self.Y_ls.flatten())
+        pr = stats.pearsonr(F0_ay, self.Y_ls.flatten())
+        return F0_ay, (round(spr[0],3),round(pr[0],3))
+    
+    def within_patch_clustering(self, channel, n_clusters, n_patch=None , to_plot=True,**kwargs):
+        """
+        return:
+            subclusters : list, list of patch (alignments)
+        """
+        self.n_patch = n_patch
+        act_patch, _, _ = self.locate_MA_seq(channel)
+        flatten_seq = np.stack([reader.one_hot(seq).flatten() for seq in act_patch])
+        print("the shape of sequence matrix {}".format(flatten_seq.shape))
+        
+        # clsutering
+        cluster_index = cluster.KMeans(n_clusters=n_clusters).fit_predict(flatten_seq)
+        # down
+        if to_plot:
+            tsne = TSNE(metric='cosine').fit(flatten_seq)
+            self.tsne=tsne
+            plt.figure(figsize=(6,5),dpi=150)
+    #         sns.set_theme(style='ticks', palette='viridis')
+            scatter_args = {"palette":'viridis'}
+            scatter_args.update(kwargs)
+            sns.scatterplot(x=tsne.embedding_[:,0], y=tsne.embedding_[:,1], 
+                            hue=cluster_index, **scatter_args);
+        
+        return act_patch, flatten_seq, cluster_index
+    
+    def activation_overview(self, **kwargs):
+        channel_spearman = []
+       
+        fig, ax = plt.subplots(1,1, figsize=(10,5), dpi=300)
+        ax.spines['top'].set_visible(False)
+        ax.spines['right'].set_visible(False)
+        
+        n_channel = self.popen.channel_ls[self.layer]
+        for i in tqdm(range(n_channel)):
+            act, (spr, pr)  = self.activation_density(i, to_print=False, feature_map=None, **kwargs);
+            channel_spearman.append(spr[0])
+        return fig, ax, np.array(channel_spearman)
+    
+    def matrix_to_seq(self, matrix):
+        """
+        return the sequence with maximum weight at each site
+        """
+        max_act_seq = ''
+        assert matrix.shape[1] == 4
+        for i in matrix.argmax(axis=1):
+            max_act_seq += ['A', 'C', 'G', 'T'][i]
+            
+        return max_act_seq
+    
+    def save_as_meme_format(self,channels:list, save_path, filter_prefix='filter', transformation='probability'):
+        motifs = []
+        success_channel = []
+        for cc in channels:
+            try:
+                region, index, patches = self.locate_MA_seq(channel=cc)
+                M = self.sequence_to_matrix(region, transformation=transformation);
+                motifs.append(M)
+                success_channel.append(cc)
+            except ValueError:
+                continue
+                
+        write_meme(success_channel, motifs ,save_path, filter_prefix)
+    
+    def gradience_scaler(self,array):
+        """
+        a scaling function to 
+        """
+        meann = np.mean(array)
+        ranges = array.max() -  array.min()
+        return (array-meann)/ranges
+    
+    def get_input_grad(self, task, focus=True, fm=None, starting_layer=None):
+        """
+        compute the gradience of Y over feature_map
+        fm : specify feature map
+        starting_layer: or None [0,3]
+        """
+        All_grad = []
+        current_layer = self.layer if starting_layer is None else starting_layer
+        model = self.load_model().to(self.popen.cuda_id)
+        model.train()
+        if fm is None:
+            fm = self.feature_map
+        
+        for start in tqdm(range(0, fm.shape[0], 64)):
+            # prepare input
+            minibatch = fm[start:start+64]
+            X = torch.as_tensor(minibatch, device=self.popen.cuda_id).float()
+            X.requires_grad=True
+
+            # forward
+            # x = model.soft_share.encoder[3](X)
+            out=X
+            for layer in model.soft_share.encoder[current_layer:]:
+                out = layer(out)
+            Z_t = torch.transpose(out, 1, 2)
+            h_prim,(c1,c2) = model.tower[task][0](Z_t)
+            out = model.tower[task][1](c2)
+
+            # auto grad
+            external_grad = torch.ones_like(out)
+            out.backward(gradient=external_grad,retain_graph=True)
+            grad = X.grad
+            
+            if focus:
+                indices = self.argmax_to_indeces(np.argmax(minibatch, axis=2))
+                grad = grad[indices].reshape(-1,256)
+            
+            All_grad.append( tensor_2_numpy(grad) )
+        
+        # concate each channel and average over input sequences
+        grad_ay = np.concatenate(All_grad, axis=0)#.mean(axis=0)
+
+        # return self.gradience_scaler(grad_ay)
+        return grad_ay  
+            
+    def argmax_to_indeces(self,index):
+        "index : of shape [batch, 256] , the result of "
+        s_index = []
+        ch_index = []
+        loc_index = []
+        for sample in range(index.shape[0]):
+            sam_loc = index[sample]
+            for ch, pos in enumerate(sam_loc):
+                s_index.append(sample)
+                ch_index.append(ch)
+                loc_index.append(pos)
+                
+        return (s_index, ch_index, loc_index)
+
+    def extract_max_seq_pattern(self, condition, n_clusters=6, n_patch=3000):
+        pattern = []
+        channel_source = []
+        for channel in np.where(condition)[0]:
+
+            try:
+                act_patch, flatten_seq, cluster_index = self.within_patch_clustering(channel, n_clusters=n_clusters, to_plot=False,n_patch=n_patch)
+            except ValueError:
+                continue
+
+            for i in range(n_clusters):
+                sub_cluster = act_patch[cluster_index==i]
+                if len(sub_cluster) > 0:
+                    matrix = self.sequence_to_matrix(sub_cluster)
+                    pattern.append(self.matrix_to_seq(matrix))
+                    channel_source.append(channel)
+            return pattern, channel_source
+
+def sum_occurrance(df, pattern):
+    pattern_occurance = []
+    for p in pattern:
+        pattern_occurance.append(np.sum([(p in utr) for utr in df.seq.values]))
+    return np.array(pattern_occurance)
+
+def generate_scramble_index(size  , N_1):
+    scramble_index=np.zeros((size,))
+
+    while scramble_index.sum() < N_1:    
+        num_ = int(N_1 - scramble_index.sum())
+        randindex = np.random.randint(0, size, size=(num_,))
+
+        for i in randindex:
+            scramble_index[i] 
+    return scramble_index    
+
+
+class merge_task_map(Maximum_activation_patch):
+    def __init__(self, popen, which_layer, n_patch , kfold_index=None, device_string='cpu'):
+        """merging all tasksa"""
+        super().__init__(popen, which_layer, n_patch , kfold_index, device_string)
+        self.old_popen = copy.copy(popen)
+        self.df_dict = {}
+        self.Y_2_task = {}
+        self.patches_ls = None
+    
+    def load_indexed_dataloader(self,task):
+        
+        if task in ['Andrev2015','muscle','pc3']:
+            self.popen.seq_col = 'utr'
+            self.popen.aux_task_columns = ['log_te']
+            self.popen.split_like = None
+            self.popen.kfold_cv = True
+        
+        loader_ls = super().load_indexed_dataloader(task)
+        self.popen = copy.copy(self.old_popen)
+        return loader_ls
+    
+    def extract_feature_map(self, which_set=0):
+        feature_map = {}
+        for task in self.popen.cycle_set:
+            feature_map[task] = super().extract_feature_map(task=task, which_set=which_set)
+            self.Y_2_task[task] = self.Y_ls
+        self.feature_map = feature_map
+        return feature_map
+    
+    def loading(self,task, which_set):
+        model = self.load_model().to(self.popen.cuda_id)
+        dataloader = self.load_indexed_dataloader(task)[which_set]
+        self.df_dict[task] = dataloader.dataset.df
+        return model, dataloader
+    
+    def locate_MA_seq(self, channel, feature_map=None, patches_ls=None):
+        """A multi-task version of MA region retrieve"""
+        regions = []
+        indeces = {}
+        patches = []
+        
+        if patches_ls is None:
+            patches_ls = self.patches_ls
+        
+        for i, task in enumerate(self.popen.cycle_set):
+            # task out : region, index, patches
+            if patches_ls is not None:
+                self.n_patch = patches_ls[i]
+                # print(f"{task} n_patch: {self.n_patch}")
+            self.df = self.df_dict[task]
+            # _region , index , patches
+            r, i, p = super().locate_MA_seq(channel, feature_map=self.feature_map[task])
+            regions.append(r)
+            indeces[task] = i
+            patches.append(p)
+            
+        return np.concatenate(regions) , indeces , np.concatenate(patches)
+    
+    def activation_density(self, channel, to_print=True, to_plot=True, feature_map=None, **kwargs):
+        all_F0 = []
+        all_spr = []
+        all_pr = []
+        for i, task in enumerate(self.popen.cycle_set):
+            # task out : region, index, patches
+            task_featmap = self.feature_map[task]
+            self.Y_ls = self.Y_2_task[task]
+            
+            if to_print:
+                print(task)
+            F0, (spr,pr) = super().activation_density(channel, to_print=to_print, 
+                                                  to_plot=False,
+                                                  feature_map = task_featmap)
+            all_F0.append(F0)
+            all_spr.append(spr)
+            all_pr.append(pr)
+        
+        # merge all activateion 
+        
+        if to_plot:
+            for F0_ay in all_F0:
+                sns.kdeplot(F0_ay, **kwargs)
+            
+        return all_F0 , (all_spr, all_pr)
+    
+    def flexible_n_patch(self,channel, qtl=0.95):
+        """
+        call this function before the all other functions will enable a flexible task-wise n_patch
+        
+        """
+        all_F0 , (all_spr, all_pr) = self.activation_density(channel,False, False,None)
+        thres = np.max([np.quantile(f0, qtl) for f0 in all_F0])
+        self.patches_ls = [np.sum(f0 > thres) for f0 in all_F0]
+        return self.patches_ls 
+        
+    def activation_overview(self):
+        pr_ay=[]
+        spr_ay=[]
+
+        # the number of convolution filters in layer 3
+        # the index is also 3 because input channel = 4 for layer 0
+        for i in tqdm(range(self.popen.channel_ls[self.layer])):
+            # 
+            _,(spr,pr) = self.activation_density(i,False,False)
+            pr_ay.append(pr)
+            spr_ay.append(spr)
+
+        # convert to ndarray
+        pr_ay = np.stack(pr_ay) 
+        spr_ay = np.stack(spr_ay) 
+        return pr_ay, spr_ay
+    
+    def save_as_meme_format(self,channels:list, save_path, filter_prefix='filter', transformation='probability', qtl=0.95, fix_patches_ls=None):
+        """
+        Save the position weight matrix as the meme-suite acceptale minimal motif format
+        """
+        
+        with open(save_path, 'w') as f:
+            f.write("MEME version 5.3.0\n\n")
+            f.write("ALPHABET= ACGT\n\n")
+            f.write("strands: + -\n\n")
+            f.write("Background letter frequencies\n")
+            f.write("A 0.25 C 0.25 G 0.25 T 0.25\n")
+
+            for cc in channels:
+                if fix_patches_ls is None:
+                    n_patches_ls = self.flexible_n_patch(cc, qtl=qtl)
+                else:
+                    n_patches_ls = fix_patches_ls
+                try:
+                    region, index, patches = self.locate_MA_seq(channel=cc, patches_ls=n_patches_ls)
+                    M = self.sequence_to_matrix(region, transformation=transformation);
+                    f.write('\n')
+                    f.write(f"MOTIF {filter_prefix}_{cc}\n")
+                    seq_len = len(region[0])
+                    f.write(f"letter-probability matrix: alength= 4 w= {seq_len} \n")
+                    for line in M.values:
+                        f.write(" "+line.__str__()[1:-1]+'\n')
+                except ValueError:
+                    continue
+                    
+            f.close()
+            print('writed')
+            
+    def get_input_grad(self, focus=True):
+        grads = {}
+        for task in self.popen.cycle_set:
+            grads[task] = super().get_input_grad(focus=True, task=task, fm=self.feature_map[task]) 
+        return grads
+
+class Maximum_activation_kmer(Maximum_activation_patch):
+    def __init__(self,popen, which_layer, n_patch, kfold_index):
+        super().__init__(popen, which_layer, n_patch, kfold_index)
+        self.virtual_pad = 0
+    
+    def compute_virtual_pad(self):
+        self.virtual_pad = 0
+        print('the virtual pad is', self.virtual_pad)
+    
+    def load_indexed_dataloader(self, task='kmer'):
+        self.csv_path = f"<DATA_DIR>/all_{self.r}mer.csv"
+        self.popen.split_like = [self.csv_path,self.csv_path]
+        self.popen.csv_path = None
+        self.popen.kfold_cv = False
+        self.popen.kfold_index = self.kfold_index
+        self.popen.pad_to = self.r
+        self.popen.seq_col = 'utr'
+        self.popen.aux_task_columns = ['rl']
+        return  reader.get_dataloader(self.popen) 
+        
+    def extract_feature_map(self):
+         
+        model, dataloader= self.loading('kmer', 1)
+        
+        for l in range(self.layer):
+            # remove padding to precisely locate
+            model.soft_share.encoder[l][0][0].padding = (0,)
+        
+        feature_map = []
+        X_ls = []
+        Y_ls = []
+        
+        model.eval()
+        with torch.no_grad():
+            for Data in tqdm(dataloader):
+                # iter each batch
+                x,y = train_val.put_data_to_cuda(Data,self.popen,False)
+                x = torch.transpose(x, 1, 2)
+#                 X_ls.append(x.numpy())
+                Y_ls.append( tensor_2_numpy(y) )
+                
+                for layer in model.soft_share.encoder[:self.layer]:
+                    out = layer(x)
+                    x = out
+                feature_map.append( tensor_2_numpy(out) )
+                
+                torch.cuda.empty_cache()
+            
+
+        feature_map_l = np.concatenate( feature_map, axis=0)
+        
+#         self.X_ls = np.concatenate(X_ls, axis=0)
+        self.Y_ls = np.concatenate(Y_ls, axis=0)
+
+        print("activation map of layer |%d|"%self.layer,feature_map_l.shape)
+#         print(self.X_ls.shape)
+        print("Y : ",self.Y_ls.shape)
+        self.feature_map = feature_map_l
+        
+        self.filters = self.get_filter_param(model)
+        
+        del model
+        self.df = pd.read_csv(self.csv_path)
+        return feature_map_l
+    
+def write_meme(channels:list, PWMs:list ,save_path, filter_prefix='filter'):
+    """
+    Save the position weight matrix as the meme-suite acceptale minimal motif format
+    """
+    assert len(channels)==len(PWMs)
+    with open(save_path, 'w') as f:
+        f.write("MEME version 5.4.1\n\n")
+        f.write("ALPHABET= ACGT\n\n")
+        f.write("strands: + -\n\n")
+        f.write("Background letter frequencies\n")
+        f.write("A 0.25 C 0.25 G 0.25 T 0.25\n")
+
+        for cc,M in zip(channels, PWMs):
+            f.write('\n')
+            f.write(f"MOTIF {filter_prefix}_{cc}\n")
+            seq_len = M.shape[0]
+            f.write(f"letter-probability matrix: alength= 4 w= {seq_len} \n")
+            for line in M.values:
+                f.write(" "+line.__str__()[1:-1]+'\n')
+
+
+        f.close()
+        print('writed to', save_path)
+        
+def extract_meme(memepath):
+    """
+    read the meme file and extract motifs to rewrite
+    """
+    with open(memepath,'r') as f:
+        all_lines = f.readlines()[9:]
+    
+    all_blocks = []
+    for i, line in enumerate(all_lines):
+        if line.startswith("MOTIF"):
+            width = re.match(r"letter-probability matrix: alength= 4 w= (\d)*",all_lines[i+1]).groups(1)
+            width = int(width[0])
+
+            all_blocks.append( all_lines[i:i+3+width])
+            
+    return all_blocks
+

model/src/mrl_te_optimization/models/popen.py

@@ -0,0 +1,150 @@
+import os,sys
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import numpy as np
+import utils
+import json
+from models import Modules
+import configparser
+import logging
+
+class Auto_popen(object):
+    def __init__(self,config_file):
+        """
+        read the config_fiel
+        """
+        # machine config path
+        self.shuffle = True
+        self.script_dir = utils.script_dir
+        self.data_dir = utils.data_dir
+        self.data_dir = '/mnt/sina/run/ml/gan/motif/MTtrans/test.csv'
+        self.log_dir = utils.log_dir
+        self.pth_dir = utils.pth_dir
+        self.set_attr_as_none(['te_net_l2','loss_fn','modual_to_fix','other_input_columns','pretrain_pth','kfold_index'])
+        self.split_like = False
+        self.loss_schema = 'constant'
+        
+        # transform to dict and convert to  specific data type
+        self.config = configparser.ConfigParser()
+        self.config.read(config_file)
+        self.config_file = config_file
+        self.config_dict = {item[0]: eval(item[1]) for item in self.config.items('DEFAULT')}
+        print(self.config_dict)
+        # assign some attr from config_dict         
+        self.set_attr_from_dict(self.config_dict.keys())
+        self.check_run_and_setting_name()                          # check run name
+        self._dataset = "_" + self.dataset if self.dataset != '' else self.dataset
+        # the saving direction
+        self.path_category = self.config_file.split('/')[-4]
+        self.vae_log_path = config_file.replace('.ini','.log')
+        
+
+        self.Resumable = False
+
+        # covariates for other input
+        self.n_covar = len(self.other_input_columns) if self.other_input_columns is not None else 0
+        
+        # generate self.model_args
+        self.get_model_config()
+
+    @property
+    def vae_pth_path(self):
+        save_to = os.path.join(self.pth_dir,self.model_type+self._dataset,self.setting_name)
+        if self.kfold_index is None:
+            pth = os.path.join(save_to, self.run_name + '-model_best.pth')
+        elif type(self.kfold_index) == int:
+            k = self.kfold_index
+            pth = os.path.join(save_to, self.run_name + f'-model_best_cv{k}.pth')
+        return pth
+    
+    @vae_pth_path.setter
+    def vae_pth_path(self, path):
+        self._vae_pth_path = path
+    
+    def set_attr_from_dict(self,attr_ls):
+        for attr in attr_ls:
+            self.__setattr__(attr,self.config_dict[attr])
+    
+    def set_attr_as_none(self,attr_ls):
+        for attr in attr_ls:
+            self.__setattr__(attr,None)
+
+    def check_run_and_setting_name(self):
+        file_name = self.config_file.split("/")[-1]
+        dir_name = self.config_file.split("/")[-2]
+        self.setting_name = dir_name
+        assert self.run_name == file_name.split(".")[0]
+        
+    def get_model_config(self):
+        """
+        assert we type in the correct model type and group them into model_args
+        """
+        
+        if self.model_type in dir(Modules):
+            self.Model_Class = eval("Modules.{}".format(self.model_type))
+        else:
+            raise NameError("not such model type")
+        
+        # conv_args define the soft-sharing part
+        conv_args = ["channel_ls","kernel_size","stride","padding_ls","diliation_ls","pad_to"]
+        self.conv_args = tuple([self.__getattribute__(arg) for arg in conv_args])
+        
+        # left args dfine the tower part in which the arguments are different among tasks
+        left_args={# Backbone models
+                    'RL_regressor':["tower_width","dropout_rate"],
+                    'RL_clf':["n_class","tower_width","dropout_rate"],
+                    'RL_gru':["tower_width","dropout_rate"],
+                    'RL_FACS': ["tower_width","dropout_rate"],
+                    'RL_hard_share':["tower_width","dropout_rate", "activation","cycle_set" ],
+                    'RL_covar_reg':["tower_width","dropout_rate", "activation", "n_covar", "cycle_set" ],
+                    'RL_covar_intercept':["tower_width","dropout_rate", "activation", "n_covar", "cycle_set" ],
+                    'RL_mish_gru':["tower_width","dropout_rate"],
+                    # GP models
+                    'GP_net': ['tower_width', 'dropout_rate', 'global_pooling', 'activation', 'cycle_set'],
+                    'Frame_GP': ['tower_width', 'dropout_rate', 'activation', 'cycle_set'],
+                    'RL_Atten': ['qk_dim', 'n_head', 'n_atten_layer', 'tower_width', 'dropout_rate', 'activation', 'cycle_set'],
+                    # Koo net
+                    'Conf_CNN' : ['pool_size'],
+                    }[self.model_type]
+        
+        self.model_args = [self.conv_args] + [self.__getattribute__(arg) for arg in left_args]
+        
+            
+    def check_experiment(self,logger):
+        """
+        check any unfinished experiment ?
+        """
+        log_save_dir = os.path.dirname(self.vae_log_path)
+        pth_save_dir = os.path.join(self.pth_dir,self.model_type+self._dataset,self.setting_name)
+        # make dirs 
+        if not os.path.exists(log_save_dir):
+            os.makedirs(log_save_dir)
+        if not os.path.exists(pth_save_dir):
+            os.makedirs(pth_save_dir)
+        
+        # check resume
+        if os.path.exists(self.vae_log_path) & os.path.exists(self.vae_pth_path):
+            self.Resumable = True
+            logger.info(' \t \t ==============<<<  Experiment detected  >>>============== \t \t \n')
+            
+    def update_ini_file(self,E,logger):
+        """
+        E  is the dict contain the things to update
+        """
+        # update the ini file
+        self.config_dict.update(E)
+        strconfig = {K: repr(V) for K,V in self.config_dict.items()}
+        self.config['DEFAULT'] = strconfig
+        
+        with open(self.config_file,'w') as f:
+            self.config.write(f)
+        
+        logger.info('   ini file updated    ')
+        
+    def chimera_weight_update(self):
+        # TODO : progressively update the loss weight between tasks
+        
+        # TODO : 1. scale the loss into the same magnitude
+        
+        # TODO : 2. update the weight by their own learning progress
+        
+        return None

model/src/mrl_te_optimization/models/reader.py

@@ -0,0 +1,512 @@
+import os
+import sys
+import copy
+sys.path.append(os.path.dirname(__file__))
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+
+import json
+import torch
+import numpy as np
+import pandas as pd
+from torch import nn
+
+from Bio import SeqIO
+from Bio.SeqRecord import SeqRecord
+
+from torch.utils.data import  DataLoader, Dataset ,random_split,IterableDataset
+from sklearn.model_selection import KFold,train_test_split
+from bucket_sampler import Bucket_Sampler
+from utils import Seq_one_hot
+
+global script_dir
+global data_dir
+
+with open(os.path.join(os.path.dirname(os.path.dirname(__file__)),"machine_configure.json"),'r') as f:
+    config = json.load(f)
+
+script_dir = config['script_dir']
+data_dir = config['data_dir']
+
+def one_hot(seq,complementary=False):
+    """
+    one_hot encoding on sequence
+    complementary: encode nucleatide into complementary one
+    """
+    # setting
+    seq = list(seq.replace("U","T"))
+    seq_len = len(seq)
+    complementary = -1 if complementary else 1
+    # compose dict
+    keys = ['A', 'C', 'G', 'T'][::complementary]
+    oh_dict = {keys[i]:i for i in range(4)}
+    # array
+    oh_array = np.zeros((seq_len,4))
+    for i,C in enumerate(seq):
+        try:
+            oh_array[i,oh_dict[C]]=1
+        except:
+            continue      # for nucleotide that are not in A C G T   
+    return oh_array 
+
+def read_bpseq(test_bpseq_path):
+    """
+    read bpseq file, extract sequence and ptable
+    """
+    with open(test_bpseq_path,'r') as f:
+        test_bpseq = f.readlines()
+        f.close()
+
+    for i in range(8):
+        if test_bpseq[i].startswith('1'):
+            start_index = i 
+
+    bp_seq = test_bpseq[start_index:]
+
+    seq = ''.join([line.strip().split(" ")[1] for line in bp_seq])
+    ptable = [int(line.strip().split(" ")[2]) for line in bp_seq]
+
+    assert len(ptable) == int(bp_seq[-1].split(" ")[0])
+
+    return seq,ptable
+
+def pad_zeros(X,pad_to):
+    """
+    zero padding at the right end of the sequence
+    """
+    if pad_to == 0:
+        # pad_to = 8*(X.shape[0]//8 + 1) + 1
+        try:
+            seq_len = X.shape[0]
+        except:
+            seq_len = X[0].shape[0]
+        pad_to = 3*(seq_len//3 + 1) + 1
+        
+    seq_len = X.shape[0] if isinstance(X, np.ndarray) else X[0].shape[0]
+    gap = pad_to - seq_len 
+    
+    #  here we change to padding ahead  , previously  nn.ZeroPad2d([0,0,0,gap])
+    pad_fn = nn.ZeroPad2d([0,0,gap,0])  #  (padding_left , padding_right , padding_top , padding_bottom )
+    # gap_array = np.zeros()
+    if isinstance(X, list):          # additional input
+        X_padded = [pad_fn(torch.tensor(X[0])), X[1]]
+    elif isinstance(X, np.ndarray):
+        X_padded = pad_fn(torch.tensor(X))
+    return X_padded
+
+def pack_seq(ds_zls:list,pad_to:int):
+    X_ts = [X for X,Y in ds_zls]
+    if pad_to == 0:
+        try:
+            max_len = np.max([X.shape[0] for X in X_ts])
+        except:
+            max_len = np.max([X[0].shape[0] for X in X_ts])
+        
+        pad_to = 3*(max_len//3 + 1) + 1
+    
+    if isinstance(X_ts[0],np.ndarray):
+        X_packed = torch.stack([pad_zeros(X=x,pad_to=pad_to) for x in X_ts])
+    else:
+        X_packed = [torch.stack([pad_zeros(X=x[0],pad_to=pad_to) for x in X_ts]), torch.tensor([x[1] for x in X_ts])]
+    Y_packed = torch.tensor(np.array([Y for X,Y in ds_zls]))
+    
+    return X_packed , Y_packed
+
+
+
+class mask_reader(Dataset):
+    def __init__(self,npy_path):
+        """
+        read the mask A549 sequence and with real sequence
+        """
+        self.data_set = np.load(npy_path)
+        self.X = self.data_set[:,0]
+        self.Y = self.data_set[:,1]
+    
+    def __len__(self):
+        return self.X.shape[0]
+    
+    def __getitem__(self,index):
+        return (self.X[index,:,:],self.Y[index,:,:])
+
+class mix_dataset(Dataset):
+    def __init__(self):
+        # dataset_ls = [mask_reader(os.path.join(data_dir,'mix_data',"mix_%s.npy"%set)) for set in ['train','val','test']]
+        raise ValueError("the dataset is not defined")
+
+class mask_dataset(Dataset):
+    def __init__(self):
+        raise ValueError("the dataset is not defined")
+
+class ribo_dataset(Dataset):
+    
+    def __init__(self,DF, pad_to, trunc_len=50, seq_col='utr',
+                      aux_columns='TE_count', other_input_columns=None):
+        """
+        Dataset to trancate sequence and return in one-hot encoding way
+        `dataset(DF,pad_to,trunc_len=50,seq_col='utr')`
+        ...DF: the dataframe contain sequence and its meta-info
+        ...pad_to: final size of the output tensor
+        ...trunc_len: maximum sequence to retain. number of nt preceding AUG
+        ...seq_col : which col of the DF contain sequence to convert
+        """
+        DF[seq_col] = DF[seq_col].astype(str)
+        self.df = DF
+        self.pad_to = pad_to
+        self.trunc_len = 0 if trunc_len is None else trunc_len
+        
+        # X and Y
+        self.seqs = self.df.loc[:,seq_col].values
+        self.other_input_columns = other_input_columns
+        self.Y = self.df.loc[:,aux_columns].values
+        
+    def __len__(self):
+        return self.df.shape[0]
+    
+    def __getitem__(self,i):
+        seq = self.seqs[i]
+        x_padded = self.seq_chunk_N_oh(seq)
+        input = x_padded
+        if self.other_input_columns is not None:
+                input = [x_padded]
+                for col in self.other_input_columns:
+                    input.append(self.df.loc[:,col].values[i]) 
+        y = self.Y[i]
+        return input,y
+    
+    def seq_chunk_N_oh(self,seq):
+        """
+        truncate the sequence and encode in one hot
+        """
+        if (len(seq) >  self.trunc_len)&(self.trunc_len >0):
+            seq = seq[-1* self.trunc_len:]
+        
+        X = one_hot(seq)
+        X = torch.tensor(X)
+
+        # X_padded = pad_zeros(X, self.pad_to)
+        
+        return X.float()
+        
+        
+class MTL_dataset(Dataset):
+    def __init__(self, DF, pad_to=100, seq_col='utr', aux_columns=None,
+                other_input_columns=None, trunc_len=None):
+        """
+        the dataset for Multi-task learning, will return sequence in one-hot encoded version, together with some auxilary task label
+        arguments:
+        ...csv_path: abs path of csv file, column `utr` should be in the csv
+        ...pad_to : maximum length of the sequences
+        ...columns : list  contains what  axuslary task label will be 
+        """
+        self.pad_to = pad_to
+        self.trunc_len = trunc_len
+        DF[seq_col] = DF[seq_col].astype(str)
+        self.df = DF     # read Df
+        self.seqs = self.df[seq_col].values       # take out all the sequence in DF
+        self.columns = aux_columns
+        self.other_input_columns = other_input_columns
+        
+        assert trunc_len == None, "MTL dataset does not support argument trunc len"
+                
+    def __len__(self):
+        return self.df.shape[0]
+    
+    def __getitem__(self,index):
+        seq = self.seqs[index]        # sequence: str of len 25~100
+        
+        X = one_hot(seq).astype(float)       # seq_oh : np array, one hot encoding sequence 
+        # X = pad_zeros(X)         # X  : torch.tensor , zero padded to 100
+        
+        if self.columns == None:
+            # which means no auxilary label is needed
+            item = X ,X
+        elif (type(self.columns) == list) & (len(self.columns)!=0):
+            # return what's in columns
+            aux_labels = self.df.loc[:,self.columns].values[index]
+            # if len(self.columns) == 1:
+            #     aux_labels = aux_labels.reshape(-1,1)
+            item = X ,aux_labels
+
+            if self.other_input_columns is not None:
+                input = []
+                for col in self.other_input_columns:
+                    input.append(self.df.loc[:,col].values[index]) 
+                item = (X, input),aux_labels
+
+        return item   
+
+class kmer_scan_dataset(Dataset):
+    def __init__(self, DF, seq_col, kmer_size, aux_columns=None, *args):
+        super().__init__()
+        self.df = DF
+        self.seqs = DF[seq_col].values
+        self.rls = DF[aux_columns].values
+        self.k = kmer_size
+        self.kmer_sets = self.create_kmer()
+        self.kmer_lookup = { kmer : i for i, kmer in enumerate(self.kmer_sets) }
+
+    def __len__(self):
+        return self.df.shape[0]
+
+    def __getitem__(self, index):
+        seq  = self.seqs[index]
+        y = self.rls[index]
+
+        x = self.scan_kmer(seq)
+        return torch.tensor(x.astype(float)), torch.tensor(y)
+    
+    def scan_kmer(self, seq):
+
+        mat = np.zeros((len(self.kmer_lookup), len(seq)-self.k))
+        for i in range(0, len(seq)-self.k):
+            seq_let = seq[i:i+self.k]
+            mat[self.kmer_lookup[seq_let], i]  = 1
+
+        return mat
+
+    
+    def create_kmer(self):
+        all_kmer = {0:['']}
+        k = 1
+        while k <= self.k:
+            k_mer = []  # 1 ; 4 ; 2: 4**2 ...
+            for source in all_kmer[k-1]:
+                k_mer += [source + base for base in ['A','C','G','T']]
+
+            assert len(k_mer) == 4**k, f"new kmers {len(k_mer)}, not equal to {4**k}"
+            all_kmer[k] = k_mer
+            k += 1
+
+        all_kmer.pop(0)
+        return all_kmer[self.k]
+
+
+
+
+class FASTA_dataset(MTL_dataset):
+    def __init__(self,DF, pad_to,trunc_len, seq_col, aux_columns, other_input_columns):
+        super().__init__(DF=DF, pad_to=pad_to, trunc_len=trunc_len, seq_col=seq_col, 
+                         aux_columns=aux_columns, other_input_columns=other_input_columns)
+
+def get_splited_dataloader(dataset_func, df_ls, ratio:list, batch_size, shuffle, pad_to, seed=42,return_dataset=False):
+    """
+    split the total dataset into train val test, and return in a DataLoader (train_loader,val_loader,test_loader) 
+    dataset : the defined <UTR_dataset>
+    ratio : the ratio of train : val : test
+    batch_size : int
+    """
+
+    #  determined set-ls
+    
+    set_ls = [dataset_func(df) for df in df_ls]
+    if return_dataset:
+        return set_ls
+    
+    if pad_to == 0 :
+        # automatic padding to `8x -1`
+        # and we will pad the sequence of similar length with bucket sampler
+        col_fn = lambda x: pack_seq(x,0)
+        sampler_ls = [{"batch_sampler":Bucket_Sampler(df,batch_size=batch_size,drop_last=True),
+                       "num_workers":4,
+                       "drop_last":False,
+                       "collate_fn":col_fn} for df in df_ls]
+        #  wrap dataset to dataloader
+        loader_ls = [DataLoader(subset,**kwargs) for subset,kwargs in zip(set_ls,sampler_ls)]
+    else:
+        col_fn = lambda x: pack_seq(x,pad_to)
+        loaderargs = {"batch_size": batch_size,
+                        "generator":torch.Generator().manual_seed(42),
+                        "drop_last":False ,
+                        "num_workers":4,
+                        "shuffle":shuffle,
+                        "collate_fn":col_fn}
+        #  wrap dataset to dataloader
+        loader_ls = [DataLoader(subset,**loaderargs) for subset in set_ls]
+        
+    if len(loader_ls) == 2:
+        # a complement of empty test set
+        loader_ls.append(None) 
+        
+    return loader_ls
+
+def split_DF(data_path,split_like,ratio, kfold_cv, kfold_index=None,seed=42):
+    
+    class _cf_data(object):
+        def __init__(self,data_path):
+            self.path = data_path
+            self.is_fasta = data_path.endswith('.fasta')
+            self.__read__()
+    
+        def __read__(self):
+            if self.is_fasta:
+                self.data = list(SeqIO.parse(self.path,'fasta'))
+            else:
+                self.data = pd.read_csv(os.path.join(data_dir,self.path),low_memory=False)
+        
+        def __len__(self):
+            return len(self.data)
+                
+        def _slice_(self, indices):
+            if isinstance(self.data, list) & self.is_fasta:
+                return np.array(self.data)[indices]
+            elif isinstance(self.data, pd.DataFrame):
+                return self.data.iloc[indices]
+            
+        
+    if kfold_cv == True:
+        full_df = pd.read_csv(os.path.join(data_dir,data_path),low_memory=False)
+        # K-fold CV : 8:1:1 for each partition
+        df_ls = KFold_df_split(full_df,kfold_index)
+        
+    if kfold_cv == 'train_val':
+        train_val, test = [_cf_data(data_path).data for data_path in split_like]
+        # K-fold CV : 8:1:1 for each partition
+        train, val = KFold_1_split(train_val,kfold_index)
+        df_ls = [train, val , test]
+    
+    elif type(split_like) == list:
+        df_ls = [_cf_data(data_path).data for data_path in split_like]
+
+    elif type(kfold_cv) == str:
+        full_df = pd.read_csv(os.path.join(data_dir,data_path),low_memory=False)
+        if kfold_cv in full_df.columns:
+            kfold_index = str(kfold_index) 
+            test = full_df.query(f"`{kfold_cv}` == @kfold_index")
+            train_val = full_df.query(f"`{kfold_cv}` != @kfold_index")
+            train, val = train_test_split(train_val, test_size=0.05)
+            df_ls = [train, val, test]
+    
+    else:
+        full_df = _cf_data(data_path)
+        # POPEN.ratio will determine train :val :test ratio
+        total_len = len(full_df)
+        lengths = [int(total_len*sub_ratio) for sub_ratio in ratio[:-1]]
+        lengths.append(total_len-sum(lengths))         # make sure the sum of length is the total len
+
+        set_ls = random_split(full_df,lengths,generator=torch.Generator().manual_seed(seed)) 
+        df_ls = [full_df._slice_(subset.indices) for subset in set_ls] # df.iloc [ idx ]
+        
+    return df_ls
+    
+def split_DF_call(data_path,split_like,ratio, kfold_cv, kfold_index=None,seed=42):
+    
+    class _cf_data(object):
+        def __init__(self,data_path):
+            self.path = data_path
+            self.is_fasta = data_path.endswith('.fasta')
+            self.__read__()
+    
+        def __read__(self):
+            if self.is_fasta:
+                self.data = list(SeqIO.parse(self.path,'fasta'))
+            else:
+                self.data = pd.read_csv(os.path.join(data_dir,self.path),low_memory=False)
+        
+        def __len__(self):
+            return len(self.data)
+                
+        def _slice_(self, indices):
+            if isinstance(self.data, list) & self.is_fasta:
+                return np.array(self.data)[indices]
+            elif isinstance(self.data, pd.DataFrame):
+                return self.data.iloc[indices]
+            
+        
+    if kfold_cv == True:
+        full_df = pd.read_csv(os.path.join(data_dir,data_path),low_memory=False)
+        # K-fold CV : 8:1:1 for each partition
+        df_ls = KFold_df_split(full_df,kfold_index)
+        
+    if kfold_cv == 'train_val':
+        train_val, test = [_cf_data(data_path).data for data_path in split_like]
+        # K-fold CV : 8:1:1 for each partition
+        train, val = KFold_1_split(train_val,kfold_index)
+        df_ls = [train, val , test]
+    
+    elif type(split_like) == list:
+        df_ls = [_cf_data(data_path).data for data_path in split_like]
+
+    elif type(kfold_cv) == str:
+        full_df = pd.read_csv(os.path.join(data_dir,data_path),low_memory=False)
+        if kfold_cv in full_df.columns:
+            kfold_index = str(kfold_index) 
+            test = full_df.query(f"`{kfold_cv}` == @kfold_index")
+            train_val = full_df.query(f"`{kfold_cv}` != @kfold_index")
+            train, val = train_test_split(train_val, test_size=0.05)
+            df_ls = [train, val, test]
+    
+    else:
+        full_df = _cf_data(data_path)
+        # POPEN.ratio will determine train :val :test ratio
+        total_len = len(full_df)
+        lengths = [int(total_len*sub_ratio) for sub_ratio in ratio[:-1]]
+        lengths.append(total_len-sum(lengths))         # make sure the sum of length is the total len
+
+        set_ls = random_split(full_df,lengths,generator=torch.Generator().manual_seed(seed)) 
+        df_ls = [full_df._slice_(subset.indices) for subset in set_ls] # df.iloc [ idx ]
+        
+    return df_ls
+
+
+def KFold_df_split(df,K,**kfoldargs):
+    """
+    split the dataset DF in a ratio of 8:1:1 , train:val:test in the framework of  K-fold CV 
+    set random seed = 43
+    arguments:
+    df : the `pd.DataFrame` object containing all data info
+    K : [0,4] , the index of subfold 
+    """
+    
+    # K-fold partition : n_splits=5
+    fold_index = list(KFold(10,shuffle=True,random_state=42).split(df))
+    train_val_index, test_index = fold_index[K]  
+    # the first 4/5 part of it is train set
+     
+    # index the df
+    train_val_df = df.iloc[train_val_index]
+    test_df = df.iloc[test_index]
+    
+    # the remaining 1/5 data will further break into val and test
+    train_df, val_df = train_test_split(train_val_df,test_size=0.05,random_state=42)
+    
+    return [train_df,val_df,test_df]   
+    
+def KFold_1_split(df,K,**kfoldargs):
+    """
+    for dataset with standard test set, where only train_val can be splited
+    """
+    fold_index = list(KFold(10,shuffle=True,random_state=42).split(df))
+    train_index, val_index = fold_index[K]  
+    # the first 4/5 part of it is train set
+     
+    # index the df
+    train = df.iloc[train_index]
+    val = df.iloc[val_index]
+
+    return train , val
+    
+def get_dataloader(POPEN):
+    """
+    wrapper
+    """
+    
+    # POPEN.csv_path = "/mnt/sina/run/ml/gan/motif/MTtrans/test.csv"
+    # print(POPEN.csv_path)
+    # print(POPEN.split_like)
+    # print(POPEN.train_test_ratio)
+    # print(POPEN.kfold_cv)
+    # print(POPEN.kfold_index)
+    df_ls = split_DF(POPEN.csv_path,POPEN.split_like,POPEN.train_test_ratio,POPEN.kfold_cv,POPEN.kfold_index,seed=42)
+    
+
+    DS_Class = eval(POPEN.dataset+"_dataset")
+         
+    dataset_func = lambda x : DS_Class(x,pad_to=POPEN.pad_to,trunc_len=POPEN.trunc_len,
+                                         seq_col=POPEN.seq_col, aux_columns=POPEN.aux_task_columns,
+                                         other_input_columns=POPEN.other_input_columns)
+    
+    loader_ls = get_splited_dataloader(dataset_func, df_ls,ratio=POPEN.train_test_ratio,
+                                        batch_size=POPEN.batch_size, shuffle=POPEN.shuffle,
+                                        pad_to=POPEN.pad_to, seed=42) # new function
+        
+    return loader_ls 

model/src/mrl_te_optimization/models/train_val.py

@@ -0,0 +1,337 @@
+import os
+import sys
+import time
+import numpy as np
+import pandas as pd
+import logging
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import utils
+import torch
+from scipy.stats import pearsonr
+from torch import optim
+from sklearn.metrics import r2_score, f1_score, roc_auc_score
+from models.ScheduleOptimizer import ScheduledOptim , find_lr
+from models.loss import Dynamic_Task_Priority as DTP
+
+def train(dataloader,model,optimizer,popen,epoch,lr=None, verbose=True):
+
+    logger = logging.getLogger("VAE")
+    loader_len = len(dataloader)       # number of iteration
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    model = model.to(device)
+    
+    model.train()
+    verbose_list=[]
+    verbose_df = pd.DataFrame()
+    for idx,data in enumerate(dataloader):
+        
+        X,Y = put_data_to_cuda(data,popen,require_grad=True)
+        optimizer.zero_grad()
+        
+        out = model(X)
+       
+        loss_dict = model.compute_loss(out,X,Y,popen)
+        loss = loss_dict['Total']
+        acc_dict = model.compute_acc(out,X,Y,popen)
+        loss_dict.update(acc_dict) # adding the acc ditc into loss dict
+         
+        loss.backward()        
+        optimizer.step()
+        
+        # ====== update lr =======
+        if type(optimizer) == ScheduledOptim:
+            lr = optimizer.update_learning_rate()      # see model.optim.py
+            loss_dict["lr"]=lr
+        elif popen.optimizer == 'Adam':
+            lr = optimizer.param_groups[0]['lr']
+        if popen.loss_schema != 'constant':
+            popen.chimerla_weight = popen.loss_schedualer._update(loss_dict)
+            for t in popen.tasks:
+                loss_dict[popen.loss_schema+"_wt_"+t] = popen.chimerla_weight[t]
+                
+        with torch.no_grad():   
+            loss_dict= utils.clean_value_dict(loss_dict)
+            verbose_list.append(loss_dict)
+       
+        # ======== verbose ========
+        # record result 5 times for a epoch
+        if verbose:
+            if idx % int(loader_len/5) == 0:
+                
+                # plot that in loss dict
+                loss_dict_keys = list(loss_dict.keys())
+            
+                train_verbose = "{:5d} / {:5d} ({:s}%):"
+                verbose_args = [idx,loader_len,str(int(idx/loader_len*100)).zfill(3)]
+                for key in loss_dict_keys:
+                    train_verbose += "\t %s:{:.7f}"%key
+                    verbose_args.append(loss_dict[key])
+                
+                # plot the cumulative mean total loss
+                short_batch_df = pd.json_normalize(verbose_list) # this will be cumulative
+                mean_total = short_batch_df.loc[:,'Total'].mean()
+                train_verbose += "\t %s:{:.7f}"%"Mean_Total"
+                verbose_args.append(mean_total)
+                
+                train_verbose = train_verbose.format(*verbose_args)                         
+            
+                logger.info(train_verbose)
+        if popen.cuda_id != torch.device('cpu'):
+            with torch.cuda.device(popen.cuda_id):
+                torch.cuda.empty_cache()
+
+def validate(dataloader,model,popen,epoch):
+
+    logger = logging.getLogger("VAE")
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    model = model.to(device)
+    model.teacher_forcing = False    # turn off teacher_forcing
+    
+    # ====== set up empty =====
+    verbose_list=[]
+    Y_ls = []
+    pred_ls = []
+    metric_dict = {}
+    # ======== evaluate =======
+    model.eval()
+    with torch.no_grad():
+        for idx,data in enumerate(dataloader):
+            X,Y = put_data_to_cuda(data,popen,require_grad=False)
+            Y_ls.append(Y.cpu().numpy())
+            out = model(X)
+            pred_ls.append(out.cpu().numpy())
+            loss_dict = model.compute_loss(out,X,Y,popen)
+            loss = loss_dict['Total']
+            acc_dict = model.compute_acc(out,X,Y,popen)
+            loss_dict.update(acc_dict) 
+            
+            loss_dict = utils.clean_value_dict(loss_dict)  # convert possible torch to single item
+            verbose_list.append(loss_dict)
+
+            if popen.cuda_id != torch.device('cpu'):   
+                with torch.cuda.device(popen.cuda_id):  
+                    torch.cuda.empty_cache()
+            
+          # # average among batch
+    
+    # ======== verbose ========
+    Y_ay = np.concatenate(Y_ls,axis=0).flatten()
+    pred_ay = np.concatenate(pred_ls,axis=0).flatten()
+
+    if popen.model_type == 'RL_clf':
+        metric_dict["F1"] = f1_score(Y_ay, pred_ay>0.5, average='binary')
+        metric_dict["AUROC"] = roc_auc_score(Y_ay, pred_ay)
+    else:
+        metric_dict[f"r2"] = r2_score(Y_ay, pred_ay)
+        metric_dict[f"pr"] = pearsonr(Y_ay, pred_ay)[0]
+    
+    verbose_df = pd.json_normalize(verbose_list)
+        
+    val_verbose = ""
+    verbose_args = []
+    verbose_dict = {key:verbose_df[key].mean() for key in verbose_df.columns}
+    verbose_dict.update(metric_dict)
+
+    for key,values in verbose_dict.items():
+        val_verbose += "\t %s:{:.7f}"%key
+        verbose_args.append(values)
+
+    val_verbose = val_verbose.format(*verbose_args)                         
+
+    logger.info(val_verbose)
+    
+    # what avg acc return : mean of  RL_Acc , Recons_Acc, Motif_Acc
+    acc_col = list(acc_dict.keys())
+    Avg_acc = np.mean(verbose_df.loc[:,acc_col].mean(axis=0))  
+    
+    # return these to save current performance
+    return (verbose_df['Total'].mean(),Avg_acc) if 'RL_loss' not in verbose_df.keys() else (verbose_df['RL_loss'].mean(),verbose_df['RL_Acc'].mean())
+
+def iter_train(loader_dict, model, optimizer, popen, epoch, verbose=True):
+
+    logger = logging.getLogger("VAE")
+    # loader_len = len(dataloader)       # number of iteration
+    all_len = [len(loader[0]) for loader in loader_dict.values()]
+    max_len = np.max(all_len)
+    n_task = len(popen.cycle_set)
+    total_len = max_len*n_task
+    
+    all_train = {task : iter(loader[0]) for task,loader in loader_dict.items()}
+    def try_next(all_train, task):
+        try:
+            data = next(all_train[task])
+            return data
+        except StopIteration:
+            all_train[task] = iter(loader_dict[task][0])
+            data = next(all_train[task])
+            return data
+    
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model = model.to(device)
+    model.train()
+    
+    verbose_list=[]
+    for idx in range(total_len):
+        task = popen.cycle_set[idx%n_task]
+        model.task = task
+        
+        data = try_next(all_train, task)
+        
+        X,Y = put_data_to_cuda(data,popen,require_grad=True)
+        optimizer.zero_grad()
+        
+        out = model(X)
+        loss_dict = model.compute_loss(out,X,Y,popen)
+        loss = loss_dict['Total']
+        acc_dict = model.compute_acc(out,X,Y,popen)
+        loss_dict.update(acc_dict) # adding the acc ditc into loss dict
+         
+        loss.backward()        
+        optimizer.step()
+        
+        # ====== update lr =======
+        if type(optimizer) == ScheduledOptim:
+            lr = optimizer.update_learning_rate()      # see model.optim.py
+            loss_dict["lr"]=lr
+        elif popen.optimizer == 'Adam':
+            lr = optimizer.param_groups[0]['lr']
+        if popen.loss_schema != 'constant':
+            popen.chimerla_weight = popen.loss_schedualer._update(loss_dict)
+            for t in popen.tasks:
+                loss_dict[popen.loss_schema+"_wt_"+t] = popen.chimerla_weight[t]
+                
+        with torch.no_grad():   
+            loss_dict= utils.clean_value_dict(loss_dict)
+            verbose_list.append(loss_dict)
+       
+        # ======== verbose ========
+        # record result 5 times for a epoch
+        if verbose:
+            if idx % int(total_len/5) == 0:
+                
+                # plot that in loss dict
+                loss_dict_keys = list(loss_dict.keys())
+            
+                train_verbose = "{:5d} / {:5d} ({:s}%):"
+                verbose_args = [idx,total_len,str(int(idx/total_len*100)).zfill(3)]
+                for key in loss_dict_keys:
+                    train_verbose += "\t %s:{:.7f}"%key
+                    verbose_args.append(loss_dict[key])
+                
+                # plot the cumulative mean total loss
+                short_batch_df = pd.json_normalize(verbose_list) # this will be cumulative
+                mean_total = short_batch_df.loc[:,'Total'].mean()
+                train_verbose += "\t %s:{:.7f}"%"Mean_Total"
+                verbose_args.append(mean_total)
+                
+                train_verbose = train_verbose.format(*verbose_args)                         
+            
+                logger.info(train_verbose)
+
+        if popen.cuda_id != torch.device('cpu'):
+            with torch.cuda.device(popen.cuda_id):
+                torch.cuda.empty_cache()
+
+def cycle_validate(loader_dict, model, optimizer, popen, epoch , which_set=1, return_=False):
+
+    logger = logging.getLogger("VAE")
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    model = model.to(device)
+    
+    # ====== set up empty =====
+    # model.loss_dict_keys = ['RL_loss', 'Recons_loss', 'Motif_loss', 'Total', 'RL_Acc', 'Recons_Acc', 'Motif_Acc']
+    verbose_list=[]
+    r2_dict = {}
+    Y_n_pred = {}
+    # ======== evaluate =======
+    
+    
+    for subset, dataloader in loader_dict.items():
+        
+        # fix 
+        model.task = subset
+        # # logger.info("        =======================|     fix      |=======================        ")
+        # model = utils.fix_parameter(model, popen.modual_to_fix[0])
+        # train(dataloader[0], model, optimizer, popen, epoch, verbose=True)
+        Y_ls = []
+        pred_ls = []
+        with torch.no_grad():
+            model.eval()
+            for idx,data in enumerate(dataloader[which_set]):
+                X,Y = put_data_to_cuda(data,popen,require_grad=False)
+                out = model(X)
+
+                Y_ls.append(Y.detach().cpu().numpy())
+                pred_ls.append(out.detach().cpu().numpy())
+
+                loss_dict = model.compute_loss(out,X,Y,popen)
+                loss_dict['%s_loss'%subset] = loss_dict['Total']
+                acc_dict = model.compute_acc(out,X,Y,popen)
+                loss_dict.update(acc_dict) 
+                
+                loss_dict = utils.clean_value_dict(loss_dict)  # convert possible torch to single item
+                verbose_list.append(loss_dict)
+                
+                if popen.cuda_id != torch.device('cpu'):
+                    with torch.cuda.device(popen.cuda_id):  
+                        torch.cuda.empty_cache()
+        
+        Y_ay = np.concatenate(Y_ls,axis=0).flatten()
+        pred_ay = np.concatenate(pred_ls,axis=0).flatten()
+        if return_:
+            Y_n_pred[subset] = (Y_ay, pred_ay)
+        r2_dict[f"{subset}_r2"] = r2_score(Y_ay, pred_ay)
+        r2_dict[f"{subset}_pr"] = pearsonr(Y_ay, pred_ay)[0]
+
+          # # average among batch
+    
+    # ======== verbose ========
+    
+    verbose_df = pd.json_normalize(verbose_list)
+        
+    val_verbose = ""
+    verbose_args = []
+    verbose_dict = {key:verbose_df[key].mean() for key in verbose_df.columns}
+    verbose_dict.update(r2_dict)
+    for key,values in verbose_dict.items():
+        val_verbose += "\t %s:{:.7f}"%key
+        verbose_args.append(values)
+        
+    val_verbose = val_verbose.format(*verbose_args)                         
+
+    logger.info(val_verbose)
+    
+    # what avg acc return : mean of  RL_Acc , Recons_Acc, Motif_Acc
+    acc_col = list(acc_dict.keys())
+    Avg_acc = np.mean(verbose_df.loc[:,acc_col].mean(axis=0))  
+    
+    if return_:
+        return verbose_dict, Y_n_pred
+    else:
+        return verbose_dict
+
+         
+def put_data_to_cuda(data,popen,require_grad=True):
+
+    X,y = data       
+    device = popen.cuda_id 
+    # for X is a list [seq, uAUG ....]
+    if popen.other_input_columns is not None:
+        X = [x_i.float().to(device) for x_i in X]
+        if require_grad:
+            for x_i in X:
+                x_i.required_grad = True
+
+    # X is not a list : seq
+    else:
+        X = X.float().to(device)
+        if require_grad:
+            X.required_grad = True  # check !!!
+    
+    Y = y.float().to(device)
+    
+    # Y = Y if X.shape == Y.shape else None  # for mask data
+    return X,Y

model/src/mrl_te_optimization/popen.py

@@ -0,0 +1,150 @@
+import os,sys
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import numpy as np
+import utils
+import json
+from models import Modules
+import configparser
+import logging
+
+class Auto_popen(object):
+    def __init__(self,config_file):
+        """
+        read the config_fiel
+        """
+        # machine config path
+        self.shuffle = True
+        self.script_dir = utils.script_dir
+        self.data_dir = utils.data_dir
+        self.log_dir = utils.log_dir
+        self.pth_dir = utils.pth_dir
+        self.set_attr_as_none(['te_net_l2','loss_fn','modual_to_fix','other_input_columns','pretrain_pth','kfold_index'])
+        self.split_like = False
+        self.loss_schema = 'constant'
+        
+        # transform to dict and convert to  specific data type
+        self.config = configparser.ConfigParser()
+        self.config.read(config_file)
+        self.config_file = config_file
+        self.config_dict = {item[0]: eval(item[1]) for item in self.config.items('DEFAULT')}
+        
+        # assign some attr from config_dict         
+        self.set_attr_from_dict(self.config_dict.keys())
+        self.check_run_and_setting_name()                          # check run name
+        self._dataset = "_" + self.dataset if self.dataset != '' else self.dataset
+        # the saving direction
+        self.path_category = self.config_file.split('/')[-4]
+        self.vae_log_path = config_file.replace('.ini','.log')
+        
+
+        self.Resumable = False
+
+        # covariates for other input
+        self.n_covar = len(self.other_input_columns) if self.other_input_columns is not None else 0
+        
+        # generate self.model_args
+        self.get_model_config()
+
+    @property
+    def vae_pth_path(self):
+        save_to = os.path.join(self.pth_dir,self.model_type+self._dataset,self.setting_name)
+        if self.kfold_index is None:
+            pth = os.path.join(save_to, self.run_name + '-model_best.pth')
+        elif type(self.kfold_index) == int:
+            k = self.kfold_index
+            pth = os.path.join(save_to, self.run_name + f'-model_best_cv{k}.pth')
+        return pth
+    
+    @vae_pth_path.setter
+    def vae_pth_path(self, path):
+        self._vae_pth_path = path
+    
+    def set_attr_from_dict(self,attr_ls):
+        for attr in attr_ls:
+            self.__setattr__(attr,self.config_dict[attr])
+    
+    def set_attr_as_none(self,attr_ls):
+        for attr in attr_ls:
+            self.__setattr__(attr,None)
+
+    def check_run_and_setting_name(self):
+        file_name = self.config_file.split("/")[-1]
+        dir_name = self.config_file.split("/")[-2]
+        self.setting_name = dir_name
+        assert self.run_name == file_name.split(".")[0]
+        # assert self.run_name == file_name.split(".")[0]
+        
+    def get_model_config(self):
+        """
+        assert we type in the correct model type and group them into model_args
+        """
+        self.model_type = 'RL_hard_share'
+        if self.model_type in dir(Modules):
+            self.Model_Class = eval("Modules.{}".format(self.model_type))
+        else:
+            raise NameError("not such model type")
+        
+        # conv_args define the soft-sharing part
+        conv_args = ["channel_ls","kernel_size","stride","padding_ls","diliation_ls","pad_to"]
+        self.conv_args = tuple([self.__getattribute__(arg) for arg in conv_args])
+        
+        # left args dfine the tower part in which the arguments are different among tasks
+        left_args={# Backbone models
+                    'RL_regressor':["tower_width","dropout_rate"],
+                    'RL_clf':["n_class","tower_width","dropout_rate"],
+                    'RL_gru':["tower_width","dropout_rate"],
+                    'RL_FACS': ["tower_width","dropout_rate"],
+                    'RL_hard_share':["tower_width","dropout_rate", "activation","cycle_set" ],
+                    'RL_covar_reg':["tower_width","dropout_rate", "activation", "n_covar", "cycle_set" ],
+                    'RL_covar_intercept':["tower_width","dropout_rate", "activation", "n_covar", "cycle_set" ],
+                    'RL_mish_gru':["tower_width","dropout_rate"],
+                    # GP models
+                    'GP_net': ['tower_width', 'dropout_rate', 'global_pooling', 'activation', 'cycle_set'],
+                    'Frame_GP': ['tower_width', 'dropout_rate', 'activation', 'cycle_set'],
+                    'RL_Atten': ['qk_dim', 'n_head', 'n_atten_layer', 'tower_width', 'dropout_rate', 'activation', 'cycle_set'],
+                    # Koo net
+                    'Conf_CNN' : ['pool_size'],
+                    }[self.model_type]
+        
+        self.model_args = [self.conv_args] + [self.__getattribute__(arg) for arg in left_args]
+        
+            
+    def check_experiment(self,logger):
+        """
+        check any unfinished experiment ?
+        """
+        log_save_dir = os.path.dirname(self.vae_log_path)
+        pth_save_dir = os.path.join(self.pth_dir,self.model_type+self._dataset,self.setting_name)
+        # make dirs 
+        if not os.path.exists(log_save_dir):
+            os.makedirs(log_save_dir)
+        if not os.path.exists(pth_save_dir):
+            os.makedirs(pth_save_dir)
+        
+        # check resume
+        if os.path.exists(self.vae_log_path) & os.path.exists(self.vae_pth_path):
+            self.Resumable = True
+            logger.info(' \t \t ==============<<<  Experiment detected  >>>============== \t \t \n')
+            
+    def update_ini_file(self,E,logger):
+        """
+        E  is the dict contain the things to update
+        """
+        # update the ini file
+        self.config_dict.update(E)
+        strconfig = {K: repr(V) for K,V in self.config_dict.items()}
+        self.config['DEFAULT'] = strconfig
+        
+        with open(self.config_file,'w') as f:
+            self.config.write(f)
+        
+        logger.info('   ini file updated    ')
+        
+    def chimera_weight_update(self):
+        # TODO : progressively update the loss weight between tasks
+        
+        # TODO : 1. scale the loss into the same magnitude
+        
+        # TODO : 2. update the weight by their own learning progress
+        
+        return None

model/src/mrl_te_optimization/script/Nupack_MFE_prediction.py

@@ -0,0 +1,48 @@
+#!/ssd/users/wergillius/.conda/envs/pytorch/bin/python
+import os
+import sys
+import numpy as np
+import pandas as pd
+import nupack
+import PATH
+import utils
+import scipy
+from tqdm import tqdm
+tqdm.pandas()
+
+my_model=nupack.Model(material='RNA')
+eGFP_seq = utils.eGFP_seq
+
+data_dir= utils.data_dir
+csv_name = sys.argv[1]
+csv_path = csv_name if os.path.exists(csv_name) else os.path.join(data_dir, csv_name)
+
+assert os.path.exists(csv_path), "csv not found"
+
+
+df = pd.read_csv(csv_path)
+seq_col = 'seq' if 'seq' in df.columns else 'utr'
+tr_col = 'rl' if 'rl' in df.columns else 'log_te'
+
+def cor_fun(col):
+    pearson_r = scipy.stats.pearsonr(df[col].values,
+                df[tr_col].values)[0]
+    return pearson_r
+
+nupck_engery_fn = lambda x: nupack.mfe(strands=[x], model=my_model)[0].energy
+df['nupack_MFE']  = df[seq_col].progress_apply(nupck_engery_fn)
+
+print("nupack_MFE : {}".format(cor_fun('nupack_MFE')))
+
+# nupck_engery_fn = lambda x: nupack.mfe(strands=[x + eGFP_seq[:10]], model=my_model)[0].energy
+# df['eGFP10_nupMFE']  = df[seq_col].progress_apply(nupck_engery_fn)
+
+# print("eGFP10_nupMFE : {}".format(cor_fun('eGFP10_nupMFE')))
+
+# nupck_engery_fn = lambda x: nupack.mfe(strands=[x + eGFP_seq[:50]], model=my_model)[0].energy
+# df['eGFP50_nupMFE']  = df[seq_col].progress_apply(nupck_engery_fn)
+
+# print("eGFP50_nupMFE : {}".format(cor_fun('eGFP50_nupMFE')))
+
+df.to_csv(csv_path, index=False)
+print(f"saved to {csv_path}")

model/src/mrl_te_optimization/script/PATH.py

@@ -0,0 +1,3 @@
+import os,sys
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+from utils import *

model/src/mrl_te_optimization/script/covar_test.py

@@ -0,0 +1,203 @@
+import os,sys
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import argparse
+
+parser = argparse.ArgumentParser('the main to train model')
+parser.add_argument('--config_file',type=str,required=True)
+parser.add_argument('--cuda',type=str,default=0,required=False)
+parser.add_argument("--kfold_index",type=int,default=1,required=False)
+args = parser.parse_args()
+
+cuda_id = args.cuda if args.cuda is not None else utils.get_config_cuda(args.config_file)
+os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_id)
+
+
+import time
+import torch
+import copy
+import utils
+from torch import optim
+import numpy as np 
+from models import reader,train_val
+from models.ScheduleOptimizer import ScheduledOptim,scheduleoptim_dict_str
+from models.popen import Auto_popen
+from models.loss import Dynamic_Task_Priority,Dynamic_Weight_Averaging
+
+
+POPEN = Auto_popen(args.config_file)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+POPEN.cuda_id = device
+
+
+POPEN.kfold_index = args.kfold_index
+if POPEN.kfold_cv:
+    if args.kfold_index is None:
+        raise NotImplementedError("please specify the kfold index to perform K fold cross validation")
+    POPEN.vae_log_path = POPEN.vae_log_path.replace(".log","_cv%d.log"%args.kfold_index)
+    #POPEN.vae_pth_path = POPEN.vae_pth_path.replace(".pth","_cv%d.pth"%args.kfold_index)
+    
+
+# Run name
+if POPEN.run_name is None:
+    run_name = POPEN.model_type + time.strftime("__%Y_%m_%d_%H:%M")
+else:
+    run_name = POPEN.run_name
+    
+# log dir
+logger = utils.setup_logs(POPEN.vae_log_path)
+logger.info(f"    ===========================| device {device}{cuda_id} |===========================    ")
+#  built model dir or check resume 
+POPEN.check_experiment(logger)
+
+#                               |=====================================|
+#                               |===========   setup  part  ==========|
+#                               |=====================================|
+# read data
+loader_set = {}          
+n_covar_dict = {}                                                                                                                                                                                                                                                                                                       
+base_path = ['cycle_train_val.csv', 'cycle_test.csv']
+base_csv = 'cycle_MTL_transfer.csv'
+for task in POPEN.cycle_set:
+    if (task in ['MPA_U', 'MPA_H', 'MPA_V', 'SubMPA_H']):
+        datapopen = Auto_popen('log/Backbone/RL_covar_intercept/3M/no_covar.ini')
+        datapopen.split_like = [path.replace('cycle', task) for path in base_path]
+        datapopen.kfold_index = args.kfold_index
+        n_covar_dict[task] = datapopen.n_covar
+
+    elif (task in ['RP_293T', 'RP_muscle', 'RP_PC3']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/3R/schedule_MTL.ini')
+        datapopen.csv_path = base_csv.replace("cycle",task)
+        datapopen.kfold_index = args.kfold_index
+        datapopen.pad_to = POPEN.pad_to
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.n_covar = POPEN.n_covar
+        n_covar_dict[task] = datapopen.n_covar
+
+    elif (task in ['pcr3', '293']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/karollus_RPs/rp_cycle.ini')
+        datapopen.csv_path = base_csv.replace("cycle",task)
+        datapopen.kfold_index = args.kfold_index
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.pad_to = POPEN.pad_to
+        datapopen.n_covar = POPEN.n_covar
+        n_covar_dict[task] = datapopen.n_covar
+    loader_set[task] = reader.get_dataloader(datapopen)
+
+POPEN.n_covar = n_covar_dict
+POPEN.get_model_config() # update model config
+    
+    
+# ===========  setup model  ===========
+# train_iter = iter(train_loader)
+# X,Y  = next(train_iter)
+# -- pretrain -- 
+if POPEN.pretrain_pth is not None:
+    # load pretran model
+    logger.info("===============================|   pretrain   |===============================")
+    logger.info(f" {POPEN.pretrain_pth}")
+    pretrain_popen = Auto_popen(os.path.join(utils.script_dir, POPEN.pretrain_pth))
+    if not os.path.exists(pretrain_popen.vae_pth_path):
+        if type(args.kfold_index) == int:
+            pretrain_popen.kfold_index = args.kfold_index
+    pretrain_model = torch.load(pretrain_popen.vae_pth_path, map_location=torch.device('cpu'))['state_dict']
+
+    
+    if POPEN.model_type == pretrain_popen.model_type:
+        # if not POPEN.Resumable:
+        #     # we only load pre-train for the first time 
+        #     # later we can resume 
+        model = pretrain_model.to(device)
+        del pretrain_model
+        
+        if (POPEN.cycle_set != pretrain_popen.cycle_set):
+            model.all_tasks = POPEN.cycle_set
+            model.tower = torch.nn.ModuleDict(
+                {POPEN.cycle_set[i] : model.tower[t]  for i, t in enumerate(pretrain_popen.cycle_set)}
+                                                )
+        
+    elif POPEN.modual_to_fix is not None:
+        # POPEN.model_type != pretrain_popen.model_type
+        
+        model = POPEN.Model_Class(*POPEN.model_args)
+        for modual in POPEN.modual_to_fix:
+            if modual in dir(pretrain_model):    
+                eval(f'model.{modual}').load_state_dict(
+                    eval(f'model.{modual}').state_dict()
+                    )
+
+        state_dict = {'epoch': 0,
+                        'validation_acc': 0,
+                        'state_dict': model.to('cpu'),
+                        'validation_loss': 0}
+        shared_pretrain_pth = POPEN.vae_pth_path.replace(f"_cv{args.kfold_index}", '')
+        if not os.path.exists(shared_pretrain_pth):
+            utils.snapshot(shared_pretrain_pth, state_dict)
+        utils.snapshot(POPEN.vae_pth_path, state_dict)
+
+        model = torch.load(POPEN.vae_pth_path, map_location=torch.device('cpu')) 
+        model = model.to(device)
+    
+# -- end2end -- 
+else:
+    Model_Class = POPEN.Model_Class  # DL_models.LSTM_AE 
+    
+    model = Model_Class(*POPEN.model_args).to(device)
+    
+if POPEN.Resumable:
+    model = utils.load_model(POPEN, model, logger)
+    
+# =========== fix parameters ===========
+if isinstance(POPEN.modual_to_fix, list):
+    for modual in POPEN.modual_to_fix:
+        model = utils.fix_parameter(model,modual)
+        model =  model.to(device)
+    logger.info(' \t \t ==============| %s fixed |============== \t \t \n'%POPEN.modual_to_fix)
+# =========== set optimizer ===========
+if POPEN.optimizer == 'Schedule':
+    optimizer = ScheduledOptim(optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                                      betas=(0.9, 0.98), 
+                                      eps=1e-09, 
+                                      weight_decay=1e-4, 
+                                      amsgrad=True),
+                           n_warmup_steps=20)
+elif type(POPEN.optimizer) == dict:
+    optimizer = eval(scheduleoptim_dict_str.format(**POPEN.optimizer))
+else:
+    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                            lr=POPEN.lr,
+                            betas=(0.9, 0.98), 
+                            eps=1e-09, 
+                            weight_decay=POPEN.l2)
+if POPEN.loss_schema == 'DTP':
+    POPEN.loss_schedualer = Dynamic_Task_Priority(POPEN.tasks,POPEN.gamma,POPEN.chimerla_weight)
+elif POPEN.loss_schema == 'DWA':
+    POPEN.loss_schedualer = Dynamic_Weight_Averaging(POPEN.tasks,POPEN.tau,POPEN.chimerla_weight)
+# =========== resume ===========
+best_loss = np.inf
+best_acc = 0
+best_epoch = 0
+previous_epoch = 0
+if POPEN.Resumable:
+    previous_epoch,best_loss,best_acc = utils.resume(POPEN, optimizer,logger)
+    epoch = previous_epoch
+    
+
+#                               |=====================================|
+#                               |==========  test  part ==========|
+#                               |=====================================|
+
+    
+
+logger.info("===============================| testing  |===============================")
+verbose_dict = train_val.cycle_validate(loader_set,model,optimizer,popen=POPEN,epoch=epoch, which_set=2)
+# matching task performance influence what to save
+
+
+
+if np.any(['r2' in key for key in verbose_dict.keys()]):
+    val_avg_acc = np.mean([values for key, values in verbose_dict.items() if 'r2' in key])
+    acc_dict = {f"cv{args.kfold_index}_{key}":values for key, values in verbose_dict.items() if 'r2' in key}
+else:
+    val_avg_acc = np.mean([values for key, values in verbose_dict.items() if 'acc' in key])
+    acc_dict = {}
+val_total_loss = verbose_dict['Total']

model/src/mrl_te_optimization/script/covar_train.py

@@ -0,0 +1,243 @@
+import os,sys
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import argparse
+
+parser = argparse.ArgumentParser('the main to train model')
+parser.add_argument('--config_file',type=str,required=True)
+parser.add_argument('--cuda',type=str,default=0,required=False)
+parser.add_argument("--kfold_index",type=int,default=1,required=False)
+args = parser.parse_args()
+
+cuda_id = args.cuda if args.cuda is not None else utils.get_config_cuda(args.config_file)
+os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_id)
+
+
+import time
+import torch
+import copy
+import utils
+from torch import optim
+import numpy as np 
+from models import reader,train_val
+from models.ScheduleOptimizer import ScheduledOptim,scheduleoptim_dict_str
+from models.popen import Auto_popen
+from models.loss import Dynamic_Task_Priority,Dynamic_Weight_Averaging
+
+
+POPEN = Auto_popen(args.config_file)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+POPEN.cuda_id = device
+
+
+POPEN.kfold_index = args.kfold_index
+if POPEN.kfold_cv:
+    if args.kfold_index is None:
+        raise NotImplementedError("please specify the kfold index to perform K fold cross validation")
+    POPEN.vae_log_path = POPEN.vae_log_path.replace(".log","_cv%d.log"%args.kfold_index)
+    #POPEN.vae_pth_path = POPEN.vae_pth_path.replace(".pth","_cv%d.pth"%args.kfold_index)
+    
+
+# Run name
+if POPEN.run_name is None:
+    run_name = POPEN.model_type + time.strftime("__%Y_%m_%d_%H:%M")
+else:
+    run_name = POPEN.run_name
+    
+# log dir
+logger = utils.setup_logs(POPEN.vae_log_path)
+logger.info(f"    ===========================| device {device}{cuda_id} |===========================    ")
+#  built model dir or check resume 
+POPEN.check_experiment(logger)
+
+#                               |=====================================|
+#                               |===========   setup  part  ==========|
+#                               |=====================================|
+# read data
+loader_set = {}          
+n_covar_dict = {}                                                                                                                                                                                                                                                                                                       
+base_path = ['cycle_train_val.csv', 'cycle_test.csv']
+base_csv = 'cycle_MTL_transfer.csv'
+for task in POPEN.cycle_set:
+    if (task in ['MPA_U', 'MPA_H', 'MPA_V', 'SubMPA_H']):
+        datapopen = Auto_popen('log/Backbone/RL_covar_intercept/3M/no_covar.ini')
+        datapopen.split_like = [path.replace('cycle', task) for path in base_path]
+        datapopen.kfold_index = args.kfold_index
+        n_covar_dict[task] = datapopen.n_covar
+
+    elif (task in ['RP_293T', 'RP_muscle', 'RP_PC3']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/3R/schedule_MTL.ini')
+        datapopen.csv_path = base_csv.replace("cycle",task)
+        datapopen.kfold_index = args.kfold_index
+        datapopen.pad_to = POPEN.pad_to
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.n_covar = POPEN.n_covar
+        n_covar_dict[task] = datapopen.n_covar
+
+    elif (task in ['pcr3', '293']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/karollus_RPs/rp_cycle.ini')
+        datapopen.csv_path = base_csv.replace("cycle",task)
+        datapopen.kfold_index = args.kfold_index
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.pad_to = POPEN.pad_to
+        datapopen.n_covar = POPEN.n_covar
+        n_covar_dict[task] = datapopen.n_covar
+    loader_set[task] = reader.get_dataloader(datapopen)
+
+POPEN.n_covar = n_covar_dict
+POPEN.get_model_config() # update model config
+    
+    
+# ===========  setup model  ===========
+# train_iter = iter(train_loader)
+# X,Y  = next(train_iter)
+# -- pretrain -- 
+if POPEN.pretrain_pth is not None:
+    # load pretran model
+    logger.info("===============================|   pretrain   |===============================")
+    logger.info(f" {POPEN.pretrain_pth}")
+    pretrain_popen = Auto_popen(os.path.join(utils.script_dir, POPEN.pretrain_pth))
+    if not os.path.exists(pretrain_popen.vae_pth_path):
+        if type(args.kfold_index) == int:
+            pretrain_popen.kfold_index = args.kfold_index
+    pretrain_model = torch.load(pretrain_popen.vae_pth_path, map_location=torch.device('cpu'))['state_dict']
+
+    
+    if POPEN.model_type == pretrain_popen.model_type:
+        # if not POPEN.Resumable:
+        #     # we only load pre-train for the first time 
+        #     # later we can resume 
+        model = pretrain_model.to(device)
+        del pretrain_model
+        
+        if (POPEN.cycle_set != pretrain_popen.cycle_set):
+            model.all_tasks = POPEN.cycle_set
+            model.tower = torch.nn.ModuleDict(
+                {POPEN.cycle_set[i] : model.tower[t]  for i, t in enumerate(pretrain_popen.cycle_set)}
+                                                )
+        
+    elif POPEN.modual_to_fix is not None:
+        # POPEN.model_type != pretrain_popen.model_type
+        
+        model = POPEN.Model_Class(*POPEN.model_args)
+        for modual in POPEN.modual_to_fix:
+            if modual in dir(pretrain_model):    
+                eval(f'model.{modual}').load_state_dict(
+                    eval(f'model.{modual}').state_dict()
+                    )
+
+        state_dict = {'epoch': 0,
+                        'validation_acc': 0,
+                        'state_dict': model.to('cpu'),
+                        'validation_loss': 0}
+        shared_pretrain_pth = POPEN.vae_pth_path.replace(f"_cv{args.kfold_index}", '')
+        if not os.path.exists(shared_pretrain_pth):
+            utils.snapshot(shared_pretrain_pth, state_dict)
+        utils.snapshot(POPEN.vae_pth_path, state_dict)
+
+        model = torch.load(POPEN.vae_pth_path, map_location=torch.device('cpu')) 
+        model = model.to(device)
+    
+# -- end2end -- 
+else:
+    Model_Class = POPEN.Model_Class  # DL_models.LSTM_AE 
+    
+    model = Model_Class(*POPEN.model_args).to(device)
+    
+if POPEN.Resumable:
+    model = utils.load_model(POPEN, model, logger)
+    
+# =========== fix parameters ===========
+if isinstance(POPEN.modual_to_fix, list):
+    for modual in POPEN.modual_to_fix:
+        model = utils.fix_parameter(model,modual)
+        model =  model.to(device)
+    logger.info(' \t \t ==============| %s fixed |============== \t \t \n'%POPEN.modual_to_fix)
+# =========== set optimizer ===========
+if POPEN.optimizer == 'Schedule':
+    optimizer = ScheduledOptim(optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                                      betas=(0.9, 0.98), 
+                                      eps=1e-09, 
+                                      weight_decay=1e-4, 
+                                      amsgrad=True),
+                           n_warmup_steps=20)
+elif type(POPEN.optimizer) == dict:
+    optimizer = eval(scheduleoptim_dict_str.format(**POPEN.optimizer))
+else:
+    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                            lr=POPEN.lr,
+                            betas=(0.9, 0.98), 
+                            eps=1e-09, 
+                            weight_decay=POPEN.l2)
+if POPEN.loss_schema == 'DTP':
+    POPEN.loss_schedualer = Dynamic_Task_Priority(POPEN.tasks,POPEN.gamma,POPEN.chimerla_weight)
+elif POPEN.loss_schema == 'DWA':
+    POPEN.loss_schedualer = Dynamic_Weight_Averaging(POPEN.tasks,POPEN.tau,POPEN.chimerla_weight)
+# =========== resume ===========
+best_loss = np.inf
+best_acc = 0
+best_epoch = 0
+previous_epoch = 0
+if POPEN.Resumable:
+    previous_epoch,best_loss,best_acc = utils.resume(POPEN, optimizer,logger)
+    
+
+#                               |=====================================|
+#                               |==========  training  part ==========|
+#                               |=====================================|
+for epoch in range(POPEN.max_epoch-previous_epoch+1):
+    epoch += previous_epoch
+    
+    #          
+    logger.info("===============================|    epoch {}   |===============================".format(epoch))
+
+
+    train_val.iter_train(loader_set,model=model,optimizer=optimizer,popen=POPEN,epoch=epoch)
+
+    #              -----------| validate |-----------   
+    logger.info("===============================| start validation |===============================")
+    verbose_dict = train_val.cycle_validate(loader_set,model,optimizer,popen=POPEN,epoch=epoch)
+    test_dict = train_val.cycle_validate(loader_set,model,optimizer,popen=POPEN,epoch=epoch, which_set=2)
+
+    if np.any(['r2' in key for key in verbose_dict.keys()]):
+        val_avg_acc = np.mean([values for key, values in verbose_dict.items() if 'r2' in key])
+    else:
+        val_avg_acc = np.mean([values for key, values in verbose_dict.items() if 'acc' in key])
+    val_total_loss = verbose_dict['Total']
+    
+    # matching task performance influence what to save
+    
+    
+    DICT ={"ran_epoch":epoch,"n_current_steps":optimizer.n_current_steps,"delta":optimizer.delta} if type(optimizer) == ScheduledOptim else {"ran_epoch":epoch}
+    POPEN.update_ini_file(DICT,logger)
+    
+    
+#    -----------| compare the result |-----------
+    if (best_loss > val_total_loss) :
+        # update best performance
+        best_loss = min(best_loss,val_total_loss)
+        best_acc = max(best_acc,val_avg_acc)
+        best_epoch = epoch
+        
+        # save
+        utils.snapshot(POPEN.vae_pth_path, {
+                    'epoch': epoch + 1,
+                    'validation_acc': val_avg_acc,
+                    # 'state_dict': model.state_dict(),
+                    'state_dict': model,
+                    'validation_loss': val_total_loss,
+                    'optimizer': optimizer.state_dict(),
+                })
+        
+        # update the popen
+        POPEN.update_ini_file({'run_name':run_name,
+                            "ran_epoch":epoch,
+                            "best_acc":best_acc},
+                            logger)
+        
+    elif (epoch - best_epoch >= 30)&((type(optimizer) == ScheduledOptim)):
+        optimizer.increase_delta()
+        
+    elif (epoch - best_epoch >= 60)&(epoch > POPEN.max_epoch/2):
+        # at the late phase of training
+        logger.info("<<<<<<<<<<< Early Stopping >>>>>>>>>>")
+        break

model/src/mrl_te_optimization/script/data_prepocessing.sh

@@ -0,0 +1,2 @@
+python ./preprocessing_MPA.py
+python ./preprocessing_RP.py

model/src/mrl_te_optimization/script/detect_changepoint_featmap.py

@@ -0,0 +1,146 @@
+import os
+import sys
+import pandas as pd
+import numpy as np
+import PATH
+import torch
+import argparse
+import seaborn as sns
+# 
+from models import reader
+from models import train_val
+from models.popen import Auto_popen
+from models import max_activation_patch as MAP
+# 
+from sklearn.linear_model import Lasso, Ridge, ElasticNet, LassoCV, RidgeCV, ElasticNetCV
+import warnings
+from scipy.cluster import hierarchy
+from matplotlib.backends.backend_pdf import PdfPages
+warnings.filterwarnings('ignore')
+
+
+#
+parser = argparse.ArgumentParser('the script to evlauate the effect of ')
+parser.add_argument("-c", "--config", type=str, required=True, help='the model config file: xxx.ini')
+parser.add_argument("-s", "--set", type=int, default=2, help='train - 0 ,val - 1, test - 2 ')
+parser.add_argument("-p", "--n_max_act", type=int, default=500, help='the number of seq')
+parser.add_argument("-k", "--kfold_cv", type=int, default=1, help='the repeat')
+parser.add_argument("-d", "--device", type=str, default='cpu', help='the device to use to extract featmap, digit or cpu')
+args = parser.parse_args()
+
+# args = parser.parse_args(["-c","/ssd/users/wergillius/Project/MTtrans/log/Backbone/RL_hard_share/3M/small_repective_filed_strides1113.ini",
+#     "-s","2",
+#     "-d", "1"])
+
+config_path = args.config
+save_path = config_path.replace(".ini", "_coef")
+config = Auto_popen(config_path)
+config.batch_size = 256
+config.shuffle = False
+config.kfold_cv = 'train_val'
+all_task = config.cycle_set
+
+task_channel_effect = {}
+task_performance = {}
+
+# path check
+if os.path.exists(config_path) and not os.path.exists(save_path):
+    os.mkdir(save_path)
+
+saved_pdf =os.path.join(save_path, 'changepoint_actmap.pdf')
+pp = PdfPages(saved_pdf)
+channel_cluster_task = {}
+for task in all_task:
+
+    # .... format featmap as data ....
+    print(f"\n\nevaluating for task: {task}")
+    # re-instance the map for each task
+    map_task = MAP.Maximum_activation_patch(popen=config, which_layer=4,
+                                      n_patch=args.n_max_act,
+                                      kfold_index=args.kfold_cv,
+                                      device_string=args.device)
+
+    # extract feature map and rl decision chain
+    featmap = map_task.extract_feature_map(task=task, which_set=args.set)
+    cum_rl_trend = map_task.cumulative_rl_decision(task=task, which_set=args.set)
+
+    # truncate the featmap and rl trend according to sequence length
+    seq_len_ls = map_task.df[config.seq_col].apply(len)
+    total_stride = np.product(map_task.strides)
+    to_stay = seq_len_ls//total_stride - 4
+    trunc_start = featmap.shape[2] - to_stay
+
+    # if max_seq_len == 50:
+    #     trunc_start = -12
+
+
+    # find the low rl sequences
+    rl_pred = cum_rl_trend[:,-1]
+    threshold = np.quantile(rl_pred, [0.05, 0.95])
+
+    low_rl = rl_pred < threshold[0]
+    high_rl = rl_pred > threshold[1]
+
+    # subset find low rl change point feature
+    lowrl_rl_chain = cum_rl_trend[low_rl]
+    lowrl_ft = featmap[low_rl]
+    lowrl_detect_region = [slice(s,None) for s in trunc_start[low_rl]]
+    changepoint_map = map_task.retrieve_featmap_at_changepoint(lowrl_ft, lowrl_rl_chain, 
+                                threshold=-1, direction='less', detect_region=lowrl_detect_region)
+    print(changepoint_map.shape)
+
+    # sample high rl feature
+    highrl_rl_chain = cum_rl_trend[high_rl]
+    highrl_ft = featmap[high_rl]
+    highrl_detect_region = [slice(s,None) for s in trunc_start[high_rl]]
+    background_map = map_task.retrieve_featmap_at_changepoint(highrl_ft, highrl_rl_chain, 
+                                threshold=0.5, direction='greater', detect_region=highrl_detect_region)
+    
+    # and then subsample
+    n_background = background_map.shape[0]
+    n_foreground = changepoint_map.shape[0]
+    if n_background > n_foreground:
+        downsample_seed = np.random.choice(np.arange(0,n_background), size=n_foreground)
+        background_map = background_map[downsample_seed]
+        n_background = background_map.shape[0]
+    
+    # concate and order two group
+    row_colors = ["#E09832"]*n_background + ["#192D48"]*n_foreground
+    all_act=np.concatenate([background_map, changepoint_map], axis=0)
+    feat_norm_act = (all_act - all_act.min(axis=0)) / (all_act.max(axis=0) - all_act.min(axis=0))
+
+
+    # Visualization
+    g=sns.clustermap(feat_norm_act, row_colors=row_colors);
+
+    # 'array', 'axis', 'calculate_dendrogram', 'calculated_linkage', 'data', 
+    # 'dendrogram', 'dependent_coord', 'independent_coord', 'label', 'linkage', 
+    # 'method', 'metric', 'plot', 'reordered_ind', 'rotate', 'shape', 'xlabel', 
+    # 'xticklabels', 'xticks', 'ylabel', 'yticklabels', 'yticks'
+
+    
+    Z_col=g.dendrogram_col.linkage
+    thres = 0.8*max(Z_col[:,2])
+    R = hierarchy.dendrogram(Z_col ,color_threshold=thres, truncate_mode=None,
+        above_threshold_color='#AAAAAA', p=10, orientation='top',ax=g.ax_col_dendrogram);
+    # R['leaves']
+    # R['ivl']
+
+    Z_row=g.dendrogram_row.linkage
+    thres = 0.8*max(Z_row[:,2])
+    R2 = hierarchy.dendrogram(Z_row ,color_threshold=thres, truncate_mode=None, orientation='left',
+        above_threshold_color='#AAAAAA', p=10, ax=g.ax_row_dendrogram);
+    g.ax_row_dendrogram.invert_yaxis()
+    g.figure.suptitle(task)
+
+    pp.savefig(g.figure)
+
+    channel_cluster_task[task] = pd.DataFrame(dict(zip(R['leaves'],R['leaves_color_list'])))
+    print("No..")
+pp.close()
+channel_cluster_df = pd.DataFrame(channel_cluster_task)
+saved_csv = os.path.join(save_path, "changepoint_channel_cluster.csv")
+channel_cluster_df.to_csv(saved_csv, index=False)
+print("==DONE==")
+print(f"result save to {saved_pdf}\n \t\t {saved_csv}")
+

model/src/mrl_te_optimization/script/download_training_data.sh

@@ -0,0 +1,40 @@
+DATA_DIR=$(cat machine_configure.json |grep data_dir|awk  -F ': "' '{print $2}'|awk -F '",' '{print $1}')
+
+if [ ! -d ${DATA_DIR} ]; then
+    echo "=============================================================="
+    echo "                   "
+    echo ${DATA_DIR} does not exist! Redirect to ./data
+    echo "                   "
+    echo "=============================================================="
+    DATA_DIR=$(pwd)/data
+    fi
+
+cd ${DATA_DIR}
+
+## download the ribosome profiling datasest 
+# RP-muscle
+wget https://raw.githubusercontent.com/zzz2010/5UTR_Optimizer/master/data/df_counts_and_len.TE_sorted.Muscle.with_annot.txt 
+# RP-293T
+wget https://raw.githubusercontent.com/zzz2010/5UTR_Optimizer/master/data/df_counts_and_len.TE_sorted.HEK_Andrev2015.with_annot.txt 
+# RP-PC3
+wget https://raw.githubusercontent.com/zzz2010/5UTR_Optimizer/master/data/df_counts_and_len.TE_sorted.pc3.with_annot.txt 
+# ref
+wget https://raw.githubusercontent.com/zzz2010/5UTR_Optimizer/master/data/gencode_v17_5utr_15bpcds.fa
+
+## download the Massively parallel report assay datasets
+# MPA_U
+wget ftp.ncbi.nlm.nih.gov/geo/samples/GSM3130nnn/GSM3130435/suppl/GSM3130435_egfp_unmod_1.csv.gz 
+# MPA_H
+wget ftp.ncbi.nlm.nih.gov/geo/samples/GSM3130nnn/GSM3130443/suppl/GSM3130443_designed_library.csv.gz 
+# MPA_V
+wget ftp.ncbi.nlm.nih.gov/geo/samples/GSM4084nnn/GSM4084997/suppl/GSM4084997_varying_length_25to100.csv.gz
+
+
+## dowlaod the yeast dataset
+wget ftp.ncbi.nlm.nih.gov/geo/samples/GSM2793nnn/GSM2793752/suppl/GSM2793752_Random_UTRs.csv.gz
+
+
+gzip -d *.gz
+
+echo "Finished ! All data ready~"
+echo "Please step to the pre-proccessing"

model/src/mrl_te_optimization/script/evaluate_channel_effect.py

@@ -0,0 +1,112 @@
+import os
+import sys
+import pandas as pd
+import numpy as np
+import PATH
+import torch
+import argparse
+# 
+from models import reader
+from models import train_val
+from models.popen import Auto_popen
+from models import max_activation_patch as MAP
+# 
+from sklearn.linear_model import LassoCV, RidgeCV, ElasticNetCV, LogisticRegressionCV
+import warnings
+warnings.filterwarnings('ignore')
+
+
+#
+parser = argparse.ArgumentParser('the script to evlauate the effect of ')
+parser.add_argument("-c", "--config", type=str, required=True, help='the model config file: xxx.ini')
+parser.add_argument("-s", "--set", type=int, default=2, help='train - 0 ,val - 1, test - 2 ')
+parser.add_argument("-p", "--n_max_act", type=int, default=500, help='the number of seq')
+parser.add_argument("-k", "--kfold_cv", type=int, default=1, help='the repeat')
+parser.add_argument("-t", "--task", type=str, default='regression', help='either regression or classification')
+parser.add_argument("-d", "--device", type=str, default='cpu', help='the device to use to extract featmap, digit or cpu')
+args = parser.parse_args()
+
+config_path = args.config
+save_path = config_path.replace(".ini", "_coef")
+config = Auto_popen(config_path)
+config.batch_size = 256
+# config.kfold_cv = 'train_val'
+all_task = config.cycle_set
+
+task_channel_effect = {}
+task_performance = {}
+
+# path check
+if os.path.exists(config_path) and not os.path.exists(save_path):
+    os.mkdir(save_path)
+
+for task in all_task:
+
+    # .... format featmap as data ....
+    print(f"\n\nevaluating for task: {task}")
+    # re-instance the map for each task
+    map_task = MAP.Maximum_activation_patch(popen=config, which_layer=4,
+                                      n_patch=args.n_max_act,
+                                      kfold_index=args.kfold_cv,
+                                      device_string=args.device)
+
+    # extract feature map and rl decision chain
+    featmap = map_task.extract_feature_map(task=task, which_set=args.set)
+    cum_rl_trend = map_task.cumulative_rl_decision(task=task, which_set=args.set)
+
+    # truncate the featmap and rl trend according to sequence length
+    max_seq_len = map_task.df[config.seq_col].apply(len).max()
+    to_stay = max_seq_len // np.product(map_task.strides) +1
+    trunc_start = featmap.shape[2] - to_stay
+    featmap = featmap[:,:,trunc_start:]
+    cum_rl_trend = cum_rl_trend[:,trunc_start:]
+
+    # construct input for linear regression
+    n_sample,n_channel,n_posi = featmap.shape
+    
+    X = featmap.reshape(n_sample,-1)
+    Y = map_task.Y_ls.flatten()
+
+    # .... regression ....
+    if args.task == 'regression':
+        L1 = LassoCV(alphas=np.linspace(2e-3, 0.1, 49))
+        L2 = RidgeCV(alphas=np.linspace(0.001, 0.101, 20))
+        elastic = ElasticNetCV(alphas=np.linspace(2e-3, 0.1, 49), n_jobs=10)    
+        models = [L1, L2, elastic]
+        model_names = ['Lasso', 'Ridge', 'Elastic']
+ 
+    else:
+        LR = LogisticRegressionCV(n_jobs=10)
+        # L1 = LogisticRegressionCV(n_jobs=10,penalty='l1', solver='saga')
+        # elastic = LogisticRegressionCV(n_jobs=10,penalty='elasticnet', solver='saga', l1_ratios=np.linspace(0.0, 0.5, 10))
+        models = [LR]
+        model_names = ['Logistic']
+ 
+
+    for model,name in zip(models, model_names):
+        print(f"\nregressing {name}..")
+        model.fit(X,Y)
+        r2 = model.score(X,Y)  # will be
+        sparsity = np.sum(model.coef_==0) / X.shape[1] *100 
+        try:
+            alpha = model.alpha_
+        except:
+            alpha = 0.0
+
+        print(f"{name} with optimal alpha {alpha:.5f}, r2/acc {r2:.3f} , zero coeff {sparsity:.1f}%")
+    
+        # save df
+        effect = model.coef_.reshape(n_channel,-1)
+        fullcoef_df = pd.DataFrame(effect, columns=[f"{name}_posi_"+str(trunc_start+i) for i in range(to_stay)])
+        fullcoef_df.to_csv( os.path.join(save_path , f"{task}_{name}_coef.csv"), index=False)
+        task_channel_effect[f"{task}_{name}"] = effect.mean(axis=1)
+
+        task_performance[f"{task}_{name}"] = [alpha, r2, sparsity, model.coef_.max(), model.coef_.min()]  
+
+
+all_effect = pd.DataFrame(task_channel_effect)
+all_effect.to_csv(os.path.join(save_path, "all_task_mean_effect.csv"), index=False)
+
+report_df = pd.DataFrame(task_performance)
+report_df.index = ['optim_alpha','r2', 'zero_pctg', 'max_coef', 'min_coef']
+report_df.to_csv(os.path.join(save_path, "regression_report.csv"), index=False)

model/src/mrl_te_optimization/script/fix_test_for_RP.py

@@ -0,0 +1,32 @@
+
+##  generating the hand craft featuress
+
+
+##
+
+test_set_tid = []
+trainval_dict = {}
+test_set_dict = {}
+for cell_line in ['muscle',"PC3","293T"]:
+
+    seed = 41
+
+    df = processed_dict[cell_line]
+    df_to_sample = df.query("`T_id` not in @test_set_tid") 
+    df_overlap = df.query("(`T_id` in @test_set_tid)")
+
+    # the number to sample reduce as the list cumulated from previous cell type
+    # what to sample 
+    n_2_sample = int(0.1*df.shape[0]) - df_overlap.shape[0]
+    sampled_subset = df_to_sample.sample(n=n_2_sample,random_state=seed)
+
+    # merge newly sampled with those in the list
+    test_set_df = sampled_subset.append(df_overlap)
+    trainval_dict[cell_line] = pd.concat([df,test_set_df]).drop_duplicates(keep=False)
+    test_set_dict[cell_line] = test_set_df
+    # adding new Tids
+    test_set_tid += sampled_subset.T_id.values.tolist() 
+
+for cell_line in ["muscle","PC3","293T"]:
+    trainval_dict[cell_line].to_csv(pj(f"RP_{cell_line}_train_val.csv"))
+    test_set_dict[cell_line].to_csv(pj(f"RP_{cell_line}_test.csv"))

model/src/mrl_te_optimization/script/generate_config_js.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+script_dir=$(pwd)
+echo "{" > machine_configure.json
+echo '"script_dir": "'${script_dir}'",' >> machine_configure.json
+echo '"data_dir": "'${script_dir}'/data",' >> machine_configure.json
+echo '"log_dir": "'${script_dir}'/log",' >> machine_configure.json
+echo '"pth_dir": "'${script_dir}'/checkpoint"' >> machine_configure.json
+echo "}" >> machine_configure.json
+
+if [ ! -d ${script_dir}/data ]; then
+    mkdir data
+    mkdir checkpoint    
+    fi

model/src/mrl_te_optimization/script/iter_test.py

@@ -0,0 +1,219 @@
+import os,sys
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import argparse
+
+parser = argparse.ArgumentParser('the main to train model')
+parser.add_argument('--config_file',type=str,required=True)
+parser.add_argument('--cuda',type=str,default=0,required=False)
+parser.add_argument("--kfold_index",type=int,default=1,required=False)
+args = parser.parse_args()
+
+cuda_id = args.cuda if args.cuda is not None else utils.get_config_cuda(args.config_file)
+os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_id)
+
+
+import time
+import torch
+import copy
+import utils
+from torch import optim
+import numpy as np 
+from models import reader,train_val
+from models.ScheduleOptimizer import ScheduledOptim,scheduleoptim_dict_str
+from models.popen import Auto_popen
+from models.loss import Dynamic_Task_Priority,Dynamic_Weight_Averaging
+
+
+POPEN = Auto_popen(args.config_file)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+POPEN.cuda_id = device
+
+
+POPEN.kfold_index = args.kfold_index
+if POPEN.kfold_cv:
+    if args.kfold_index is None:
+        raise NotImplementedError("please specify the kfold index to perform K fold cross validation")
+    POPEN.vae_log_path = POPEN.vae_log_path.replace(".log","_cv%d.log"%args.kfold_index)
+    #POPEN.vae_pth_path = POPEN.vae_pth_path.replace(".pth","_cv%d.pth"%args.kfold_index)
+    
+
+# Run name
+if POPEN.run_name is None:
+    run_name = POPEN.model_type + time.strftime("__%Y_%m_%d_%H:%M")
+else:
+    run_name = POPEN.run_name
+    
+# log dir
+logger = utils.setup_logs(POPEN.vae_log_path)
+logger.info(f"    ===========================| device {device}{cuda_id} |===========================    ")
+#  built model dir or check resume 
+POPEN.check_experiment(logger)
+
+#                               |=====================================|
+#                               |===========   setup  part  ==========|
+#                               |=====================================|
+# read data
+loader_set = {}          
+n_covar_dict = {}                                                                                                                                                                                                                                                                                                       
+base_path = ['cycle_train_val.csv', 'cycle_test.csv']
+base_csv = 'cycle_MTL_transfer.csv'
+for task in POPEN.cycle_set:
+    if (task in ['MPA_U', 'MPA_H', 'MPA_V', 'SubMPA_H']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/3M/schedule_lr.ini')
+        datapopen.split_like = [path.replace('cycle', task) for path in base_path]
+        datapopen.kfold_index = args.kfold_index
+        n_covar_dict[task] = datapopen.n_covar
+
+    elif (task in ['RP_293T', 'RP_muscle', 'RP_PC3']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/3R/schedule_MTL.ini')
+        datapopen.csv_path = base_csv.replace("cycle",task)
+        datapopen.kfold_index = args.kfold_index
+        datapopen.pad_to = POPEN.pad_to
+        datapopen.aux_task_columns = POPEN.aux_task_columns
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.n_covar = POPEN.n_covar
+
+    elif (task in ['pcr3', '293']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/karollus_RPs/rp_cycle.ini')
+        datapopen.csv_path = base_csv.replace("cycle",task)
+        datapopen.kfold_index = args.kfold_index
+        datapopen.aux_task_columns = POPEN.aux_task_columns
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.pad_to = POPEN.pad_to
+        datapopen.n_covar = POPEN.n_covar
+    datapopen.shuffle = False
+    loader_set[task] = reader.get_dataloader(datapopen)
+
+POPEN.n_covar = n_covar_dict
+POPEN.get_model_config() # update model config
+    
+# ===========  setup model  ===========
+# train_iter = iter(train_loader)
+# X,Y  = next(train_iter)
+# -- pretrain -- 
+if POPEN.pretrain_pth is not None:
+    # load pretran model
+    logger.info("===============================|   pretrain   |===============================")
+    logger.info(f" {POPEN.pretrain_pth}")
+    pretrain_popen = Auto_popen(os.path.join(utils.script_dir, POPEN.pretrain_pth))
+    if not os.path.exists(pretrain_popen.vae_pth_path):
+        if type(args.kfold_index) == int:
+            pretrain_popen.kfold_index = args.kfold_index
+    pretrain_model = torch.load(pretrain_popen.vae_pth_path, map_location=torch.device('cpu'))['state_dict']
+
+    
+    if POPEN.model_type == pretrain_popen.model_type:
+        # if not POPEN.Resumable:
+        #     # we only load pre-train for the first time 
+        #     # later we can resume 
+        model = pretrain_model.to(device)
+        del pretrain_model
+        
+        if (POPEN.cycle_set != pretrain_popen.cycle_set):
+            model.all_tasks = POPEN.cycle_set
+            model.tower = torch.nn.ModuleDict(
+                {POPEN.cycle_set[i] : model.tower[t]  for i, t in enumerate(pretrain_popen.cycle_set)}
+                                                )
+        
+    elif POPEN.modual_to_fix is not None:
+        # POPEN.model_type != pretrain_popen.model_type
+        
+        model = POPEN.Model_Class(*POPEN.model_args)
+        for modual in POPEN.modual_to_fix:
+            if modual in dir(pretrain_model):    
+                eval(f'model.{modual}').load_state_dict(
+                    eval(f'model.{modual}').state_dict()
+                    )
+
+        state_dict = {'epoch': 0,
+                        'validation_acc': 0,
+                        'state_dict': model.to('cpu'),
+                        'validation_loss': 0}
+        shared_pretrain_pth = POPEN.vae_pth_path.replace(f"_cv{args.kfold_index}", '')
+        if not os.path.exists(shared_pretrain_pth):
+            utils.snapshot(shared_pretrain_pth, state_dict)
+        utils.snapshot(POPEN.vae_pth_path, state_dict)
+
+        model = torch.load(POPEN.vae_pth_path, map_location=torch.device('cpu')) 
+        model = model.to(device)
+    
+# -- end2end -- 
+elif POPEN.model_type == "CrossStitch_Model":
+    backbone = {}
+    for t in POPEN.tasks:
+        task_popen = Auto_popen(POPEN.backbone_config[t])
+        task_model = task_popen.Model_Class(*task_popen.model_args)
+        utils.load_model(task_popen,task_model,logger)
+        backbone[t] = task_model.to(device)
+    POPEN.model_args = [backbone] + POPEN.model_args
+    model = POPEN.Model_Class(*POPEN.model_args).to(device)
+else:
+    Model_Class = POPEN.Model_Class  # DL_models.LSTM_AE 
+    model = Model_Class(*POPEN.model_args).to(device)
+    
+if POPEN.Resumable:
+    model = utils.load_model(POPEN, model, logger)
+    
+# =========== fix parameters ===========
+if isinstance(POPEN.modual_to_fix, list):
+    for modual in POPEN.modual_to_fix:
+        model = utils.fix_parameter(model,modual)
+        model =  model.to(device)
+    logger.info(' \t \t ==============| %s fixed |============== \t \t \n'%POPEN.modual_to_fix)
+# =========== set optimizer ===========
+if POPEN.optimizer == 'Schedule':
+    optimizer = ScheduledOptim(optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                                      betas=(0.9, 0.98), 
+                                      eps=1e-09, 
+                                      weight_decay=1e-4, 
+                                      amsgrad=True),
+                           n_warmup_steps=20)
+elif type(POPEN.optimizer) == dict:
+    optimizer = eval(scheduleoptim_dict_str.format(**POPEN.optimizer))
+else:
+    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                            lr=POPEN.lr,
+                            betas=(0.9, 0.98), 
+                            eps=1e-09, 
+                            weight_decay=POPEN.l2)
+if POPEN.loss_schema == 'DTP':
+    POPEN.loss_schedualer = Dynamic_Task_Priority(POPEN.tasks,POPEN.gamma,POPEN.chimerla_weight)
+elif POPEN.loss_schema == 'DWA':
+    POPEN.loss_schedualer = Dynamic_Weight_Averaging(POPEN.tasks,POPEN.tau,POPEN.chimerla_weight)
+# =========== resume ===========
+best_loss = np.inf
+best_acc = 0
+best_epoch = 0
+previous_epoch = 0
+epoch = 0
+if POPEN.Resumable:
+    previous_epoch,best_loss,best_acc = utils.resume(POPEN, optimizer,logger)
+    
+
+#                               |=====================================|
+#                               |==========  training  part ==========|
+#                               |=====================================|
+epoch += previous_epoch
+
+
+
+#              -----------| validate |-----------   
+logger.info("===============================| test |===============================")
+verbose_dict = train_val.cycle_validate(loader_set,model,optimizer,popen=POPEN,epoch=epoch, which_set=2)
+# matching task performance influence what to save
+
+if np.any(['r2' in key for key in verbose_dict.keys()]):
+    val_avg_acc = np.mean([values for key, values in verbose_dict.items() if 'r2' in key])
+    acc_dict = {f"cv{args.kfold_index}_{key}":values for key, values in verbose_dict.items() if 'r2' in key}
+else:
+    val_avg_acc = np.mean([values for key, values in verbose_dict.items() if 'acc' in key])
+    acc_dict = {}
+val_total_loss = verbose_dict['Total']
+
+# DICT ={"ran_epoch":epoch,"n_current_steps":optimizer.n_current_steps,"delta":optimizer.delta} if type(optimizer) == ScheduledOptim else {"ran_epoch":epoch}
+# POPEN.update_ini_file(DICT,logger)
+    
+    
+
+
+# POPEN.update_ini_file(acc_dict,logger)

model/src/mrl_te_optimization/script/iter_train.py

@@ -0,0 +1,251 @@
+import os,sys
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import argparse
+
+parser = argparse.ArgumentParser('the main to train model')
+parser.add_argument('--config_file',type=str,required=True)
+parser.add_argument('--cuda',type=str,default=0,required=False)
+parser.add_argument("--kfold_index",type=int,default=1,required=False)
+args = parser.parse_args()
+
+cuda_id = args.cuda if args.cuda is not None else utils.get_config_cuda(args.config_file)
+os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_id)
+
+
+import time
+import torch
+import copy
+import utils
+from torch import optim
+import numpy as np 
+from models import reader,train_val
+from models.ScheduleOptimizer import ScheduledOptim,scheduleoptim_dict_str
+from models.popen import Auto_popen
+from models.loss import Dynamic_Task_Priority,Dynamic_Weight_Averaging
+
+
+POPEN = Auto_popen(args.config_file)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+POPEN.cuda_id = device
+
+
+POPEN.kfold_index = args.kfold_index
+if POPEN.kfold_cv:
+    if args.kfold_index is None:
+        raise NotImplementedError("please specify the kfold index to perform K fold cross validation")
+    POPEN.vae_log_path = POPEN.vae_log_path.replace(".log","_cv%d.log"%args.kfold_index)
+    #POPEN.vae_pth_path = POPEN.vae_pth_path.replace(".pth","_cv%d.pth"%args.kfold_index)
+    
+
+# Run name
+if POPEN.run_name is None:
+    run_name = POPEN.model_type + time.strftime("__%Y_%m_%d_%H:%M")
+else:
+    run_name = POPEN.run_name
+    
+# log dir
+logger = utils.setup_logs(POPEN.vae_log_path)
+logger.info(f"    ===========================| device {device}{cuda_id} |===========================    ")
+#  built model dir or check resume 
+POPEN.check_experiment(logger)
+
+#                               |=====================================|
+#                               |===========   setup  part  ==========|
+#                               |=====================================|
+# read data
+loader_set = {}                                                                                                                                                                                                                                                                                                                 
+base_path = ['cycle_train_val.csv', 'cycle_test.csv']
+base_csv = copy.deepcopy(POPEN.csv_path) #'cycle_MTL_transfer.csv'
+for task in POPEN.cycle_set:
+    if (task in ['MPA_U', 'MPA_H', 'MPA_V', 'SubMPA_H']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/3M/schedule_lr.ini')
+        datapopen.split_like = [path.replace('cycle', task) for path in base_path]
+        datapopen.kfold_index = args.kfold_index
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.n_covar = POPEN.n_covar
+
+    elif (task in ['RP_293T', 'RP_muscle', 'RP_PC3']):
+        # base_csv = 'cycle_protein_coding.csv'
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/3R/schedule_MTL.ini')
+        datapopen.csv_path = base_csv.replace("cycle",task)
+        datapopen.kfold_index = args.kfold_index
+        datapopen.kfold_cv = POPEN.kfold_cv
+        datapopen.pad_to = POPEN.pad_to
+        datapopen.aux_task_columns = POPEN.aux_task_columns
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.n_covar = POPEN.n_covar
+
+    elif (task in ['pcr3', '293']):
+        datapopen = Auto_popen('log/Backbone/RL_hard_share/karollus_RPs/rp_cycle.ini')
+        datapopen.csv_path = base_csv.replace("cycle",task)
+        datapopen.kfold_index = args.kfold_index
+        datapopen.kfold_cv = POPEN.kfold_cv
+        datapopen.aux_task_columns = POPEN.aux_task_columns
+        datapopen.other_input_columns = POPEN.other_input_columns
+        datapopen.pad_to = POPEN.pad_to
+        datapopen.n_covar = POPEN.n_covar
+    
+
+    loader_set[task] = reader.get_dataloader(datapopen)
+
+    
+# ===========  setup model  ===========
+# train_iter = iter(train_loader)
+# X,Y  = next(train_iter)
+# -- pretrain -- 
+if POPEN.pretrain_pth is not None:
+    # load pretran model
+    logger.info("===============================|   pretrain   |===============================")
+    logger.info(f" {POPEN.pretrain_pth}")
+    pretrain_popen = Auto_popen(os.path.join(utils.script_dir, POPEN.pretrain_pth))
+    if not os.path.exists(pretrain_popen.vae_pth_path):
+        if type(args.kfold_index) == int:
+            pretrain_popen.kfold_index = args.kfold_index
+    pretrain_model = torch.load(pretrain_popen.vae_pth_path, map_location=torch.device('cpu'))['state_dict']
+
+    
+    if POPEN.model_type == pretrain_popen.model_type:
+        # if not POPEN.Resumable:
+        #     # we only load pre-train for the first time 
+        #     # later we can resume 
+        model = pretrain_model.to(device)
+        del pretrain_model
+        
+        if (POPEN.cycle_set != pretrain_popen.cycle_set):
+            model.all_tasks = POPEN.cycle_set
+            model.tower = torch.nn.ModuleDict(
+                {POPEN.cycle_set[i] : model.tower[t]  for i, t in enumerate(pretrain_popen.cycle_set)}
+                                                )
+        
+    elif POPEN.modual_to_fix is not None:
+        # POPEN.model_type != pretrain_popen.model_type
+        
+        model = POPEN.Model_Class(*POPEN.model_args)
+        for modual in POPEN.modual_to_fix:
+            if modual in dir(pretrain_model):    
+                eval(f'model.{modual}').load_state_dict(
+                    eval(f'model.{modual}').state_dict()
+                    )
+
+        state_dict = {'epoch': 0,
+                        'validation_acc': 0,
+                        'state_dict': model.to('cpu'),
+                        'validation_loss': 0}
+        shared_pretrain_pth = POPEN.vae_pth_path.replace(f"_cv{args.kfold_index}", '')
+        if not os.path.exists(shared_pretrain_pth):
+            utils.snapshot(shared_pretrain_pth, state_dict)
+        utils.snapshot(POPEN.vae_pth_path, state_dict)
+
+        model = torch.load(POPEN.vae_pth_path, map_location=torch.device('cpu')) 
+        model = model.to(device)
+    
+# -- end2end -- 
+elif POPEN.model_type == "CrossStitch_Model":
+    backbone = {}
+    for t in POPEN.tasks:
+        task_popen = Auto_popen(POPEN.backbone_config[t])
+        task_model = task_popen.Model_Class(*task_popen.model_args)
+        utils.load_model(task_popen,task_model,logger)
+        backbone[t] = task_model.to(device)
+    POPEN.model_args = [backbone] + POPEN.model_args
+    model = POPEN.Model_Class(*POPEN.model_args).to(device)
+else:
+    Model_Class = POPEN.Model_Class  # DL_models.LSTM_AE 
+    model = Model_Class(*POPEN.model_args).to(device)
+    
+if POPEN.Resumable:
+    model = utils.load_model(POPEN, model, logger)
+    
+# =========== fix parameters ===========
+if isinstance(POPEN.modual_to_fix, list):
+    for modual in POPEN.modual_to_fix:
+        model = utils.fix_parameter(model,modual)
+        model =  model.to(device)
+    logger.info(' \t \t ==============| %s fixed |============== \t \t \n'%POPEN.modual_to_fix)
+# =========== set optimizer ===========
+if POPEN.optimizer == 'Schedule':
+    optimizer = ScheduledOptim(optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                                      betas=(0.9, 0.98), 
+                                      eps=1e-09, 
+                                      weight_decay=1e-4, 
+                                      amsgrad=True),
+                           n_warmup_steps=20)
+elif type(POPEN.optimizer) == dict:
+    optimizer = eval(scheduleoptim_dict_str.format(**POPEN.optimizer))
+else:
+    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                            lr=POPEN.lr,
+                            betas=(0.9, 0.98), 
+                            eps=1e-09, 
+                            weight_decay=POPEN.l2)
+if POPEN.loss_schema == 'DTP':
+    POPEN.loss_schedualer = Dynamic_Task_Priority(POPEN.tasks,POPEN.gamma,POPEN.chimerla_weight)
+elif POPEN.loss_schema == 'DWA':
+    POPEN.loss_schedualer = Dynamic_Weight_Averaging(POPEN.tasks,POPEN.tau,POPEN.chimerla_weight)
+# =========== resume ===========
+best_loss = np.inf
+best_acc = 0
+best_epoch = 0
+previous_epoch = 0
+if POPEN.Resumable:
+    previous_epoch,best_loss,best_acc = utils.resume(POPEN, optimizer,logger)
+    
+
+#                               |=====================================|
+#                               |==========  training  part ==========|
+#                               |=====================================|
+for epoch in range(POPEN.max_epoch-previous_epoch+1):
+    epoch += previous_epoch
+    
+    #          
+    logger.info("===============================|    epoch {}   |===============================".format(epoch))
+
+
+    train_val.iter_train(loader_set,model=model,optimizer=optimizer,popen=POPEN,epoch=epoch)
+
+#              -----------| validate |-----------   
+    logger.info("===============================| start validation |===============================")
+    verbose_dict = train_val.cycle_validate(loader_set,model,optimizer,popen=POPEN,epoch=epoch, which_set=1)
+    test_dict = train_val.cycle_validate(loader_set,model,optimizer,popen=POPEN,epoch=epoch, which_set=2)
+
+    if np.any(['r2' in key for key in verbose_dict.keys()]):
+        val_avg_acc = np.mean([values for key, values in verbose_dict.items() if 'r2' in key])
+    else:
+        val_avg_acc = np.mean([values for key, values in verbose_dict.items() if 'acc' in key])
+    val_total_loss = verbose_dict['Total']
+    
+    
+    DICT ={"ran_epoch":epoch,"n_current_steps":optimizer.n_current_steps,"delta":optimizer.delta} if type(optimizer) == ScheduledOptim else {"ran_epoch":epoch}
+    POPEN.update_ini_file(DICT,logger)
+    
+    
+#    -----------| compare the result |-----------
+    if (best_loss > val_total_loss) :
+        # update best performance
+        best_loss = min(best_loss,val_total_loss)
+        best_acc = max(best_acc,val_avg_acc)
+        best_epoch = epoch
+        
+        # save
+        utils.snapshot(POPEN.vae_pth_path, {
+                    'epoch': epoch + 1,
+                    'validation_acc': val_avg_acc,
+                    # 'state_dict': model.state_dict(),
+                    'state_dict': model,
+                    'validation_loss': val_total_loss,
+                    'optimizer': optimizer.state_dict(),
+                })
+        
+        # update the popen
+        POPEN.update_ini_file({'run_name':run_name,
+                            "ran_epoch":epoch,
+                            "best_acc":best_acc},
+                            logger)
+        
+    elif (epoch - best_epoch >= 30)&((type(optimizer) == ScheduledOptim)):
+        optimizer.increase_delta()
+        
+    elif (epoch - best_epoch >= 60)&(epoch > POPEN.max_epoch/2):
+        # at the late phase of training
+        logger.info("<<<<<<<<<<< Early Stopping >>>>>>>>>>")
+        break

model/src/mrl_te_optimization/script/main_train.py

@@ -0,0 +1,201 @@
+import os,sys
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+import argparse
+
+parser = argparse.ArgumentParser('the main to train model')
+parser.add_argument('--config_file',type=str,required=True)
+parser.add_argument('--cuda',type=int,default=None,required=False)
+parser.add_argument("--kfold_index",type=int,default=None,required=False)
+args = parser.parse_args()
+
+cuda_id = args.cuda
+os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_id)
+
+
+import time
+import torch
+import utils
+from torch import optim
+import numpy as np 
+from models import Modules,reader,train_val
+from models.ScheduleOptimizer import ScheduledOptim,scheduleoptim_dict_str
+from models.popen import Auto_popen
+from models.loss import Dynamic_Task_Priority,Dynamic_Weight_Averaging
+
+
+POPEN = Auto_popen(args.config_file)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+torch.set_num_interop_threads(4)
+POPEN.cuda_id = device
+
+
+POPEN.kfold_index = args.kfold_index
+if POPEN.kfold_cv:
+    if args.kfold_index is None:
+        raise NotImplementedError("please specify the kfold index to perform K fold cross validation")
+    POPEN.vae_log_path = POPEN.vae_log_path.replace(".log","_cv%d.log"%args.kfold_index)
+    #POPEN.vae_pth_path = POPEN.vae_pth_path.replace(".pth","_cv%d.pth"%args.kfold_index)
+    
+
+# Run name
+if POPEN.run_name is None:
+    run_name = POPEN.model_type + time.strftime("__%Y_%m_%d_%H:%M")
+else:
+    run_name = POPEN.run_name
+    
+# log dir
+logger = utils.setup_logs(POPEN.vae_log_path)
+logger.info(f"  	 	 ==============<<< device used: {device}:{cuda_id}  >>>============== 	 	 ")
+#  built model dir or check resume 
+POPEN.check_experiment(logger)
+#                               |=====================================|
+#                               |===========   setup  part  ==========|
+#                               |=====================================|
+# read data
+train_loader,val_loader,test_loader = reader.get_dataloader(POPEN)
+# ===========  setup model  ===========
+# train_iter = iter(train_loader)
+# X,Y  = next(train_iter)
+# -- pretrain -- 
+if POPEN.pretrain_pth is not None:
+    # load pretran model
+    pretrain_popen = Auto_popen(os.path.join(utils.script_dir, POPEN.pretrain_pth))
+    pretrain_model = pretrain_popen.Model_Class(*pretrain_popen.model_args)
+
+    if not os.path.exists(pretrain_popen.vae_pth_path):
+        if type(args.kfold_index) == int:
+            pretrain_popen.kfold_index = args.kfold_index
+    pretrain_model = utils.load_model(pretrain_popen,pretrain_model,logger)
+    
+
+    
+    if POPEN.Model_Class == pretrain_popen.Model_Class:
+        # if not POPEN.Resumable:
+        #     # we only load pre-train for the first time 
+        #     # later we can resume 
+        model = pretrain_model.to(device)
+        del pretrain_model
+    elif POPEN.modual_to_fix is not None:
+        # POPEN.model_type != pretrain_popen.model_type
+        
+        model = POPEN.Model_Class(*POPEN.model_args)
+        for modual in POPEN.modual_to_fix:
+            if modual in dir(pretrain_model):    
+                eval(f'model.{modual}').load_state_dict(
+                    eval(f'model.{modual}').state_dict()
+                    )
+
+        state_dict = {'epoch': 0,
+                        'validation_acc': 0,
+                        'state_dict': model.to('cpu'),
+                        'validation_loss': np.inf}
+        shared_pretrain_pth = POPEN.vae_pth_path.replace(f"_cv{args.kfold_index}", '')
+        if not os.path.exists(shared_pretrain_pth):
+            utils.snapshot(shared_pretrain_pth, state_dict)
+        utils.snapshot(POPEN.vae_pth_path, state_dict)
+
+        model = torch.load(POPEN.vae_pth_path, map_location=torch.device('cpu'))['state_dict']
+        model = model.to(device)
+
+    
+# -- end2end -- 
+elif POPEN.model_type == "CrossStitch_Model":
+    backbone = {}
+    for t in POPEN.tasks:
+        task_popen = Auto_popen(POPEN.backbone_config[t])
+        task_model = task_popen.Model_Class(*task_popen.model_args)
+        task_model = utils.load_model(task_popen,task_model,logger)
+        backbone[t] = task_model.to(device)
+    POPEN.model_args = [backbone] + POPEN.model_args
+    model = POPEN.Model_Class(*POPEN.model_args).to(device)
+else:
+    Model_Class = POPEN.Model_Class  # DL_models.LSTM_AE 
+    model = Model_Class(*POPEN.model_args).to(device)
+    
+if POPEN.Resumable:
+    model = utils.load_model(POPEN, model, logger)
+    
+# =========== fix parameters ===========
+if isinstance(POPEN.modual_to_fix, list):
+    for modual in POPEN.modual_to_fix:
+        model = utils.fix_parameter(model,modual)
+    logger.info(' \t \t ==============<<< %s part is fixed>>>============== \t \t \n'%POPEN.modual_to_fix)
+# =========== set optimizer ===========
+if POPEN.optimizer == 'Schedule':
+    optimizer = ScheduledOptim(optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                                      betas=(0.9, 0.98), 
+                                      eps=1e-09, 
+                                      weight_decay=1e-4, 
+                                      amsgrad=True),
+                           n_warmup_steps=20)
+elif type(POPEN.optimizer) == dict:
+    optimizer = eval(scheduleoptim_dict_str.format(**POPEN.optimizer))
+else:
+    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
+                            lr=POPEN.lr,
+                            betas=(0.9, 0.98), 
+                            eps=1e-09, 
+                            weight_decay=POPEN.l2)
+if POPEN.loss_schema == 'DTP':
+    POPEN.loss_schedualer = Dynamic_Task_Priority(POPEN.tasks,POPEN.gamma,POPEN.chimerla_weight)
+elif POPEN.loss_schema == 'DWA':
+    POPEN.loss_schedualer = Dynamic_Weight_Averaging(POPEN.tasks,POPEN.tau,POPEN.chimerla_weight)
+# =========== resume ===========
+best_loss = np.inf
+best_acc = 0
+best_epoch = 0
+previous_epoch = 0
+if POPEN.Resumable:
+    previous_epoch,best_loss,best_acc = utils.resume(POPEN, optimizer,logger)
+    
+
+#                               |=====================================|
+#                               |==========  training  part ==========|
+#                               |=====================================|
+for epoch in range(POPEN.max_epoch-previous_epoch+1):
+    epoch += previous_epoch
+    
+    #           ----------| train |----------
+    logger.info("===============================|    epoch {}   |===============================".format(epoch))
+    train_val.train(dataloader=train_loader,model=model,optimizer=optimizer,popen=POPEN,epoch=epoch)
+           
+    #         -----------| validate |-----------
+    
+    if epoch % POPEN.config_dict['setp_to_check'] == 0:
+        logger.info("===============================| start validation |===============================")
+        val_total_loss,val_avg_acc = train_val.validate(val_loader,model,popen=POPEN,epoch=epoch)
+        _,_ = train_val.validate(test_loader,model,popen=POPEN,epoch=epoch)
+        
+        DICT ={"ran_epoch":epoch,"n_current_steps":optimizer.n_current_steps,"delta":optimizer.delta} if type(optimizer) == ScheduledOptim else {"ran_epoch":epoch}
+        POPEN.update_ini_file(DICT,logger)
+        
+    #    -----------| compare the result |-----------
+        if (best_loss > val_total_loss): #| (best_acc < val_avg_acc):
+            # update best performance
+            best_loss = min(best_loss,val_total_loss)
+            best_acc = max(best_acc,val_avg_acc)
+            best_epoch = epoch
+            
+            # save
+            utils.snapshot(POPEN.vae_pth_path, {
+                        'epoch': epoch + 1,
+                        'validation_acc': val_avg_acc,
+                        'state_dict': model.to('cpu'),
+                        'validation_loss': val_total_loss,
+                        # 'optimizer': optimizer.state_dict(),
+                    })
+            
+            # update the popen
+            POPEN.update_ini_file({'run_name':run_name,
+                                "ran_epoch":epoch,
+                                "best_acc":best_acc,
+                                "cuda_id":cuda_id},
+                                logger)
+            
+        elif (epoch - best_epoch >= 30)&((type(optimizer) == ScheduledOptim)):
+            optimizer.increase_delta()
+            
+        elif (epoch - best_epoch >= 60)&(epoch > POPEN.max_epoch/2):
+            # at the late phase of training
+            logger.info("<<<<<<<<<<< Early Stopping >>>>>>>>>>")
+            break

model/src/mrl_te_optimization/script/preprocessing_MPA.py

@@ -0,0 +1,115 @@
+import os
+import sys
+import torch
+import PATH
+import utils
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from matplotlib import pyplot as plt
+
+from sklearn.preprocessing import StandardScaler
+
+
+# ========== 1. preprocessing MPA datasets ==========
+
+# read csv and load in to a dictionary
+Ex_data_dir = utils.data_dir
+csv_path_ls = [os.path.join(Ex_data_dir,csv) for csv in ['GSM3130435_egfp_unmod_1.csv','GSM3130443_designed_library.csv','GSM4084997_varying_length_25to100.csv']]
+for path in csv_path_ls:
+    assert os.path.exists(path), f"The file {path} is not properly downloaded"
+
+df_dict = {
+  csv_path.split("_")[-1].replace(".csv","") : pd.read_csv(csv_path,low_memory=False)  for csv_path in csv_path_ls
+           }
+# align the columns name across datasets
+df_dict['library'].rename({'total':'total_reads'},axis=1,inplace=True)
+df_dict['25to100']['r13'] = np.zeros((df_dict['25to100'].shape[0],))
+
+
+
+# - MPA_U - 
+df = df_dict['1']
+df.sort_values('total_reads', inplace=True, ascending=False)
+df.reset_index(inplace=True, drop=True)
+
+
+## train_val test spliting
+# the most abundant 5'UTR is used as test set
+test_df = df.iloc[:20000]
+train_val_df = df.iloc[20000:280000]
+
+
+
+# - MPA_H - 
+df = df_dict['library']
+# taking the natrual utr and sort by reads count
+human_lib = df[(df['library'] == 'human_utrs') | (df['library'] == 'snv')]
+human_lib = human_lib.sort_values('total_reads', ascending=False).reset_index(drop=True)
+
+# the top 25k abundant reads as test set
+sub = human_lib.iloc[:25000]
+remaining = human_lib.iloc[25000:]
+
+
+# - MPA_V - 
+df1 = df_dict['25to100']
+# take the random seqs
+df = df1[df1['set']=='random']
+
+## Filter out UTRs with too few less reads
+df=df[df['total_reads']>=10]
+
+df['utr100'] = 75*'N' +df['utr']
+df['utr100'] = df['utr100'].str[-100:]
+df.sort_values('total_reads', inplace=True, ascending=False)
+df.reset_index(inplace=True, drop=True)
+
+## some natural sequence
+human = df1[df1['set']=='human']
+## Filter out UTRs with too few less reads
+human=human[human['total_reads']>=10]
+human['utr100'] = 75*'N' +human['utr']
+human['utr100'] = human['utr100'].str[-100:]
+human.sort_values('total_reads', inplace=True, ascending=False)
+human.reset_index(inplace=True, drop=True)
+
+e_test = pd.DataFrame(columns=df.columns)
+for i in range(25,101):
+    tmp = df[df['len']==i]
+    tmp.sort_values('total_reads', inplace=True, ascending=False)
+    tmp.reset_index(inplace=True, drop=True)
+    e_test = e_test.append(tmp.iloc[:100])
+
+subhuman = pd.DataFrame(columns=human.columns)
+for i in range(25,101):
+    tmp = human[human['len']==i]
+    tmp.sort_values('total_reads', inplace=True, ascending=False)
+    tmp.reset_index(inplace=True, drop=True)
+    subhuman = subhuman.append(tmp.iloc[:100])
+
+e_train = pd.concat([df, e_test, e_test]).drop_duplicates(keep=False)
+
+vleng_test = pd.concat([e_test,subhuman])
+bins = np.arange(24, 105, 20)
+labels = [ '25-44' , '45-64', '65-84', '85-100']
+vleng_test['rng'] = pd.cut(vleng_test['len'], bins=bins)
+
+# saving all
+e_test.to_csv(os.path.join(Ex_data_dir,"MPA_V_test.csv"),index=False)
+e_train.to_csv(os.path.join(Ex_data_dir,"MPA_V_train_val.csv"),index=False)
+
+sub.to_csv(os.path.join(Ex_data_dir,"MPA_H_test.csv"),index=False)
+remaining.to_csv(os.path.join(Ex_data_dir,"MPA_H_train_val.csv"),index=False)
+
+# sub sample MPA-H
+remaining.sample(frac=0.1).to_csv(os.path.join(Ex_data_dir,"SubMPA_H_train_val.csv"),index=False)
+sub.to_csv(os.path.join(Ex_data_dir,"SubMPA_H_test.csv"),index=False) 
+
+test_df.to_csv(os.path.join(Ex_data_dir,"MPA_U_test.csv"),index=False)
+train_val_df.to_csv(os.path.join(Ex_data_dir,"MPA_U_train_val.csv"),index=False)
+
+
+
+print("The preprocssing for MPA tasks is Finished !!")
+print(f"The files are saved to {utils.data_dir}")

model/src/mrl_te_optimization/script/preprocessing_RP.py

@@ -0,0 +1,65 @@
+import os
+import numpy as np
+import pandas as pd
+import PATH
+import utils
+from Bio import SeqIO
+from scipy import sparse
+
+
+pj = lambda x: os.path.join(utils.data_dir,x)
+
+
+# read in the sequence 
+fa = SeqIO.parse(pj('gencode_v17_5utr_15bpcds.fa'),'fasta')
+tx_seq=dict()
+for seq_record in fa:
+    tx=seq_record.id
+    seq=seq_record.seq
+    if len(seq)<30: #skip when it too short
+        continue
+    if "ATG" not in seq:
+        continue
+    if tx not in tx_seq or len(seq)>len(tx_seq[tx]):
+        tx_seq[tx]=seq
+
+# txt records TE and rkpm
+RP_data_path = {}
+RP_data_path['muscle'] = pj('df_counts_and_len.TE_sorted.Muscle.with_annot.txt')
+RP_data_path['PC3'] = pj('df_counts_and_len.TE_sorted.pc3.with_annot.txt')
+RP_data_path['293T'] = pj('df_counts_and_len.TE_sorted.HEK_Andrev2015.with_annot.txt')
+
+# read the hand crafted features
+# (57415, 33)
+feat_mat = pd.read_csv("util/Zhang_et_al_features.csv",index_col=0)
+feat_mat.insert(0, 'seq', [tx_seq[x].upper().__str__() for x in feat_mat.index.values])
+
+
+processed_dict = {}
+for cell_line, csv_path in RP_data_path.items():
+
+    assert os.path.exists(csv_path), f"The raw data file for {cell_line} does not exist!"
+    
+    print(f"processing {cell_line}...\n")
+    RP_raw_data = pd.read_table(csv_path, sep=' ')
+    RP_raw_data.loc[:,'T_id'] = RP_raw_data.index.values
+    # adding sequence in 
+    RP_raw_data = RP_raw_data.query('rpkm_rnaseq >5 & rpkm_riboseq > 0.1')
+    RP_raw_data['log_te'] = np.log(RP_raw_data.te.values)
+
+    RP_feat_merge = RP_raw_data.merge(feat_mat,left_on=['T_id'],right_index=True,suffixes=["",""])
+    RP_feat_merge.sort_values('rpkm_rnaseq', ascending=False, inplace=True)
+    
+    # drop duplicated UTRs
+    RP_raw_dedup = RP_feat_merge.drop_duplicates(RP_feat_merge.columns[17:], keep='first')
+    RP_raw_dedup['utr'] = RP_raw_dedup['seq'].apply(lambda x: x[-216:-16]) # max len 200
+    RP_raw_dedup['utr_len'] = RP_raw_dedup.utr.apply(len)
+    RP_raw_dedup.query('`utr_len`>30')
+    RP_raw_dedup = RP_raw_dedup.drop_duplicates(['utr'], keep='first') 
+    
+    # save them
+    processed_dict[cell_line] = RP_raw_dedup
+    processed_dict[cell_line].to_csv(pj(f"RP_{cell_line}_MTL_transfer.csv"))
+
+print("The preprocssing for RP tasks is Finished !!")
+print(f"The files are saved to {utils.data_dir}")

model/src/mrl_te_optimization/script/run-TF-modisco.py

@@ -0,0 +1,99 @@
+import os
+import sys
+import pandas as pd
+import numpy as np
+import PATH
+import torch
+import argparse
+import h5py
+# 
+from subprocess import DEVNULL, STDOUT, check_call
+from models import train_val
+from models.popen import Auto_popen
+from models import max_activation_patch as MAP
+# 
+from sklearn.linear_model import Lasso, Ridge, ElasticNet, LassoCV, RidgeCV, ElasticNetCV
+import warnings
+warnings.filterwarnings('ignore')
+
+
+#
+parser = argparse.ArgumentParser('the script to evlauate the effect of ')
+parser.add_argument("-c", "--config", type=str, required=True, help='the model config file: xxx.ini')
+parser.add_argument("-s", "--set", type=int, default=2, help='train - 0 ,val - 1, test - 2 ')
+parser.add_argument("-p", "--n_max_act", type=int, default=500, help='the number of seq')
+parser.add_argument("-k", "--kfold_cv", type=int, default=1, help='the repeat')
+parser.add_argument("-d", "--device", type=str, default='cpu', help='the device to use to extract featmap, digit or cpu')
+args = parser.parse_args()
+
+config_path = args.config
+save_path = config_path.replace(".ini", "_coef")
+config = Auto_popen(config_path)
+config.batch_size = 256
+config.kfold_cv = 'train_val'
+all_task = config.cycle_set
+
+task_channel_effect = {}
+task_performance = {}
+
+# path check
+if os.path.exists(config_path) and not os.path.exists(save_path):
+    os.mkdir(save_path)
+
+for task in all_task:
+
+    # .... format featmap as data ....
+    print(f"\n\nevaluating for task: {task}")
+    # re-instance the map for each task
+    map_task = MAP.Maximum_activation_patch(popen=config, which_layer=4,
+                                      n_patch=args.n_max_act,
+                                      kfold_index=args.kfold_cv,
+                                      device_string=args.device)
+
+    # extract feature map and rl decision chain
+    # get X
+    model, dataloader= map_task.loading(task=task, which_set=args.set)
+    max_seq_len = map_task.df[config.seq_col].apply(len).max()
+
+    X_ls = []
+    for Data in dataloader:
+        # iter each batch
+        x,y = train_val.put_data_to_cuda(Data, map_task.popen,False) 
+        X_ls.append(x.detach().cpu().numpy())
+    X = np.concatenate(X_ls, axis=0)
+    X = np.transpose(X, (0,2,1))
+    
+    attribute = map_task.get_input_grad(task=task, focus=False, fm=X, starting_layer=0)
+    X = X[:, :, -1*max_seq_len:]
+    attribute = attribute[:,:, -1*max_seq_len:]
+
+    x_npz_path = os.path.join(save_path, f"{task}_set{args.set}_x.npz")
+    grad_npz_path = os.path.join(save_path, f"{task}_set{args.set}_grad.npz")
+    np.savez(x_npz_path, X)
+    np.savez(grad_npz_path, attribute)
+
+    
+    
+    for n_seqlet in range(500, 4000, 500):
+        outfile = os.path.join(save_path, f"{task}_n{n_seqlet}_tfmodisco.h5")
+        command = f"modisco motifs -s {x_npz_path} -a {grad_npz_path} -n {n_seqlet} -o {outfile}"
+        
+        check_call(command.split(), stdout=DEVNULL, stderr=STDOUT)
+
+        with h5py.File(outfile, "r") as f:
+            print(f.keys())
+            if 'pos_patterns' in f.keys():
+                pos_motif = f['pos_patterns'] 
+                n_pos = len(pos_motif.keys())
+            else:
+                n_pos = -1
+
+            if 'neg_patterns' in f.keys():
+                neg_motif = f['neg_patterns'] 
+                n_neg = len(neg_motif.keys())
+            else:
+                n_neg = -1
+            f.close()
+
+        print(f"discover {n_pos} pos pattern and {n_neg} neg patterns for n = {n_seqlet}")
+    print(f'done for {task}')

model/src/mrl_te_optimization/script/train_kmer_classifer.py

@@ -0,0 +1,239 @@
+import os, sys
+# os.environ["CUDA_VISIBLE_DEVICES"] = str(3)
+import pytorch_lightning as pl
+from torch.nn import functional as F
+import PATH, utils
+import torch
+from torch import nn, optim
+from models import reader
+from models.popen import Auto_popen
+import pandas as pd
+import numpy as np
+from collections import OrderedDict
+import torchmetrics 
+from torch.utils.data import DataLoader
+from pytorch_lightning import callbacks 
+from sklearn.model_selection import train_test_split
+import argparse
+
+################
+if __name__ == '__main__':
+    global_seed = int(sys.argv[6])
+    torch.manual_seed(global_seed)
+    torch.cuda.manual_seed_all(global_seed)
+
+
+###### data ######
+def get_kmer_input_shape(csv_name, kernel_size):
+    MPA_U_test = pd.read_csv(os.path.join(utils.data_dir, f"{csv_name}_test.csv"))
+
+    # dataset
+    test_DS = reader.kmer_scan_dataset(MPA_U_test, seq_col='utr', 
+                                       kmer_size=kernel_size, aux_columns='rl')
+    return np.multiply(*test_DS[0][0].shape)
+
+def get_kmer_dls(csv_name, kernel_size, seed, seq_col, label_col):
+    
+    print("seed = %s"%seed)
+    train_df, val_df, test_df = reader.split_DF(csv_name, None, [0.8,0.1,0.1], kfold_cv=True, kfold_index=seed,seed=43)
+    
+    # dataset
+    train_DS = reader.kmer_scan_dataset(train_df, seq_col=seq_col, kmer_size=kernel_size, aux_columns=label_col)
+    val_DS = reader.kmer_scan_dataset(val_df, seq_col=seq_col, kmer_size=kernel_size, aux_columns=label_col)
+    test_DS = reader.kmer_scan_dataset(test_df, seq_col=seq_col, kmer_size=kernel_size, aux_columns=label_col)
+
+    # dataloader
+    train_dl = DataLoader(train_DS, batch_size = 64, shuffle=True)
+    val_dl = DataLoader(val_DS, batch_size = 64, shuffle=False)
+    test_dl = DataLoader(test_DS, batch_size = 64, shuffle=False)
+    return train_dl, val_dl, test_dl
+
+##################
+# PyTorch Light model #
+class mlp_models(pl.LightningModule):
+    def __init__(self, dims):
+        super().__init__()
+        self.dims = dims
+        n_layer = len(dims) - 1 
+        self.train_r2 = torchmetrics.R2Score()
+        self.val_r2 = torchmetrics.R2Score()
+        # add layers
+        nns = []
+        i = 1
+        for in_dim , out_dim in zip(dims[:-1], dims[1:]):
+            nns.append( (f"Linear_{i}", nn.Linear(in_dim, out_dim)) )
+            if i < n_layer:
+                nns +=  [(f"BN_{i}", nn.BatchNorm1d(out_dim)), (f"act_{i}", nn.Mish()) ]
+            i += 1
+    
+        self.network = nn.Sequential(OrderedDict(nns))
+
+    def forward(self, x):
+        return self.network(x)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        b = x.shape[0]
+        x = x.view(b, -1).float()
+        y = y.view(b,).float()
+        y_hat = self.network(x).view(b,)
+        loss = F.mse_loss(y_hat, y.float())
+        r2 = self.train_r2(y_hat, y)
+        self.log('train_loss', loss)
+        self.log('train_acc', self.train_r2)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        b = x.shape[0]
+        x = x.view(b, -1).float()
+        y = y.view(b,).float()
+        y_hat = self.network(x).view(b,)
+        
+        loss = F.mse_loss(y_hat, y.float())
+        r2 = self.val_r2(y_hat, y)
+        self.log('val_loss', loss)
+        self.log('val_acc', self.val_r2)
+
+# PyTorch Light model #
+class rnn_models(pl.LightningModule):
+    def __init__(self, k, hidden=128):
+        super().__init__()
+        self.save_hyperparameters()
+        self.train_F1 = torchmetrics.F1Score()
+        self.train_AUROC = torchmetrics.AUROC()
+        self.train_ACC = torchmetrics.Accuracy()
+
+        self.val_F1 = torchmetrics.F1Score()
+        self.val_AUROC = torchmetrics.AUROC()
+        self.val_ACC = torchmetrics.Accuracy()
+
+        self.test_F1 = torchmetrics.F1Score()
+        self.test_AUROC = torchmetrics.AUROC()
+        self.test_ACC = torchmetrics.Accuracy()
+       
+        tower = { f"GRU_layer" : nn.GRU(input_size=4**k, hidden_size=hidden,
+                                      num_layers=2,batch_first=True),
+                 f"fc_out" : nn.Linear(hidden, 1), 
+                }
+        self.sigmod = nn.Sigmoid()
+        self.tower = nn.ModuleDict(tower)
+
+    def forward(self, x):
+        
+        x = x.transpose(1,2) # B C L -> B L C
+        h_prim,(c1,c2) = self.tower['GRU_layer'](x)
+        out = self.tower['fc_out'](c2)
+        return self.sigmod(out)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        b = x.shape[0]
+        x = x.float()
+        y = y.view(b,)
+        y_hat = self.forward(x).view(b,)
+        
+        loss = F.binary_cross_entropy(y_hat, y.float())
+        F1 = self.train_F1(y_hat, y.long())
+        auroc = self.train_AUROC(y_hat, y.long())
+        acc = self.train_ACC(y_hat, y.long())
+        self.log('train_loss', loss)
+        self.log('train_f1', self.train_F1)
+        self.log('train_AUROC', self.train_AUROC)
+        self.log('train_ACC', self.train_ACC)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        b = x.shape[0]
+        x = x.float()
+        y = y.view(b,)
+        y_hat = self.forward(x).view(b,)
+        
+        loss = F.binary_cross_entropy(y_hat, y.float())
+        self.val_F1(y_hat, y.long())
+        self.val_AUROC(y_hat, y.long())
+        self.val_ACC(y_hat, y.long())
+        self.log('val_loss', loss)
+        self.log('val_F1', self.val_F1)
+        self.log('val_AUROC', self.val_AUROC)
+        self.log('val_ACC', self.val_ACC)
+    
+    def test_step(self, test_batch, batch_idx):
+        x, y = test_batch
+        b = x.shape[0]
+        x = x.float()
+        y = y.view(b,)
+        y_hat = self.forward(x).view(b,)
+        
+        loss = F.binary_cross_entropy(y_hat, y.float())
+        self.test_F1(y_hat, y.long())
+        self.test_AUROC(y_hat, y.long())
+        self.test_ACC(y_hat, y.long())
+        self.log('test_loss', loss)
+        self.log('test_F1', self.test_F1)
+        self.log('test_AUROC', self.test_AUROC)
+        self.log('test_ACC', self.test_ACC)
+        
+################
+if __name__ == '__main__':
+
+    csv_name = sys.argv[1]
+    kmer_size = int(sys.argv[2])
+    hidden = int(sys.argv[3])
+    seq_col = sys.argv[4]
+    label_col = sys.argv[5]
+
+    
+    
+    train_dl, val_dl, test_dl = get_kmer_dls(csv_name, kmer_size, global_seed, seq_col, label_col)
+    
+
+    ####  hyper-params  ####
+    Input_length = np.multiply(*train_dl.dataset[0][0].shape)
+    # dims = [Input_length] + hidden + [1]
+    
+    model = rnn_models(kmer_size, hidden)
+    default_root_dir="/data/users/wergillius/UTR_VAE/pth/Kmer_Alan"
+    # default_root_dir="/ssd/users/wergillius/Project/MTtrans/evaluation/Kmer_results"
+    data_name = os.path.basename(csv_name).split("_")[0]
+    log_dir = os.path.join(default_root_dir, f"{data_name}_K{kmer_size}H{hidden}_sd{global_seed}")
+    ################
+
+    # training
+    trainer = pl.Trainer(accelerator='gpu',devices=1, auto_select_gpus=False, # 
+                         default_root_dir=log_dir,  
+                         limit_train_batches=0.5, max_epochs=600, 
+                         #plugins=pl.plugins.DDPPlugin(find_unused_parameters=False),
+                         callbacks=[
+                            callbacks.ModelCheckpoint(monitor="val_loss",save_top_k=1),
+                            callbacks.EarlyStopping(monitor="val_F1", mode="min", patience=15)                                    
+                         ])
+
+   
+    trainer.fit(model, train_dl, val_dl)
+
+    ckpt_dir = os.path.join(trainer.log_dir,'checkpoints')
+    ckpt_path = os.path.join(ckpt_dir, [file for file in os.listdir(ckpt_dir) if file.endswith('.ckpt')][0])
+    print('model saved to : %s\n'%ckpt_path)
+    # trainer.validate(model, test_dl)
+
+    saved_model = rnn_models.load_from_checkpoint(ckpt_path)
+    trainer.validate(saved_model, val_dl)
+    trainer.test(saved_model, test_dl)
+
+# debugging args:
+    # "/data/users/wergillius/UTR_VAE/Alan_dataset/AlanAll_binary_10pctg.csv",
+    #                  "3",
+    #                  "64",
+    #                  "seq",
+    #                  "Binary_10pc",
+    #                  "1"

model/src/mrl_te_optimization/script/train_kmer_models.py

@@ -0,0 +1,167 @@
+import os, sys
+# os.environ["CUDA_VISIBLE_DEVICES"] = str(3)
+import pytorch_lightning as pl
+from torch.nn import functional as F
+import PATH, utils
+import torch
+from torch import nn, optim
+from models import reader
+from models.popen import Auto_popen
+import pandas as pd
+import numpy as np
+from collections import OrderedDict
+import torchmetrics 
+from torch.utils.data import DataLoader
+from pytorch_lightning import callbacks 
+
+
+###### data ######
+def get_kmer_input_shape(csv_name, kernel_size):
+    MPA_U_test = pd.read_csv(os.path.join(utils.data_dir, f"{csv_name}_test.csv"))
+
+    # dataset
+    test_DS = reader.kmer_scan_dataset(MPA_U_test, seq_col='utr', 
+                                       kmer_size=kernel_size, aux_columns='rl')
+    return np.multiply(*test_DS[0][0].shape)
+
+def get_kmer_dls(csv_name, kernel_size, seed):
+    train_val = pd.read_csv(os.path.join(utils.data_dir, f"{csv_name}_train_val.csv"))
+    MPA_U_test = pd.read_csv(os.path.join(utils.data_dir, f"{csv_name}_test.csv"))
+
+    MPA_U_val = train_val.sample(frac=0.1, random_state=seed)
+    MPA_U_train = pd.concat([train_val, MPA_U_val]).drop_duplicates(keep=False)
+
+    # dataset
+    train_DS = reader.kmer_scan_dataset(MPA_U_train, seq_col='utr', kmer_size=kernel_size, aux_columns='rl')
+    val_DS = reader.kmer_scan_dataset(MPA_U_val, seq_col='utr', kmer_size=kernel_size, aux_columns='rl')
+    test_DS = reader.kmer_scan_dataset(MPA_U_test, seq_col='utr', kmer_size=kernel_size, aux_columns='rl')
+
+    # dataloader
+    train_dl = DataLoader(train_DS, batch_size = 64, shuffle=True)
+    val_dl = DataLoader(test_DS, batch_size = 64, shuffle=False)
+    test_dl = DataLoader(test_DS, batch_size = 64, shuffle=False)
+    return train_dl, val_dl, test_dl
+
+##################
+# PyTorch Light model #
+class mlp_models(pl.LightningModule):
+    def __init__(self, dims):
+        super().__init__()
+        self.dims = dims
+        n_layer = len(dims) - 1 
+        self.train_r2 = torchmetrics.R2Score()
+        self.val_r2 = torchmetrics.R2Score()
+        # add layers
+        nns = []
+        i = 1
+        for in_dim , out_dim in zip(dims[:-1], dims[1:]):
+            nns.append( (f"Linear_{i}", nn.Linear(in_dim, out_dim)) )
+            if i < n_layer:
+                nns +=  [(f"BN_{i}", nn.BatchNorm1d(out_dim)), (f"act_{i}", nn.Mish()) ]
+            i += 1
+    
+        self.network = nn.Sequential(OrderedDict(nns))
+
+    def forward(self, x):
+        return self.network(x)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        b = x.shape[0]
+        x = x.view(b, -1).float()
+        y = y.view(b,).float()
+        y_hat = self.network(x).view(b,)
+        loss = F.mse_loss(y_hat, y.float())
+        r2 = self.train_r2(y_hat, y)
+        self.log('train_loss', loss)
+        self.log('train_acc', self.train_r2)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        b = x.shape[0]
+        x = x.view(b, -1).float()
+        y = y.view(b,).float()
+        y_hat = self.network(x).view(b,)
+        
+        loss = F.mse_loss(y_hat, y.float())
+        r2 = self.val_r2(y_hat, y)
+        self.log('val_loss', loss)
+        self.log('val_acc', self.val_r2)
+
+# PyTorch Light model #
+class rnn_models(pl.LightningModule):
+    def __init__(self, k):
+        super().__init__()
+        
+        self.train_r2 = torchmetrics.R2Score()
+        self.val_r2 = torchmetrics.R2Score()
+
+       
+        tower = { f"GRU_layer" : nn.GRU(input_size=4**k, hidden_size=128,
+                                      num_layers=2,batch_first=True),
+                 f"fc_out" : nn.Linear(128, 1), 
+                }
+    
+        self.tower = nn.ModuleDict(tower)
+
+    def forward(self, x):
+        
+        x = x.transpose(1,2) # B C L -> B L C
+        h_prim,(c1,c2) = self.tower['GRU_layer'](x)
+        out = self.tower['fc_out'](c2)
+        return out
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        b = x.shape[0]
+        x = x.float()
+        y = y.view(b,).float()
+        y_hat = self.forward(x).view(b,)
+        
+        loss = F.mse_loss(y_hat, y)
+        r2 = self.train_r2(y_hat, y)
+        self.log('train_loss', loss)
+        self.log('train_acc', self.train_r2)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        b = x.shape[0]
+        x = x.float()
+        y = y.view(b,).float()
+        y_hat = self.forward(x).view(b,)
+        
+        loss = F.mse_loss(y_hat, y)
+        self.val_r2(y_hat, y)
+        self.log('val_loss', loss)
+        self.log('val_acc', self.val_r2)
+        
+################
+if __name__ == '__main__':
+
+    csv_name = sys.argv[1]
+    kmer_size = sys.argv[2]
+    hidden = sys.argv[3]
+
+    ####  hyper-params  ####
+    Input_length = np.multiply(*train_DS[0][0].shape)
+    hidden = [512]
+    dims = [Input_length] + hidden + [1]
+    ################
+
+    # training
+    trainer = pl.Trainer(gpus=1, num_processes=8,
+                         default_root_dir="/ssd/users/wergillius/Project/MTtrans/evaluation/Kmer_results",
+                         limit_train_batches=0.5, max_epochs=6, 
+                         callbacks=[callbacks.EarlyStopping(monitor="val_loss", mode="min", patience=5)])
+    trainer.fit(model, train_dl, val_dl)
+    trainer.test(test_dl)

model/src/mrl_te_optimization/script/train_oneLayer_CNN.py

@@ -0,0 +1,114 @@
+import os, sys
+import pytorch_lightning as pl
+from torch.nn import functional as F
+import PATH, utils
+import torch
+from torch import nn, optim
+from models import reader
+from models.popen import Auto_popen
+import pandas as pd
+import numpy as np
+from collections import OrderedDict
+import torchmetrics 
+from torch.utils.data import DataLoader
+from pytorch_lightning import callbacks 
+
+
+###### data ######
+def get_dls(csv_name, kernel_size, seed):
+    train_val = pd.read_csv(os.path.join(utils.data_dir, f"{csv_name}_train_val.csv"))
+    MPA_U_test = pd.read_csv(os.path.join(utils.data_dir, f"{csv_name}_test.csv"))
+
+    MPA_U_val = train_val.sample(frac=0.1, random_state=seed)
+    MPA_U_train = pd.concat([train_val, MPA_U_val]).drop_duplicates(keep=False)
+
+    # dataset
+    train_DS = reader.MTL_dataset(MPA_U_train, seq_col='utr', aux_columns=['rl'])
+    val_DS = reader.MTL_dataset(MPA_U_val, seq_col='utr', aux_columns=['rl'])
+    test_DS = reader.MTL_dataset(MPA_U_test, seq_col='utr', aux_columns=['rl'])
+
+    # dataloader
+    train_dl = DataLoader(train_DS, batch_size = 64, shuffle=True)
+    val_dl = DataLoader(test_DS, batch_size = 64, shuffle=False)
+    test_dl = DataLoader(test_DS, batch_size = 64, shuffle=False)
+    return train_dl, val_dl, test_dl
+##################
+
+
+# PyTorch Light model #
+class Onelayer_CNN(pl.LightningModule):
+    def __init__(self, kernel_size):
+        super().__init__()
+
+        self.train_r2 = torchmetrics.R2Score()
+        self.val_r2 = torchmetrics.R2Score()
+
+        self.conv_layer = nn.Sequential(nn.Conv1d(4, 256, kernel_size),
+                                    nn.BatchNorm1d(256),
+                                    nn.Mish())
+        
+        tower = { f"GRU_layer" : nn.GRU(input_size=256, hidden_size=128,
+                                      num_layers=2,batch_first=True),
+                 f"fc_out" : nn.Linear(128, 1), 
+                }
+
+        self.tower = nn.ModuleDict(tower)
+
+    def forward(self, x):
+        if x.shape[1] != 4:
+            x = x.transpose(1,2)
+        Z = self.conv_layer(x)
+        Z_t = torch.transpose(Z, 1, 2)
+        h_prim,(c1,c2) = self.tower['GRU_layer'](Z_t)
+        out = self.tower['fc_out'](c2)
+        return out
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        b = x.shape[0]
+        x = x.float()
+        y = y.view(b,).float()
+        y_hat = self.forward(x).view(b,)
+        
+        loss = F.mse_loss(y_hat, y)
+        r2 = self.train_r2(y_hat, y)
+        self.log('train_loss', loss)
+        self.log('train_acc', self.train_r2)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        b = x.shape[0]
+        x = x.float()
+        y = y.view(b,).float()
+        y_hat = self.forward(x).view(b,)
+        loss = F.mse_loss(y_hat, y)
+        self.val_r2(y_hat, y)
+        self.log('val_loss', loss)
+        self.log('val_acc', self.val_r2)
+
+################
+if __name__ == '__main__':
+    
+    ####  hyper-params  ####
+    csv_name = sys.argv[1]
+    kernel_size = sys.argv[2]
+    seed = 41
+    
+    ################
+    train_dl, val_dl, test_dl = get_dls(csv_name, kernel_size, seed)
+    
+    model = Onelayer_CNN(int(kernel_size))
+
+
+    # training
+    trainer = pl.Trainer(gpus=1, num_processes=8,
+                         default_root_dir="/ssd/users/wergillius/Project/MTtrans/evaluation/one_layer_logs",
+                         limit_train_batches=0.5, max_epochs=60, 
+                         callbacks=[callbacks.EarlyStopping(monitor="val_loss", mode="min", patience=5)])
+    trainer.fit(model, train_dl, val_dl)
+    trainer.test(test_dl)

model/src/mrl_te_optimization/util.py

@@ -0,0 +1,384 @@
+import os
+import operator
+import pickle
+import time
+import seaborn as sns
+import matplotlib.pyplot as plt
+import pandas as pd
+import numpy as np
+import math
+import itertools
+import random
+import os
+from polyleven import levenshtein
+import operator
+import pickle
+import time
+import tensorflow as tf
+
+tf.random.set_seed(35)
+np.random.seed(35)
+
+rna_vocab = {"A":0,
+             "C":1,
+             "G":2,
+             "U":3,
+             "*":4}
+
+rev_rna_vocab = {v:k for k,v in rna_vocab.items()}
+
+def recover_seq(samples, rev_charmap=rev_rna_vocab):
+    """Convert samples to strings and save to log directory."""
+    
+    char_probs = samples
+    argmax = np.argmax(char_probs, 2)
+    seqs = []
+    for line in argmax:
+        s = "".join(rev_charmap[d] for d in line)
+        s = s.replace('*','')
+        seqs.append(s)
+    return seqs 
+
+def file_to_list(file_name,size):
+    data = []
+    with open(file_name,'r') as f:
+        lines = f.readlines()
+        
+        for seq in lines:
+            seq_ = seq.replace('\n','')
+            data.append(seq_)
+            # if len(seq) == size:
+            #     data.append(seq)
+
+    return data
+
+def one_hot_encode(seq):
+    mapping = dict(zip("ACGT", range(4)))
+    seq = seq.replace('U','T')    
+    seq2 = [mapping[i] for i in seq]
+    return np.eye(4)[seq2]
+
+def encode(seqs):
+    return np.reshape([one_hot_encode(seqs[i]) for i in range(len(seqs))],(np.array(seqs).shape[0],-1))
+
+min_len = None
+nuc_dict = {'a':[1.0,0.0,0.0,0.0],'c':[0.0,1.0,0.0,0.0],'g':[0.0,0.0,1.0,0.0], 
+                'u':[0.0,0.0,0.0,1.0], 't':[0.0,0.0,0.0,1.0], 
+                'n':[0.0,0.0,0.0,0.0], 'x':[1/4,1/4,1/4,1/4]}
+
+def encode_seq_framepool(seq, max_len=128):
+    # print(seq)
+
+    length = len(seq)
+    if max_len > 0 and min_len is None:
+        padding_needed = max_len - length
+        seq = "N"*padding_needed + seq
+    if min_len is not None:
+        if len(seq) < min_len:
+            seq = "N"*(min_len - len(seq)) + seq
+
+        if len(seq) > min_len:
+            seq = seq[(len(seq) - min_len):]
+    seq = seq.lower()
+    one_hot = np.array([nuc_dict[x] for x in seq]) # get stacked on top of each other
+    
+    return one_hot
+
+def list_to_file(filename,list):
+    with open(filename + ".txt", 'w') as f:
+        for element in list[:-1]:
+            f.write(element+"\n")
+        f.write(list[-1])       
+
+def tsne_natural(file_name, length, key='seq'):
+    df = pd.read_csv(file_name)
+    seqs = np.array(df[key]).tolist()
+
+    selected_seqs = []
+
+    for i in range(len(seqs)):
+        seq = seqs[i]
+        seq = seq.upper()
+        if seq not in selected_seqs and len(seq) == length:
+            selected_seqs.append(seq)
+
+    return selected_seqs
+
+def read_real(file_name, UTR_LEN, key='seq', all= True, samples= 128 ):
+    df = pd.read_csv(file_name)
+    seqs = np.array(df[key]).tolist()
+
+    selected_seqs = []
+
+    for i in range(len(seqs)):
+        if len(seqs[i]) < (UTR_LEN + 1) and len(seqs[i]) > int(UTR_LEN/2):
+            seqs[i] = seqs[i].upper()
+            if seqs[i] not in selected_seqs:
+                selected_seqs.append(seqs[i])
+
+    if all:
+        return selected_seqs
+    
+    else:
+        indices = []
+
+        for i in range(len(selected_seqs)):
+
+            indices.append(i)
+
+        samples = np.random.choice(len(indices),samples,replace=False)
+
+        chosen = []
+
+        for i in range(len(samples)):
+            chosen.append(selected_seqs[samples[i]])
+            
+        return chosen
+    
+def random_sample(length):
+    rna_vocab = {"A":0,
+             "C":1,
+             "G":2,
+             "T":3}
+
+    rev_rna_vocab = {v:k for k,v in rna_vocab.items()}
+
+    mapping = dict(zip([0,1,2,3],"ACGT"))
+
+    sample = ''
+    for i in range(length):
+        r = random.random()
+        if r < 0.6:
+            sample += random.choice(['C','G'])
+        else:
+            sample += random.choice(['A','T'])
+    # rsample = [random.randrange(4) for i in range(length)]
+    # rsample = [mapping[i] for i in rsample]
+    # string = ''
+    # for neuc in rsample:
+    #     string += neuc
+
+    return sample 
+
+def random_data(length, size):
+    samples = []
+    for i in range(size):
+        samples.append(random_sample(length))
+
+    return samples
+
+def motif_count(gens,motifs_path='./../../data/motifs.csv'):
+    motifs = pd.read_csv(motifs_path)
+    seqs = motifs['seq']
+    scores = motifs['spr']
+    gen_positive = 0
+    gen_negative = 0
+    for i in range(len(seqs)):
+        for j in range(len(gens)):
+            if seqs[i] in gens[j]:
+                if scores[i] < 0:
+                    gen_negative += 1
+                else:
+                    gen_positive += 1
+
+    return gen_positive, gen_negative
+
+def tsne_gen(path,length=128,BATCH_SIZE=64,DIM=40):
+    wgan = tf.keras.models.load_model(path)
+
+    selected = []
+    while len(selected) < BATCH_SIZE:
+
+        noise = tf.Variable(tf.random.normal(shape=[64,DIM]))
+
+        sequences_init = wgan(noise)
+
+        gen_seqs_init = sequences_init.numpy().astype('float')
+
+        seqs_gen_init = recover_seq(gen_seqs_init, rev_rna_vocab)
+
+        for i in range(len(seqs_gen_init)):
+            if len(seqs_gen_init[i]) == length:
+                selected.append(seqs_gen_init[i])
+
+    return selected[:BATCH_SIZE]
+
+def generate_data(path,BATCH_SIZE=64,UTR_LEN=128,DIM=40):
+    wgan = tf.keras.models.load_model(path)
+
+    gens = []
+
+    while len(gens) < BATCH_SIZE:
+
+        noise = tf.Variable(tf.random.normal(shape=[BATCH_SIZE,DIM]))
+
+        sequences_init = wgan(noise)
+
+        gen_seqs_init = sequences_init.numpy().astype('float')
+
+        seqs_gen_init = recover_seq(gen_seqs_init, rev_rna_vocab)
+
+        for i in range(len(seqs_gen_init)):
+            if len(seqs_gen_init[i]) > int(UTR_LEN/2) and len(seqs_gen_init[i])< UTR_LEN+1:
+                if seqs_gen_init[i] not in gens:
+                    gens.append(seqs_gen_init[i])
+                    if len(gens) == BATCH_SIZE:
+                        break
+
+    return gens[:BATCH_SIZE]
+
+def gc_percentage(seq):
+    count = 0.0
+    for char in seq:
+        if char == 'C' or char == 'G':
+            count +=1
+
+    return float(count/len(seq))
+
+def get_gc_content(data):
+    gc_content = []
+    for seq in data:
+        seq.replace('\n','')
+        seq.replace('*','')
+        gc = gc_percentage(seq)
+        gc_content.append(gc)
+
+    return gc_content
+
+def get_gc_content_many(data):
+
+    collection = []
+    gc_contents = []
+    for seq in data:
+        seq = seq.upper()
+        seq.replace('\n','')
+        seq.replace('*','')
+        gc = gc_percentage(seq)
+        gc_contents.append(gc)
+    
+    return gc_contents
+
+def get_4mers():
+    neucs = ['A','C','G','T']
+    
+    mers = [p for p in itertools.product(neucs, repeat=4)]
+    for i in range(len(mers)):
+        mers[i] = mers[i][0] + mers[i][1] + mers[i][2] + mers[i][3]
+
+    return mers
+
+def get_4mer_dic(seqs):
+    
+    _4mers = get_4mers()
+    length = 0
+    dics = []
+
+    for seq in seqs:
+        dic = {}
+        for item in _4mers:
+            dic[item] = 0
+
+        # Iterate With the Sliding Window
+        length = len(seq)
+        limit = length - 4
+        for i in range(limit):
+            mer = seq[i:i+4]
+            dic[mer] += 1
+
+        dics.append(dic)
+    
+    return dics
+
+def euclidean_kmer(item, ref, mers):
+    dist = 0
+    for mer in mers:
+        diff = item[mer] - ref[mer]
+        dist += math.pow(diff,2)
+
+    return dist
+
+def euclidean_kmer_all(item,refs,mers):
+    dists = []
+    for ref in refs:
+        dists.append(euclidean_kmer(item,ref,mers))
+
+    return min(dists)
+
+def hamming_dist(src, target):
+    
+    dists = []
+    for i in range(len(src)):
+        smallest = np.inf
+        for j in range(len(target)):
+            dist = levenshtein(src[i],target[j])
+            if dist > 0 and dist < smallest:
+                smallest = dist
+
+        dists.append(dist)
+
+    return np.array(dists)
+
+
+
+
+
+
+
+
+
+def one_hot_motif(seq,length=128,complementary=False):
+    """
+    one_hot encoding on sequence
+    complementary: encode nucleatide into complementary one
+    """
+    
+    if length == -1:
+        length = len(seq)
+    
+    # seq = str(seq)
+    # setting
+    seq = list(seq.replace("U","T"))
+    seq_len = len(seq)
+    complementary = -1 if complementary else 1
+    # compose dict
+    keys = ['A', 'C', 'G', 'T'][::complementary]
+    oh_dict = {keys[i]:i for i in range(4)}
+    # array
+    oh_array = np.zeros((length,4),dtype=np.double)
+    for i,C in enumerate(seq):
+        try:
+            oh_array[i,oh_dict[C]]=1
+        except:
+            continue      # for nucleotide that are not in A C G T   
+    return oh_array
+
+def one_hot_all_motif(seqs):
+    # length = np.max([len(seq) for seq in seqs])
+    return [one_hot_motif(seqs[i], length = 128) for i in range(len(seqs))]
+
+
+
+
+def one_hot(seq):
+    convert = False
+    if isinstance(seq, tf.Tensor):
+        seq = seq.numpy().astype(str)
+        convert = True
+
+    num_seqs = len(seq)
+    seq_len = len(seq[0])
+    seqindex = {'A':0, 'C':1, 'G':2, 'T':3, 'a':0, 'c':1, 'g':2, 't':3}
+    seq_vec = np.zeros((num_seqs,seq_len,4), dtype='bool')
+    for i in range(num_seqs):
+        thisseq = seq[i]
+        for j in range(seq_len):
+            try:
+                seq_vec[i,j,seqindex[thisseq[j]]] = 1
+            except:
+                pass
+    
+    if convert:
+        seq_vec = tf.convert_to_tensor(seq_vec,dtype=tf.float32)
+
+
+    return seq_vec

model/src/mrl_te_optimization/utils.py

@@ -0,0 +1,209 @@
+from sklearn.preprocessing import OneHotEncoder
+import logging
+import numpy as np
+import pandas as pd
+import os
+import json
+import re
+import torch
+import collections
+# from models import reader
+from models.ScheduleOptimizer import ScheduledOptim
+
+print(os.path.dirname(__file__))
+
+# ====================|   some path   |=======================
+global script_dir
+global data_dir
+global log_dir
+global pth_dir
+# global cell_lines
+
+global egfp_seq
+
+with open(os.path.join(os.path.dirname(__file__),"machine_configure.json"),'r') as f:
+    config = json.load(f)   
+
+script_dir = config['script_dir']
+data_dir = config['data_dir']
+log_dir = config['log_dir']
+pth_dir = config['pth_dir']
+
+
+
+# =====================| one hot encode |=======================
+
+class Seq_one_hot(object):
+    def __init__(self,seq_type='nn',seq_len=100):
+        """
+        initiate the sequence one hot encoder
+        """
+        self.seq_len=seq_len
+        self.seq_type =seq_type
+        self.enable_encoder()
+        
+    def enable_encoder(self):
+        if self.seq_type == 'nn':
+            self.encoder = OneHotEncoder(sparse=False)
+            self.encoder.drop_idx_ = None
+            self.encoder.categories_ = [np.array(['A', 'C', 'G', 'T'], dtype='<U1')]*self.seq_len
+
+    def discretize_seq(self,data):
+        """
+        discretize sequence into character
+        argument:
+        ...data: can be dataframe with UTR columns , or can be single string
+        """
+        if type(data) is pd.DataFrame:
+            return np.stack(data.UTR.apply(lambda x: list(x)))
+        elif type(data) is str:
+            return np.array(list(data))
+    
+    def transform(self,data,flattern=True):
+        """
+        One hot encode
+        argument:
+        data : is a 2D array
+        flattern : True
+        """
+        X = self.encoder.transform(data)                             # 400 for each seq
+        X_M = np.stack([seq.reshape(self.seq_len,4) for seq in X])   # i.e 100*4
+        return X if flattern else X_M
+    
+    def d_transform(self,data,flattern=True):
+        """
+        discretize data and put into transform
+        """
+        X = self.discretize_seq(data)
+        return self.transform(X,flattern)
+
+
+# =====================|   logger       |=======================
+
+def setup_logs(vae_log_path,level=None):
+    """
+
+    :param save_dir:  the directory to set up logs
+    :param type:  'model' for saving logs in 'logs/cpc'; 'imp' for saving logs in 'logs/imp'
+    :param run_name:
+    :return:logger
+    """
+    # initialize logger
+    logger = logging.getLogger("VAE")
+    logger.setLevel(logging.INFO)
+    if level=='warning':
+        logger.setLevel(logging.WARNING)
+
+    # create the logging file handler
+    log_file = os.path.join(vae_log_path)
+    fh = logging.FileHandler(log_file)
+
+    # create the logging console handler
+    ch = logging.StreamHandler()
+
+    # format
+    formatter = logging.Formatter("%(asctime)s - %(message)s")
+    fh.setFormatter(formatter)
+
+    # add handlers to logger object
+    logger.addHandler(fh)
+    logger.addHandler(ch)
+
+    return logger
+
+def clean_value_dict(dict):
+    """
+    deal with verbose dict where the values maybe torch object, extact the item and return clean dict
+    """
+    clean_dict={}
+    for k,v in dict.items():
+        
+        try:
+            v = v.item()
+        except:
+            v = v
+        clean_dict[k] = v
+    return clean_dict
+
+def fix_parameter(model,modual_to_fix,fix_or_unfix=False):
+    """
+    for a given model, fix part of the parameter to fine-tuning / transfering 
+    args:
+    model : `nn.Modual`,initiated model instance
+    modual_to_fix : str, define which part of the model will not update by gradient 
+                    e.g. "soft_share" then 
+    """
+    
+    fix_part = eval("model."+modual_to_fix)   # e.g. model.shoft_share
+     
+    for param in fix_part.parameters():
+            param.requires_grad = fix_or_unfix
+    
+    return model
+
+def unfix_parameter(model,modual_to_fix,fix_or_unfix=False):
+    return fix_parameter(model,modual_to_fix,fix_or_unfix=True)
+
+def snapshot(vae_pth_path, state):
+    logger = logging.getLogger("VAE")
+    # torch.save can save any object
+    # dict type object in our cases
+    torch.save(state, vae_pth_path)
+    logger.info("Snapshot saved to {}\n".format(vae_pth_path))
+
+
+def load_model(popen,model,logger=None):
+    
+    info = lambda x: print(x) if logger==None else logger.info(x)
+    # popen.vae_pth_path = '/mnt/sina/run/ml/gan/dev/git/UTRGAN/src/mrl_optimization/script/checkpoint/RL_hard_share_MTL/3M/small_repective_filed_strides1113-model_best_cv1.pth'
+    checkpoint = torch.load(popen.vae_pth_path, map_location=torch.device('cpu')) 
+    if isinstance(checkpoint['state_dict'], collections.OrderedDict):
+            # optimizer.load_state_dict(checkpoint['optimizer'])
+        model.load_state_dict(checkpoint['state_dict'])
+    else:
+        model = checkpoint['state_dict']
+    
+    info(' \t \t ==============<<< encoder load from >>>============== \t \t ')
+    info(" \t"+popen.vae_pth_path)
+    
+    return model
+    
+def get_config_cuda(config_file):
+    with open(config_file,'r') as f:
+        lines = f.read_lines()
+        for line in lines:
+            if "cuda_id =" in line:
+                device = line.split("=")[1].strip()
+                break
+    device = int(device) if device.isdigit() else device
+    return device
+
+def resume(popen,optimizer,logger):
+    """
+    for a experiment, check whether it;s a new run, and create dir 
+    """
+    #run_name = model_stype + time.strftime("__%Y_%m_%d_%H:%M"))
+    
+    if popen.Resumable:
+        
+        checkpoint = torch.load(popen.vae_pth_path, map_location=torch.device('cpu'))   # xx-model-best.pth
+        previous_epoch = checkpoint['epoch']
+        previous_loss = checkpoint['validation_loss']
+        previous_acc = checkpoint['validation_acc']
+        
+        
+        # very important
+        if (type(optimizer) == ScheduledOptim):
+            optimizer.n_current_steps = popen.n_current_steps
+            optimizer.delta = popen.delta
+        
+        logger.info(" \t \t ========================================================= \t \t ")
+        logger.info(' \t \t ==============<<< Resume from checkpoint>>>============== \t \t \n')
+        logger.info(" \t"+popen.vae_pth_path+'\n')
+        logger.info(" \t \t ========================================================= \t \t \n")
+        
+        return previous_epoch,previous_loss,previous_acc
+        
+    
+egfp_seq = "atgggcgaattaagtaagggcgaggagctgttcaccggggtggtgcccatcctggtcgagctggacggcgacgtaaacggccacaagttcagcgtgtccggcgagggcgagggcgatgccacctacggcaagctgaccctgaagttcatctgcaccaccggcaagctgcccgtgccctggcccaccctcgtgaccaccctgacctacggcgtgcagtgcttcagccgctaccccgaccacatgaagcagcacgacttcttcaagtccgccatgcccgaaggctacgtccaggagcgcaccatcttct"
+eGFP_seq = egfp_seq.upper()

requirements.txt

@@ -0,0 +1,62 @@
+# UTRGAN on the OneScience base environment
+# Python 3.11; SCNet BW DCU
+#
+# IMPORTANT RUNTIME DIFFERENCE
+# ----------------------------
+# OneScience is pinned to DTK 25.04 / DAS1.7:
+#   torch==2.5.1+das.opt1.dtk25042
+#   tensorflow==2.18.0+das.opt1.dtk25042
+# UTRGAN ReLU/model inference failed there with "Failure when generating HSACO".
+# The validated UTRGAN environment therefore uses DTK 26.04 / DAS1.8 and must
+# load `compiler/dtk/26.04` (or an equivalent /opt/dtk-26.04 runtime).
+#
+# Install the active entries without replacing their dependencies:
+#   python -m pip install --no-deps -r requirements.txt
+# Recheck torch/TensorFlow/HIP/DCU visibility immediately after installation.
+
+# Required replacements for the OneScience DTK 26.04 framework wheels.
+tensorflow @ https://download.sourcefind.cn:65024/file/4/tensorflow/DAS1.8/tensorflow-2.18.0+das.opt1.dtk2604-cp311-cp311-manylinux_2_28_x86_64.whl
+torch @ https://download.sourcefind.cn:65024/file/4/pytorch/DAS1.8/torch-2.5.1+das.opt1.dtk2604-cp311-cp311-manylinux_2_28_x86_64.whl
+
+# Required compatibility/additional packages.
+# OneScience includes tf-keras without a strict version; 2.18.0 is the tested
+# version for loading the released legacy H5 files with TF_USE_LEGACY_KERAS=1.
+tf-keras==2.18.0
+polyleven==0.9.0
+
+# Additional packages needed only by the repository's G4/analysis utilities.
+# XGBoost is held below 2.3 because the released G4 JSON files use its old
+# pre-1.6 JSON format, whose removal was announced for XGBoost 2.3.
+xgboost==2.1.4
+ViennaRNA==2.7.0
+logomaker
+ruptures
+cliffs-delta
+
+# NUPACK is used only by script/Nupack_MFE_prediction.py. NUPACK 4 is an
+# externally licensed/downloaded Python package and is not installed from this
+# requirements file. Install it separately from https://www.nupack.org/ only
+# when that optional preprocessing script is needed.
+# nupack  # external optional dependency; no normal PyPI pin
+
+# Provided by OneScience; complete UTRGAN dependency inventory follows.
+# Keep these commented to avoid reinstalling generic wheels over OneScience.
+# numpy==1.26.3
+# pandas>=2.2.2
+# scipy==1.14.1
+# h5py>=3.7.0
+# matplotlib
+# seaborn
+# tqdm>=4.60.0
+# requests
+# biopython==1.84
+# scikit-learn>=1.2.2,<=1.6.0
+# einops>=0.7.0
+# pytorch-lightning==2.0.6
+# torchmetrics
+
+# Supplied transitively by the validated DTK TensorFlow stack; do not adjust
+# these independently because doing so can break TensorFlow compatibility.
+# keras==3.15.1
+# protobuf
+# ml-dtypes

scripts/optimize_te_mrl.py

@@ -0,0 +1,424 @@
+
+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)}")

scripts/predict.py

@@ -0,0 +1,213 @@
+"""Generate 5' UTR candidates and rank them with FramePool and MTtrans."""
+
+import argparse
+import json
+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"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Generate UTRGAN candidates and rank by MRL and TE."
+    )
+    parser.add_argument("--num-candidates", type=int, default=1024)
+    parser.add_argument("--batch-size", type=int, default=128)
+    parser.add_argument("--seed", type=int, default=33)
+    parser.add_argument("--device", choices=("dcu", "cpu"), default="dcu")
+    parser.add_argument("--device-id", default="0")
+    parser.add_argument(
+        "--output-dir",
+        default=str(PROJECT_ROOT / "outputs" / "pretrained_batch_ranking"),
+    )
+    return parser.parse_args()
+
+
+def configure_runtime(args):
+    os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
+    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
+    if args.device == "cpu":
+        os.environ["HIP_VISIBLE_DEVICES"] = "-1"
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+    else:
+        os.environ["HIP_VISIBLE_DEVICES"] = args.device_id
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id
+    for import_root in (MODEL_ROOT, MODULE_ROOT):
+        if str(import_root) not in sys.path:
+            sys.path.insert(0, str(import_root))
+
+
+def main():
+    args = parse_args()
+    if args.num_candidates < 1 or args.batch_size < 1:
+        raise ValueError("--num-candidates and --batch-size must be positive")
+    configure_runtime(args)
+
+    import numpy as np
+    import pandas as pd
+    import tensorflow as tf
+    import torch
+
+    import framepool
+    import util
+
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    generator_path = PROJECT_ROOT / "weight" / "checkpoint_3000.h5"
+    framepool_path = PROJECT_ROOT / "weight" / "utr_model_combined_residual_new.h5"
+    mttrans_path = (
+        PROJECT_ROOT
+        / "weight"
+        / "mttrans"
+        / "RL_hard_share_MTL"
+        / "3R"
+        / "schedule_MTL-model_best_cv1.pth"
+    )
+    for path in (generator_path, framepool_path, mttrans_path):
+        if not path.is_file():
+            raise FileNotFoundError(path)
+
+    tf_device = "/GPU:0" if args.device == "dcu" else "/CPU:0"
+    torch_device = torch.device("cuda:0" if args.device == "dcu" else "cpu")
+    if args.device == "dcu":
+        tf_gpus = tf.config.list_physical_devices("GPU")
+        if not tf_gpus:
+            raise RuntimeError("TensorFlow did not detect a DCU")
+        if not torch.cuda.is_available():
+            raise RuntimeError("PyTorch did not detect a DCU")
+        for gpu in tf_gpus:
+            try:
+                tf.config.experimental.set_memory_growth(gpu, True)
+            except RuntimeError:
+                pass
+
+    # Loading on CPU avoids device-side random-initializer kernels; inference
+    # is explicitly placed on the requested device below.
+    with tf.device("/CPU:0"):
+        generator = tf.keras.models.load_model(generator_path, compile=False)
+        mrl_model = framepool.load_framepool(str(framepool_path))
+    generator.trainable = False
+    mrl_model.trainable = False
+
+    checkpoint = torch.load(
+        mttrans_path, map_location="cpu", weights_only=False
+    )
+    te_model = checkpoint["state_dict"].to(torch_device)
+    te_model.eval()
+
+    np.random.seed(args.seed)
+    tf.random.set_seed(args.seed)
+    torch.manual_seed(args.seed)
+    if args.device == "dcu":
+        torch.cuda.manual_seed_all(args.seed)
+
+    noise = np.random.RandomState(args.seed).normal(
+        size=(args.num_candidates, 40)
+    ).astype(np.float32)
+
+    generated_batches = []
+    with tf.device(tf_device):
+        for start in range(0, args.num_candidates, args.batch_size):
+            stop = min(start + args.batch_size, args.num_candidates)
+            generated_batches.append(
+                generator(tf.convert_to_tensor(noise[start:stop]), training=False).numpy()
+            )
+    generated = np.concatenate(generated_batches, axis=0)
+    if generated.shape != (args.num_candidates, 128, 5):
+        raise RuntimeError(f"Unexpected generator shape: {generated.shape}")
+    if not np.isfinite(generated).all():
+        raise RuntimeError("Generator output contains NaN/Inf")
+
+    sequences = list(util.recover_seq(generated, util.rev_rna_vocab))
+    mrl_scores = []
+    with tf.device(tf_device):
+        for start in range(0, len(sequences), args.batch_size):
+            chunk = sequences[start : start + args.batch_size]
+            encoded = np.asarray(
+                [util.encode_seq_framepool(seq) for seq in chunk],
+                dtype=np.float32,
+            )
+            prediction = mrl_model(tf.convert_to_tensor(encoded), training=False)
+            mrl_scores.extend(tf.reshape(prediction, (-1,)).numpy().tolist())
+
+    te_scores = []
+    with torch.inference_mode():
+        for start in range(0, len(sequences), args.batch_size):
+            chunk = sequences[start : start + args.batch_size]
+            encoded = np.asarray(util.one_hot_all_motif(chunk), dtype=np.float32)
+            encoded = torch.from_numpy(encoded).transpose(1, 2).to(torch_device)
+            prediction = te_model(encoded)
+            te_scores.extend(prediction.reshape(-1).cpu().numpy().tolist())
+
+    mrl_scores = np.asarray(mrl_scores, dtype=np.float32)
+    te_scores = np.asarray(te_scores, dtype=np.float32)
+    if not np.isfinite(mrl_scores).all() or not np.isfinite(te_scores).all():
+        raise RuntimeError("MRL/TE scores contain NaN/Inf")
+
+    table = pd.DataFrame(
+        {
+            "candidate_id": [
+                f"UTRGAN_{index + 1:05d}" for index in range(len(sequences))
+            ],
+            "sequence": sequences,
+            "length": [len(sequence) for sequence in sequences],
+            "mrl_score": mrl_scores,
+            "te_score": te_scores,
+        }
+    )
+    table["is_duplicate"] = table.duplicated("sequence", keep="first")
+    table["mrl_rank"] = table["mrl_score"].rank(
+        method="first", ascending=False
+    ).astype(int)
+    table["te_rank"] = table["te_score"].rank(
+        method="first", ascending=False
+    ).astype(int)
+    unique = table.drop_duplicates("sequence", keep="first").copy()
+
+    table.to_csv(output_dir / "all_candidates_scores.csv", index=False)
+    unique.sort_values("mrl_score", ascending=False).to_csv(
+        output_dir / "ranked_by_mrl.csv", index=False
+    )
+    unique.sort_values("te_score", ascending=False).to_csv(
+        output_dir / "ranked_by_te.csv", index=False
+    )
+    np.save(output_dir / "generator_probabilities.npy", generated)
+
+    summary = {
+        "requested_candidates": args.num_candidates,
+        "generated_candidates": len(table),
+        "unique_sequences": len(unique),
+        "duplicate_sequences": int(table["is_duplicate"].sum()),
+        "generator_shape": list(generated.shape),
+        "generator_probability_max_error": float(
+            np.max(np.abs(generated.sum(axis=-1) - 1.0))
+        ),
+        "length_min": int(table["length"].min()),
+        "length_max": int(table["length"].max()),
+        "mrl_min": float(mrl_scores.min()),
+        "mrl_max": float(mrl_scores.max()),
+        "mrl_mean": float(mrl_scores.mean()),
+        "te_min": float(te_scores.min()),
+        "te_max": float(te_scores.max()),
+        "te_mean": float(te_scores.mean()),
+        "tensorflow_version": tf.__version__,
+        "torch_version": torch.__version__,
+        "torch_hip": torch.version.hip,
+        "device": args.device,
+        "seed": args.seed,
+    }
+    (output_dir / "summary.json").write_text(
+        json.dumps(summary, indent=2), encoding="utf-8"
+    )
+    print(json.dumps(summary, indent=2))
+    print("UTRGAN_PRETRAINED_BATCH_MRL_TE_RANKING_PASS")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,398 @@
+import os
+import sys
+from pathlib import Path
+
+# Keep the official ``src`` package importable after adopting the standard
+# conf/model/scripts/weight layout.
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+MODEL_ROOT = PROJECT_ROOT / "model"
+if str(MODEL_ROOT) not in sys.path:
+    sys.path.insert(0, str(MODEL_ROOT))
+
+# The released H5 files use the legacy Keras serialization format.
+os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
+
+import numpy as np
+import pandas as pd
+import pdb
+import tensorflow as tf
+# from IPython.display import clear_output
+import time
+from src.gan.lib import models
+from src.gan.lib import utils
+import socket
+import datetime
+from tqdm import tqdm
+import random
+import matplotlib.pyplot as plt
+import argparse
+
+import tensorflow.keras.backend as K
+from tensorflow.keras.optimizers import Adam
+
+tf.compat.v1.enable_eager_execution()
+
+
+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=5)
+parser.add_argument('-mil', type=int, required=False ,default=64)
+parser.add_argument('-mxl', type=int, required=False ,default=128)
+parser.add_argument('-dim', type=int, required=False ,default=40)
+parser.add_argument('-gpu', type=str, required=False ,default='-1')
+parser.add_argument('-e', '--epochs', type=int, default=4000,
+                    help='Number of complete passes over the filtered UTRdb dataset.')
+parser.add_argument('--output-dir', type=str,
+                    default=str(PROJECT_ROOT / 'outputs' / 'train'),
+                    help='Directory for logs, generated samples, and checkpoints.')
+args = parser.parse_args()
+
+MODEL_NAME = 'WGAN-TF2'
+RUN_ROOT = Path(args.output_dir).expanduser().resolve()
+OUTPUT_PATH = RUN_ROOT / 'outputs' / MODEL_NAME
+TRAIN_LOGDIR = RUN_ROOT / 'tensorboard' / MODEL_NAME / 'train_data'
+OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
+TRAIN_LOGDIR.mkdir(parents=True, exist_ok=True)
+file_writer = tf.summary.create_file_writer(str(TRAIN_LOGDIR))
+
+if args.gpu == '-1':
+    device = 'cpu'
+else:
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
+    device = 'cuda'
+
+def plot(x, y, logdir, name, xlabel=None, ylabel=None, title=None):
+
+    plt.plot(x,y,'-')
+    plt.title(title)
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+    plt.savefig(logdir+'/'+name+'.png')
+    plt.clf()
+
+
+def plot_valid(x1, y1, x2, y2, logdir, name, xlabel=None, ylabel=None, title=None):
+
+    plt.plot(x2,y2,'-',color='tab:blue')
+    plt.plot(x1,y1,'-',color='tab:orange')
+
+    plt.title(title)
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+
+    plt.savefig(logdir+'/'+name+'.png')
+    plt.clf()    
+
+def gradient_penalty_loss( y_true, y_pred, discriminator):
+  """
+  Computes gradient penalty based on prediction and weighted real / fake samples
+  """
+  alpha = K.random_uniform((DIM, 1, 1))
+  averaged_samples = (alpha * y_pred) + ((1 - alpha) * y_true)
+
+  gradients = K.gradients(y_pred, averaged_samples)[0]
+  # compute the euclidean norm by squaring ...
+  gradients_sqr = K.square(gradients)
+  #   ... summing over the rows ...
+  gradients_sqr_sum = K.sum(gradients_sqr,
+                            axis=np.arange(1, len(gradients_sqr.shape)))
+  #   ... and sqrt
+  gradient_l2_norm = K.sqrt(gradients_sqr_sum)
+  # compute lambda * (1 - ||grad||)^2 still for each single sample
+  gradient_penalty = K.square(1 - gradient_l2_norm)
+  # return the mean as loss over all the batch samples
+  return K.mean(gradient_penalty)
+
+def log(samples_dir=False,suff=None):
+    stamp = datetime.date.strftime(datetime.datetime.now(), "%Y.%m.%d-%Hh%Mm%Ss") + "_{}".format(socket.gethostname())
+    full_logdir = os.path.join(str(RUN_ROOT / "logs"), stamp)
+    if suff:
+        full_logdir = full_logdir + suff
+    os.makedirs(full_logdir, exist_ok=True)
+    if samples_dir: os.makedirs(os.path.join(full_logdir, "samples"), exist_ok=True)
+    log_dir = "{}:{}".format(socket.gethostname(), full_logdir)
+    
+    return full_logdir, 0
+
+data_path = args.d
+data_utr = pd.read_csv(data_path)
+UTRdf = data_utr['seq'].to_numpy()
+
+seqs = []
+
+UTR_LEN = args.mxl
+
+for i in range(len(UTRdf)):
+    if len(UTRdf[i]) < UTR_LEN+1 and len(UTRdf[i]) > int(UTR_LEN/2):
+        if UTRdf[i] not in seqs:
+            seqs.append(UTRdf[i])
+
+sequences = np.array(seqs)
+
+sequences = [x.upper() for x in sequences]
+
+rna_vocab = {"A":0,
+             "C":1,
+             "G":2,
+             "U":3,
+             "*":4}
+
+rev_rna_vocab = {v:k for k,v in rna_vocab.items()}
+
+def one_hot_encode(seq, SEQ_LEN=UTR_LEN):
+    mapping = dict(zip("ACGT*", range(5)))    
+    seq2 = [mapping[i] for i in seq]
+    if len(seq2) < SEQ_LEN:
+        extra = [np.eye(5)[4]] * (SEQ_LEN - len(seq2))
+        return np.vstack([np.eye(5)[seq2] , extra])
+    return np.eye(5)[seq2]
+
+def one_hot_encode_2(seq, SEQ_LEN=UTR_LEN):
+    mapping = dict(zip("ACGT", range(4)))    
+    seq2 = [mapping[i] for i in seq]
+    return np.eye(4)[seq2]    
+
+ohe_sequences = np.asarray([one_hot_encode(x) for x in sequences])
+
+
+BATCH_SIZE = args.bs # Batch size
+ITERS = args.epochs # How many complete dataset iterations to train for
+SEQ_LEN = UTR_LEN # Sequence length in characters
+DIM = args.dim # Model dimensionality.
+CRITIC_ITERS = 5 # How many critic iterations per generator iteration. 
+LAMBDA = 10 # Gradient penalty lambda hyperparameter.
+# README defines ``-lr 5`` as 1e-5.  The upstream ``e**args.lr`` expression
+# is inconsistent with that documented interface.
+LR = np.power(10.0, -args.lr)
+
+
+LAMBDA = 10 # For gradient penalty
+
+CURRENT_EPOCH = 1 # Epoch start from
+SAVE_EVERY_N_EPOCH = 50 # Save checkpoint at every n epoch
+
+MIN_LR = 0.000001 # Minimum value of learning rate
+DECAY_FACTOR=1.00004 # learning rate decay factor
+'''
+Set seed for reproducibility
+'''
+seed = 35
+np.random.seed(seed)
+# tf.set_random_seed(seed)
+tf.random.set_seed(seed)
+
+logdir, checkpoint_baseline = log(samples_dir=True)
+logdir2 = ''
+
+'''
+Build GAN
+'''
+model_type = "resnet"
+data_enc_dim = 5
+data_size = SEQ_LEN * data_enc_dim
+# data_size = 256
+gen_layers = 3
+disc_layers = 3
+lmbda = 10. #lipschitz penalty hyperparameter.
+
+SAMPLE_SIZE = 128
+
+N_CHANNELS = DIM
+
+G = models.resnet_g2(DIM,N_CHANNELS,SEQ_LEN,5,res_layers=gen_layers)
+D = models.resnet_d2(N_CHANNELS,SEQ_LEN,5,res_layers=disc_layers)
+
+G.summary()
+D.summary()
+
+D_optimizer = Adam(learning_rate=LR, beta_1=0.5, beta_2=0.99)
+G_optimizer = Adam(learning_rate=LR, beta_1=0.5, beta_2=0.99)
+
+EPOCHs = ITERS
+
+@tf.function
+def WGAN_GP_train_d_step(real_sample, batch_size, step):
+
+    noise = tf.random.normal([batch_size, DIM])
+    epsilon = tf.random.uniform(shape=[batch_size, 1, 1], minval=0, maxval=1)
+    ###################################
+    # Train D
+    ###################################
+    with tf.GradientTape(persistent=True) as d_tape:
+        with tf.GradientTape() as gp_tape:
+            fake_sample = G([noise], training=True)
+            fake_sample_mixed = epsilon * tf.dtypes.cast(real_sample, tf.float32) + ((1 - epsilon) * fake_sample)
+            fake_mixed_pred = D([fake_sample_mixed], training=True)
+            
+        # Compute gradient penalty
+        grads = gp_tape.gradient(fake_mixed_pred, fake_sample_mixed)
+        grad_norms = tf.sqrt(tf.reduce_sum(tf.square(grads), axis=[1,2])) # Originally axis=[1,2]
+        # grad_norms = tf.norm(grads, axis=[1,2])
+        gradient_penalty = tf.reduce_mean(tf.square(grad_norms - 1.))
+        
+        fake_pred = D([fake_sample], training=True)
+        real_pred = D([real_sample], training=True)
+        
+        D_loss = tf.reduce_mean(fake_pred) - tf.reduce_mean(real_pred) + LAMBDA * gradient_penalty
+    # Calculate the gradients for discriminator
+    D_gradients = d_tape.gradient(D_loss,D.trainable_variables)
+    # Apply the gradients to the optimizer
+    D_optimizer.apply_gradients(zip(D_gradients,D.trainable_variables))
+    # Write loss values to tensorboard
+    if step % 10 == 0:
+        with file_writer.as_default():
+            tf.summary.scalar('D_loss', tf.reduce_mean(D_loss), step=step)
+
+    return D_loss, gradient_penalty
+
+@tf.function
+def WGAN_GP_train_g_step(real_sample, batch_size, step):
+
+    noise = tf.random.normal([batch_size, DIM])
+    ###################################
+    # Train G
+    ###################################
+    with tf.GradientTape() as g_tape:
+        fake_sample = G([noise], training=True)
+        fake_pred = D([fake_sample], training=True)
+        G_loss = -tf.reduce_mean(fake_pred)
+
+    G_gradients = g_tape.gradient(G_loss,
+                                            G.trainable_variables)
+    # Apply the gradients to the optimizer
+    G_optimizer.apply_gradients(zip(G_gradients,
+                                                G.trainable_variables))
+    # Write loss values to tensorboard
+    if step % 10 == 0:
+        with file_writer.as_default():
+            tf.summary.scalar('G_loss', G_loss, step=step)
+
+    return G_loss, noise
+
+checkpoint_path = os.path.join(str(RUN_ROOT), "checkpoints", "tensorflow", MODEL_NAME)
+
+ckpt = tf.train.Checkpoint(generator=G,
+                        discriminator=D,
+                        G_optimizer=G_optimizer,
+                        D_optimizer=D_optimizer)
+
+ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=40)
+
+def generate_and_save_samples(model, epoch, test_input, figure_size=(12,6), subplot=(3,6), save=True, is_flatten=False):
+    '''
+        Generate samples and plot it.
+    '''
+    predictions = model.predict(test_input)
+    utils.save_samples(logdir, predictions, epoch, rev_rna_vocab, annotated=False)
+
+
+'''
+load data
+'''
+
+Train = True
+validate = True
+
+data = ohe_sequences
+
+
+if validate:
+    split = len(data) // 10
+    # print(split)
+    valid_data = data[:split]
+    train_data = data[split:]
+    if len(train_data) == 1: train_data = train_data[0]
+    if len(valid_data) == 1: valid_data = valid_data[0]
+else:
+    train_data = data
+
+
+train_data = train_data.astype('float32')
+valid_data = valid_data.astype('float32')
+
+train_seqs = tf.data.Dataset.from_tensor_slices(train_data).shuffle(len(train_data)).batch(BATCH_SIZE)
+valid_seqs = tf.data.Dataset.from_tensor_slices(valid_data).shuffle(len(valid_data)).batch(split)
+
+
+plot_iter = 10
+
+current_learning_rate = LR
+trace = True
+n_critic_count = 0
+
+d_losses = []
+g_losses = []
+gradient_penalties = []
+iterations = 0
+iteration_numbers = []
+iteration_numbers_valid = []
+d_validation_losses = []
+
+random_name = time.strftime("%Y%m%d-%H%M%S")
+
+gen_iters = 0
+
+if Train:
+    sample_noise = tf.random.normal([BATCH_SIZE, DIM])
+    generate_and_save_samples(G, 0, [sample_noise], figure_size=(12,6), subplot=(3,6), save=False, is_flatten=False)
+
+    pbar = tqdm(range(EPOCHs))
+    for epoch in pbar:
+        start = time.time()
+
+        tdataset = train_seqs.enumerate()
+
+        for step, tdata in tdataset.as_numpy_iterator():
+            current_batch_size = tdata.shape[0]
+
+            d_loss, gp = WGAN_GP_train_d_step(tdata, batch_size=tf.constant(current_batch_size, dtype=tf.int64), step=tf.constant(step, dtype=tf.int64))
+            n_critic_count += 1
+            if n_critic_count >= CRITIC_ITERS: 
+                g_loss, noise = WGAN_GP_train_g_step(tdata, batch_size= tf.constant(current_batch_size, dtype=tf.int64), step=tf.constant(step, dtype=tf.int64))
+                gen_iters += 1
+                n_critic_count = 0
+            
+            if step % 10 == 0:
+                print ('.', end='')
+
+        
+        if epoch % SAVE_EVERY_N_EPOCH == 0 and epoch != 0:
+            ckpt_save_path = ckpt_manager.save()
+            utils.save_checkpoints(logdir,G,epoch)
+
+        if epoch % 50 == 0:    
+            
+            generate_and_save_samples(G, epoch, [sample_noise], figure_size=(12,6), subplot=(3,6), save=True, is_flatten=False)
+
+        os.system('clear')
+
+        iteration_numbers.append(iterations)
+        g_losses.append(-g_loss)
+        d_losses.append(-d_loss)
+        gradient_penalties.append(gp)
+        # plot(iteration_numbers, d_losses, logdir, 'discriminator_loss', xlabel="Iteration", ylabel="Discriminator Cost")
+        # plot(iteration_numbers, g_losses, logdir, 'generator_loss', xlabel="Iteration", ylabel="Generator Cost")
+        # plot(iteration_numbers, gradient_penalties, logdir, 'gradient_penalty', xlabel="Iteration", ylabel="Gradient Penalty")
+
+        iterations+=1
+
+        iteration_numbers_valid.append(iterations)
+        fake_sample_valid = G([noise], training=True)
+        fake_pred_valid = D([fake_sample_valid], training=True)
+        real_pred_valid = D([tf.convert_to_tensor(list(valid_seqs.as_numpy_iterator())[0])], training=True)
+        
+        D_loss_valid = tf.reduce_mean(fake_pred_valid) - tf.reduce_mean(real_pred_valid)
+        d_validation_losses.append(-D_loss_valid)
+
+        # plot_valid(iteration_numbers_valid, d_validation_losses, iteration_numbers, d_losses, logdir, 'validation_loss', xlabel="Iteration", ylabel="D Validation Loss")
+
+    ckpt_save_path = ckpt_manager.save()
+    print ('Saving checkpoint for epoch {} at {}'.format(EPOCHs,
+                                                            ckpt_save_path))
+
+
+print("####################################################")
+print(f"############### Gen Iterations : {gen_iters} #####")
+print("####################################################")