analysis/Hsu.et.al.git/src/unirep.py

"""
The trained 1900-dimensional mLSTM babbler.
From the UniRep Paper [Alley2019]
"""

import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
import tensorflow_addons as tfa
import tensorflow_probability as tfp
import numpy as np
import pandas as pd
import sys
sys.path.append('../')
from utils import aa_seq_to_int, int_to_aa
from utils.unirep_utils import bucketbatchpad
import os

# Helpers
def tf_get_shape(tensor):
    static_shape = tensor.shape.as_list()
    dynamic_shape = tf.unstack(tf.shape(tensor))
    dims = [s[1] if s[0] is None else s[0]
            for s in zip(static_shape, dynamic_shape)]
    return dims

def sample_with_temp(logits, t):
    """
    Takes temperature between 0 and 1 -> zero most conservative, 1 most liberal. Samples.
    """
    t_adjusted = logits / t  # broadcast temperature normalization
    softed = tf.nn.softmax(t_adjusted)
    
    # Make a categorical distribution from the softmax and sample
    return tfp.distributions.Categorical(probs=softed).sample()

def initialize_uninitialized(sess):
    """
    from https://stackoverflow.com/questions/35164529/in-tensorflow-is-there-any-way-to-just-initialize-uninitialised-variables
    """
    global_vars = tf.global_variables()
    is_not_initialized = sess.run([tf.is_variable_initialized(var) for var in global_vars])
    not_initialized_vars = [v for (v, f) in zip(global_vars, is_not_initialized) if not f]
    if len(not_initialized_vars):
        sess.run(tf.variables_initializer(not_initialized_vars))


# Setup to initialize from the correctly named model files.
class mLSTMCell1900(tf.nn.rnn_cell.RNNCell):

    def __init__(self,
                 num_units,
                 model_path="./",
                 wn=True,
                 scope='mlstm',
                 var_device='cpu:0',
                 ):
        # Really not sure if I should reuse here
        super(mLSTMCell1900, self).__init__()
        self._num_units = num_units
        self._model_path = model_path
        self._wn = wn
        self._scope = scope
        self._var_device = var_device

    @property
    def state_size(self):
        # The state is a tuple of c and h
        return (self._num_units, self._num_units)

    @property
    def output_size(self):
        # The output is h
        return (self._num_units)

    def zero_state(self, batch_size, dtype):
        c = tf.zeros([batch_size, self._num_units], dtype=dtype)
        h = tf.zeros([batch_size, self._num_units], dtype=dtype)
        return (c, h)

    def call(self, inputs, state):
        # Inputs will be a [batch_size, input_dim] tensor.
        # Eg, input_dim for a 10-D embedding is 10
        nin = inputs.get_shape()[1].value

        # Unpack the state tuple
        c_prev, h_prev = state
        with tf.variable_scope(self._scope):
            wx_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_wx:0.npy"))
            wh_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_wh:0.npy"))
            wmx_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_wmx:0.npy"))
            wmh_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_wmh:0.npy"))
            b_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_b:0.npy"))
            gx_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_gx:0.npy"))
            gh_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_gh:0.npy"))
            gmx_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_gmx:0.npy"))
            gmh_init = np.load(os.path.join(self._model_path, "rnn_mlstm_mlstm_gmh:0.npy"))        
            wx = tf.get_variable(
                "wx", initializer=wx_init)
            wh = tf.get_variable(
                "wh", initializer=wh_init)
            wmx = tf.get_variable(
                "wmx", initializer=wmx_init)
            wmh = tf.get_variable(
                "wmh", initializer=wmh_init)
            b = tf.get_variable(
                "b", initializer=b_init)
            if self._wn:
                gx = tf.get_variable(
                    "gx", initializer=gx_init)
                gh = tf.get_variable(
                    "gh", initializer=gh_init)
                gmx = tf.get_variable(
                    "gmx", initializer=gmx_init)
                gmh = tf.get_variable(
                    "gmh", initializer=gmh_init)

        if self._wn:
            wx = tf.nn.l2_normalize(wx, axis=0) * gx
            wh = tf.nn.l2_normalize(wh, axis=0) * gh
            wmx = tf.nn.l2_normalize(wmx, axis=0) * gmx
            wmh = tf.nn.l2_normalize(wmh, axis=0) * gmh
        m = tf.matmul(inputs, wmx) * tf.matmul(h_prev, wmh)
        z = tf.matmul(inputs, wx) + tf.matmul(m, wh) + b
        i, f, o, u = tf.split(z, 4, 1)
        i = tf.nn.sigmoid(i)
        f = tf.nn.sigmoid(f)
        o = tf.nn.sigmoid(o)
        u = tf.tanh(u)
        c = f * c_prev + i * u
        h = o * tf.tanh(c)
        return h, (c, h)

