scmtg 0.0.1__tar.gz

scmtg-0.0.1/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.1
+Name: scmtg
+Version: 0.0.1
+Summary: scMTG learns generative Markov transitions for single-cell temporal dynamics
+Home-page: https://github.com/liuq-lab/scMTG
+Author: Xuejian Cui
+Author-email: cuixj@hit.edu.cn
+License: MIT
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.8.0
+Requires-Dist: tensorflow-gpu==2.6.0
+Requires-Dist: keras==2.6.0

scmtg-0.0.1/scMTG/__init__.py

@@ -0,0 +1,3 @@
+__version__ = '0.0.1'
+from .scMTG import *
+from . util import *

scmtg-0.0.1/scMTG/loss.py

@@ -0,0 +1,132 @@
+import numpy as np
+import tensorflow as tf
+import tensorflow.keras.backend as K
+
+
+def _nan2zero(x):
+    return tf.where(tf.math.is_nan(x), tf.zeros_like(x), x)
+
+def _nan2inf(x):
+    return tf.where(tf.math.is_nan(x), tf.zeros_like(x)+np.inf, x)
+
+def _nelem(x):
+    nelem = tf.reduce_sum(tf.cast(~tf.math.is_nan(x), tf.float32))
+    return tf.cast(tf.where(tf.equal(nelem, 0.), 1., nelem), x.dtype)
+
+
+def _reduce_mean(x):
+    nelem = _nelem(x)
+    x = _nan2zero(x)
+    return tf.divide(tf.reduce_sum(x), nelem)
+
+
+def mse_loss(y_true, y_pred):
+    ret = tf.square(y_pred - y_true)
+
+    return _reduce_mean(ret)
+
+
+def poisson_loss(y_true, y_pred):
+    y_pred = tf.cast(y_pred, tf.float32)
+    y_true = tf.cast(y_true, tf.float32)
+    nelem = _nelem(y_true)
+    y_true = _nan2zero(y_true)
+
+    # last term can be avoided since it doesn't depend on y_pred
+    # however keeping it gives a nice lower bound to zero
+    ret = y_pred - y_true*tf.math.log(y_pred+1e-10) + tf.math.lgamma(y_true+1.0)
+
+    return tf.divide(tf.reduce_sum(ret), nelem)
+
+class NB(object):
+    def __init__(self, theta=None, masking=False, scope='nbinom_loss/',
+                 scale_factor=1.0, debug=False):
+
+        self.eps = 1e-10
+        self.scale_factor = scale_factor
+        self.debug = debug
+        self.scope = scope
+        self.masking = masking
+        self.theta = theta
+
+    def loss(self, y_true, y_pred, mean=True):
+        scale_factor = self.scale_factor
+        eps = self.eps
+
+        with tf.name_scope(self.scope):
+            y_true = tf.cast(y_true, tf.float32)
+            y_pred = tf.cast(y_pred, tf.float32) * scale_factor
+
+            if self.masking:
+                nelem = _nelem(y_true)
+                y_true = _nan2zero(y_true)
+
+            # Clip theta
+            theta = tf.minimum(self.theta, 1e6)
+
+            t1 = tf.math.lgamma(1/(theta+eps)+eps) + tf.math.lgamma(y_true+1.0) - tf.math.lgamma(y_true+1/(theta+eps)+eps) #lgamma ln(gamma(x))
+            t2 = (1/(theta+eps)+y_true) * tf.math.log(1.0 + y_pred*theta) - (y_true * tf.math.log(theta*y_pred+eps))
+
+            if self.debug:
+                assert_ops = [
+                        tf.verify_tensor_all_finite(y_pred, 'y_pred has inf/nans'),
+                        tf.verify_tensor_all_finite(t1, 't1 has inf/nans'),
+                        tf.verify_tensor_all_finite(t2, 't2 has inf/nans')]
+
+                tf.summary.histogram('t1', t1)
+                tf.summary.histogram('t2', t2)
+
+                with tf.control_dependencies(assert_ops):
+                    final = t1 + t2
+
+            else:
+                final = t1 + t2
+
+            final = _nan2inf(final)
+
+            if mean:
+                if self.masking:
+                    final = tf.divide(tf.reduce_sum(final), nelem)
+                else:
+                    final = tf.reduce_mean(final)
+
+        return final
+
+class ZINB(NB):
+    def __init__(self, pi, ridge_lambda=0.0, scope='zinb_loss/', **kwargs):
+        super().__init__(scope=scope, **kwargs)
+        self.pi = pi
+        self.ridge_lambda = ridge_lambda
+
+    def loss(self, y_true, y_pred, mean=True):
+        scale_factor = self.scale_factor
+        eps = self.eps
+
+        with tf.name_scope(self.scope):
+            nb_case = super().loss(y_true, y_pred, mean=False) - tf.math.log(1.0-self.pi+eps)
+
+            y_true = tf.cast(y_true, tf.float32)
+            y_pred = tf.cast(y_pred, tf.float32) * scale_factor
+            theta = tf.minimum(self.theta, 1e6)
+
+            zero_nb = tf.pow((1/theta)/(y_pred+1/theta), 1/theta)
+            zero_case = -tf.math.log(self.pi + ((1.0-self.pi)*zero_nb)+eps)
+            result = tf.where(tf.less(y_true, 1e-8), zero_case, nb_case)
+            ridge = self.ridge_lambda*tf.square(self.pi)
+            result += ridge
+
+            if mean:
+                if self.masking:
+                    result = _reduce_mean(result)
+                else:
+                    result = tf.reduce_mean(result)
+
+            result = _nan2inf(result)
+
+            if self.debug:
+                tf.summary.histogram('nb_case', nb_case)
+                tf.summary.histogram('zero_nb', zero_nb)
+                tf.summary.histogram('zero_case', zero_case)
+                tf.summary.histogram('ridge', ridge)
+
+        return result

