text: #!/usr/bin/python
  # -*- coding: utf-8 -*-
  
  import tensorflow as tf
  import matplotlib
  matplotlib.use('Agg')
  import os
  import time
  import matplotlib.pyplot as plt
  import numpy as np
  from matplotlib import cm
  from estimators.utilities import Generator, Critic
  from estimators.utilities import add_random_input, add_random_labels
  from estimators.utilities import set_learning_rate, set_global_step, rescale
  from estimators.utilities import gradient_step, sc_summary, save_generated_cells
  from MulticoreTSNE import MulticoreTSNE as TSNE
  tsne = TSNE(n_jobs=20)
  
  
  class cscGAN:
      """
      Contains the class for the conditional scGAN (cscGAN).
      Methods include the creation of the graph, the training of the model,
       the validation and generation of the cells.
      """
      def __init__(self, train_files, valid_files, genes_no, clusters_no,
                   scaling, scale_value, max_steps, batch_size, latent_dim,
                   gen_layers, output_lsn, gene_cond_type, critic_layers,
                   optimizer, lambd, beta1, beta2, decay, alpha_0, alpha_final):
          """
          Constructor for the cscGAN.
  
          Parameters
          ----------
          train_files : list
              List of TFRecord files used for training.
          valid_files : list
              List of TFRecord files used for validation.
          genes_no : int
              Number of genes in the expression matrix.
          clusters_no : int
              Number of clusters.
          scaling : str
              Method used to scale the data, see the scaling method of the
              GeneMatrix class in preprocessing/process_raw.py for more details.
          scale_value : int, float
              Parameter of the scaling function.
          max_steps : int
              Number of steps in the (outer) training loop.
          batch_size : int
              Batch size used for the training.
          latent_dim : int
              Dimension of the latent space used from which the input noise
               of the generator is sampled.
          gen_layers : list
              List of integers corresponding to the number of neurons of each
              layer of the generator.
          output_lsn : int, None
              Parameter of the LSN layer at the output of the critic
               (i.e. total number of counts per generated cell).
              If set to None, the layer won't be added in the generator.
          gene_cond_type : str
              Conditional normalization layers used in the generator.
               Can be either "batchnorm" or "layernorm".
              If anything else, it won't be added in the model
              (there will be no conditioning in the generation).
          critic_layers : list
              List of integers corresponding to the number of neurons of each
              layer of the critic.
          optimizer : str
              Optimizer used in the training.
              Can be "AMSGrad" for AMSGrad.
              If anything else, Adam will be used.
          lambd : float
              Regularization hyper-parameter to be used with the gradient
               penalty in the WGAN loss.
          beta1 : float
              Exponential decay for the first-moment estimates.
          beta2 : float
              Exponential decay for the second-moment estimates.
          decay : str
              If True, uses an exponential decay of the learning rate.
          alpha_0 : float
              Initial learning rate value.
          alpha_final : float
              Final value of the learning rate if the decay is set to True.
          """
  
          # read the parameters
          self.clusters_no = clusters_no
          self.latent_dim = latent_dim
          self.lambd = lambd
          self.gen_cond_type = gene_cond_type
          self.critic_layers = critic_layers
          self.gen_layers = gen_layers
          self.output_lsn = output_lsn
          self.optimizer = optimizer
          self.beta1 = beta1
          self.beta2 = beta2
          self.lr_decay = decay
          self.alpha_0 = alpha_0
          self.alpha_final = alpha_final
          self.batch_size = batch_size
          self.max_steps = max_steps
          self.scaling = scaling
          self.scale_value = scale_value
          self.train_files = train_files
          self.valid_files = valid_files
          self.genes_no = genes_no
  
          # prepare input pipeline for training
          self.train_cells, self.train_cells_clusters = self.make_input_fn(
              self.train_files)
  
          # prepare input pipeline for validation
          self.test_cells, self.test_cells_clusters = self.make_input_fn(
              self.valid_files)
  
          # module parameters
          self.generator = None
          self.critic_real = None
          self.critic_fake = None
          self.gradient_penalty = None
          self.gen_loss = None
          self.critic_loss = None
          self.global_step = None
          self.incr_global_step = None
          self.learning_rate = None
          self.critic_train = None
          self.critic_grads_and_vars = None
          self.gen_train = None
          self.gen_grads_and_vars = None
          self.model_train = None
          self.output_tensor = None
          self.build_model()
  
          # add visualization
          self.visualization()
  
          # the total number of all trainable parameters
          self.parameter_count = tf.reduce_sum(
              [tf.reduce_prod(tf.shape(v)) for v in tf.trainable_variables()])
  
      def make_input_fn(self, file_paths, epochs=None):
          """
          Function that loads the TFRecords files and creates the placeholders
          for the data inputs.
  
          Parameters
          ----------
          file_paths : list
              List of TFRecord files from which to read from.
          epochs : int
              Integer specifying the number of times to read through the dataset.
              If None, cycles through the dataset forever.
              NOTE - If specified, creates a variable that must be initialized,
              so call tf.local_variables_initializer() and run the op in a session.
              Default is None.
  
          Returns
          -------
          features : Tensor
              Tensor containing a batch of cells (vector of expression levels).
          cluster : Tensor
              Tensor containing (a batch of) the cluster indexes of the
              corresponding cells.
          """
  
          feature_map = {'scg': tf.SparseFeature(index_key='indices',
                                                 value_key='values',
                                                 dtype=tf.float32,
                                                 size=self.genes_no),
                         'cluster_int': tf.FixedLenFeature(1, tf.int64)}
  
          options = tf.python_io.TFRecordOptions(
              tf.python_io.TFRecordCompressionType.GZIP)
  
          batched_features = tf.contrib.learn.read_batch_features(
              file_pattern=file_paths,
              batch_size=self.batch_size,
              features=feature_map,
              reader=lambda: tf.TFRecordReader(
                  options=options),
              num_epochs=epochs)
  
          sgc = batched_features['scg']
  
          sparse = tf.sparse_reshape(sgc, (self.batch_size, self.genes_no))
  
          dense = tf.sparse_tensor_to_dense(sparse)
  
          cluster = tf.squeeze(tf.to_int32(batched_features['cluster_int']))
  
          features = tf.reshape(dense, (self.batch_size, self.genes_no))
  
          return features, cluster
  
      def build_model(self):
          """
          Method that initializes the cscGAN model, creates the graph and
          defines the loss and optimizer.
  
          Returns
          -------
  
          """
          # training or inference (used for the batch normalization)
          is_training = tf.placeholder(dtype=tf.bool, name='is_training')
  
          clusters_ratios = tf.placeholder(dtype=tf.float32,
                                           shape=(1, self.clusters_no),
                                           name='clusters_ratios')
  
          z_input = add_random_input(self.batch_size, self.latent_dim)
          input_clusters = add_random_labels(clusters_ratios, self.batch_size)
  
          # create generator
          self.generator = Generator.create_cond_generator(
              z_input=z_input,
              batch_size=self.batch_size,
              latent_dim=self.latent_dim,
              output_cells_dim=self.genes_no,
              var_scope='generator',
              gen_layers=self.gen_layers,
              output_lsn=self.output_lsn,
              gen_cond_type=self.gen_cond_type,
              clusters_no=self.clusters_no,
              input_clusters=input_clusters,
              is_training=is_training,
              clusters_ratios=clusters_ratios,
              reuse=None)
  
          # Critic with real cells as input
          with tf.name_scope('real_critic'):
              self.critic_real = Critic.create_cond_critic(
                  xinput=self.train_cells,
                  input_clusters=self.train_cells_clusters,
                  var_scope="critic",
                  critic_layers=self.critic_layers,
                  clusters_no=self.clusters_no,
                  reuse=None)
  
          # Critic with generated cells as input (shares weights with critic_real)
          with tf.name_scope('fake_critic'):
              self.critic_fake = \
                  Critic.create_cond_critic(
                      xinput=self.generator.fake_outputs,
                      input_clusters=self.generator.input_clusters,
                      var_scope="critic",
                      critic_layers=self.critic_layers,
                      clusters_no=self.clusters_no, reuse=True)
  
          # Disc loss
          with tf.name_scope('critic_loss'):
              critic_loss_wgan = tf.reduce_mean(self.critic_fake.dist) \
                                 - tf.reduce_mean(self.critic_real.dist)
  
              # The following lines implement the gradient penalty term
              alpha = tf.random_uniform(shape=[self.batch_size, 1],
                                        minval=0., maxval=1.)
  
              generator_interpolates = \
                  Generator.create_cond_generator(
                      z_input=z_input,
                      batch_size=self.batch_size,
                      latent_dim=self.latent_dim,
                      output_cells_dim=self.genes_no,
                      var_scope='generator',
                      gen_layers=self.gen_layers,
                      output_lsn=self.output_lsn,
                      gen_cond_type=self.gen_cond_type,
                      clusters_no=self.clusters_no,
                      input_clusters=self.train_cells_clusters,
                      is_training=is_training,
                      clusters_ratios=clusters_ratios,
                      reuse=True)
  
              differences = generator_interpolates.fake_outputs - self.train_cells
  
              interpolates = self.train_cells + (alpha * differences)
  
              with tf.name_scope('help_critic'):
                  critic_interpolates = \
                      Critic.create_cond_critic(
                          xinput=interpolates,
                          input_clusters=self.train_cells_clusters,
                          var_scope="critic",
                          critic_layers=self.critic_layers,
                          clusters_no=self.clusters_no,
                          reuse=True)
  
              gradients = tf.gradients(critic_interpolates.dist, [interpolates])[0]
              slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients),
                                             reduction_indices=[1]))
              self.gradient_penalty = tf.reduce_mean((slopes - 1) ** 2)
              critic_loss_wgan += self.lambd * self.gradient_penalty
              self.critic_loss = critic_loss_wgan
  
          # gen loss
          with tf.name_scope('generator_loss'):
              self.gen_loss = -tf.reduce_mean(self.critic_fake.dist)
  
          # add global step
          self.global_step, self.incr_global_step = set_global_step()
  
          # add decaying learning rate
          self.learning_rate = set_learning_rate(self.alpha_0,
                                                 self.alpha_final,
                                                 self.global_step,
                                                 self.max_steps)
  
          # training the critic
          with tf.name_scope("critic_train"):
              critic_params = [var for var in tf.trainable_variables()
                               if var.name.startswith('critic')]
              self.critic_train, self.critic_grads_and_vars =\
                  gradient_step(critic_params,
                                training_loss=self.critic_loss,
                                learning_rate=self.learning_rate,
                                beta1=self.beta1,
                                beta2=self.beta2,
                                optimizer=self.optimizer)
  
          # training the generator
          with tf.name_scope("generator_train"):
              with tf.control_dependencies(
                      [self.critic_train] +
                      tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
                  gen_params = [var for var in tf.trainable_variables()
                                if var.name.startswith('gen')]
                  self.gen_train, self.gen_grads_and_vars = gradient_step(
                      gen_params,
                      training_loss=self.gen_loss,
                      learning_rate=self.learning_rate,
                      beta1=self.beta1,
                      beta2=self.beta2,
                      optimizer=self.optimizer)
  
          self.model_train = tf.group(self.incr_global_step, self.gen_train)
          self.critic_train = tf.group(self.critic_train)
  
      def visualization(self):
          """
          Method creating the placeholders to log the values for monitoring
          with the Tensorboard.
  
          Returns
          -------
  
          """
  
          # histograms and mean of real and generated cells
          sc_summary("single_cell_fake", self.generator.fake_outputs)
          sc_summary("single_cell_real", self.train_cells)
  
          # loss functions visualization
          tf.summary.scalar("gen_loss", self.gen_loss)
          tf.summary.scalar("critic_loss", self.critic_loss)
  
          tf.summary.scalar("Penalty", self.gradient_penalty)
          tf.summary.scalar("learning_rate", self.learning_rate)
  
          tf.summary.histogram("Distance_fake", self.critic_fake.dist)
          tf.summary.histogram("Distance_real", self.critic_real.dist)
  
          # visualize  trainable variables
          for var in tf.trainable_variables():
              tf.summary.histogram(var.op.name + "/values", var)
  
          # visualize gradients
          for grad, var in self.gen_grads_and_vars + self.critic_grads_and_vars:
              if grad is not None:
                  tf.summary.histogram(var.op.name + "/gradients", grad)
  
      def training(self, exp_folder, clusters_ratios, checkpoint=None,
                   progress_freq=1000, summary_freq=200, save_freq=5000,
                   validation_freq=500, critic_iter=5, valid_cells_no=500):
          """
          Method that trains the cscGAN.
  
          Parameters
          ----------
          exp_folder : str
              Path where TF will write the logs, save the model, the t-SNE plots etc.
          clusters_ratios : dict
              Dictionary containing the different cluster names and their
              ratio in the data.
          checkpoint : str, None
              Path to the checkpoint to start from, or None to start training
               from scratch.
              Default is None.
          progress_freq : int
              Period (in steps) between displays of the losses values on the
              standard output.
              Default is 1000.
          summary_freq : int
              Period (in steps) between logs for the Tensorboard.
              Default is 200.
          save_freq : int
              Period (in steps) between saves of the model.
              Default is 5000.
          validation_freq : int
              Period (in steps) between validation measures are computed
              (e.g. t-SNE plots).
              Default is 500.
          critic_iter : int
              Number of training iterations of the critic (inner loop) for each
              iteration on the generator (outer loop).
              Default is 5.
          valid_cells_no :
              Number of cells in the validation set.
              Default is 500.
  
          Returns
          -------
  
          """
          exp_name = exp_folder.split('/')[-1]
  
          # Transform the cluster ratios dictionary into an ordered list
          clusters_ratios = [value for (key, value) in sorted(clusters_ratios.items())]
          clusters_ratios.sort(reverse=True)
          clusters_ratios = np.reshape(clusters_ratios, (1, len(clusters_ratios)))
  
          # Number of different models to keep (each time a model is saved,
          #  it will overwrite the oldest)
          saver = tf.train.Saver(max_to_keep=1)
  
          train_supervisor = tf.train.Supervisor(logdir=exp_folder,
                                                 save_summaries_secs=0,
                                                 saver=None)
  
          start = time.time()
  
          # Start the TF session
          with train_supervisor.managed_session() as sess:
  
              train_feed_dict = {self.generator.is_training: True,
                                 self.generator.clusters_ratios: clusters_ratios}
  
              print("Parameter Count is [ %d ]." % (sess.run(self.parameter_count)))
  
              # load checkpoint and instantiate the start_step accordingly
              if checkpoint is not None:
                  print("Loading model from checkpoint....")
                  checkpoint = tf.train.latest_checkpoint(checkpoint)
                  saver.restore(sess, checkpoint)
                  start_step = sess.run(train_supervisor.global_step)
              else:
                  start_step = 0
  
              critic_fetch = {"train": self.critic_train}
  
              # Outer loop, one step for the generator and several for the critic
              for step in range(start_step, self.max_steps):
  
                  # small utility function to perform tasks at defined intervals
                  def should(freq):
                      return freq > 0 and \
                             ((step + 1) % freq == 0 or
                              step == self.max_steps - 1)
  
                  # Inner loop, for each generator step, several critic steps
                  if step > 0:
                      for i_critic in range(critic_iter):
                          sess.run(fetches=critic_fetch, feed_dict=train_feed_dict)
  
                  model_fetches = {"train": self.model_train}
  
                  # Add the corresponding summary tensors to the fetches
                  if should(summary_freq):
                      model_fetches["summary"] = train_supervisor.summary_op
  
                  if should(progress_freq):
                      model_fetches["gen_loss"] = self.gen_loss
                      model_fetches["critic_loss"] = self.critic_loss
  
                  results = sess.run(model_fetches, feed_dict=train_feed_dict)
  
                  # Update the summaries for Tensorboard
                  if should(summary_freq):
                      print("Recording summary ...")
                      train_supervisor.summary_writer.add_summary(
                          results["summary"], step)
  
                  # Launch the validation steps
                  if should(validation_freq):
                      self.validation(sess, valid_cells_no,
                                      exp_folder, step, clusters_ratios)
  
                  # Print out the progresses
                  if should(progress_freq):
                      rate = (step + 1) / (time.time() - start)
                      remaining = (self.max_steps - (step + 1)) / rate
                      print("[ %s ] Step number %d ." % (exp_name, step))
                      print("[ %s ] Running rate  %0.3f steps/sec."
                            % (exp_name, rate))
                      print("[ %s ] Estimated remaining time  %d m"
                            % (exp_name, remaining // 60))
                      print("[ %s ] Critic batch loss %0.3f"
                            % (exp_name, results["critic_loss"]))
                      print("[ %s ] Generator batch loss %0.f"
                            % (exp_name, results["gen_loss"]))
  
                  # Save the model
                  if should(save_freq):
                      saver.save(sess, os.path.join(exp_folder, "model"),
                                 global_step=step)
  
      def read_valid_cells(self, sess, cells_no):
          """
          Method that reads a given number of cells from the validation set.
  
          Parameters
          ----------
          sess : Session
              The TF Session in use.
          cells_no : int
              Number of validation cells to read.
  
          Returns
          -------
          real_cells : numpy array
              Matrix with the required amount of validation cells.
          real_clusters : list
              List containing the corresponding cluster indexes.
          """
  
          batches_no = int(np.ceil(cells_no // self.batch_size))
  
          real_cells = []
          real_clusters = []
          for i_batch in range(batches_no):
              test_inputs, test_clusters = sess.run(
                  [self.test_cells, self.test_cells_clusters])
              real_cells.append(test_inputs)
              real_clusters.append(test_clusters)
  
          real_cells = np.array(real_cells, dtype=np.float32)
          real_cells = real_cells.reshape((-1, self.test_cells.shape[1]))
  
          real_cells = rescale(real_cells, scaling=self.scaling,
                               scale_value=self.scale_value)
  
          return real_cells, real_clusters
  
      def generate_cells(self, cells_no, clusters_ratios=None,
                         sess=None, save_path=None, checkpoint=None):
          """
          Method that generate cells from the current model.
  
          Parameters
          ----------
          cells_no : int or list
              Numbers of cells per cluster to be generated.
              If the clusters_ratios are provided, should be an int (total number of cells).
              If cluster_ratios is None, should be a list of number of cells per cluster.
          clusters_ratios : numpy array
              List containing the different cluster ratios to use for
              the conditional generation.
              Default is None.
          sess : Session
              The TF Session in use.
              If None, a Session is created.
              Default is None.
          save_path : str
              Path in which to write the generated cells.
              If None, the cells are only returned and not written.
              Default is None.
          checkpoint : str /None
              Path to the checkpoint from which to load the model.
              If None, uses the current model loaded in the session.
              Default is None.
  
          Returns
          -------
          fake_cells : Numpy array
              2-D Array with the gene expression matrix of the generated cells.
          fake_labels : Numpy array
              Array containing the cluster index of the generated cells.
          """
  
          if sess is None:
              sess = tf.Session()
  
          if checkpoint is not None:
              saver = tf.train.Saver()
              saver.restore(sess, tf.train.latest_checkpoint(checkpoint))
  
          fake_cells = np.empty((0, self.genes_no), dtype=np.float32)
          fake_labels = np.empty([0, 1], dtype=np.int32)
  
          if clusters_ratios is None and len(cells_no) > 1:
  
              for cluster, cells_per_cluster in enumerate(cells_no):
                  if int(cells_per_cluster) == 0:
                      continue
                  clusters_ratios = np.zeros((1, self.clusters_no), dtype=np.float)
                  clusters_ratios[0][cluster] = 1
  
                  fc, fl = self.generate_cells(sess=sess, checkpoint=checkpoint,
                                               cells_no=int(cells_per_cluster),
                                               clusters_ratios=clusters_ratios)
  
                  fake_cells = np.append(fake_cells, fc, axis=0)
                  fake_labels = np.append(fake_labels, fl)
  
          else:
  
              batches_no = int(np.ceil(cells_no / self.batch_size))
  
              clusters_ratios_ph = self.generator.clusters_ratios
              fake_labels_tensor = self.generator.input_clusters
              is_training = self.generator.is_training
              fake_cells_tensor = self.generator.fake_outputs
  
              eval_feed_dict = {is_training: False,
                                clusters_ratios_ph: clusters_ratios}
  
              for i_batch in range(batches_no):
                  fc, fl = sess.run([fake_cells_tensor, fake_labels_tensor],
                                    feed_dict=eval_feed_dict)
                  fake_cells = np.append(fake_cells, fc, axis=0)
                  fake_labels = np.append(fake_labels, fl)
  
              fake_labels = fake_labels[0:cells_no]
              fake_cells = fake_cells[0:cells_no]
  
          rescale(fake_cells, scaling=self.scaling, scale_value=self.scale_value)
  
          if save_path is not None:
              save_generated_cells(fake_cells, save_path, fake_labels)
  
          return fake_cells, fake_labels
  
      def validation(self, sess, cells_no, exp_folder,
                     train_step, clusters_ratios):
          """
          Method that initiates some validation steps of the current model.
  
          Parameters
          ----------
          sess : Session
              The TF Session in use.
          cells_no : int
              Number of cells to use for the validation step.
          exp_folder : str
              Path to the job folder in which the outputs will be saved.
          train_step : int
              Index of the current training step.
          clusters_ratios : list
              List containing the different cluster ratios to use for the
              conditional generation.
  
          Returns
          -------
  
          """
  
          print("Find tSNE embedding for the generated and the validation cells")
          self.generate_tSNE_image(sess, cells_no, exp_folder,
                                   train_step, clusters_ratios)
  
      def generate_tSNE_image(self, sess, cells_no, exp_folder,
                              train_step, clusters_ratios):
          """
          Generates and saves a t-SNE plot with real and simulated cells
  
          Parameters
          ----------
          sess : Session
              The TF Session in use.
          cells_no : int
              Number of cells to use for the real and simulated cells (each)
              used for the plot.
          exp_folder : str
              Path to the job folder in which the outputs will be saved.
          train_step : int
              Index of the current training step.
          clusters_ratios : list
              List containing the different cluster ratios to use for the
              conditional generation.
  
          Returns
          -------
  
          """
  
          tnse_logdir = os.path.join(exp_folder, 'TSNE')
          if not os.path.isdir(tnse_logdir):
              os.makedirs(tnse_logdir)
  
          # generate fake cells
          fake_cells, fake_clusters = self.generate_cells(
              checkpoint=None,
              cells_no=cells_no,
              clusters_ratios=clusters_ratios,
              sess=sess)
  
          valid_cells, valid_clusters = self.read_valid_cells(sess, cells_no)
  
          real_cells_clusters = np.array(valid_clusters, dtype=np.float32).flatten()
          fake_cells_clusters = np.array(fake_clusters, dtype=np.float32).flatten()
  
          embedded_cells = tsne.fit_transform(
              np.concatenate((valid_cells, fake_cells), axis=0))
          embedded_cells_real = embedded_cells[0:real_cells_clusters.shape[0], :]
          embedded_cells_fake = embedded_cells[real_cells_clusters.shape[0]:, :]
  
          colormap = cm.nipy_spectral
          colors = [colormap(i) for i in np.linspace(0, 1, self.clusters_no)]
  
          plt.clf()
          plt.figure(figsize=(16, 12))
  
          for i in range(self.clusters_no):
              mask = real_cells_clusters[:] == i
  
              plt.scatter(embedded_cells_real[mask, 0],
                          embedded_cells_real[mask, 1],
                          c=colors[i], marker='*',
                          label='real_' + str(i))
  
          for i in range(self.clusters_no):
              mask = fake_cells_clusters[:] == i
              plt.scatter(embedded_cells_fake[mask, 0],
                          embedded_cells_fake[mask, 1],
                          c=colors[i], marker='o',
                          label='fake_' + str(i))
  
          plt.grid(True)
          plt.legend(loc='lower left',
                     numpoints=1, ncol=3,
                     fontsize=8, bbox_to_anchor=(0, 0))
          plt.savefig(tnse_logdir + '/step_' + str(train_step) + '.jpg')
          plt.close()
  