Module dimdrop.models

class Autoencoder:
    """
    A deep autoencoder model as baseline for other autoencoder based
    dimensionality reduction methods.

    The defaults are set to the parameters explained in a paper of
    Geoffrey Hinton.

    Parameters
    ----------
    in_dim : int
        The input dimension
    out_dim : int
        The output dimension
    layer_sizes : array of int, optional
        The sizes of the layers of the network, is mirrored over encoder and
        decoder parts, default `[2000, 1000, 500]`
    lr : float, optional
        The learning rate of the network, default `0.01`
    log : boolean, optional
        Whether log-transformation should be performed, default `False`
    scale : boolean, optional
        Whether scaling (making values within [0,1]) should be performed,
        default `True`
    batch_size : int, optional
        The batch size of the network, default `100`
    patience : int, optional
        The amount of epochs without improvement before the network stops
        training, default `3`
    epochs : int, optional
        The maximum amount of epochs, default `1000`
    regularizer : keras regularizer, optional
        A regularizer to use for the middle layer of the autoencoder.
        `None` or instance of `dimdrop.regularizers.KMeansRegularizer`,
        `dimdrop.regularizers.GMMRegularizer`,
        `dimdrop.regularizers.TSNERegularizer`.
    pretrain_method : string, optional
        The pretrain method to use. `None`, `'rbm'` or `'stacked'`
    decay : bool, optional
        Whether to decay the learning rate during training, default `True`.
    verbose : int, optional
        The verbosity of the network, default `0`

    Attributes
    ----------
    model :  keras model
        The autoencoder model
    encoder : keras model
        The encoder model
    layers : array of keras layers
        The layers of the network
    data_transform : Transform object
        The transformation to apply on the data before using it

    References
    ----------
    - G E Hinton and R R Salakhutdinov. Reducing the dimensionality of data
      with neural networks. *Science*, 313(5786):504–507, July 2006.
    """

    def __init__(
        self,
        in_dim,
        out_dim,
        layer_sizes=[2000, 1000, 500],
        lr=0.01,
        scale=True,
        log=False,
        batch_size=100,
        patience=3,
        epochs=1000,
        regularizer=None,
        pretrain_method='rbm',
        decay=True,
        verbose=0
    ):
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.layer_sizes = layer_sizes
        self.lr = lr
        self.data_transform = Transform(scale, log)
        self.batch_size = batch_size
        self.patience = patience
        self.epochs = epochs
        self.regularizer = regularizer
        self.pretrain_method = pretrain_method
        self.decay = decay
        self.verbose = verbose
        self.__pretrainers = {'rbm': self.__pretrain_rbm,
                              'stacked': self.__pretrain_stacked}
        self.__init_network()

    def __init_network(self):
        activation = 'sigmoid' if self.pretrain_method == 'rbm' else 'relu'
        self.layers = [Dense(
            self.layer_sizes[0],
            activation=activation,
            input_shape=(self.in_dim,)
        )]
        self.layers += [Dense(size, activation=activation)
                        for size in self.layer_sizes[1:]]
        self.layers += [Dense(
            self.out_dim,
            activity_regularizer=self.regularizer
        )] if self.regularizer else [Dense(self.out_dim)]
        self.layers += [Dense(size, activation=activation)
                        for size in self.layer_sizes[::-1]]
        self.layers += [Dense(self.in_dim)]

        self.model = Sequential(
            self.layers
        )
        self.model.compile(
            loss='mse',
            optimizer=Adam(lr=self.lr, decay=self.lr /
                           self.epochs if self.decay else 0.0)
        )
        self.encoder = Sequential(
            self.layers[:len(self.layers) // 2]
        )

    def __pretrain_rbm(self, data):
        rbms = [BernoulliRBM(
            batch_size=self.batch_size,
            learning_rate=self.lr,
            n_components=num,
            n_iter=20,
            verbose=self.verbose
        ) for num in self.layer_sizes + [self.out_dim]]
        current = data
        for i, rbm in enumerate(rbms):
            if self.verbose:
                print('Training RBM {}/{}'.format(i + 1, len(rbms)))
            rbm.fit(current)
            current = rbm.transform(current)
            dec_layer = self.layers[len(self.layers) - 1 - i]
            enc_layer = self.layers[i]
            dec_layer.set_weights(
                [rbm.components_, dec_layer.get_weights()[1]])
            enc_layer.set_weights(
                [np.transpose(rbm.components_), enc_layer.get_weights()[1]])

    def __pretrain_stacked(self, data):
        num_layers = len(self.layers)
        cur_data = data
        early_stopping = EarlyStopping(monitor="loss", patience=self.patience)
        input_shapes = [(el,) for el in [self.in_dim] + self.layer_sizes]
        for i in range(num_layers // 2):
            if self.verbose:
                print('Training Stack {}/{}'.format(i+1, num_layers // 2))

            stack = Sequential([
                Dropout(0.2, input_shape=input_shapes[i]),
                self.layers[i],
                Dropout(0.2),
                self.layers[num_layers - i - 1]
            ])

            input_layer = Input(shape=input_shapes[i])
            encoder_layer = self.layers[i](input_layer)
            stack_encode = Model(input_layer, encoder_layer)

            stack.compile(
                loss='mse',
                optimizer=Adam(lr=self.lr, decay=self.lr /
                               self.epochs if self.decay else 0.0)
            )

            stack.fit(
                cur_data,
                cur_data,
                epochs=self.epochs,
                callbacks=[early_stopping],
                batch_size=self.batch_size,
                verbose=self.verbose
            )
            cur_data = stack_encode.predict(cur_data)

    def fit(self, data):
        """
        Fit the given data to the model.

        Parameters
        ----------
        data : array
            Array of training samples where each sample is of size `in_dim`
        """
        data = self.data_transform(data)

        if self.pretrain_method:
            self.__pretrainers[self.pretrain_method](data)

        early_stopping = EarlyStopping(monitor='loss', patience=self.patience)
        callbacks = [early_stopping]
        if self.regularizer:
            self.regularizer.init_fit(self.encoder, data)
            callbacks.append(self.regularizer)
        self.model.fit(data, data, epochs=self.epochs, callbacks=callbacks,
                       batch_size=self.batch_size, verbose=self.verbose)

    def transform(self, data):
        """
        Transform the given data

        Parameters
        ----------
        data : array
            Array of samples to be transformed, where each sample is of size
            `in_dim`

        Returns
        -------
        array
            Transformed samples, where each sample is of size `out_dim`
        """
        data = self.data_transform(data)
        return self.encoder.predict(data, verbose=self.verbose)

    def fit_transform(self, data):
        """
        Fit the given data to the model and return its transformation

        Parameters
        ----------
        data : array
            Array of training samples where each sample is of size `in_dim`

        Returns
        -------
        array
            Transformed samples, where each sample is of size `out_dim`
        """
        self.fit(data)
        return self.transform(data)

class DEC(Autoencoder):
    """
    Deep Embedded Clustering model

    References
    ----------
    * Junyuan Xie, Ross B. Girshick, and Ali Farhadi. Unsupervised deep
      embedding for clustering analysis. *CoRR*, abs/1511.06335, 2015.
    """

    def __init__(
            self,
            in_dim,
            out_dim,
            k,
            layer_sizes=[500, 500, 2000],
            epochs=1000,
            lr=0.1,
            scale=True,
            log=False,
            batch_size=256,
            patience=3,
            tol=0.01,
            decay=True,
            verbose=0):
        super().__init__(
            in_dim,
            out_dim,
            layer_sizes=layer_sizes,
            lr=lr,
            scale=scale,
            log=log,
            batch_size=batch_size,
            patience=patience,
            epochs=epochs,
            regularizer=None,
            pretrain_method='stacked',
            decay=decay,
            verbose=verbose
        )
        self.k = k
        self.tol = tol

    def fit(self, data):
        super().fit(data)
        data = self.data_transform(data)
        clustering_layer = ClusteringLayer(
            self.k, name='clustering')(self.encoder.output)

        if self.verbose:
            print('Initializing cluster centers')

        kmeans = KMeans(n_clusters=self.k, n_init=20)
        y_pred = kmeans.fit_predict(self.encoder.predict(data))
        self.clustering_model = Model(inputs=self.encoder.input,
                                      outputs=clustering_layer)

        self.clustering_model.get_layer(name='clustering').set_weights(
            [kmeans.cluster_centers_])

        self.clustering_model.compile(
            optimizer=Adam(self.lr, decay=self.lr / self.epochs),
            loss='kld'
        )
        sequence = DECSequence(data, self.clustering_model, self.batch_size)

        early_stopping = EarlyStopping(monitor='loss', patience=self.patience)
        if self.verbose:
            print('Clustering optimization')
        self.clustering_model.fit_generator(
            sequence,
            math.ceil(data.shape[0] / self.batch_size),
            epochs=self.epochs,
            callbacks=[early_stopping],
            verbose=self.verbose
        )

    def soft_cluster_assignments(self, data):
        """Get the soft cluster assignments of the clustering layer of the DEC
        network for the input data.

        Parameters
        ----------
        data : array
            The input data

        Returns
        -------
        array of cluster assignments
        """
        return np.apply_along_axis(
            np.argmax,
            1,
            self.clustering_model.predict(data)
        )

class DeepCluster(Autoencoder):
    """
    DeepCluster model

    References
    ----------
    * Kai Tian, Shuigeng Zhou, and Jihong Guan. Deepcluster: A general
      clustering framework based on deep learning. In Michelangelo Ceci,
      Jaakko Hollmén, Ljupčo Todorovski, Celine Vens, and Sašo Džeroski,
      editors, *Machine Learning and Knowledge Discovery in Databases*,
      pages 809–825, Cham, 2017. Springer International Publishing.
    """

    def __init__(
            self,
            in_dim,
            out_dim,
            k,
            layer_sizes=[2000, 1000, 500],
            lr=0.01,
            scale=True,
            log=False,
            batch_size=100,
            patience=3,
            epochs=1000,
            regularizer='kmeans',
            regularizer_weight=0.5,
            decay=True,
            verbose=0
    ):
        if regularizer == 'kmeans':
            regularizer_obj = KMeansRegularizer(
                np.random.rand(k, out_dim), weight=regularizer_weight)
        elif regularizer == 'gmm':
            # TODO add gmm cluster regularizer
            regularizer_obj = None
        else:
            raise AttributeError
        super().__init__(
            in_dim,
            out_dim,
            layer_sizes=layer_sizes,
            lr=lr,
            scale=scale,
            log=log,
            batch_size=batch_size,
            patience=patience,
            epochs=epochs,
            regularizer=regularizer_obj,
            pretrain_method=None,
            decay=decay,
            verbose=verbose
        )

    def get_cluster_centers(self):
        """Get the cluster centers generated by the last fit operation.

        Returns:
        --------
        Array of cluster centers
        """
        return self.regularizer.cluster_centers

    def get_cluster_assignments(self):
        """Get the cluster assignments for the data of the last fit operation

        Returns:
        --------
        Array of cluster assignments
        """
        return self.regularizer.cluster_assignments

class ParametricTSNE:
    """
    Implementation of the parametric variant of t-distributed neighborhood
    embedding.

    Parameters
    ----------
    in_dim : int
        The input dimension
    out_dim : int
        The output dimension
    layer_sizes : array, optional
        sizes of each layer in the neural network, default is the structure
        proposed in the original paper, namely: `[500, 2000, 2000]`
    lr : float, optional
        The learning rate of the network, default `0.01`
    log : boolean, optional
        Whether log-transformation should be performed, default `False`
    scale : boolean, optional
        Whether scaling (making values within [0,1]) should be performed,
        default `True`
    batch_size : int, optional
        The batch size of the network, default `100`
    pretrain : int, optional
        Whether to perform pretraining using Restricted Boltzmann Machines,
        default `False`
    perplexity : int, optional
        Perplexity parameter in the t-SNE formula, controls how many neighbors
        are considered in the local neighborhood, default `30`
    tol : float, optional
        Tolerance of the perplexity, default `1e-5`
    patience : int, optional
        The amount of epochs without improvement before fitting will stop
        early, default `3`
    epochs : int, optional
        Maximum amount of epochs, default `1000`
    decay : bool, optional
        Whether to decay the learning rate during training, default `True`.
    verbose : int, optional
        Controls the verbosity of the model, default `0`

    Attributes
    ----------
    model : keras Sequential model
        The neural network
    layers : keras layers
        The layers of the neural network
    data_transform : Transform object
        The transformation to apply on data before using it

    References
    ----------
    - Laurens van der Maaten. Learning a parametric embedding by preserving
      local structure. In David van Dyk and Max Welling, editors, *Proceedings
      of the Twelth International Conference on Artificial Intelligence and
      Statistics*, volume 5 of *Proceedings of Machine Learning Research*,
      pages 384–391, Hilton Clearwater Beach Resort, Clearwater Beach, Florida
      USA, 16–18 Apr 2009. PMLR.
    """

    def __init__(
            self,
            in_dim,
            out_dim,
            layer_sizes=[500, 500, 2000],
            lr=0.01,
            log=False,
            scale=True,
            batch_size=100,
            pretrain=False,
            perplexity=30,
            tol=1e-5,
            patience=3,
            epochs=1000,
            decay=True,
            verbose=0):
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.layer_sizes = layer_sizes
        self.lr = lr
        self.data_transform = Transform(scale, log)
        self.batch_size = batch_size
        self.pretrain = pretrain
        self.perplexity = perplexity
        self.tol = tol
        self.patience = patience
        self.epochs = epochs
        self.verbose = verbose
        self.decay = decay
        self.__init_network()

    def __init_network(self):
        if self.pretrain:
            self.rbms = [BernoulliRBM(
                batch_size=self.batch_size,
                learning_rate=self.lr,
                n_components=num,
                n_iter=20,
                verbose=self.verbose
            ) for num in self.layer_sizes + [self.out_dim]]
            activation = 'sigmoid'
        else:
            activation = 'relu'
        self.layers = []
        for i, num in enumerate(self.layer_sizes):
            if i == 0:
                self.layers.append(Dense(
                    num,
                    activation=activation,
                    input_shape=(self.in_dim,)
                ))
            else:
                self.layers.append(Dense(num, activation=activation))
        self.layers.append(Dense(self.out_dim))

        self.model = Sequential(self.layers)

        optimizer = SGD(lr=self.lr, decay=self.lr /
                        self.epochs if self.decay else 0.0)
        loss = TSNELoss(self.in_dim, self.batch_size)

        self.model.compile(
            optimizer=optimizer,
            loss=loss
        )

        if self.verbose:
            self.model.summary()

    def __pretrain(self, data):
        current = data
        for i, rbm in enumerate(self.rbms):
            if self.verbose:
                print('Training RBM {}/{}'.format(i + 1, len(self.rbms)))
            rbm.fit(current)
            current = rbm.transform(current)

            self.layers[i].set_weights(
                [np.transpose(rbm.components_), rbm.intercept_hidden_])

    def fit(self, data):
        """
        Fit the given data to the model.

        Parameters
        ----------
        data : array
            Array of training samples where each sample is of size `in_dim`
        """
        # make data length be a multiple of batch size
        data = data[:(data.shape[0] // self.batch_size) * self.batch_size]

        data = self.data_transform(data)

        if self.pretrain:
            if self.verbose:
                print('Pretraining network')
            self.__pretrain(data)

        early_stopping = EarlyStopping(monitor='loss', patience=self.patience)

        P = compute_joint_probabilities(
            data,
            batch_size=self.batch_size,
            d=self.out_dim,
            perplexity=self.perplexity,
            tol=self.tol,
            verbose=self.verbose
        )
        y_train = P.reshape(data.shape[0], -1)

        self.model.fit(
            data,
            y_train,
            epochs=self.epochs,
            callbacks=[early_stopping],
            batch_size=self.batch_size,
            shuffle=False,
            verbose=self.verbose
        )

    def transform(self, data):
        """
        Transform the given data

        Parameters
        ----------
        data : array
            Array of samples to be transformed, where each sample is of size
            `in_dim`

        Returns
        -------
        array
            Transformed samples, where each sample is of size `out_dim`
        """
        data = self.data_transform(data)
        return self.model.predict(data, verbose=self.verbose)

    def fit_transform(self, data):
        """
        Fit the given data to the model and return its transformation

        Parameters
        ----------
        data : array
            Array of training samples where each sample is of size `in_dim`

        Returns
        -------
        array
            Transformed samples, where each sample is of size `out_dim`
        """
        self.fit(data)
        return self.transform(data)

class VASC:
    """
    VASC: variational autoencoder for scRNA-seq datasets

    Parameters
    ----------
    in_dim : int
        The input dimension
    out_dim : int
        The output dimension
    layer_sizes : array
        sizes of each layer in the neural network, default is the structure
        proposed in the original paper, namely: `[512, 128, 32]`
    epochs: int, optional
        Maximum number of epochs, default `5000`
    patience: int, optional
        The amount of epochs without improvement before the network stops
        training, default `3`
    batch_size: int, optional
        The batch size for stochastic optimization, default `100`
    lr : float, optional
        The learning rate of the network, default `0.01`
    var: boolean, optional
        Whether to estimate the variance parameters, default `False`
    log: boolean, optional
        Whether log-transformation should be performed, default `False`
    scale: boolean, optional
        Whether scaling (making values within [0,1]) should be performed,
        default `True`
    decay : bool, optional
        Whether to decay the learning rate during training, default `True`.
    verbose : int, optional
        Controls the verbosity of the model, default `0`

    Attributes
    ----------
    ae : keras model
        Partial model used for generating the embbeddings
    vae : keras model
        The variational autoencoder

    References
    ----------
    * Dongfang Wang and Jin Gu. Vasc: dimension reduction and visualization of
      single cell rna sequencing data by deep variational autoencoder.
      *bioRxiv*, 2017.
    """

    def __init__(
            self,
            in_dim,
            out_dim,
            layer_sizes=[512, 128, 32],
            epochs=1000,
            patience=3,
            batch_size=100,
            lr=0.01,
            var=False,
            log=False,
            scale=True,
            decay=True,
            verbose=0):
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.layer_sizes = layer_sizes
        self.epochs = epochs
        self.patience = patience
        self.batch_size = batch_size
        self.lr = lr
        self.var = var
        self.data_transform = Transform(scale, log)
        self.verbose = verbose
        self.decay = decay
        self.__init_network()

    def __init_network(self):
        in_dim = self.in_dim
        expr_in = Input(shape=(self.in_dim,))

        # The first part of model to recover the expr.
        h0 = Dropout(0.5)(expr_in)
        # Encoder layers
        last_layer = h0
        for i, layer_size in enumerate(self.layer_sizes):
            if i == 0:
                last_layer = Dense(
                    units=layer_size,
                    name='encoder_{}'.format(i + 1),
                    kernel_regularizer=regularizers.l1(0.01)
                )(last_layer)
            else:
                last_layer = Dense(
                    units=layer_size,
                    name='encoder_{}'.format(i + 1)
                )(last_layer)
            last_layer = Activation('relu')(last_layer)

        z_mean = Dense(units=self.out_dim, name='z_mean')(last_layer)
        z_log_var = None
        if self.var:
            z_log_var = Dense(units=2, name='z_log_var')(last_layer)
            z_log_var = Activation('softplus')(z_log_var)

        # sampling new samples
            z = Lambda(sampling, output_shape=(
                self.out_dim,))([z_mean, z_log_var])
        else:
            z = Lambda(sampling, output_shape=(self.out_dim,))([z_mean])

        # Decoder layers
        last_layer = z
        for i, layer_size in enumerate(self.layer_sizes[::-1]):
            last_layer = Dense(
                units=layer_size, name='decoder_{}'.format(i + 1))(last_layer)
            last_layer = Activation('relu')(last_layer)

        expr_x = Dense(units=self.in_dim, activation='sigmoid')(
            last_layer)

        expr_x_drop = Lambda(lambda x: -x ** 2)(expr_x)
        expr_x_drop_p = Lambda(lambda x: K.exp(x))(expr_x_drop)
        expr_x_nondrop_p = Lambda(lambda x: 1-x)(expr_x_drop_p)
        expr_x_nondrop_log = Lambda(lambda x: K.log(x+1e-20))(expr_x_nondrop_p)
        expr_x_drop_log = Lambda(lambda x: K.log(x+1e-20))(expr_x_drop_p)
        expr_x_drop_log = Reshape(target_shape=(
            self.in_dim, 1))(expr_x_drop_log)
        expr_x_nondrop_log = Reshape(target_shape=(
            self.in_dim, 1))(expr_x_nondrop_log)
        logits = concatenate(
            [expr_x_drop_log, expr_x_nondrop_log], axis=-1)

        temp_in = Input(shape=(self.in_dim,))
        temp_ = RepeatVector(2)(temp_in)

        temp_ = Permute((2, 1))(temp_)
        samples = Lambda(gumbel_softmax, output_shape=(
            self.in_dim, 2,))([logits, temp_])
        samples = Lambda(lambda x: x[:, :, 1])(samples)
        samples = Reshape(target_shape=(self.in_dim,))(samples)

        out = multiply([expr_x, samples])

        vae = Model(inputs=[expr_in, temp_in], outputs=out)

        loss_func = VAELoss(in_dim, z_log_var, z_mean)

        opt = RMSprop(lr=self.lr, decay=self.lr /
                      self.epochs if self.decay else 0.0)
        vae.compile(optimizer=opt, loss=loss_func)

        ae = Model(inputs=[expr_in, temp_in], outputs=[
                   z_mean, z, samples, out])

        self.vae = vae
        self.ae = ae

        if self.verbose:
            self.vae.summary()

    def fit(self, data):
        """
        Fit the given data to the model.

        Parameters
        ----------
        data : array
            Array of training samples where each sample is of size `in_dim`
        """
        data = self.data_transform(data)

        tau_in = np.ones(data.shape, dtype='float32')

        early_stopping = EarlyStopping(monitor='loss', patience=self.patience)

        self.vae.fit(
            [data, tau_in],
            data,
            epochs=self.epochs,
            batch_size=self.batch_size,
            shuffle=True,
            verbose=self.verbose,
            callbacks=[early_stopping]
        )

    def fit_transform(self, data):
        """
        Fit the given data to the model and return its transformation

        Parameters
        ----------
        data : array
            Array of training samples where each sample is of size `in_dim`

        Returns
        -------
        array
            Transformed samples, where each sample is of size `out_dim`
        """
        self.fit(data)
        return self.transform(data)

    def transform(self, data):
        """
        Transform the given data

        Parameters
        ----------
        data : array
            Array of samples to be transformed, where each sample is of size
            `in_dim`

        Returns
        -------
        array
            Transformed samples, where each sample is of size `out_dim`
        """
        data = self.data_transform(data)
        return self.ae.predict([data, np.ones(data.shape, dtype='float32')])[0]