scmtg-0.0.1/scMTG/model.py

@@ -0,0 +1,278 @@
+import os
+import tensorflow as tf
+import tensorflow.keras.backend as K
+from keras.layers import Lambda, Layer
+
+MeanAct = lambda x: tf.clip_by_value(tf.nn.softplus(x), 0, 1e4)
+DispAct = lambda x: tf.clip_by_value(tf.nn.softplus(x), 1e-4, 1e4)
+
+
+class BaseFullyConnectedNet(tf.keras.Model):
+    """ Encoder network.
+    """
+    def __init__(self, input_dim, z_dim, output_dim, model_name, nb_units=[256], concat_every_fcl=False, batchnorm=False, dropout=True, last_relu=False):  
+        super(BaseFullyConnectedNet, self).__init__()
+        self.input_layer = tf.keras.layers.Input((input_dim,))
+        self.input_dim = input_dim
+        self.z_dim = z_dim
+        self.output_dim = output_dim
+        self.model_name = model_name
+        self.nb_units = nb_units
+        self.concat_every_fcl = concat_every_fcl
+        self.batchnorm = batchnorm
+        self.dropout = dropout
+        self.last_relu = last_relu
+        self.all_layers = []
+        """ Builds the FC stacks. """
+        for i in range(len(nb_units)):
+            fc_layer = tf.keras.layers.Dense(
+                units = self.nb_units[i],
+                activation = None,
+#                 kernel_regularizer = tf.keras.regularizers.L2(2.5e-5)
+            )  
+            norm_layer = tf.keras.layers.BatchNormalization()
+            dropout_layer = tf.keras.layers.Dropout(0.1)
+            act_layer = tf.keras.layers.LeakyReLU(alpha=0.2)
+#             act_layer = tf.keras.layers.ReLU()
+            self.all_layers.append([fc_layer, norm_layer, dropout_layer, act_layer])
+        fc_layer = tf.keras.layers.Dense(
+            units = self.output_dim,
+            activation = None,
+#             kernel_regularizer = tf.keras.regularizers.L2(2.5e-5),
+#             activity_regularizer = tf.keras.regularizers.L1(2.5e-3)
+        )
+        self.all_layers.append([fc_layer, None, None, None])
+        
+        self.out = self.call(self.input_layer)
+
+    def call(self, inputs, training=True):
+        """ Return the output of the Generator.
+        Args:
+            inputs: tensor with shape [batch_size, input_dim]
+        Returns:
+            Output of Generator.
+            float32 tensor with shape [batch_size, output_dim]
+        """
+        y = inputs[:,self.z_dim:]
+        for i, layers in enumerate(self.all_layers[:-1]):
+            # Run inputs through the sublayers.
+            fc_layer, norm_layer, dropout_layer, act_layer = layers
+            with tf.name_scope("%s_layer_%d" % (self.model_name, i+1)):
+                x = fc_layer(inputs) if i==0 else fc_layer(x)
+                if self.batchnorm:
+                    x = norm_layer(x)
+                if self.dropout:
+                    x = dropout_layer(x)
+                x = act_layer(x)
+                if self.concat_every_fcl:
+                    x = tf.keras.layers.concatenate([x,y],axis=1)
+        fc_layer, _, _, _ = self.all_layers[-1]
+        with tf.name_scope("%s_layer_ouput" % self.model_name):
+            output = fc_layer(x)
+            # No activation func at last layer
+            if self.last_relu:
+#                 output = tf.keras.layers.ReLU()(output)
+                output = tf.keras.activations.tanh(output)
+        return output
+    
+class Decoder2(tf.keras.Model):
+    """ Decoder network.
+    """
+    def __init__(self, input_dim, z_dim, output_dim, model_name, nb_units=[256], batchnorm=False, dropout=True, last_relu=False):  
+        super(Decoder2, self).__init__()
+        self.input_layer = tf.keras.layers.Input((input_dim,))
+        self.input_dim = input_dim
+        self.z_dim = z_dim
+        self.output_dim = output_dim
+        self.model_name = model_name
+        self.nb_units = nb_units
+        self.batchnorm = batchnorm
+        self.dropout = dropout
+        self.last_relu = last_relu
+        self.all_layers = []
+        """ Builds the FC stacks. """
+        for i in range(len(nb_units)):
+            fc_layer = tf.keras.layers.Dense(
+                units = self.nb_units[i],
+                activation = None,
+#                 kernel_regularizer = tf.keras.regularizers.L2(2.5e-5)
+            )  
+            norm_layer = tf.keras.layers.BatchNormalization()
+            dropout_layer = tf.keras.layers.Dropout(0.1)
+            act_layer = tf.keras.layers.LeakyReLU(alpha=0.2)
+#             act_layer = tf.keras.layers.ReLU()
+            self.all_layers.append([fc_layer, norm_layer, dropout_layer, act_layer])
+            
+        disp_layer = tf.keras.layers.Dense(
+            units=self.output_dim, 
+            activation=DispAct, 
+#             kernel_initializer='glorot_uniform', 
+            name='dispersion'
+        )
+        mean_layer = tf.keras.layers.Dense(
+            units=self.output_dim, 
+            activation=MeanAct, 
+#             kernel_initializer='glorot_uniform', 
+            name='mean'
+        )
+        self.all_layers.append([disp_layer, mean_layer])
+        
+        self.disp, self.mean = self.call(self.input_layer)
+
+    def call(self, inputs, training=True):
+        """ Return the output of the Generator.
+        Args:
+            inputs: tensor with shape [batch_size, input_dim]
+        Returns:
+            Output of Generator.
+            float32 tensor with shape [batch_size, output_dim]
+        """
+        for i, layers in enumerate(self.all_layers[:-1]):
+            # Run inputs through the sublayers.
+            fc_layer, norm_layer, dropout_layer, act_layer = layers
+            with tf.name_scope("%s_layer_%d" % (self.model_name, i+1)):
+                x = fc_layer(inputs) if i==0 else fc_layer(x)
+                if self.batchnorm:
+                    x = norm_layer(x)
+                if self.dropout:
+                    x = dropout_layer(x)
+                x = act_layer(x)
+        disp_layer, mean_layer = self.all_layers[-1]
+        with tf.name_scope("%s_layer_output" % self.model_name):
+            disp = disp_layer(x)
+            mean = mean_layer(x)
+            # No activation func at last layer
+            if self.last_relu:
+                mean = tf.keras.layers.ReLU()(mean)
+        return disp, mean
+    
+class Generator(tf.keras.Model):
+    """Generator network.
+    """
+    def __init__(self, input_dim, z_dim, output_dim, model_name, nb_units=[256], concat_every_fcl=False, batchnorm=False, dropout=True, last_relu=False):  
+        super(Generator, self).__init__()
+        self.input_layer = tf.keras.layers.Input((input_dim,))
+        self.input_dim = input_dim
+        self.z_dim = z_dim
+        self.output_dim = output_dim
+        self.model_name = model_name
+        self.nb_units = nb_units
+        self.concat_every_fcl = concat_every_fcl
+        self.batchnorm = batchnorm
+        self.dropout = dropout
+        self.last_relu = last_relu
+        self.scale_factor = tf.Variable(initial_value=3.0, trainable=True)
+        self.all_layers = []
+        
+        """Builds the FC stacks."""
+        for i in range(len(self.nb_units) + 1):
+            units = self.output_dim if i == len(nb_units) else self.nb_units[i]
+            fc_layer = tf.keras.layers.Dense(
+                units = units,
+                activation = None,
+#                 kernel_regularizer = tf.keras.regularizers.L2(2.5e-5)
+            )   
+            norm_layer = tf.keras.layers.BatchNormalization()
+            dropout_layer = tf.keras.layers.Dropout(0.1)
+            self.all_layers.append([fc_layer, norm_layer, dropout_layer])
+        self.out = self.call(self.input_layer)
+
+    def call(self, inputs, training=True):
+        """Return the output of the Generator.
+        Args:
+            inputs: tensor with shape [batch_size, z_dim + nb_classes]
+        Returns:
+            Output of Generator.
+            float32 tensor with shape [batch_size, output_dim]
+        """
+        y = inputs[:,self.z_dim:]
+        for i, layers in enumerate(self.all_layers[:-1]):
+            # Run inputs through the sublayers.
+            fc_layer, norm_layer, dropout_layer = layers
+            with tf.name_scope("%s_g_layer_%d" % (self.model_name,i)):
+                x = fc_layer(inputs) if i==0 else fc_layer(x)
+#                 if i==len(self.nb_units):
+#                     x = norm_layer(x)
+                if self.batchnorm:
+                    x = norm_layer(x)
+                if self.dropout:
+                    x = dropout_layer(x)
+                x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)
+#                 x = tf.keras.layers.ReLU()(x)
+                if self.concat_every_fcl:
+                    x = tf.keras.layers.concatenate([x,y],axis=1)
+        fc_layer, norm_layer, dropout_layer = self.all_layers[-1]
+        with tf.name_scope("%s_g_layer_output"%self.model_name):
+            output = fc_layer(x)
+            # No activation func at last layer
+            if self.last_relu:
+#                 output = tf.keras.layers.ReLU()(output)
+#                 output = tf.keras.activations.tanh(output)
+                output = self.scale_factor * tf.keras.activations.tanh(output)
+        return output
+    
+class Discriminator(tf.keras.Model):
+    """Discriminator network.
+    """
+    def __init__(self, input_dim, z_dim, output_dim, model_name, nb_units=[256], batchnorm=False, dropout=True, last_relu=False):  
+        super(Discriminator, self).__init__()
+        self.input_layer = tf.keras.layers.Input((input_dim,))
+        self.input_dim = input_dim
+        self.z_dim = z_dim
+        self.output_dim = output_dim
+        self.model_name = model_name
+        self.nb_units = nb_units
+        self.batchnorm = batchnorm
+        self.dropout = dropout
+        self.last_relu = last_relu
+        self.all_layers = []
+        """Builds the FC stacks."""
+        for i in range(len(nb_units) + 1):
+            units = self.output_dim if i == len(nb_units) else self.nb_units[i]
+            fc_layer = tf.keras.layers.Dense(
+                units = units,
+                activation = None,
+#                 kernel_regularizer = tf.keras.regularizers.L2(2.5e-5)
+            )   
+            norm_layer = tf.keras.layers.BatchNormalization()
+            dropout_layer = tf.keras.layers.Dropout(0.1)
+            self.all_layers.append([fc_layer, norm_layer, dropout_layer])
+            
+        self.out = self.call(self.input_layer)
+
+    def call(self, inputs, training=True):
+        """Return the output of the Discriminator network.
+        Args:
+            inputs: tensor with shape [batch_size, input_dim]
+        Returns:
+            Output of Discriminator.
+            float32 tensor with shape [batch_size, 1]
+        """
+        fc_layer, norm_layer, dropout_layer = self.all_layers[0]
+        with tf.name_scope("%s_d_layer_0" % self.model_name):
+            x = fc_layer(inputs)
+            if self.dropout:
+                x = dropout_layer(x)
+            x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)
+#             x = tf.keras.layers.ReLU()(x)
+            
+        for i, layers in enumerate(self.all_layers[1:-1]):
+            # Run inputs through the sublayers.
+            fc_layer, norm_layer, dropout_layer = layers
+            with tf.name_scope("%s_d_layer_%d" % (self.model_name,i+1)):
+                x = fc_layer(x)
+#                 if i==len(self.nb_units):
+#                     x = norm_layer(x)
+                if self.batchnorm:
+                    x = norm_layer(x)
+                if self.dropout:
+                    x = dropout_layer(x)
+#                 x = tf.keras.activations.tanh(x)
+                x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)
+#                 x = tf.keras.layers.ReLU()(x)
+        fc_layer, norm_layer, dropout_layer = self.all_layers[-1]
+        with tf.name_scope("%s_d_layer_output" % self.model_name):
+            output = fc_layer(x)
+            if self.last_relu:
+                output = tf.keras.layers.ReLU()(output)
+        return output