class mLSTMCell(tf.nn.rnn_cell.RNNCell):

    def __init__(self,
                 num_units,
                 wx_init=tf.orthogonal_initializer(),
                 wh_init=tf.orthogonal_initializer(),
                 wmx_init=tf.orthogonal_initializer(),
                 wmh_init=tf.orthogonal_initializer(),
                 b_init=tf.orthogonal_initializer(),
                 gx_init=tf.ones_initializer(),
                 gh_init=tf.ones_initializer(),
                 gmx_init=tf.ones_initializer(),
                 gmh_init=tf.ones_initializer(),
                 wn=True,
                 scope='mlstm',
                 var_device='cpu:0',
                 ):
        # Really not sure if I should reuse here
        super(mLSTMCell, self).__init__()
        self._num_units = num_units
        self._wn = wn
        self._scope = scope
        self._var_device = var_device
        self._wx_init = wx_init
        self._wh_init = wh_init
        self._wmx_init = wmx_init
        self._wmh_init = wmh_init
        self._b_init = b_init
        self._gx_init = gx_init
        self._gh_init = gh_init
        self._gmx_init = gmx_init
        self._gmh_init = gmh_init

    @property
    def state_size(self):
        # The state is a tuple of c and h
        return (self._num_units, self._num_units)

    @property
    def output_size(self):
        # The output is h
        return (self._num_units)

    def zero_state(self, batch_size, dtype):
        c = tf.zeros([batch_size, self._num_units], dtype=dtype)
        h = tf.zeros([batch_size, self._num_units], dtype=dtype)
        return (c, h)

    def call(self, inputs, state):
        # Inputs will be a [batch_size, input_dim] tensor.
        # Eg, input_dim for a 10-D embedding is 10
        nin = inputs.get_shape()[1].value

        # Unpack the state tuple
        c_prev, h_prev = state
        with tf.variable_scope(self._scope):
            wx = tf.get_variable(
                "wx", initializer=self._wx_init)
            wh = tf.get_variable(
                "wh", initializer=self._wh_init)
            wmx = tf.get_variable(
                "wmx", initializer=self._wmx_init)
            wmh = tf.get_variable(
                "wmh", initializer=self._wmh_init)
            b = tf.get_variable(
                "b", initializer=self._b_init)
            if self._wn:
                gx = tf.get_variable(
                    "gx", initializer=self._gx_init)
                gh = tf.get_variable(
                    "gh", initializer=self._gh_init)
                gmx = tf.get_variable(
                    "gmx", initializer=self._gmx_init)
                gmh = tf.get_variable(
                    "gmh", initializer=self._gmh_init)

        if self._wn:
            wx = tf.nn.l2_normalize(wx, dim=0) * gx
            wh = tf.nn.l2_normalize(wh, dim=0) * gh
            wmx = tf.nn.l2_normalize(wmx, dim=0) * gmx
            wmh = tf.nn.l2_normalize(wmh, dim=0) * gmh
        m = tf.matmul(inputs, wmx) * tf.matmul(h_prev, wmh)
        z = tf.matmul(inputs, wx) + tf.matmul(m, wh) + b
        i, f, o, u = tf.split(z, 4, 1)
        i = tf.nn.sigmoid(i)
        f = tf.nn.sigmoid(f)
        o = tf.nn.sigmoid(o)
        u = tf.tanh(u)
        c = f * c_prev + i * u
        h = o * tf.tanh(c)
        return h, (c, h)

