ga.py

'''
Author: LI Min
'''

import numpy as np
import random
from xor_solver import XORSolver
from functools import reduce
import matplotlib.pyplot as plt
plt.switch_backend('agg')



class GAforXOR(object):
    '''
    Fitnesss: Encoding Scucess Rate
    Secondary Metric: Activated Scan Chain Rate.

    Every instances: matrix A, and matrix P;
    Matrix A refers to XORNet, matrix P refers to ANDNet;
    So far, we do the GA for matrix A, and keep matrix P as the same as EDT.
    Args: 
        num_sc: The number of scan chain to be encoded;
        num_ctrl: The number of encoding bits; This parameter should be set close to (num_sc * specified percentage).
        num_generation: The number of generation to be evolved.
        num_pop: The number of instances, in other words, how many A's and P's.
        num_parent: The number of parents selected to mate, in other words, select top-N instances in the population;
        num_crossover: The number of children, generated by crossovering from parents;
        num_mutation: The number of children to do mutation;
        mutation_rate: The rate of mutation;
        connection_rate: The initial connection configuration of matrix A. E.g., 0.2, then 20% of A being 1's, other elements being 0's;
        power_limit: The initial configuration of matrix P. E.g., 2, then each row of matrix has 2 1's, meaning that the average activating rate of scan chain is 25%.
        freq_sc_file: The frequency statistic of scan chain. We use it to generate test data;
        specified_percentage: The specified scan chain rate when generating test data;
        num_test: The number of test data to be generated. Out of concern of cost, we use a smaller number of test data.

    Another idea:
        X in b[10000, 415]    B in b[10000, 40]   A in b[40 * 415].
        minimize |X - BA|. Just like matrix factorization.
    Another idea:
        Graph based approcah. Analyize the joint distribution of scan chains... 

    '''

    def __init__(self, num_sc=415, num_ctrl=25, num_generation=10, num_pop=20, num_parent=5, num_crossover=15, num_mutation=5, mutation_rate=0.05, connection_percentage=0.2, power_lit=1, freq_sc_file='data/freq_sc.npy', specified_percentage=0.1, num_test=100, beta=0.0, power_uppper=0.5):
        self.num_sc = num_sc
        self.num_ctrl = num_ctrl
        self.num_generation = num_generation
        self.num_pop = num_pop
        ### To Do: change the number into percentage
        self.num_parent = num_parent
        self.num_crossover = num_crossover
        self.num_mutation = num_mutation
        self.mutation_rate = mutation_rate
        self.connection_percentage = connection_percentage
        self.power_limit = power_lit   # The power upper limit; 1-50%, 2-25%, 3-33%...
        self.freq_sc = np.load(freq_sc_file)
        self.test_data = self.initialize_testdata(specified_percentage, num_test)
        self.pop = self.initialize_pop()   # contains (pop_A, pop_P)
        # save a random xor matrix
        self.save_xor_original()
        self.generation_idx = 0
        self.fitness_history = {}
        self.encode_history = {}
        self.act_history = {}
        self.uns_history = {}
        self.beta = beta # The coefficient for fitness function
        self.power_upper = power_uppper


        

    def initialize_testdata(self, specified_percentage, num_test):
        # test_data = np.zeros((num_test, self.num_sc))
        # for (i, row) in enumerate(test_data):
        #     generate_row = np.random.choice(self.num_sc, size=int(self.num_sc *  specified_percentage), replace=False, p=self.freq_sc)
        #     test_data[i][generate_row] = 1
        test_data = np.load('data/mlb_sc.npy')
        test_data = test_data.astype(dtype=bool)
        return test_data

    def initialize_pop(self):
        pop_A = np.random.choice(2, size=(self.num_pop, self.num_sc * self.power_limit, self.num_ctrl), p=[1-self.connection_percentage, self.connection_percentage]).astype(dtype=bool)
        # Another way to generate population
        pop_A = np.zeros((self.num_pop, self.num_sc * self.power_limit, self.num_ctrl))
        for i_p in range(np.shape(pop_A)[0]):
            for i_r in range(np.shape(pop_A)[1]):
                k_base = 2 if random.random() < 0.7 else 3
                k = k_base + np.random.geometric(0.4)
                idx = np.random.choice(self.num_ctrl, size=k, replace=False)
                np.put(pop_A[i_p][i_r], idx, 1)
                
        pop_P = np.zeros((self.num_sc, self.num_sc * self.power_limit))
        for (i, row) in enumerate(pop_P):
            row[(i*self.power_limit):(i*self.power_limit+self.power_limit)] = 1
        pop_P = np.repeat(np.expand_dims(pop_P, axis=0), repeats=self.num_pop, axis=0).astype(dtype=bool)

        return pop_A, pop_P

    def cal_pop_fitness(self):
        fitness_pop = []
        encode_pop = []
        act_pop = []
        uns_pop = []
        for i in range(self.num_pop):
            encode_success_i, act_i, violate_power = self.xor_solving(i)
            uns_i = np.shape(self.test_data)[0] - int(encode_success_i * (np.shape(self.test_data)[0] + 1)) + violate_power
            # fitness_i = encode_success_i - self.beta * act_i  # maximize encoding success rate and minimize the activated percentage.
            fitness_i = - int(uns_i)
            fitness_pop.append(fitness_i)
            encode_pop.append(encode_success_i)
            act_pop.append(act_i)
            uns_pop.append(uns_i)
            # print('A index:', i)
        self.fitness_history[self.generation_idx] = fitness_pop
        self.encode_history[self.generation_idx] = encode_pop
        self.act_history[self.generation_idx] = act_pop
        self.uns_history[self.generation_idx] = uns_pop
        # self.generation_idx += 1

    def xor_solving(self, i):
        A = self.pop[0][i]
        P = self.pop[1][i]
        # A being boolean matrix
        encoded_count = 0.0
        total = 0.0
        activated_rate_acculmulate = 0.0
        violate_power = 0
        for (i, cube) in enumerate(self.test_data):
            total += 1
            P_hat = P[cube.astype(dtype=bool)]
            idx = np.sum(P_hat, axis=0).astype(dtype=bool)
            A_hat = A[idx]
            b_hat = np.ones(np.sum(idx)).astype(dtype=bool)
            equation = XORSolver(A_hat, b_hat)
            equation.gaussian_elimination()
            if equation.status:
                encoded_count += 1
                activated_rate = self.calculate_activated_rate(A, equation.x)
                activated_rate_acculmulate += activated_rate
                if activated_rate >= self.power_limit:
                    violate_power += 1
                
        return encoded_count/(total + 1.0), activated_rate_acculmulate/(encoded_count + 1.0), violate_power

    def calculate_activated_rate(self, A, x):
        b = np.zeros(np.shape(A)[0]).astype(dtype=bool)
        for (i, A_i) in enumerate(A):
            if np.sum(A_i) == 0:
                continue
            x_valid = x[A_i]
            b[i] = reduce(np.logical_xor, x_valid)
        return np.sum(b) / len(b)
    

    def select_mating_pool(self):
        parents = np.empty((self.num_parent, self.num_sc, self.num_ctrl))
        ind = np.argpartition(self.fitness_history[self.generation_idx], -self.num_parent)[-self.num_parent:]
        parents = self.pop[0][ind]
        return parents

    def crossover(self, parents):
        # There are different types of crossover. For XORNet, more flexible solution is to crossover under the granularity of each row of A.
        # kind of like Uniform Crossover
        offspring = np.empty((self.num_crossover, self.num_sc * self.power_limit, self.num_ctrl))
        # The source of each row of offspring coming from
        source_crossover = np.random.choice(2, size=(self.num_crossover, self.num_sc * self.power_limit)).astype(dtype=bool)
        for k in range(self.num_crossover):
            parents_idx = np.random.choice(self.num_parent, size=2, replace=False)
            offspring[k][source_crossover[k]] = parents[parents_idx[0]][source_crossover[k]]
            offspring[k][np.invert(source_crossover[k])] = parents[parents_idx[1]][np.invert(source_crossover[k])]
        return offspring.astype(dtype=bool)

    def mutation(self, offspring):
        # Mutation. For XORNet, mutation happends elemently in A.
        selected_idx = np.random.choice(self.num_crossover, size=self.num_mutation, replace=False)
        # Mutation points
        mu_idx = np.random.choice(2, size=(self.num_mutation, self.num_sc * self.power_limit, self.num_ctrl), p=[1-self.mutation_rate, self.mutation_rate]).astype(dtype=bool)
        offspring[selected_idx][mu_idx] = np.invert(offspring[selected_idx][mu_idx])

        return offspring.astype(dtype=bool)


    
    def GALoop(self):
        # new_pop = np.empty((self.num_pop, self.num_sc, self.num_ctrl))
        self.cal_pop_fitness()
        parents = self.select_mating_pool()
        offspring = self.crossover(parents)
        offspring = self.mutation(offspring)
        self.pop[0][:self.num_parent] = parents
        self.pop[0][-self.num_crossover:] = offspring
        self.generation_idx += 1
    
    def GA(self):
        for i in range(self.num_generation):
            print('###### No. {} generation ####'.format(i))
            self.GALoop()
            print('max ', i, ' :',  np.max(self.fitness_history[i]))
            print('average: ', i, ':', np.average(self.fitness_history[i]))
    
    def visulization(self):
        # fig, axs = plt.subplots(3)
        # for (key, values) in self.fitness_history.items():
        #     axs[0,].plot([key] * len(values), values, '.', color='k')
        #     axs[0].plot(key, np.max(values), '*', color='r')
        #     axs[0].plot(key, np.average(values), 'o', color='b')
        # axs[0].set(ylabel='Fitness(%.1f)' % self.beta)
        # for (key, values) in self.encode_history.items():
        #     axs[1].plot([key] * len(values), values, '.', color='k')
        #     axs[1].plot(key, np.max(values), '*', color='r')
        #     axs[1].plot(key, np.average(values), 'o', color='b')
        # axs[1].set(ylabel='ESR')
        # for (key, values) in self.act_history.items():
        #     axs[2].plot([key] * len(values), values, '.', color='k')
        #     axs[2].plot(key, np.min(values), '*', color='r')
        #     axs[2].plot(key, np.average(values), 'o', color='b')
        # axs[2].set(xlabel='# Generation', ylabel='AP')
        fig, ax = plt.subplots()
        for (key, values) in self.uns_history.items():
            ax.plot([key] * len(values), values, '.', color='k')
            ax.plot(key, np.max(values), '*', color='r')
            ax.plot(key, np.average(values), 'o', color='b')
        ax.set(ylabel='UNS')

        # plt.title('GA for Testing')
        plt.savefig('figs/GA_uns.pdf')

    
    def save_xor(self):
        # save the xor network and and net
        best_idx = np.argmax(self.fitness_history[self.num_generation-1])
        np.save('checkpoint/GA_XOR_best.npy', self.pop[0][best_idx])
        np.save('checkpoint/GA_AND_best.npy', self.pop[1][best_idx])
    
    def save_xor_original(self):
        # randomn save the original xor network from 1st generation
        select_idx = random.randint(0, self.num_pop-1)
        np.save('checkpoint/GA_XOR_orig.npy', self.pop[0][select_idx])
        np.save('checkpoint/GA_XOR_orig.npy', self.pop[1][select_idx])

        


        
        





def main():
    ga = GAforXOR()
    ga.GA()
    ga.visulization()
    ga.save_xor()

if __name__ == '__main__':
    main()