scmtg-0.0.1/scMTG/scMTG.py

@@ -0,0 +1,295 @@
+import tensorflow as tf
+from .model import Decoder2, Generator, Discriminator, BaseFullyConnectedNet
+import numpy as np
+import pandas as pd
+from .util import Sequential_sampler
+from .loss import NB
+import dateutil.tz
+import datetime
+import sys
+import copy
+import os
+import json
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+class scMTG(object):
+    """Markov Transition Generative model for time-series single cell analysis.
+    """
+    def __init__(self, params, timestamp=None, random_seed=None):
+        super(scMTG, self).__init__()
+        self.params = params
+        self.timestamp = timestamp
+        if random_seed is not None:
+            tf.keras.utils.set_random_seed(random_seed)
+
+        #initilize the shared autoencoder (encoder + decoder)
+        self.encoder = BaseFullyConnectedNet(input_dim=params['e_dim'],z_dim=params['z_dim'],output_dim = params['z_dim'], 
+                                        model_name='e_net', nb_units=params['e_units'], last_relu=False)
+        self.decoder = Decoder2(input_dim=params['z_dim'],z_dim=params['z_dim'],output_dim = params['e_dim'],
+                               model_name='d_net', nb_units=params['d_units'], last_relu=False)
+
+        #initilize the T-1 Markov generators and T-1 discriminators
+        self.generators = [Generator(input_dim=params['noise_dim']+params['z_dim'],z_dim=params['z_dim'],
+                                     output_dim=params['z_dim'],model_name='g_net_%d'%i,
+                                     nb_units=params['gen_units'], concat_every_fcl=False, last_relu=True) 
+                           for i in range(params['nb_time']-1)]
+        self.discriminators = [Discriminator(input_dim=params['z_dim'],z_dim=params['z_dim'],output_dim = 1,
+                                             model_name='d_net_%d'%i, nb_units=params['dis_units'], last_relu=False) 
+                               for i in range(params['nb_time']-1)]
+        lr_schedule1 = tf.keras.optimizers.schedules.ExponentialDecay(params['lr'], decay_steps=100000, decay_rate=0.9, staircase=True)
+        lr_schedule2 = tf.keras.optimizers.schedules.ExponentialDecay(params['lr']/10.0, decay_steps=100000, decay_rate=0.9, staircase=True)
+        self.ae_optimizer = tf.keras.optimizers.Adam(lr_schedule1, beta_1=0.5, beta_2=0.9)
+        self.e_optimizer = tf.keras.optimizers.Adam(lr_schedule2, beta_1=0.5, beta_2=0.9)
+        self.g_optimizer = tf.keras.optimizers.Adam(lr_schedule1, beta_1=0.5, beta_2=0.9)
+        self.d_optimizer = tf.keras.optimizers.Adam(lr_schedule1, beta_1=0.5, beta_2=0.9)
+
+        self.initialize_nets()
+
+        if self.timestamp is None:
+            now = datetime.datetime.now(dateutil.tz.tzlocal())
+            self.timestamp = now.strftime('%Y%m%d_%H%M%S')
+            
+        self.best_path = "{}/{}/best_model/".format(
+            params['output_dir'], self.timestamp)
+
+        if self.params['save_model'] and not os.path.exists(self.best_path):
+            os.makedirs(self.best_path)
+        
+        self.save_dir = "{}/{}".format(
+            params['output_dir'], self.timestamp)
+
+        if self.params['save_res'] and not os.path.exists(self.save_dir):
+            os.makedirs(self.save_dir)   
+
+        self.ckpt = tf.train.Checkpoint(encoder = self.encoder,
+                                   decoder = self.decoder,
+                                   generators = self.generators,
+                                   discriminators = self.discriminators,
+                                   ae_optimizer = self.ae_optimizer,
+                                   e_optimizer = self.e_optimizer,
+                                   g_optimizer = self.g_optimizer,
+                                   d_optimizer = self.d_optimizer)
+
+    def get_config(self):
+        return {
+                "params": self.params,
+        }
+
+    def initialize_nets(self, print_summary = True):
+        """Initialize all the networks in CausalEGM."""
+
+        self.encoder(np.zeros((1, self.params['e_dim'])))
+        self.decoder(np.zeros((1, self.params['z_dim'])), 1.0)
+        [self.generators[i](np.zeros((1, self.params['z_dim']+self.params['noise_dim']))) 
+            for i in range(self.params['nb_time']-1)]
+        [self.discriminators[i](np.zeros((1, self.params['z_dim']))) 
+            for i in range(self.params['nb_time']-1)]
+        if print_summary:
+            print(self.encoder.summary())
+            print(self.decoder.summary())
+            print([self.generators[i].summary() for i in range(self.params['nb_time']-1)])
+            print([self.discriminators[i].summary() for i in range(self.params['nb_time']-1)])
+
+    @tf.function
+    def train_ae_step(self, data_series):
+        """train shared AE.
+        """  
+        with tf.GradientTape(persistent=True) as tape:
+            embed_series = tf.map_fn(lambda item:self.encoder(item), data_series)
+            
+            disps, means = [], []
+            for i in range(len(embed_series)):
+                disp, mean = self.decoder(embed_series[i])
+                disps.append(disp)
+                means.append(mean)
+            zinb = NB(theta=tf.concat(disps, axis=0), debug=False)
+            loss_rec = zinb.loss(tf.reshape(data_series, [-1, data_series.shape[-1]]), tf.concat(means, axis=0), mean=True)
+    
+        # Calculate the gradients
+        gradients = tape.gradient(loss_rec, self.encoder.trainable_variables+self.decoder.trainable_variables)
+        # Apply the gradients to the optimizer
+        self.ae_optimizer.apply_gradients(zip(gradients, self.encoder.trainable_variables+self.decoder.trainable_variables))
+        return loss_rec
+
+    @tf.function
+    def train_gen_step(self, data_series):
+        """train generators.
+        """  
+        noise = tf.random.normal(shape=(self.params['nb_time']-1,data_series.shape[1],self.params['noise_dim']),mean=0.,stddev=self.params['sd'])
+        with tf.GradientTape(persistent=True) as gen_tape:
+            embed_series = tf.map_fn(lambda item:self.encoder(item) ,data_series)
+
+            #data from time point 1,2,...,T-1
+            data_previous = tf.concat([embed_series[:-1],noise],axis=-1)
+
+            #generated data for time point 2,...,T
+            data_gen = tf.TensorArray(tf.float32, size=data_previous.shape[0])
+            for i in range(data_previous.shape[0]):
+                data_gen = data_gen.write(i, self.generators[i](data_previous[i]))
+            data_gen=data_gen.stack()
+
+            dz_gen = tf.TensorArray(tf.float32, size=data_gen.shape[0])
+            for i in range(data_gen.shape[0]):
+                dz_gen = dz_gen.write(i, self.discriminators[i](data_gen[i]))
+            dz_gen=dz_gen.stack()
+
+            loss_g = -tf.reduce_mean(dz_gen)
+            loss_td = tf.reduce_mean((data_gen-embed_series[:-1])**2)
+            loss_g_total = loss_g + self.params['beta']*loss_td
+
+        # Calculate the gradients
+        g_gradients = gen_tape.gradient(loss_g_total, sum([item.trainable_variables for item in self.generators], []))
+        self.g_optimizer.apply_gradients(zip(g_gradients, sum([item.trainable_variables for item in self.generators], [])))
+        
+        e_gradients = gen_tape.gradient(loss_g_total, self.encoder.trainable_variables)
+        self.e_optimizer.apply_gradients(zip(e_gradients, self.encoder.trainable_variables))
+        return loss_g, loss_td, loss_g_total
+
+    @tf.function
+    def train_disc_step(self, data_series):
+        """train discriminators.
+        """  
+        epsilon_z = tf.random.uniform(shape=(self.params['nb_time']-1,1,1),minval=0., maxval=1.)
+        noise = tf.random.normal(shape=(self.params['nb_time']-1,data_series.shape[1],self.params['noise_dim']),mean=0.,stddev=self.params['sd'])
+        with tf.GradientTape(persistent=True) as disc_tape:
+            embed_series = tf.map_fn(lambda item:self.encoder(item) ,data_series)
+            data_previous = tf.concat([embed_series[:-1],noise],axis=-1)
+
+            data_gen = tf.TensorArray(tf.float32, size=self.params['nb_time']-1)
+            for i in range(self.params['nb_time']-1):
+                data_gen = data_gen.write(i, self.generators[i](data_previous[i]))
+            data_gen=data_gen.stack()
+            
+            data_true = embed_series[1:]
+
+            dz_gen = tf.TensorArray(tf.float32, size=self.params['nb_time']-1)
+            for i in range(self.params['nb_time']-1):
+                dz_gen = dz_gen.write(i, self.discriminators[i](data_gen[i]))
+            dz_gen=dz_gen.stack()
+
+            dz_true = tf.TensorArray(tf.float32, size=self.params['nb_time']-1)
+            for i in range(self.params['nb_time']-1):
+                dz_true = dz_true.write(i, self.discriminators[i](data_true[i]))
+            dz_true = dz_true.stack()
+            loss_d = tf.reduce_mean(dz_gen)-tf.reduce_mean(dz_true)
+            
+            #gradient penalty for z
+            data_hat = epsilon_z*data_gen+(1-epsilon_z)*data_true
+
+            dz_hat = tf.TensorArray(tf.float32, size=self.params['nb_time']-1)
+            for i in range(self.params['nb_time']-1):
+                dz_hat = dz_hat.write(i, self.discriminators[i](data_hat[i]))
+            dz_hat=dz_hat.stack()
+            
+            grads = tf.gradients(dz_hat,data_hat)[0]
+            grad_norms = tf.sqrt(tf.reduce_sum(tf.square(grads), axis=[1, 2]))
+            loss_gp = tf.reduce_mean(tf.square(grad_norms - 1))
+
+            loss_d_total = loss_d + self.params['alpha']*loss_gp
+
+        # Calculate the gradients
+        d_gradients = disc_tape.gradient(loss_d_total, sum([item.trainable_variables for item in self.discriminators], []))
+        self.d_optimizer.apply_gradients(zip(d_gradients, sum([item.trainable_variables for item in self.discriminators], [])))
+        return loss_d, loss_gp, loss_d_total
+
+    def train(self, data, weights, batch_size=32, n_iter=100000, batches_per_verbose=500, verbose=1):
+        if self.params['save_res']:
+            f_params = open('{}/params.txt'.format(self.save_dir),'w')
+            f_params.write(str(self.params))
+            f_params.close()
+        self.data_sampler = Sequential_sampler(data=data, weights=weights, batch_size=batch_size)
+        
+        best_loss = float('inf')
+        best_ckpt = None
+        best_batch_idx = 0
+        loss_pd = []
+        for batch_idx in range(n_iter+1):
+            batch_data_series = self.data_sampler.next_batch()
+
+            #train autoencoders
+            loss_rec = self.train_ae_step(batch_data_series)
+
+            #control the update frequency of autoencoder vs GAN
+            if batch_idx % self.params['ae_gan_freq'] == 0:
+                #update G once and update D multiple times for WGAN-GP
+                for _ in range(self.params['g_d_freq']):
+                    batch_data_series = self.data_sampler.next_batch()
+                    loss_d, loss_gp, loss_d_total = self.train_disc_step(batch_data_series)
+                batch_data_series = self.data_sampler.next_batch()
+                loss_g, loss_td, loss_g_total = self.train_gen_step(batch_data_series)
+            
+            loss_total = loss_rec + loss_g_total + loss_d_total
+            loss_pd.append([batch_idx, loss_total.numpy(), loss_rec.numpy(), loss_gp.numpy(), loss_d.numpy(), loss_g.numpy(), loss_td.numpy()])
+            if loss_total < best_loss:
+                best_loss = loss_total
+                best_ckpt = self.ckpt
+                best_batch_idx = batch_idx
+                self.ckpt.save('{}/{}/best_model'.format(self.best_path, batch_idx))
+            if batch_idx - best_batch_idx >= 10000 and batch_idx >= 50000:
+                print("Early stop at {} and best at {}!".format(batch_idx, best_batch_idx))
+                break
+                
+            if batch_idx % batches_per_verbose == 0:
+                loss_contents = '''Iteration [%d] : loss_t [%.4f], loss_rec [%.4f], loss_gp [%.4f], loss_d [%.4f], loss_g [%.4f] loss_td [%.4f]''' \
+                %(batch_idx, loss_total, loss_rec, loss_gp, loss_d, loss_g, loss_td)
+                if verbose:
+                    print(loss_contents)
+        loss_pd = pd.DataFrame(loss_pd, columns=['batch_idx', 'loss_total', 'loss_rec', 'loss_gp', 'loss_d', 'loss_g', 'loss_td'])
+        loss_pd.to_csv(self.save_dir+'/loss_pd.csv', sep = '\t', index = False)
+#         self.ckpt.save('{}/best_model_at_{}'.format(self.best_path, batch_idx))
+        self.ckpt.restore(tf.train.latest_checkpoint('{}/best_model_at_{}'.format(self.best_path, best_batch_idx)))
+        self.evaluate(self.data_sampler.load_all(), 'best')
+    
+    def evaluate(self, data_series, batch_idx):
+        embed_series = [self.encoder.predict(item) for item in data_series]
+        data_previous = [np.concatenate([data, np.random.normal(0.,self.params['sd'],size=(data.shape[0],self.params['noise_dim']))],axis=1) for data in embed_series[:-1]] #contain T-1 time points
+        data_gen_z = [self.generators[i].predict(data) for i,data in enumerate(data_previous)]
+        data_gen_org = [self.decoder.predict(item)[-1] for item in data_gen_z]
+        
+        np.savez('{}/data_embed_at_{}.npz'.format(self.save_dir, batch_idx),embed_series)
+        np.savez('{}/data_gen_at_{}.npz'.format(self.save_dir, batch_idx),data_gen_z)
+        np.savez('{}/data_gen_org_at_{}.npz'.format(self.save_dir, batch_idx),data_gen_org)
+
+    def thresholding_(self, trans_mtx):
+        trans_mtx2 = trans_mtx / np.sum(trans_mtx, axis=1, keepdims=True)
+        trans_mtx2[np.isnan(trans_mtx2)] = 0.0
+
+        trans_mtx0 = trans_mtx2[np.where(np.sum(trans_mtx2, axis=1)>0)[0]]
+        thresh = min(trans_mtx0[i].max() for i in range(trans_mtx0.shape[0]))
+#         print(thresh)
+
+        trans_mtx3 = trans_mtx2.copy()
+        trans_mtx3[trans_mtx3<thresh] = 0.0
+        trans_mtx3 = trans_mtx3 / np.sum(trans_mtx3, axis=1, keepdims=True)
+        trans_mtx3[np.isnan(trans_mtx3)] = 0.0
+        return trans_mtx3
+    
+    def compute_trans_mat(self, times=1, n_noise=1000, n_chunk=1000, random_seed=1, thresholding=True, save_mtx=False):
+        data_series = self.data_sampler.load_all()
+        embed_series = [self.encoder.predict(item) for item in data_series]
+        
+        tf.random.set_seed(random_seed)
+        noises = tf.random.normal(shape=(n_noise, self.params['noise_dim']), mean=0.0, stddev=self.params['sd'])
+
+        embed_data = embed_series[times]
+        trans_mtx = []
+        const = -0.5 / (self.params['sd'] ** 2)
+        for embed_data0 in embed_data:
+            embed_gen1 = self.generators[times].predict(tf.concat([tf.tile(tf.reshape(embed_data0,(1,-1)),[n_noise,1]),noises], axis=-1))
+            trans2 = []
+            for i in range(n_noise//n_chunk):
+                trans2.append(tf.reduce_mean(tf.math.exp(const*tf.reduce_sum((tf.reshape(tf.tile(embed_series[times+1], [1,n_chunk]), [embed_series[times+1].shape[0],-1,embed_series[times+1].shape[1]])-embed_gen1[i*n_chunk:(i+1)*n_chunk])**2,axis=-1)),axis=-1))
+            trans_mtx.append(tf.reduce_mean(trans2, axis=0))
+        trans_mtx = np.array(trans_mtx)
+        
+        if thresholding:
+            trans_mtx = self.thresholding_(trans_mtx)
+            
+        if save_mtx:
+            np.savez('{}/trans_mtx_{}.npz'.format(self.save_dir, times),t1=trans_mtx)
+        return trans_mtx
+
+    
+    

scmtg-0.0.1/scMTG/util.py

@@ -0,0 +1,108 @@
+import numpy as np
+import os
+import math
+import pandas as pd
+import scipy
+import anndata as ad
+import scib
+
+
+def _logistic(x, L, k, center=0):
+    return L / (1 + np.exp(-k * (x - center)))
+
+def _gen_logistic(p, sup, inf, center, width):
+    return inf + _logistic(p, L=sup - inf, k=4.0 / width, center=center)
+
+def beta(p, beta_max=1.7, beta_min=0.3, beta_center=0.25, beta_width=0.5):
+    return _gen_logistic(p, beta_max, beta_min, beta_center, beta_width)
+
+def delta(a, delta_max=1.7, delta_min=0.3, delta_center=0.1, delta_width=0.2):
+    return _gen_logistic(a, delta_max, delta_min, delta_center, delta_width)
+
+def growth_rate(adata, proliferation_key="proliferation", apoptosis_key="apoptosis", delta_t=1.0, 
+                beta_max=1.7, beta_min=0.3, beta_center=0.25, beta_width=0.5, 
+                delta_max=1.7, delta_min=0.3, delta_center=0.1, delta_width=0.2):
+    birth = beta(adata.obs[proliferation_key].values, 
+                 beta_max=beta_max, beta_min=beta_min, beta_center=beta_center, beta_width=beta_width)
+    death = delta(adata.obs[apoptosis_key].values, 
+                  delta_max=delta_max, delta_min=delta_min, delta_center=delta_center, delta_width=delta_width)
+    gr = np.exp((birth - death) * delta_t)
+    return gr
+
+class Sequential_sampler(object):
+    def __init__(self, data, weights=None, batch_size=32, random_seed=123):
+        np.random.seed(random_seed)
+        self.data = [np.array(item, dtype='float32') for item in data]
+        self.nb_time = len(self.data)
+        self.batch_size = batch_size
+        self.sample_sizes = [item.shape[0] for item in self.data]
+        if weights is None:
+            self.weights = [np.ones(item, dtype='float32') / item for item in self.sample_sizes]
+        else:
+            self.weights = [np.array(item, dtype='float32') / np.sum(item) for item in weights]
+        self.idx_gens = [self.create_idx_generator(sample_size=item, time_idx=i) for i,item in enumerate(self.sample_sizes)]
+        
+    def create_idx_generator(self, sample_size, time_idx, random_seed=123):
+        while True:
+#             indices = np.random.choice(sample_size, size=self.batch_size, replace=True, p=self.weights[time_idx])
+            indices = np.random.choice(sample_size, size=self.batch_size, replace=False, p=self.weights[time_idx])
+            yield indices
+#             indices = np.random.choice(sample_size, size=3*self.batch_size, replace=False, p=self.weights[time_idx])
+#             np.random.shuffle(indices)
+#             yield indices[:self.batch_size]
+
+    def next_batch(self):
+        indexes = [next(item) for item in self.idx_gens]
+        return np.stack([item[indexes[i],:] for i,item in enumerate(self.data)])
+    
+    def load_all(self):
+        return self.data
+
+def interpolate(p0, p1, tmap, interp_frac, size, seed=1):
+    p0 = p0.toarray() if scipy.sparse.isspmatrix(p0) else p0
+    p1 = p1.toarray() if scipy.sparse.isspmatrix(p1) else p1
+    p0 = np.asarray(p0, dtype=np.float64)
+    p1 = np.asarray(p1, dtype=np.float64)
+    tmap = np.asarray(tmap, dtype=np.float64)
+    if p0.shape[1] != p1.shape[1]:
+        raise ValueError("Unable to interpolate. Number of genes do not match")
+    if p0.shape[0] != tmap.shape[0] or p1.shape[0] != tmap.shape[1]:
+        raise ValueError("Unable to interpolate. Tmap size is {}, expected {}"
+                         .format(tmap.shape, (len(p0), len(p1))))
+    I = len(p0);
+    J = len(p1)
+    a = np.power(tmap.sum(axis=0), 1. - interp_frac)
+    a[a==0] = np.finfo(float).eps
+    p = tmap / a
+    p = p.flatten(order='C')
+    p = p / p.sum()
+    np.random.seed(seed)
+    choices = np.random.choice(I * J, p=p, size=size)
+    return np.asarray([p0[i // J] * (1 - interp_frac) + p1[i % J] * interp_frac for i in choices], dtype=np.float64)
+    
+def cal_metrics(gen_data, real_data):
+    mmd_value = cal_mmd(gen_data, real_data)
+    
+    x, y = np.mean(real_data, axis=0), np.mean(gen_data, axis=0)
+    pearson_corr, _ = scipy.stats.pearsonr(x, y)
+    spearman_corr, _ = scipy.stats.spearmanr(x, y)
+    
+    lisi_value = cal_lisi(gen_data, real_data)
+
+    return mmd_value, pearson_corr, spearman_corr, lisi_value
+
+def cal_mmd(f_of_X, f_of_Y):
+    loss = 0.0
+    delta = f_of_X - f_of_Y
+    mmd_value = np.mean((delta[:-1] * delta[1:]).sum(1))
+    return mmd_value
+
+def cal_lisi(gen_data, real_data):
+    adata_scib = ad.AnnData(np.concatenate((gen_data, real_data), axis=0))
+    adata_scib.obs['batch'] = pd.Categorical(['gen']*len(gen_data) + ['true']*len(real_data))
+    lisi_value = scib.me.ilisi_graph(adata_scib, batch_key="batch", type_="full")
+    return lisi_value
+  
+    
+    
+    

scmtg-0.0.1/scmtg.egg-info/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.1
+Name: scmtg
+Version: 0.0.1
+Summary: scMTG learns generative Markov transitions for single-cell temporal dynamics
+Home-page: https://github.com/liuq-lab/scMTG
+Author: Xuejian Cui
+Author-email: cuixj@hit.edu.cn
+License: MIT
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.8.0
+Requires-Dist: tensorflow-gpu==2.6.0
+Requires-Dist: keras==2.6.0

scmtg-0.0.1/scmtg.egg-info/SOURCES.txt

@@ -0,0 +1,11 @@
+setup.py
+scMTG/__init__.py
+scMTG/loss.py
+scMTG/model.py
+scMTG/scMTG.py
+scMTG/util.py
+scmtg.egg-info/PKG-INFO
+scmtg.egg-info/SOURCES.txt
+scmtg.egg-info/dependency_links.txt
+scmtg.egg-info/requires.txt
+scmtg.egg-info/top_level.txt

scmtg-0.0.1/scmtg.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

scmtg-0.0.1/scmtg.egg-info/requires.txt

@@ -0,0 +1,2 @@
+tensorflow-gpu==2.6.0
+keras==2.6.0

scmtg-0.0.1/scmtg.egg-info/top_level.txt

@@ -0,0 +1 @@
+scMTG

scmtg-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

scmtg-0.0.1/setup.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from setuptools import setup, find_packages
+
+setup(name='scmtg',
+      version='0.0.1',
+      packages=find_packages(),
+      description='scMTG learns generative Markov transitions for single-cell temporal dynamics',
+      long_description='',
+
+      author='Xuejian Cui',
+      author_email='cuixj@hit.edu.cn',
+      url="https://github.com/liuq-lab/scMTG",
+      python_requires='>=3.8.0',
+      license='MIT',
+
+      classifiers=[
+          'Development Status :: 4 - Beta',
+          'Intended Audience :: Science/Research',
+          'License :: OSI Approved :: MIT License',
+          'Programming Language :: Python :: 3.8',
+          'Operating System :: MacOS :: MacOS X',
+          'Operating System :: Microsoft :: Windows',
+          'Operating System :: POSIX :: Linux',
+          'Topic :: Scientific/Engineering :: Bio-Informatics',
+     ],
+     
+    install_requires=[
+        'tensorflow-gpu==2.6.0',
+        'keras==2.6.0',
+    ]
+     )