class mLSTMCellStackNPY(tf.nn.rnn_cell.RNNCell):

    def __init__(self,
                 num_units=256,
                 num_layers=4,
                 dropout=None,
                 res_connect=False,
                 wn=True,
                 scope='mlstm_stack',
                 var_device='cpu:0',
                 model_path="./"
                 ):
        # Really not sure if I should reuse here
        super(mLSTMCellStackNPY, self).__init__()
        self._model_path=model_path
        self._num_units = num_units
        self._num_layers = num_layers
        self._dropout = dropout
        self._res_connect = res_connect
        self._wn = wn
        self._scope = scope
        self._var_device = var_device
        bs = "rnn_mlstm_stack_mlstm_stack" # base scope see weight file names
        join = lambda x: os.path.join(self._model_path, x)
        layers = [mLSTMCell(
            num_units=self._num_units,
            wn=self._wn,
            scope=self._scope + str(i),
            var_device=self._var_device,
            wx_init=np.load(join(bs + "{0}_mlstm_stack{1}_wx:0.npy".format(i,i))),
            wh_init=np.load(join(bs + "{0}_mlstm_stack{1}_wh:0.npy".format(i,i))),
            wmx_init=np.load(join(bs + "{0}_mlstm_stack{1}_wmx:0.npy".format(i,i))),
            wmh_init=np.load(join(bs + "{0}_mlstm_stack{1}_wmh:0.npy".format(i,i))),
            b_init=np.load(join(bs + "{0}_mlstm_stack{1}_b:0.npy".format(i,i))),
            gx_init=np.load(join(bs + "{0}_mlstm_stack{1}_gx:0.npy".format(i,i))),
            gh_init=np.load(join(bs + "{0}_mlstm_stack{1}_gh:0.npy".format(i,i))),
            gmx_init=np.load(join(bs + "{0}_mlstm_stack{1}_gmx:0.npy".format(i,i))),
            gmh_init=np.load(join(bs + "{0}_mlstm_stack{1}_gmh:0.npy".format(i,i)))      
                 ) for i in range(self._num_layers)]
        if self._dropout:
            layers = [
                tf.nn.rnn_cell.DropoutWrapper(
                    layer, output_keep_prob=1-self._dropout) for layer in layers[:-1]] + layers[-1:]
        self._layers = layers

    @property
    def state_size(self):
        # The state is a tuple of c and h
        return (
            tuple(self._num_units for _ in range(self._num_layers)), 
            tuple(self._num_units for _ in range(self._num_layers))
            )

    @property
    def output_size(self):
        # The output is h
        return (self._num_units)

    def zero_state(self, batch_size, dtype):
        c_stack = tuple(tf.zeros([batch_size, self._num_units], dtype=dtype) for _ in range(self._num_layers))
        h_stack = tuple(tf.zeros([batch_size, self._num_units], dtype=dtype) for _ in range(self._num_layers))
        return (c_stack, h_stack)

    def call(self, inputs, state):
        # Inputs will be a [batch_size, input_dim] tensor.
        # Eg, input_dim for a 10-D embedding is 10

        # Unpack the state tuple
        c_prev, h_prev = state
        
        new_outputs = []
        new_cs = []
        new_hs = []
        for i, layer in enumerate(self._layers):
            if i == 0:
                h, (c,h_state) = layer(inputs, (c_prev[i],h_prev[i]))
            else:
                h, (c,h_state) = layer(new_outputs[-1], (c_prev[i],h_prev[i]))
            new_outputs.append(h)
            new_cs.append(c)
            new_hs.append(h_state)
        
        if self._res_connect:
            # Make sure number of layers does not affect the scale of the output
            scale_factor = tf.constant(1 / float(self._num_layers))
            final_output = tf.scalar_mul(scale_factor,tf.add_n(new_outputs))
        else:
            final_output = new_outputs[-1]

        return final_output, (tuple(new_cs), tuple(new_hs))

    
