ConvGeNCode/library/generators/convGAN.py

import numpy as np
import matplotlib.pyplot as plt

from library.interfaces import GanBaseClass
from library.dataset import DataSet

from keras.layers import Dense, Input, Multiply, Flatten, Conv1D, Reshape
from keras.models import Model
from keras import backend as K
from tqdm import tqdm

import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.layers import Lambda

from library.NNSearch import NNSearch

import warnings
warnings.filterwarnings("ignore")



def repeat(x, times):
    return [x for _i in range(times)]

def create01Labels(totalSize, sizeFirstHalf):
    labels = repeat(np.array([1,0]), sizeFirstHalf)
    labels.extend(repeat(np.array([0,1]), totalSize - sizeFirstHalf))
    return np.array(labels)

class ConvGAN(GanBaseClass):
    """
    This is a toy example of a GAN.
    It repeats the first point of the training-data-set.
    """
    def __init__(self, n_feat, neb=5, gen=None, neb_epochs=10, withMajorhoodNbSearch=False, debug=False):
        self.isTrained = False
        self.n_feat = n_feat
        self.neb = neb
        self.nebInitial = neb
        self.genInitial = gen
        self.gen = gen if gen is not None else self.neb
        self.neb_epochs = 10
        self.loss_history = None
        self.debug = debug
        self.minSetSize = 0
        self.conv_sample_generator = None
        self.maj_min_discriminator = None
        self.withMajorhoodNbSearch = withMajorhoodNbSearch
        self.cg = None
        self.canPredict = True

        if self.neb is not None and self.gen is not None and self.neb > self.gen:
            raise ValueError(f"Expected neb <= gen but got neb={neb} and gen={gen}.")

    def reset(self, dataSet):
        """
        Resets the trained GAN to an random state.
        """
        self.isTrained = False

        if dataSet is not None:
            nMinoryPoints = dataSet.data1.shape[0]
            if self.nebInitial is None:
                self.neb = nMinoryPoints
            else:
                self.neb = min(self.nebInitial, nMinoryPoints)
        else:
            self.neb = self.nebInitial

        self.gen = self.genInitial if self.genInitial is not None else self.neb

        ## instanciate generator network and visualize architecture
        self.conv_sample_generator = self._conv_sample_gen()

        ## instanciate discriminator network and visualize architecture
        self.maj_min_discriminator = self._maj_min_disc()

        ## instanciate network and visualize architecture
        self.cg = self._convGAN(self.conv_sample_generator, self.maj_min_discriminator)

        if self.debug:
            print(f"neb={self.neb}, gen={self.gen}")

            print(self.conv_sample_generator.summary())
            print('\n')
            
            print(self.maj_min_discriminator.summary())
            print('\n')

            print(self.cg.summary())
            print('\n')

    def train(self, dataSet, discTrainCount=5):
        """
        Trains the GAN.

        It stores the data points in the training data set and mark as trained.

        *dataSet* is a instance of /library.dataset.DataSet/. It contains the training dataset.
        We are only interested in the first *maxListSize* points in class 1.
        """
        if dataSet.data1.shape[0] <= 0:
            raise AttributeError("Train: Expected data class 1 to contain at least one point.")

        # Store size of minority class. This is needed during point generation.
        self.minSetSize = dataSet.data1.shape[0]

        # Precalculate neighborhoods
        self.nmbMin = NNSearch(self.neb).fit(haystack=dataSet.data1)
        if self.withMajorhoodNbSearch:
            self.nmbMaj = NNSearch(self.neb).fit(haystack=dataSet.data0, needles=dataSet.data1)
        else:
            self.nmbMaj = None

        # Do the training.
        self._rough_learning(dataSet.data1, dataSet.data0, discTrainCount)
        
        # Neighborhood in majority class is no longer needed. So save memory.
        self.nmbMaj = None
        self.isTrained = True

    def generateDataPoint(self):
        """
        Returns one synthetic data point by repeating the stored list.
        """
        return (self.generateData(1))[0]


    def generateData(self, numOfSamples=1):
        """
        Generates a list of synthetic data-points.

        *numOfSamples* is a integer > 0. It gives the number of new generated samples.
        """
        if not self.isTrained:
            raise ValueError("Try to generate data with untrained Re.")

        ## roughly claculate the upper bound of the synthetic samples to be generated from each neighbourhood
        synth_num = (numOfSamples // self.minSetSize) + 1

        ## generate synth_num synthetic samples from each minority neighbourhood
        synth_set=[]
        for i in range(self.minSetSize):
            synth_set.extend(self._generate_data_for_min_point(i, synth_num))

        ## extract the exact number of synthetic samples needed to exactly balance the two classes
        synth_set = np.array(synth_set[:numOfSamples]) 

        return synth_set

    def predictReal(self, data):
        prediction = self.maj_min_discriminator.predict(data)
        return np.array([x[0] for x in prediction])

    # ###############################################################
    # Hidden internal functions
    # ###############################################################

    # Creating the GAN
    def _conv_sample_gen(self):
        """
        the generator network to generate synthetic samples from the convex space
        of arbitrary minority neighbourhoods
        """

        ## takes minority batch as input
        min_neb_batch = Input(shape=(self.n_feat,))

        ## reshaping the 2D tensor to 3D for using 1-D convolution,
        ## otherwise 1-D convolution won't work.
        x = tf.reshape(min_neb_batch, (1, self.neb, self.n_feat), name=None)
        ## using 1-D convolution, feature dimension remains the same
        x = Conv1D(self.n_feat, 3, activation='relu')(x)
        ## flatten after convolution
        x = Flatten()(x)
        ## add dense layer to transform the vector to a convenient dimension
        x = Dense(self.neb * self.gen, activation='relu')(x)

        ## again, witching to 2-D tensor once we have the convenient shape
        x = Reshape((self.neb, self.gen))(x)
        ## row wise sum
        s = K.sum(x, axis=1)
        ## adding a small constant to always ensure the row sums are non zero.
        ## if this is not done then during initialization the sum can be zero.
        s_non_zero = Lambda(lambda x: x + .000001)(s)
        ## reprocals of the approximated row sum
        sinv = tf.math.reciprocal(s_non_zero)
        ## At this step we ensure that row sum is 1 for every row in x.
        ## That means, each row is set of convex co-efficient
        x = Multiply()([sinv, x])
        ## Now we transpose the matrix. So each column is now a set of convex coefficients
        aff=tf.transpose(x[0])
        ## We now do matrix multiplication of the affine combinations with the original
        ## minority batch taken as input. This generates a convex transformation
        ## of the input minority batch
        synth=tf.matmul(aff, min_neb_batch)
        ## finally we compile the generator with an arbitrary minortiy neighbourhood batch
        ## as input and a covex space transformation of the same number of samples as output
        model = Model(inputs=min_neb_batch, outputs=synth)
        opt = Adam(learning_rate=0.001)
        model.compile(loss='mean_squared_logarithmic_error', optimizer=opt)
        return model

    def _maj_min_disc(self):
        """
        the discriminator is trained intwo phase:
        first phase:  while training GAN the discriminator learns to differentiate synthetic
                      minority samples generated from convex minority data space against
                      the borderline majority samples
        second phase: after the GAN generator learns to create synthetic samples,
                      it can be used to generate synthetic samples to balance the dataset
                      and then rettrain the discriminator with the balanced dataset
        """

        ## takes as input synthetic sample generated as input stacked upon a batch of
        ## borderline majority samples
        samples = Input(shape=(self.n_feat,))
        
        ## passed through two dense layers
        y = Dense(250, activation='relu')(samples)
        y = Dense(125, activation='relu')(y)
        y = Dense(75, activation='relu')(y)
        
        ## two output nodes. outputs have to be one-hot coded (see labels variable before)
        output = Dense(2, activation='sigmoid')(y)
        
        ## compile model
        model = Model(inputs=samples, outputs=output)
        opt = Adam(learning_rate=0.0001)
        model.compile(loss='binary_crossentropy', optimizer=opt)
        return model

    def _convGAN(self, generator, discriminator):
        """
        for joining the generator and the discriminator
        conv_coeff_generator-> generator network instance
        maj_min_discriminator -> discriminator network instance
        """
        ## by default the discriminator trainability is switched off.
        ## Thus training the GAN means training the generator network as per previously
        ## trained discriminator network.
        discriminator.trainable = False

        ## input receives a neighbourhood minority batch
        ## and a proximal majority batch concatenated
        batch_data = Input(shape=(self.n_feat,))
        
        ##- print(f"GAN: 0..{self.neb}/{self.gen}..")

        ## extract minority batch
        min_batch = Lambda(lambda x: x[:self.neb])(batch_data)
        
        ## extract majority batch
        maj_batch = Lambda(lambda x: x[self.gen:])(batch_data)
        
        ## pass minority batch into generator to obtain convex space transformation
        ## (synthetic samples) of the minority neighbourhood input batch
        conv_samples = generator(min_batch)
        
        ## concatenate the synthetic samples with the majority samples
        new_samples = tf.concat([conv_samples, maj_batch],axis=0)
        ##- new_samples = tf.concat([conv_samples, conv_samples, conv_samples, conv_samples],axis=0)
        
        ## pass the concatenated vector into the discriminator to know its decisions
        output = discriminator(new_samples)
        ##- output = Lambda(lambda x: x[:2 * self.gen])(output)
        
        ## note that, the discriminator will not be traied but will make decisions based
        ## on its previous training while using this function
        model = Model(inputs=batch_data, outputs=output)
        opt = Adam(learning_rate=0.0001)
        model.compile(loss='mse', optimizer=opt)
        return model

    # Create synthetic points
    def _generate_data_for_min_point(self, index, synth_num):
        """
        generate synth_num synthetic points for a particular minoity sample
        synth_num -> required number of data points that can be generated from a neighbourhood
        data_min -> minority class data
        neb -> oversampling neighbourhood
        index -> index of the minority sample in a training data whose neighbourhood we want to obtain
        """

        runs = int(synth_num / self.neb) + 1
        synth_set = []
        for _run in range(runs):
            batch = self.nmbMin.getNbhPointsOfItem(index)
            synth_batch = self.conv_sample_generator.predict(batch, batch_size=self.neb)
            synth_set.extend(synth_batch)

        return synth_set[:synth_num]



    # Training
    def _rough_learning(self, data_min, data_maj, discTrainCount):
        generator = self.conv_sample_generator
        discriminator = self.maj_min_discriminator
        GAN = self.cg
        loss_history = [] ## this is for stroring the loss for every run
        step = 0
        minSetSize = len(data_min)

        labels = tf.convert_to_tensor(create01Labels(2 * self.gen, self.gen))
        nLabels = 2 * self.gen

        for neb_epoch_count in range(self.neb_epochs):
            if discTrainCount > 0:
                for n in range(discTrainCount):
                    for min_idx in range(minSetSize):
                        ## generate minority neighbourhood batch for every minority class sampls by index
                        min_batch_indices = self.nmbMin.neighbourhoodOfItem(min_idx)
                        min_batch = self.nmbMin.getPointsFromIndices(min_batch_indices)
                        ## generate random proximal majority batch
                        maj_batch = self._BMB(data_maj, min_batch_indices)

                        ## generate synthetic samples from convex space
                        ## of minority neighbourhood batch using generator
                        conv_samples = generator.predict(min_batch, batch_size=self.neb)
                        ## concatenate them with the majority batch
                        concat_sample = tf.concat([conv_samples, maj_batch], axis=0)

                        ## switch on discriminator training
                        discriminator.trainable = True
                        ## train the discriminator with the concatenated samples and the one-hot encoded labels
                        discriminator.fit(x=concat_sample, y=labels, verbose=0, batch_size=nLabels)
                        ## switch off the discriminator training again
                        discriminator.trainable = False

            for min_idx in range(minSetSize):
                ## generate minority neighbourhood batch for every minority class sampls by index
                min_batch_indices = self.nmbMin.neighbourhoodOfItem(min_idx)
                min_batch = self.nmbMin.getPointsFromIndices(min_batch_indices)
                ## generate random proximal majority batch
                maj_batch = self._BMB(data_maj, min_batch_indices)

                ## generate synthetic samples from convex space
                ## of minority neighbourhood batch using generator
                conv_samples = generator.predict(min_batch, batch_size=self.neb)
                ## concatenate them with the majority batch
                concat_sample = tf.concat([conv_samples, maj_batch], axis=0)

                ## switch on discriminator training
                discriminator.trainable = True
                ## train the discriminator with the concatenated samples and the one-hot encoded labels
                discriminator.fit(x=concat_sample, y=labels, verbose=0, batch_size=nLabels)
                ## switch off the discriminator training again
                discriminator.trainable = False

                ## use the GAN to make the generator learn on the decisions
                ## made by the previous discriminator training
                ##- print(f"concat sample shape: {concat_sample.shape}/{labels.shape}")
                gan_loss_history = GAN.fit(concat_sample, y=labels, verbose=0, batch_size=nLabels)

                ## store the loss for the step
                loss_history.append(gan_loss_history.history['loss'])

                step += 1
                if self.debug and (step % 10 == 0):
                    print(f"{step} neighbourhood batches trained; running neighbourhood epoch {neb_epoch_count}")

            if self.debug:
                print(f"Neighbourhood epoch {neb_epoch_count + 1} complete")

        if self.debug:
            run_range = range(1, len(loss_history) + 1)
            plt.rcParams["figure.figsize"] = (16,10)
            plt.xticks(fontsize=20)
            plt.yticks(fontsize=20)
            plt.xlabel('runs', fontsize=25)
            plt.ylabel('loss', fontsize=25)
            plt.title('Rough learning loss for discriminator', fontsize=25)
            plt.plot(run_range, loss_history)
            plt.show()

        self.conv_sample_generator = generator
        self.maj_min_discriminator = discriminator
        self.cg = GAN
        self.loss_history = loss_history



    ## convGAN
    def _BMB(self, data_maj, min_idxs):

        ## Generate a borderline majority batch
        ## data_maj -> majority class data
        ## min_idxs -> indices of points in minority class
        ## gen -> convex combinations generated from each neighbourhood

        if self.nmbMaj is not None:
            return self.nmbMaj.neighbourhoodOfItemList(min_idxs, maxCount=self.gen)
        else:
            return tf.convert_to_tensor(
                data_maj[np.random.randint(len(data_maj), size=self.gen)]
                )