class babbler1900():

    def __init__(self,
                 model_path="./pbab_weights",
                 batch_size=256
                 ):
        self._rnn_size = 1900
        self._vocab_size = 26
        self._embed_dim = 10
        self._wn = True
        self._shuffle_buffer = 10000
        self._model_path = model_path
        self._batch_size = batch_size
        self._batch_size_placeholder = tf.placeholder(tf.int32, shape=[], name="batch_size")
        self._minibatch_x_placeholder = tf.placeholder(
            tf.int32, shape=[None, None], name="minibatch_x")
        self._initial_state_placeholder = (
            tf.placeholder(tf.float32, shape=[None, self._rnn_size]),
            tf.placeholder(tf.float32, shape=[None, self._rnn_size])
        )
        self._minibatch_y_placeholder = tf.placeholder(
            tf.int32, shape=[None, None], name="minibatch_y")
        # Batch size dimensional placeholder which gives the
        # Lengths of the input sequence batch. Used to index into
        # The final_hidden output and select the stop codon -1
        # final hidden for the graph operation.
        self._seq_length_placeholder = tf.placeholder(
            tf.int32, shape=[None], name="seq_len")
        self._temp_placeholder = tf.placeholder(tf.float32, shape=[], name="temp")
        rnn = mLSTMCell1900(self._rnn_size, model_path=model_path, wn=self._wn)
        zero_state = rnn.zero_state(self._batch_size, tf.float32)
        single_zero = rnn.zero_state(1, tf.float32)
        mask = tf.sign(self._minibatch_y_placeholder)  # 1 for nonpad, zero for pad
        inverse_mask = 1 - mask  # 0 for nonpad, 1 for pad

        total_padded = tf.reduce_sum(inverse_mask)

        pad_adjusted_targets = (self._minibatch_y_placeholder - 1) + inverse_mask

        embed_matrix = tf.get_variable(
            "embed_matrix", dtype=tf.float32,
            initializer=np.load(os.path.join(self._model_path, "embed_matrix:0.npy"))
        )
        embed_cell = tf.nn.embedding_lookup(embed_matrix, self._minibatch_x_placeholder)
        self._output, self._final_state = tf.nn.dynamic_rnn(
            rnn,
            embed_cell,
            initial_state=self._initial_state_placeholder,
            swap_memory=True,
            parallel_iterations=1
        )
        
        # If we are training a model on top of the rep model, we need to access
        # the final_hidden rep from output. Recall we are padding these sequences
        # to max length, so the -1 position will not necessarily be the right rep.
        # to get the right rep, I will use the provided sequence length to index.
        # Subtract one for the last place
        indices = self._seq_length_placeholder - 1
        self._top_final_hidden = tf.gather_nd(self._output,
                tf.stack([tf.range(tf_get_shape(self._output)[0],
                    dtype=tf.int32), indices], axis=1))
        fmask = tf.cast(mask, tf.float32)[:, :, None]
        self._avg_hidden = tf.reduce_sum(fmask * self._output,
                axis=1) / tf.reduce_sum(fmask, axis=1)
        # LEFTOFF self._output is a batch size, seq_len, num_hidden.
        # I want to average along num_hidden, but I'll have to figure out how to mask out
        # the dimensions along sequence_length which are longer than the given sequence.
        flat = tf.reshape(self._output, [-1, self._rnn_size])
        if os.path.exists(os.path.join(self._model_path, "fully_connected_weights:0.npy")):
            weights_name="fully_connected_weights"
            bias_name="fully_connected_biases"
        else:
            weights_name="dense_kernel"
            bias_name="dense_bias"
        weights_init = tf.constant_initializer(
                np.load(os.path.join(self._model_path, f"{weights_name}:0.npy")))
        bias_init = tf.constant_initializer(
                np.load(os.path.join(self._model_path, f"{bias_name}:0.npy")))
        self.dense_layer = tf.keras.layers.Dense(self._vocab_size-1,
                activation=None, kernel_initializer=weights_init,
                bias_initializer=bias_init)
        logits_flat = self.dense_layer(flat)
        seqlen = tf_get_shape(self._minibatch_x_placeholder)[1]
        self._logits = tf.reshape(
            logits_flat, [batch_size, seqlen, self._vocab_size-1])
        self.batch_losses = tfa.seq2seq.sequence_loss(
            self._logits,
            tf.cast(pad_adjusted_targets, tf.int32),
            tf.cast(mask, tf.float32),
            average_across_batch=False
        )
        self._loss = tf.reduce_mean(self.batch_losses)
        self._sample = sample_with_temp(self._logits, self._temp_placeholder)
        with tf.Session() as sess:
            self._zero_state = sess.run(zero_state)
            self._single_zero = sess.run(single_zero)

     
    def get_rep(self,seq):
        """
        Input a valid amino acid sequence, 
        outputs a tuple of average hidden, final hidden, final cell representation arrays.
        Unfortunately, this method accepts one sequence at a time and is as such quite
        slow.
        """
        with tf.Session() as sess:
            initialize_uninitialized(sess)
            # Strip any whitespace and convert to integers with the correct coding
            int_seq = aa_seq_to_int(seq.strip())[:-1]
            # Final state is a cell_state, hidden_state tuple. Output is
            # all hidden states
            final_state_, hs = sess.run(
                [self._final_state, self._output], feed_dict={
                    self._batch_size_placeholder: 1,
                    self._minibatch_x_placeholder: [int_seq],
                    self._initial_state_placeholder: self._zero_state}
            )

        final_cell, final_hidden = final_state_
        # Drop the batch dimension so it is just seq len by
        # representation size
        final_cell = final_cell[0]
        final_hidden = final_hidden[0]
        hs = hs[0]
        avg_hidden = np.mean(hs, axis=0)
        return avg_hidden, final_hidden, final_cell

    def get_babble(self, seed, length=250, temp=1):
        """
        Return a babble at temperature temp (on (0,1] with 1 being the noisiest)
        starting with seed and continuing to length length.
        Unfortunately, this method accepts one sequence at a time and is as such quite
        slow.

        """
        with tf.Session() as sess:
            initialize_uninitialized(sess)
            int_seed = aa_seq_to_int(seed.strip())[:-1]
        
            # No need for padding because this is a single element
            seed_samples, final_state_ = sess.run(
                [self._sample, self._final_state], 
                feed_dict={
                    self._minibatch_x_placeholder: [int_seed],
                    self._initial_state_placeholder: self._zero_state, 
                    self._batch_size_placeholder: 1,
                    self._temp_placeholder: temp
                }
            )
            # Just the actual character prediction
            pred_int = seed_samples[0, -1] + 1
            seed = seed + int_to_aa[pred_int]
        
            for i in range(length - len(seed)):
                pred_int, final_state_ = sess.run(
                    [self._sample, self._final_state], 
                    feed_dict={
                        self._minibatch_x_placeholder: [[pred_int]],
                        self._initial_state_placeholder: final_state_, 
                        self._batch_size_placeholder: 1,
                        self._temp_placeholder: temp
                    }
                )
                pred_int = pred_int[0, 0] + 1
                seed = seed + int_to_aa[pred_int]
        return seed        
        
    def get_rep_ops(self):
        """
        Return tensorflow operations for the final_hidden state and placeholder.
        """
        return self._top_final_hidden, self._avg_hidden, self._minibatch_x_placeholder, self._batch_size_placeholder, self._seq_length_placeholder, self._initial_state_placeholder
        
    def get_babbler_ops(self):
        """
        Return tensorflow operations for 
        the logits, masked loss, minibatch_x placeholder, minibatch y placeholder, batch_size placeholder, initial_state placeholder
        Use if you plan on using babbler1900 as an initialization for another babbler, 
        eg for fine tuning the babbler to babble a differenct distribution.
        """
        return self._logits, self._loss, self._minibatch_x_placeholder, self._minibatch_y_placeholder, self._batch_size_placeholder, self._initial_state_placeholder

    def dump_weights(self,sess,dir_name="./1900_weights"):
        """
        Saves the weights of the model in dir_name in the format required 
        for loading in this module. Must be called within a tf.Session
        For which the weights are already initialized.
        """
        vs = tf.trainable_variables()
        for v in vs:
            name = v.name
            value = sess.run(v)
            np.save(os.path.join(dir_name,name.replace('/', '_') + ".npy"), np.array(value))
            


    def format_seq(self,seq,stop=False):
        """
        Takes an amino acid sequence, returns a list of integers in the codex of the babbler.
        Here, the default is to strip the stop symbol (stop=False) which would have 
        otherwise been added to the end of the sequence. If you are trying to generate
        a rep, do not include the stop. It is probably best to ignore the stop if you are
        co-tuning the babbler and a top model as well.
        """
        if stop:
            int_seq = aa_seq_to_int(seq.strip())
        else:
            int_seq = aa_seq_to_int(seq.strip())[:-1]
        return int_seq


    def bucket_batch_pad(self,filepath, upper=2000, lower=50, interval=10):
        """
        Read sequences from a filepath, batch them into buckets of similar lengths, and
        pad out to the longest sequence.
        Upper, lower and interval define how the buckets are created.
        Any sequence shorter than lower will be grouped together, as with any greater 
        than upper. Interval defines the "walls" of all the other buckets.
        WARNING: Define large intervals for small datasets because the default behavior
        is to repeat the same sequence to fill a batch. If there is only one sequence
        within a bucket, it will be repeated batch_size -1 times to fill the batch.
        """
        self._bucket_upper = upper
        self._bucket_lower = lower
        self._bucket_interval = interval
        self._bucket = [self._bucket_lower + (i * self._bucket_interval) for i in range(int(self._bucket_upper / self._bucket_interval))]
        self._bucket_batch =  bucketbatchpad(
                    batch_size=self._batch_size,
                    pad_shape=([None]),
                    window_size=self._batch_size,
                    bounds=self._bucket,
                    path_to_data=filepath,
                    shuffle_buffer=self._shuffle_buffer,
                    repeat=None
        ).make_one_shot_iterator().get_next()
        return self._bucket_batch

    def split_to_tuple(self, seq_batch):
        """
        NOTICE THAT BY DEFAULT THIS STRIPS THE LAST CHARACTER.
        IF USING IN COMBINATION WITH format_seq then set stop=True there.
        Return a list of batch, target tuples.
        The input (array-like) should
        look like 
        1. . . . . . . . sequence_length
        .
        .
        .
        batch_size
        """
        q = None
        num_steps = seq_batch.shape[1]
        # Minibatches should start at zero index and go to -1
        # Don't even try to get what is happenning here its a brainfuck and
        # probably inefficient
        xypairs = [
            (seq_batch[:, :-1][:, idx:idx + num_steps], seq_batch[:, idx + 1:idx + num_steps + 1]) for idx in np.arange(len(seq_batch[0]))[0:-1:num_steps]
        ]
        if q:
            for e in xypairs:
                q.put(e)
        else:
            return xypairs[0]

    def is_valid_seq(self, seq, max_len=2000):
        """
        True if seq is valid for the babbler, False otherwise.
        """
        l = len(seq)
        ## important: added X here, not in original unirep code
        valid_aas = "MRHKDESTNQCUGPAVIFYWLOX" 
        if (l < max_len) and set(seq) <= set(valid_aas):
            return True
        else:
            return False