text:
from core.individual import Individual
from .types import DefaultConfig
import secrets
import os
import math
import sys
import random
import gc
import pickle
import subprocess
from anytree import search, RenderTree
import copy
import json
import numpy as np
from utils.operators import IfThenElse, Assignment, Mul, Sub, Sum, Termination
from generators.tree import generate_operator_node, generate_random_expression, generate_individual_from_seed
from utils.fitness import tcp_variant_fitness_write_switch
from multiprocessing.pool import ThreadPool as Pool
from anytree import PreOrderIter
import time
from dotenv import load_dotenv
from prompt_toolkit import print_formatted_text, HTML
import psutil

PROCWIFITCP = "wifi-tcp"

load_dotenv()

BASE_NS3_PATH = os.getenv('BASE_NS3_PATH')
CPUS = os.getenv('CPUS')

pool_size = int(CPUS)

class Incubator:

    def __init__(
                    self,
                    config=DefaultConfig,
                    variables=None,
                    pop_size=2,
                    generations=20,
                    const_termination_range=[-20, 20],
                    fitness=None,
                    generator=generate_random_expression,
                    max_depth=10,
                    max_width=10,
                    max_branch_depth=3,
                    value_bottom=-100,
                    value_upper=100,
                    alpha_var=35,
                    save_individual=False,
                    max_mutations=10,
                    max_code_lines=100,
                    sdep_k=2,
                    sdep_p=30,
                    sdep_t=0
                ):

        self.DefaultConfig = DefaultConfig
        self.pop_size = pop_size
        self.generations = generations
        self.population = []
        self.fitness = fitness
        self.variables = variables
        self.current_generation = 1
        self.hall_of_fame = []
        self.generator = generator
        self.max_depth = max_depth
        self.max_width = max_width
        self.max_branch_depth = max_branch_depth
        self.value_bottom = value_bottom
        self.value_upper = value_upper
        self.alpha_var = alpha_var
        self.save_individual=False
        self.max_mutations = max_mutations
        self.max_code_lines = max_code_lines
        self.current_gen = []
        # SDEP: targeted extinction
        self.sdep_k = sdep_k # how many generation ind are not improving
        self.sdep_p = sdep_p # probability of individual is extincted
        self.sdep_t = sdep_t # threshold regarding elite member
        self.sdep_last_best_fitness = 0
        self.sdep_current_stagnation = 0
        self.sdep_fresh_counter = 0
        self.last_best_individual_info = None
        self.last_best_individual = None
        self.ROOT_DIR = os.path.abspath(os.curdir)

    # @jit
    def init_population(self, generator, pickles=[]):
        loaded_pickle = 0
        for idx, p in enumerate(pickles):
            try:
                with open(p, 'rb') as f:
                    indiv = pickle.load(f)
                    indiv.id = idx
                    self.population.append(indiv)
                    loaded_pickle += 1
            except Exception as e:
                print("Cannot load pickle, try next one")

        for i in range(self.pop_size-loaded_pickle):
            variables = copy.deepcopy(self.variables)
            # self.population.append(generator.generate_individual_from_seed(variables=variables))
            indiv = generator.generate_individual_from_seed(
                variables=variables,
                wildcard_codes=self.DefaultConfig.WILD_CARD_CODE,
                max_depth=self.max_depth,
                max_width=self.max_width,
                max_branch_depth=self.max_branch_depth,
                alpha_var_gen=self.alpha_var,
                value_bottom=self.value_bottom,
                value_upper=self.value_upper
            )
            indiv.id = i + loaded_pickle
            self.population.append(indiv)

    # @jit
    def calculate_fitness(self, generator):
        codes = []
        for idx, individual in enumerate(self.population):
            codes.append(individual.render_code())

        #print("codes", codes)

        switch_case_lines = []
        for idx, code in enumerate(codes):
            switch_case_lines.append("case {}:\n{{".format(idx))

            for line in code:
                switch_case_lines.append(line)

            switch_case_lines.append("\n\tbreak;}")

        #print("switch", switch_case_lines)

        tcp_variant_fitness_write_switch(switch_case_lines)

        os.chdir(BASE_NS3_PATH)

        configure_command = 'CXXFLAGS="-Wno-error -Wno-unused-variable" ./waf configure --disable-python'
        build_command = 'CXXFLAGS="-Wno-error -Wno-unused-variable" ./waf build'

        retry_counter = 1
        while retry_counter >= 0:
            try:
                print('configuring...')
                start = time.time()
                out = subprocess.check_output(configure_command, shell=True, timeout=20)
                end = time.time()
                print("configured in {:.3f} seconds".format(end - start))

                print('building...')
                start = time.time()
                try:
                    out = subprocess.check_output(build_command, shell=True, timeout=500)
                except Exception as e:
                    print("IMPOSSIBLE TO BUILD", e)
                    retry_counter -= 1
                    continue
                end = time.time()
                print("built in {:.3f} seconds".format(end - start))
                break
            except:
                print('error on configuring or building...')
                retry_counter -= 1
                continue

        if retry_counter == -1:
            # sys.exit("CRITICAL ERROR! CANNOT CONTINUE THE PROCESS, EXPORT A SNAPSHOT TO NOT LOSE THE PROGRESS")
            print("CRITICAL ERROR! CANNOT CONTINUE THE PROCESS, RETRY GENERATION")
            os.chdir(self.ROOT_DIR) # changing current directory
            subprocess.call('rm -Rf ./snapshots_pickles/{self.current_generation}_gen', shell=True)
            subprocess.call('rm -Rf ./snapshots/{self.current_generation}_gen', shell=True)
            self.hall_of_fame = self.hall_of_fame[:-1]
            self.current_generation -= 1
            os.chdir(self.ROOT_DIR)
            pickles = [f.path for f in os.scandir('./snapshots_pickles/{}_gen/'.format(self.current_generation))]
            self.init_population(generator, pickles)
            self.check_config_build()
            # subprocess.call('mv ./tmp/extract/snapshots/{}_gen.json init', shell=True)

        # print("calculating fitness...")
        #print("[{}] has fitness {}".format(0, self.population[0].max_fitness(0, self.fitness)))

        pool = Pool(pool_size)

        for idx, individual in enumerate(self.population):
            pool.apply_async(self.multiprocessing_fitness, (idx, individual))

        pool.close()
        pool.join()
        # subprocess.call(['killall', '-r', 'my_pattern'])

    def multiprocessing_fitness(self, idx, individual):
        #for idx, individual in enumerate(self.population):
            # print(">>>>>>>>>>>FITNESS>>>>>>>>>>>")
        # print("[{}] calculating...".format(idx))
        start = time.time()
        fit = individual.max_fitness(idx, self.fitness, self.max_code_lines)
        print("[{}] \thas fitness\t {:.2f}".format(idx, fit), end='')
        end = time.time()
        print(" calculated in {:.1f} \tseconds".format(end - start))

        self.current_gen.append({
            'gen': self.current_generation,
            'id': idx,
            'fitness': fit,
            'path': "snapshots/{}_gen/{}.cc".format(self.current_generation, idx),
            'time': (end - start)
        })

        path_folder_out_gen = os.path.join(sys.path[0], "snapshots")
        path_file_out_gen = "{}/current_gen.json".format(path_folder_out_gen)
        with open(path_file_out_gen, 'w') as out_gen:
            json.dump(self.current_gen, out_gen)

    # @jit
    def tournament_selection(self, k=35, s=15):
        '''
        modified implementation of tournament selection as described in the referenced link.
        ref: https://en.wikipedia.org/wiki/Tournament_selection
        s is the number of selected individual from the tournament to perform crossover
        '''
        # selected individuals
        selected_individuals = []
        # pick k random individuals
        if len(selected_individuals) >= k:
            while len(selected_individuals) < k:
                # picking random individual
                indiv = self.population[secrets.randbelow(self.pop_size)]
                # check if already picked
                if indiv in selected_individuals: continue
                # append individual
                selected_individuals.append(indiv)
        else:
            selected_individuals = self.population

        best_individuals = []
        min_fitness_individual, max_fitness_individual = self.find_min_max_fitness(selected_individuals)
        # max_fitness_individual = None
        # find the best s individuals
        for idx, individual in enumerate(selected_individuals):
            if len(best_individuals) >= s \
                and individual.fitness > min_fitness_individual.fitness:

                best_individuals = [i for i in best_individuals if i != min_fitness_individual]
                min_fitness_individual, _ = self.find_min_max_fitness(best_individuals)
                continue

            if len(best_individuals) < s:
                if individual.fitness > min_fitness_individual.fitness:
                    min_fitness_individual = individual

                elif individual.fitness >= max_fitness_individual.fitness:
                    max_fitness_individual = individual

                best_individuals.append(individual)
        result = [ind for ind in best_individuals if ind.fitness > 5.0]
        return result if len(result) else best_individuals

    # @jit
    def find_min_max_fitness(self, individuals):
        min_fitness = min(individual.fitness for individual in individuals)
        max_fitness = max(individual.fitness for individual in individuals)
        min_individual = [individual for individual in individuals if individual.fitness == min_fitness][0]
        max_individual = [individual for individual in individuals if individual.fitness == max_fitness][0]
        return min_individual, max_individual

    # @jit
    def take_snapshot(self):
        path_folder = os.path.join(sys.path[0], "snapshots", "{}_gen".format(self.current_generation))
        path_file = "{}.cc"

        path_pickles = os.path.join(sys.path[0], "snapshots_pickles", "{}_gen".format(self.current_generation))
        path_file_pickles = "{}.pickle"

        # assure folder existence
        if not os.path.exists(path_folder):
            os.makedirs(path_folder)

        if not os.path.exists(path_pickles):
            os.makedirs(path_pickles)

        # output individuals

        for idx, individual in enumerate(self.population):
            try:
                # print("snapshot", path_folder + "/" + path_file.format(idx))
                individual.save_to_file(folder=path_folder, file=path_file.format(idx))
            except Exception as e:
                print("not exported individual, but still continuing")
        print_formatted_text(HTML('<aaa fg="black" bg="ansigreen">{} SNAPSHOTS EXPORTED TO: {}</aaa>'.format(len(self.population), path_folder)))

        # output pickles
        for idx, individual in enumerate(self.population):
            try:
                # print("pickle", path_pickles + "/" + path_file_pickles.format(idx))
                path = os.path.join(path_pickles, path_file_pickles.format(idx))
                with open(path, 'wb') as f:
                    pickle.dump(individual, f, protocol=pickle.HIGHEST_PROTOCOL)
                f.close()
            except Exception as e:
                print("not exported individual pickle, but still continuing")
        print_formatted_text(HTML('<aaa fg="black" bg="ansigreen">{} PICKLES EXPORTED TO: {}</aaa>'.format(len(self.population), path_pickles)))

    def add_hall_of_fame(self, add_to_elite=False):
        self.population.sort(key=lambda x:x.fitness, reverse=True)
        if int(self.population[0].fitness) > 0:
            self.hall_of_fame.append({"id": self.population[0].id, "fitness": self.population[0].fitness, "gen": self.current_generation})
            if add_to_elite == True:
                self.population[0].is_elite = True
            self.last_best_individual_info = {"id": self.population[0].id, "fitness": self.population[0].fitness, "gen": self.current_generation}
            self.last_best_individual = copy.deepcopy(self.population[0])
            return self.population[0] # return best individual
        elif self.current_generation > 1:
            return self.last_best_individual
        return None

    # @jit
    def crossover(self, best_individuals, there_is_elite, elite=None):
        # count how many offsprings to generate
        offsprings = self.DefaultConfig.TOURNAMENT['k'] - len(best_individuals)
        print("offsprings to generate are {}".format(offsprings))
        # reset population to best individuals
        self.population = []
        for b in best_individuals:
            self.population.append(copy.deepcopy(b))
        # create parent pool set
        parent_pool = set(best_individuals)
        if there_is_elite and elite is not None:
            parent_pool.add(copy.deepcopy(elite))

        place_for_elite = 1 if there_is_elite else 0
        counter_idx = len(self.population)+place_for_elite

        # generate new offsprings by
        while len(self.population) < self.pop_size-place_for_elite-self.sdep_fresh_counter: # -1 because elite individual is added afterwards
            # pick 2 random parents
            parents = []
            if len(parent_pool) > 0:
                for i in range(2):
                    parents.append(random.sample(parent_pool, 1)[0])

            variables = copy.deepcopy(self.variables)
            # self.population.append(generator.generate_individual_from_seed(variables=variables))
            indiv = generate_individual_from_seed(
                variables=variables,
                wildcard_codes=self.DefaultConfig.WILD_CARD_CODE,
                max_depth=self.max_depth,
                max_width=self.max_width,
                max_branch_depth=self.max_branch_depth,
                alpha_var_gen=self.alpha_var,
                value_bottom=self.value_bottom,
                value_upper=self.value_upper
            )
            indiv.id = counter_idx
            counter_idx += 1

            if len(parents) == 0:
                new_parent = generate_individual_from_seed(
                    variables=variables,
                    wildcard_codes=self.DefaultConfig.WILD_CARD_CODE,
                    max_depth=self.max_depth,
                    max_width=self.max_width,
                    max_branch_depth=self.max_branch_depth,
                    alpha_var_gen=self.alpha_var,
                    value_bottom=self.value_bottom,
                    value_upper=self.value_upper
                )
                new_parent.id = counter_idx
                counter_idx += 1
                parents.append(new_parent)

            print(RenderTree(parents[0].root))
            # convert nodes to sets
            parent_one_nodes = parents[0].root.descendants
            parent_two_nodes = indiv.root.descendants

            # parent_three_nodes = [n for n in PreOrderIter(parents[1].root.children[0])]
            # parent_two_nodes = search.findall(parents[1].root)
            # init new individual
            child = copy.deepcopy(parents[0])
            # child = copy.deepcopy(indiv)
            # pick random part of the tree
            if len(parent_one_nodes) > 1:
                subtree_one = copy.deepcopy(random.sample(parent_one_nodes, 1)[0])
                subtree_one.parent = None
                # perform crossover
                # first operation
                first_child_branch = random.sample(child.root.descendants, 1)[0]
                # first_child_branch.children = [subtree_one]
                subtree_one.parent = first_child_branch

            if len(parent_two_nodes) > 1:
                subtree_two = copy.deepcopy(random.sample(parent_two_nodes, 1)[0])
                subtree_two.parent = None
                # perform crossover
                # second operation
                second_child_branch = random.sample(indiv.root.descendants, 1)[0]
                subtree_two.parent = second_child_branch

            # add child to population
            self.population.append(child)

    def rand_int(self, start, end):
        return random.randint(start, end-1)

    def extinct(self, y_prob=0):
        return True if self.rand_int(0, 1000) < y_prob*10 else False

    def mutate(self, y_prob=50.0):
        return True if self.rand_int(0, 1000) < y_prob*10 else False

    def mutate_switch_branches(self, y_prob=30):
        return True if self.rand_int(0, 1000) < y_prob*10 else False

    # @jit
    def mutation(self):
        for idx, individual in enumerate(self.population):
            if not isinstance(individual, Individual):
                print("ERROR added something different as individual")
                quit()
            counter = self.rand_int(1, self.max_mutations)
            while counter > 0:
                if self.mutate(y_prob=self.DefaultConfig.MUTATION['truncate_node']) and not individual.is_elite:
                    # truncate mutation
                    nodes = [node for node in PreOrderIter(individual.root)]
                    selected_node = nodes[secrets.randbelow(len(nodes))]
                    # selected_node.parent.children = []
                    selected_node.parent = None
                    counter -= 1
                if counter > 0 and \
                    self.mutate(y_prob=self.DefaultConfig.MUTATION['operator_flip']) and not individual.is_elite:
                    # operator flipping
                    if (len(individual.root.children) >= 1):

                        variables = copy.deepcopy(individual.variables)
                        random_node = random.sample(individual.root.children, 1)[0]

                        print("----------------------------------------BEFORE----------------------------------------")
                        print(individual.render_code())
                        print("----------------------------------------MUTATE----------------------------------------")

                        new_node = generate_operator_node(
                                        random_node.parent,
                                        variables=variables,
                                        wildcard_codes=self.DefaultConfig.WILD_CARD_CODE,
                                        max_depth=self.max_depth,
                                        max_width=self.max_width,
                                        max_branch_depth=self.max_branch_depth,
                                        alpha_var_gen=self.alpha_var,
                                        value_bottom=self.value_bottom,
                                        value_upper=self.value_upper
                                    )
                        for child in random_node.children:
                            child.parent = new_node
                        random_node.parent = None
                        del random_node
                        gc.collect()
                        # random_node.children = []
                        print(individual.render_code())
                        print("--------------------------------------------------------------------------------------")
                    counter -= 1
                # switch branches
                if counter > 0 and \
                    self.mutate_switch_branches(y_prob=self.DefaultConfig.MUTATION['switch_branches']):
                    for node in PreOrderIter(individual.root):
                        if isinstance(node, IfThenElse):
                            tmp = copy.deepcopy(node.exp_f)
                            node.exp_f = copy.deepcopy(node.exp_t)
                            node.exp_t = tmp
                            break
                    counter -= 1
                # switch compatible operators
                # for instance sum with multiplication
                if counter > 0 and \
                    self.mutate(y_prob=self.DefaultConfig.MUTATION['switch_exp']):
                    for node in PreOrderIter(individual.root):
                        if isinstance(node, Assignment):
                            for oper in PreOrderIter(node.exp):
                                if self.mutate(y_prob=20):
                                    if isinstance(oper, Mul):
                                        new_oper = Sum(oper.nums)
                                        new_oper.parent = oper.parent
                                        for child in new_oper.children:
                                            child.parent = new_oper
                                    elif isinstance(oper, Sum):
                                        new_oper = Mul(oper.nums)
                                        new_oper.parent = oper.parent
                                        for child in new_oper.children:
                                            child.parent = new_oper
                                    elif isinstance(oper, Sub):
                                        new_oper = Mul(oper.nums)
                                        new_oper.parent = oper.parent
                                        for child in new_oper.children:
                                            child.parent = new_oper
                                    break
                    counter -= 1

    def sdep_f(self, x):
        '''
                    x^(1/2)
        f(x) = 10 --------------
                    1 + e^(-x)
        '''
        return 10 * ((x**(1/2))/(1+math.e**(-x)))

    def sdep(self):
        '''
        Stagnation-driven extinction protocol: Targeted Extinction
        https://www.mdpi.com/2076-3417/11/8/3461

        Notes:
            individuals are extincted with probability (sdep_p+sdep_f(sdep_current_stagnation-sdep_k))
            if their fitness is under a threshold calculated as elite_fitness*(1-sdep_t) if generations
            are stagnant for sdep_k generations.
        '''
        print("[SDEP] running...")
        self.sdep_fresh_counter = 0
        extinction_probability = 0
        # if generation are stagnant more than the activation threshold
        if self.sdep_current_stagnation >= self.sdep_k:
            x = self.sdep_current_stagnation-self.sdep_k
            extinction_probability = self.sdep_p + self.sdep_f(x)

            extinct_individuals = []

            for ind in self.population:
                if ind.fitness <= self.sdep_last_best_fitness*(1-self.sdep_t) and \
                    self.extinct(y_prob=extinction_probability):
                    extinct_individuals.append(ind)

            self.sdep_fresh_counter = len(extinct_individuals)

            for ind in extinct_individuals:
                self.population.remove(ind)
        print("[SDEP] FRESH INDIVIDUALS: {} \t STAGNATION: {} \t EXTINCTION P: {:.2f}".format(
                                                self.sdep_fresh_counter,
                                                self.sdep_current_stagnation,
                                                extinction_probability
                                            ))

    def check_config_build(self):
        codes = []
        for idx, individual in enumerate(self.population):
            codes.append(individual.render_code())

        #print("codes", codes)

        switch_case_lines = []
        for idx, code in enumerate(codes):
            switch_case_lines.append("case {}:\n{{".format(idx))

            for line in code:
                switch_case_lines.append(line)

            switch_case_lines.append("\n\tbreak;}")

        #print("switch", switch_case_lines)

        tcp_variant_fitness_write_switch(switch_case_lines)

        os.chdir(BASE_NS3_PATH)

        configure_command = 'CXXFLAGS="-Wno-error -Wno-unused-variable" ./waf configure --disable-python'
        build_command = 'CXXFLAGS="-Wno-error -Wno-unused-variable" ./waf build'

        retry_counter = 1
        while retry_counter >= 0:
            try:
                print('configuring...')
                start = time.time()
                out = subprocess.check_output(configure_command, shell=True, timeout=20)
                end = time.time()
                print("configured in {:.3f} seconds".format(end - start))

                print('building...')
                start = time.time()
                try:
                    out = subprocess.check_output(build_command, shell=True, timeout=500)
                except Exception as e:
                    print("IMPOSSIBLE TO BUILD", e)
                    retry_counter -= 1
                    continue
                end = time.time()
                print("built in {:.3f} seconds".format(end - start))
                break
            except:
                print('error on configuring or building...')
                retry_counter -= 1
                continue

        if retry_counter == -1:
            print("CRITICAL ERROR! CANNOT CONTINUE THE PROCESS, RETRY GENERATION")
            os.chdir(self.ROOT_DIR) # changing current directory
            subprocess.call('rm -Rf ./snapshots_pickles/{self.current_generation}_gen', shell=True)
            subprocess.call('rm -Rf ./snapshots/{self.current_generation}_gen', shell=True)
            self.hall_of_fame = self.hall_of_fame[:-1]
            print("Retry generating population...")
            return False
        return True

    def run(self, generator, pickles=[]):
        # writing that simulation is running
        path_termination_flag = os.path.join(sys.path[0], 'termination_flag')
        with open(path_termination_flag, 'w') as termination_flag:
            termination_flag.write('RUNNING')

        self.ROOT_DIR = os.path.abspath(os.curdir)
        # init population giving a generator
        self.init_population(generator, pickles)

        for i in range(self.generations):
            # take snapshot of generated individuals
            self.take_snapshot()

            #self.mutation()
            #break

            self.calculate_fitness(generator)

            elite_individual = copy.deepcopy(self.add_hall_of_fame(add_to_elite=True))

            if elite_individual is not None:
                if self.sdep_last_best_fitness != elite_individual.fitness:
                    self.sdep_current_stagnation = 0
                else:
                    self.sdep_current_stagnation += 1
                self.sdep_last_best_fitness = elite_individual.fitness
            else:
                self.sdep_current_stagnation += 1

            # apply SDEP
            self.sdep()

            # write out the hall of fame
            path_folder_hall = os.path.join(sys.path[0], "snapshots")
            path_file_hall = "{}/hall_of_fame.json".format(path_folder_hall)
            with open(path_file_hall, 'w') as out_hall:
                json.dump(self.hall_of_fame, out_hall)

            print("[GENERATION {}/{}] with population:".format(self.current_generation, self.pop_size))

            for idx, ind in enumerate(self.population):
                ind.id = idx
                print("id:{}, fitness:{:.2f}, variables: {}".format(ind.id,ind.fitness, ind.variables.variables_name()))

            selected = self.tournament_selection(
                k=self.DefaultConfig.TOURNAMENT["k"],
                s=self.DefaultConfig.TOURNAMENT["s"])

            # print("[GENERATION {}/{}] HALL OF FAME: {}".format(self.current_generation, self.pop_size, [str(ind) for ind in self.hall_of_fame]))
            print_formatted_text(HTML(
                "<aaa fg=\"white\" bg=\"red\">[GENERATION {}/{}] HALL OF FAME: {}</aaa>".format(self.current_generation, self.generations, [str(ind) for ind in self.hall_of_fame])
            ))

            print("{} SELECTED INDIVIDUALS".format(len(selected)))

            for idx, ind in enumerate(selected):
                print("id:{}, fitness:{:.2f}".format(ind.id, ind.fitness))

            # self.population = []
            there_is_elite = elite_individual is not None
            if len(selected) >= 2:
                self.crossover(selected, there_is_elite=there_is_elite, elite=elite_individual)
                # self.fix_not_valid_crossover()

            self.mutation()

            # fixing individual
            for ind in self.population:
                variables = copy.deepcopy(self.variables)
                # updating variables
                ind.update_variable_registry(variables)

            counter_idx = len(self.population)
            # inject fresh individuals from extinction
            for i in range(self.sdep_fresh_counter):
                variables = copy.deepcopy(self.variables)
                indiv = generate_individual_from_seed(
                    variables=variables,
                    wildcard_codes=self.DefaultConfig.WILD_CARD_CODE,
                    max_depth=self.max_depth,
                    max_width=self.max_width,
                    max_branch_depth=self.max_branch_depth,
                    alpha_var_gen=self.alpha_var,
                    value_bottom=self.value_bottom,
                    value_upper=self.value_upper
                )
                indiv.id = counter_idx
                counter_idx += 1

                # add indiv to population
                self.population.append(copy.deepcopy(indiv))
                print("FRESH INDIVIDUAL CODE: {}".format(indiv.render_code()))

            if len(selected) == 0:
                print("GOT NOT AVAILABLE INDIVIDUALS")
                self.population = []
                self.init_population(generator, pickles)

            # cleaning ID
            for idx, individual in enumerate(self.population):
                individual.id = idx

            if elite_individual is not None:
                # print("[ELITE INDIVIDUAL] id:{}, fitness:{}".format(elite_individual.id, elite_individual.fitness))
                print_formatted_text(HTML(
                    "<aaa fg=\"black\" bg=\"yellow\">[ELITE INDIVIDUAL] id:{}, fitness:{}</aaa>".format(elite_individual.id, elite_individual.fitness)
                ))
                elite_individual.id = len(self.population)
                self.population.append(elite_individual)

            self.current_generation += 1
            self.current_gen = []
            gc.collect()

            for proc in psutil.process_iter():
                # check whether the process name matches
                if proc.name() == PROCWIFITCP:
                    proc.kill()

        print("FINAL POPULATION")
        print([ind.fitness for ind in self.population])
        print("HALL OF FAME:")
        print([str(ind) for ind in self.hall_of_fame])

        # write out the hall of fame
        path_folder_hall = os.path.join(sys.path[0], "snapshots")
        path_file_hall = "{}/hall_of_fame.json".format(path_folder_hall)
        with open(path_file_hall, 'w') as out_hall:
            json.dump(self.hall_of_fame, out_hall)
        # writing that simulation has stopped
        with open(path_termination_flag, 'w') as termination_flag:
            termination_flag.write('STOP')

    def kill(self):
        sys.exit()

max_stars_repo_name: carbogninalberto/pyGENP

max_issues_repo_name: carbogninalberto/pyGENP

max_forks_repo_name: carbogninalberto/pyGENP


text:
# 必要なライブラリをインポート
import numpy as np
import pandas as pd
import warnings
warnings.simplefilter('ignore')
import datetime
from pandas_datareader import data as pdr
import matplotlib.pyplot as plt
import japanize_matplotlib
%matplotlib inline

# 平均を丸める期間を自分で選択できるタイプ
# 引数spanに平均を計算する期間を入力する
def SPG(company_code, start, end, span=1):
    # 平均算出のため、指定範囲より多めにデータを取っておく
    start = (datetime.datetime.strptime(start, '%Y-%m-%d')-datetime.timedelta(days=80)).strftime('%Y-%m-%d')
    df = pdr.StooqDailyReader(company_code, start=start, end=end).read().sort_index()

    # 各種平均計算
    date = df.index
    price = df['Close']

    def wma(w):
        weight = np.arange(len(w)) + 1
        wma = np.sum(weight * w) / weight.sum()
        return wma

    df['sma'] = price.rolling(window=span).mean()
    df['wma'] = price.rolling(window=span).apply(wma,raw=True).round(1)
    df['ema'] = price.ewm(span=span, adjust=False).mean()

    # 指定された期間に範囲を切り取る
    start = (datetime.datetime.strptime(start, '%Y-%m-%d')+datetime.timedelta(days=80)).strftime('%Y-%m-%d')
    df = df[(start <= df.index) & (df.index <= end)]
    date = df.index
    price = df['Close']

    plt.figure(figsize=(30,15))

    # 指定された丸め期間で計算された各種平均をプロット
    plt.subplot(2,1,1)
    plt.title(f'{span}日',color='black',backgroundcolor='white',size=30,loc='center')
    plt.plot(date,price,label='Close',color='black')
    plt.plot(date,df['sma'],label='sma',color='red')
    plt.plot(date,df['wma'],label='wma',color='blue')
    plt.plot(date,df['ema'],label='ema',color='green')
    plt.legend()

    # 取引量
    plt.subplot(2,1,2)
    plt.title('Volume',color='black',backgroundcolor='white',size=30,loc='center')
    plt.bar(date,df['Volume'],label='Volume',color='grey')
    plt.legend()

# テスト
# S&P 500：1547.jp
SPG('1547.jp', '2021-03-10', '2022-03-10', span=10)

ext: py

reason: import numpy


text:
import json
import os

import numpy
import thulac
import psutil
from sklearn.feature_extraction.text import TfidfVectorizer as TFIDF
import joblib
from sklearn.ensemble import GradientBoostingClassifier
import re

#training score :  0.738

class CompOpinion():
    def __init__(self, cwd=".", tfidf='statement_tfidf.model', gbt='statement_som_gbt.model'):
        patterns1 = "本院认为"
        self.ps1 = re.compile(patterns1)
        patterne1 = '综上'
        patterne2 = '依照'
        self.pe1 = re.compile(patterne1)
        self.pe2 = re.compile(patterne2)

    def cut_text(self, alltext):
        count = 0
        train_text = []
        for text in alltext:
            count += 1
            if count % 2000 == 0:
                print(count)
            train_text.append(self.cut.cut(text, text=True))

        return train_text

    def print_mem(self):
        process = psutil.Process(os.getpid())  # os.getpid()
        memInfo = process.memory_info()
        return '{:.4f}G'.format(1.0 * memInfo.rss / 1024 /1024 /1024)

    def checkResult(self, statement):
        train_data = self.cut_text([statement])
        vec = self.tfidf.transform(train_data)
        yp = self.gbt.predict(vec)
        return yp[0]

    def getOpinion(self, sentences):
        tap = False
        cr = []
        for dic in sentences:
            if tap and (self.pe1.search(dic['sentence']) or self.pe2.search(dic['sentence'])):
                break

            if self.ps1.search(dic['sentence']):
                tap = True

            if tap:
                cr.append(dic['sentence'])

        # print(cr)
        return cr

max_stars_repo_name: ShenDezhou/CAIL

max_issues_repo_name: ShenDezhou/CAIL

max_forks_repo_name: ShenDezhou/CAIL


text:
import logging

import numpy as np
import torch
import torch.optim as optim

INFTY = 1e20

class DKNN_PGD(object):
    """
    Implement gradient-based attack on DkNN with L-inf norm constraint.
    The loss function is the same as the L-2 attack, but it uses PGD as an
    optimizer.
    """

    def __init__(self, dknn):
        self.dknn = dknn
        self.device = dknn.device
        self.layers = dknn.layers
        self.guide_reps = {}
        self.thres = None
        self.coeff = None

    def __call__(self, x_orig, label, guide_layer, m, epsilon=0.1,
                 max_epsilon=0.3, max_iterations=1000, num_restart=1,
                 rand_start=True, thres_steps=100, check_adv_steps=100,
                 verbose=True):

        # make sure we run at least once
        if num_restart < 1:
            num_restart = 1

        # if not using randomized start, no point in doing more than one start
        if not rand_start:
            num_restart = 1

        label = label.cpu().numpy()
        batch_size = x_orig.size(0)
        min_, max_ = x_orig.min(), x_orig.max()
        # initialize adv to the original
        x_adv = x_orig.detach()
        best_num_nn = np.zeros((batch_size, ))

        # set coefficient of guide samples
        self.coeff = torch.zeros((x_orig.size(0), m))
        self.coeff[:, :m // 2] += 1
        self.coeff[:, m // 2:] -= 1

        for i in range(num_restart):

            # initialize perturbation
            delta = torch.zeros_like(x_adv)
            if rand_start:
                delta.uniform_(- max_epsilon, max_epsilon)
            delta.requires_grad_()

            for iteration in range(max_iterations):
                x = torch.clamp(x_orig + delta, min_, max_)

                # adaptively choose threshold and guide samples every
                # <thres_steps> iterations
                with torch.no_grad():
                    if iteration % thres_steps == 0:
                        thres = self.dknn.get_neighbors(x)[0][0][:, -1]
                        self.thres = torch.tensor(thres).to(self.device).view(
                            batch_size, 1)
                        self.find_guide_samples(
                            x, label, m=m, layer=guide_layer)

                reps = self.dknn.get_activations(x, requires_grad=True)
                loss = self.loss_function(reps)
                loss.backward()
                # perform update on delta
                with torch.no_grad():
                    delta -= epsilon * delta.grad.detach().sign()
                    delta.clamp_(- max_epsilon, max_epsilon)

                if (verbose and iteration % (np.ceil(max_iterations / 10)) == 0):
                    print('    step: %d; loss: %.3f' %
                          (iteration, loss.cpu().detach().numpy()))

                if ((iteration + 1) % check_adv_steps == 0 or
                        iteration == max_iterations):
                    with torch.no_grad():
                        # check if x are adversarial. Only store adversarial
                        # examples if they have a larger number of wrong
                        # neighbors than orevious
                        is_adv, num_nn = self.check_adv(x, label)
                        for j in range(batch_size):
                            if is_adv[j] and num_nn[j] > best_num_nn[j]:
                                x_adv[j] = x[j]
                                best_num_nn[j] = num_nn[j]

            with torch.no_grad():
                is_adv, _ = self.check_adv(x_adv, label)
            if verbose:
                print('number of successful adv: %d/%d' %
                      (is_adv.sum(), batch_size))

        return x_adv

    def check_adv(self, x, label):
        """Check if label of <x> predicted by <dknn> matches with <label>"""
        output = self.dknn.classify(x)
        num_nn = output.max(1)
        y_pred = output.argmax(1)
        is_adv = (y_pred != label).astype(np.float32)
        return is_adv, num_nn

    def loss_function(self, reps):
        """Returns the loss averaged over the batch (first dimension of x) and
        L-2 norm squared of the perturbation
        """

        batch_size = reps[self.layers[0]].size(0)
        adv_loss = torch.zeros(
            (batch_size, len(self.layers)), device=self.device)
        # find squared L-2 distance between original samples and their
        # adversarial examples at each layer
        for l, layer in enumerate(self.layers):
            rep = reps[layer].view(batch_size, 1, -1)
            dist = ((rep - self.guide_reps[layer])**2).sum(2)
            fx = self.thres - dist
            Fx = torch.max(torch.tensor(0., device=self.device),
                           self.coeff.to(self.device) * fx).sum(1)
            adv_loss[:, l] = Fx

        return adv_loss.mean()

    def find_guide_samples(self, x, label, m=100, layer='relu1'):
        """Find k nearest neighbors to <x> that all have the same class but not
        equal to <label>
        """
        num_classes = self.dknn.num_classes
        x_train = self.dknn.x_train
        y_train = self.dknn.y_train
        batch_size = x.size(0)
        nn = torch.zeros((m, ) + x.size()).transpose(0, 1)
        D, I = self.dknn.get_neighbors(
            x, k=x_train.size(0), layers=[layer])[0]

        for i, (d, ind) in enumerate(zip(D, I)):
            mean_dist = np.zeros((num_classes, ))
            for j in range(num_classes):
                mean_dist[j] = np.mean(
                    d[np.where(y_train[ind] == j)[0]][:m // 2])
            mean_dist[label[i]] += INFTY
            nearest_label = mean_dist.argmin()
            nn_ind = np.where(y_train[ind] == nearest_label)[0][:m // 2]
            nn[i, m // 2:] = x_train[ind[nn_ind]]
            nn_ind = np.where(y_train[ind] == label[i])[0][:m // 2]
            nn[i, :m // 2] = x_train[ind[nn_ind]]

        # initialize self.guide_reps if empty
        if not self.guide_reps:
            guide_rep = self.dknn.get_activations(
                nn[0], requires_grad=False)
            for l in self.layers:
                # set a zero tensor before filling it
                size = (batch_size, ) + guide_rep[l].view(m, -1).size()
                self.guide_reps[l] = torch.zeros(size, device=self.device)

        # fill self.guide_reps
        for i in range(batch_size):
            guide_rep = self.dknn.get_activations(
                nn[i], requires_grad=False)
            self.guide_reps[layer][i] = guide_rep[layer].view(
                m, -1).detach()

max_stars_repo_name: chawins/entangle-rep

max_issues_repo_name: chawins/entangle-rep

max_forks_repo_name: chawins/entangle-rep


text:
import numpy as np
import os

from . import misc

class LCurve(object):
    """Light curve holder class"""

    def __init__(self, t, r, re, dt=None, fexp=None):
        """Initialize LCurve from array r and optional t, re, dt

        Parameters:
            t:  an array containing the time axis
            r:  an array containing the count rate
            re: an array containing the measurement errors.
            dt: time sampling. If not given, dt = min(diff(t))
            fexp: fraction exposure array. If not given, fexp=np.ones_like(t)
        """

        # check input arrays #
        if not (len(t) == len(r) == len(re)):
            raise ValueError('arrays t, r, re do not match')

        # time sampling #
        if dt is None:
            dt = np.min(np.diff(t))

        # fraction exposure #
        if fexp is None:
            fexp = np.ones_like(t)

        # is the light curve evenly sampled? #
        iseven = np.all(np.isclose(np.diff(t), dt))

        # global variables #
        self.time = np.array(t)
        self.rate = np.array(r)
        self.rerr = np.array(re)
        self.fexp = fexp
        self.dt = dt
        self.iseven = iseven
        self.nt = len(t)

    def __repr__(self):
        return ('<LCurve :: nt({}) :: dt({})>').format(self.nt, self.dt)

    def make_even(self, fill=np.nan):
        """Make the light curve even in time, filling gaps with fill

        Parameters:
            fill: value to use in gaps.

        Returns:
            a new LCurve object

        """

        if self.iseven:
            return self

        # make sure time axis can be made even #
        itime = np.round((self.time - self.time[0]) / self.dt)
        if not np.allclose(itime - np.array(itime, np.int), 0):
            raise ValueError('time axis cannot be made even')

        # do work #
        t_new = np.arange(np.int(itime[-1]) + 1) * self.dt + self.time[0]
        idx = np.in1d(t_new, self.time)
        r, re = [np.zeros_like(t_new) + fill for i in range(2)]
        f = np.zeros_like(t_new)
        r[idx]  = self.rate
        re[idx] = self.rerr
        f[idx]  = self.fexp

        # return a new LCurve object #
        return LCurve(t_new, r, re, self.dt, f)

    def rebin(self, factor, error='norm', min_exp=0.0):
        """Rebin the light curve to so new_dt = dt*factor

        Parameters:
            factor: rebinning factor. dt_new = factor * dt
            error: error type (poiss|norm).
                If poiss: rerr = sqrt(rate*dt)/dt, otherwise,
                errors are summed quadratically
            min_exp: minimum fractional exposure to leave [0-1]

        Return:
            new binned LCurve

        """

        # check input error type #
        if not error in ['poiss', 'norm']:
            raise ValueError('error need to be poiss|norm')

        # make lc evenly sampled, so we bin arrays easily #
        iseven = self.iseven
        lc = self.make_even()

        # new sampling time and length #
        factor = np.int(factor)
        dt_new = lc.dt * factor
        nt_new = lc.nt//factor
        nt_scal = nt_new * factor

        # pre-binning #
        t  = lc.time[:nt_scal].reshape(nt_new, factor)
        r  = lc.rate[:nt_scal].reshape(nt_new, factor)
        re = lc.rerr[:nt_scal].reshape(nt_new, factor)
        f  = lc.fexp[:nt_scal].reshape(nt_new, factor)

        # rescale the rates to pre-fexp counts/bin #
        c  = r  * (f * lc.dt)
        ce = re * (f * lc.dt)

        # do binning #
        t  = np.mean(t, 1)
        c  = np.nansum(c, 1)
        if error == 'poiss':
            ce = np.sqrt(c)
            ce[ce==0] = np.nanmean(ce[ce!=0])
        else:
            ce = np.nansum(ce**2, 1)**0.5

        f  = np.mean(f, 1)
        fs = np.array(f)
        it = (fs != 0)
        fs[~it] = np.nan
        r  = c /(dt_new * fs)
        re = ce/(dt_new * fs)

        # # leave nan values if original lc had nan (i.e it was even)
        if iseven:
            it = np.ones_like(it) == 1

        # # filter on fracexp if needed #
        if min_exp > 0:
            it[f < min_exp] = False

        # return a new LCurve object #
        return LCurve(t[it], r[it], re[it], dt_new, f[it])

    def interp_small_gaps(self, maxgap=None, noise='poiss', seed=None):
        """Interpolate small gaps in the lightcurve if the gap
            is <maxgap; applying noise if requested

        Parameters:
            maxgap: the maximum length of a gap to be interpolated
            noise: poiss|norm|None
            seed: random seen if noise is requested

        """
        from itertools import groupby

        if not self.iseven:
            raise ValueError('lc is not even; make even first')

        # random seed if noise is needed #
        if noise is not None:
            np.random.seed(seed)

        # find gap lengths in the data #
        maxn = self.nt if maxgap is None else maxgap
        iarr = [list(i[1]) for i in groupby(np.arange(self.nt),
                    lambda ix:np.isfinite(self.rate[ix]))]
        # indices of non-finite segments #
        iinf = iarr[(1 if np.isfinite(iarr[0][0]) else 0)::2]
        # length of each non-finite segment #
        iinf = [i for i in iinf if len(i)<=maxn]
        iinf = [j for i in iinf for j in i]

        # interpolate all values then keep only those with length<maxn #
        idx = np.isfinite(self.rate)
        y  = np.interp(self.time, self.time[idx], self.rate[idx])
        ye = np.zeros_like(y)
        me = np.mean(self.rerr[idx])
        if noise is None:
            # no noise requested; the value is not altered from the interp
            # while the error is the average of all errors
            ye += me

        elif noise == 'poiss':
            # apply noise to counts/bin then convert back to counts/s
            yp = np.random.poisson(y*self.dt)
            y  = yp / self.dt
            ye = np.sqrt(yp) / self.dt
            # reset points where y=0 (hence ye=0)
            ye[yp == 0] = me

        elif noise == 'norm':
            y  += np.random.randn(len(y)) * me
            ye += me

        # now update fill in the gaps with length<maxn #
        self.rate[iinf] = y[iinf]
        self.rerr[iinf] = ye[iinf]

    @staticmethod
    def sync(lc_list, tbase=None):
        """Synchronize a list of arrays or LCurve's
        lc_list: a list of arrays or a list of LCurve objects.
            if arrays, the shape is (nlc, 3 (or 4 for fexp), nt).
            The 3 is for (time, rate, rerr)
        tbase: time array to use for reference.
            If not given, use the intersection of all t
        """
        if not isinstance(lc_list, (list, np.ndarray)):
            raise ValueError('lc_list must be a list')
        if isinstance(lc_list[0], LCurve):
            data = [np.array([l.time, l.rate, l.rerr]) for l in lc_list]
        else:
            # consider if fexp is not given
            data = [np.array(l) for l in lc_list]

        if tbase is None:
            tbase = data[0][0]
            for d in data[1:]:
                tbase = tbase[np.in1d(tbase, d[0])]

        data = [d[:, np.in1d(d[0], tbase)] for d in data]
        return data

    @staticmethod
    def read_fits_file(fits_file, **kwargs):
        """Read LCurve from fits file

        Parameters:
            fits_file: name of the fits file

        Keywords:
            min_exp: minimum fractional exposure to allow. Default 0.0 for all
            rate_tbl: name or number of hdu that contains lc data. Default: RATE
            rate_col: name or number of rate column. Default: RATE
            time_col: name or number of time column. Default: TIME
            rerr_col: name or number of rerr column. Default: ERROR
            fexp_col: name or number of the fracexp column. Default: FRACEXP
            gti_table: name or number of gti extension hdu. Default: GTI
            dt_key: name of dt keyword in header. Default: TIMEDEL
            gti_skip: how many seconds to skip at the gti boundaries. Default: 0
            verbose.

        Returns:
            ldata (shape: 4,nt containing, time, rate, rerr, fexp), dt
        """

        # pyfits #
        import astropy.io.fits as pyfits

        # default parameters #
        min_exp  = kwargs.get('min_exp', 0.)
        rate_tbl = kwargs.get('rate_tbl', 'RATE')
        rate_col = kwargs.get('rate_col', 'RATE')
        time_col = kwargs.get('time_col', 'TIME')
        rerr_col = kwargs.get('rerr_col', 'ERROR')
        fexp_col = kwargs.get('fexp_col', 'FRACEXP')
        gti_tbl  = kwargs.get('gti_tbl' , 'GTI')
        dt_key   = kwargs.get('dt_key', 'TIMEDEL')
        #dt = pyfits.open(fits_file)[1].header['timedel']
        gti_skip = kwargs.get('gti_skip', 0.0)
        verbose  = kwargs.get('verbose', False)

        # does file exist? #
        if not os.path.exists(fits_file):
            raise ValueError('file {} does not exist'.format(fits_file))

        # read file #
        with pyfits.open(fits_file) as fs:

            # lc data #
            data = fs[rate_tbl].data
            ldata = np.array([  data.field(time_col),
                                data.field(rate_col),
                                data.field(rerr_col)], dtype=np.double)

            # start time and time sampling #
            t0 = (fs[rate_tbl].header['TSTART'] if 'TSTART' in
                        fs[rate_tbl].header.keys() else 0.0)
            dt = (fs[rate_tbl].header[dt_key] if dt_key in
                        fs[rate_tbl].header.keys() else None)
            if not dt is None: t0 += dt/2

            # if the time-axis offset, correct it #
            if t0/ldata[0, 1] > 1e5:
                ldata[0] += t0

            # gti #
            try:
                ghdu = fs[gti_tbl]
                lgti  = np.array([ghdu.data.field(0), ghdu.data.field(1)],
                                dtype=np.double)
            except KeyError:
                if verbose:
                    print('No GTI found in {}'.format(fits_file))
                lgti = np.array([[ldata[0, 0]], [ldata[0, -1]]])

            # fractional exposure #
            try:
                lfracexp = data.field(fexp_col)
            except KeyError:
                if verbose:
                    print('cannot read fracexp_col in {}'.format(fits_file))
                lfracexp = np.ones_like(ldata[0])

            # apply gti #
            igti  = ldata[0] < 0
            for gstart, gstop in lgti.T:
                igti = igti | ( (ldata[0] >= (gstart+gti_skip)) &
                                (ldata[0] <= (gstop -gti_skip)) )
            igood = igti & (lfracexp >= min_exp) & (np.isfinite(ldata[0]))
            ldata = np.vstack([ldata, lfracexp])
            ldata = ldata[:, igood]

        return ldata, dt

    @staticmethod
    def read_pn_lcurve(fits_file, **kwargs):
        """Read pn lcurve fits_file created with xmmlc_lc.
            This sets values relevant to PN and calls read_fits_file

        Parameters:
            fits_file: name of the files file

        Keywords:
            See @LCurve.read_fits_file

        Returns:
            LCurve object
        """

        # set values relevant to XMM-PN files #
        kwargs.setdefault('min_exp' , 0.7)
        kwargs.setdefault('gti_tbl' , 2)

        data, dt = LCurve.read_fits_file(fits_file, **kwargs)
        return LCurve(data[0], data[1], data[2], dt, data[3])

    @staticmethod
    def read_pca_lcurve(fits_file, **kwargs):
        """Read pca lcurve fits_file.
            This sets values relevant to PCA and calls read_fits_file

        Parameters:
            fits_file: name of the files file

        Keywords:
            See @LCurve.read_fits_file

        Returns:
            LCurve object
        """

        # set values relevant to XMM-PN files #
        kwargs.setdefault('min_exp' , 0.99)
        kwargs.setdefault('gti_tbl' , 'STDGTI')

        data, dt = LCurve.read_fits_file(fits_file, **kwargs)
        return LCurve(data[0], data[1], data[2], dt, data[3])

    @staticmethod
    def read_nu_lcurve(fits_file, **kwargs):
        """Read nustar lcurve fits_file.
            This sets values relevant to NUSTAR and calls read_fits_file

        Parameters:
            fits_file: name of the files file

        Keywords:
            See @LCurve.read_fits_file

        Returns:
            LCurve object
        """

        # set values relevant to XMM-PN files #
        kwargs.setdefault('min_exp' , 0.1)
        kwargs.setdefault('gti_tbl' , 'GTI')
        kwargs.setdefault('gti_skip', 3.0)

        data, dt = LCurve.read_fits_file(fits_file, **kwargs)
        return LCurve(data[0], data[1], data[2], dt, data[3])

    @staticmethod
    def read_xis_lcurve(fits_file, **kwargs):
        """Read suzaku xis lcurve fits_file.
            This sets values relevant to NUSTAR and calls read_fits_file

        Parameters:
            fits_file: name of the files file

        Keywords:
            See @LCurve.read_fits_file

        Returns:
            LCurve object
        """

        # set values relevant to XIS files #
        kwargs.setdefault('min_exp' , 0.1)
        kwargs.setdefault('gti_tbl' , 'GTI')

        data, dt = LCurve.read_fits_file(fits_file, **kwargs)
        return LCurve(data[0], data[1], data[2], dt, data[3])

    @staticmethod
    def read_ni_lcurve(fits_file, **kwargs):
        """Read nicer lcurve fits_file.
            This sets values relevant to NICER and calls read_fits_file

        Parameters:
            fits_file: name of the files file

        Keywords:
            See @LCurve.read_fits_file

        Returns:
            LCurve object
        """

        # set values relevant to NICER files #
        kwargs.setdefault('min_exp' , 0.99)
        kwargs.setdefault('gti_tbl' , 'GTI')
        #kwargs.setdefault('gti_skip', 3.0)

        data, dt = LCurve.read_fits_file(fits_file, **kwargs)
        return LCurve(data[0], data[1], data[2], dt, data[3])

    @staticmethod
    def calculate_psd(rate, dt, norm='var', **kwargs):
        """Calculate raw psd from a list of light curves.

        Parameters:
            rate: array or list of arrays of lcurve rates
            dt: time bin width of the light curve
            norm: psd normalization: var|rms|leahy

        Keywords:
            rerr: array or list of errors on rate. If not give,
                assume, poisson noise.
            bgd: array or list of background rates. In this case,
                rate above is assumed background subtracted.
            taper: apply Hanning tapering before calculating the psd
                see p388 Bendat & Piersol; the psd need to be multiplied
                by 8/3 to componsate for the reduced variance.

        Return:
            freq, rpsd, nois.
        """

        # check input #
        if not isinstance(rate[0], (np.ndarray, list)):
            rate = [rate]

        if not norm in ['var', 'rms', 'leahy']:
            raise ValueError('norm need to be var|rms|leahy')

        # rerr and bgd; for estimating noise level #
        rerr = kwargs.get('rerr', None)
        bgd  = kwargs.get('bgd', 0.0)
        if not isinstance(bgd, (np.ndarray, list)):
            bgd = [bgd for r in rate]
        if rerr is None:
            # err is sqrt of number of counts/bin
            rerr = [np.sqrt((r+b)/dt) for r,b in zip(rate, bgd)]

        # tapering ? #
        taper = kwargs.get('taper', False)
        if taper:
            rate = [(r-r.mean()) * np.hanning(len(r)) + r.mean() for r in rate]

        # fft; remove the 0-freq and the nyquist #
        freq = [np.fft.rfftfreq(len(r), dt)[1:-1] for r in rate]
        rfft = [np.fft.rfft(r)[1:-1] for r in rate]
        mean = [np.mean(r) for r in rate]

        # normalize psd #
        expo = {'var':0, 'leahy':1, 'rms':2}
        rpsd = [(2.*dt / (len(r) * mu**expo[norm])) * np.abs(rf)**2
                    for r,rf,mu in zip(rate, rfft, mean)]

        # renormalize rpsd if tapering has been applied #
        if taper:
            rpsd = [r * 8/3 for r in rpsd]

        ## ------ noise level ------- ##
        # noise level is: 2*(mu+bgd)/(mu^2) for RMS normalization; eqn A2, Vaughan+03
        # This the special case of poisson noise light curves.
        # Generally: noise = <e^2>/(mu^2 fq_nyq)
        # where <e^2> is the averaged squared error in the measurements
        # which for poisson case: e = sqrt((mu+bgd)*dt)/dt = sqrt((mu+bgd)/dt)
        # --> <e^2> = (mu+bgd)/dt
        # fq_nyq: is the Nyquist frequency: fq_nyq = 1/(2*dt)
        # ==> for poisson case: noise = 2*(mu+bgd)/mu^2
        # For other normalization, we need to renormalize accordingly
        ## -------------------------- ##
        fnyq = 0.5/dt
        nois = [ff*0+np.mean(re**2)/(fnyq*mu**expo[norm])
                    for ff,re,mu in zip(freq, rerr, mean)]

        # flattern lists #
        _c = np.concatenate
        freq = _c(freq)
        isort = np.argsort(freq)
        freq = freq[isort]
        rpsd = _c(rpsd)[isort]
        nois = _c(nois)[isort]

        return freq, rpsd, nois

    @staticmethod
    def bin_psd(freq, rpsd, fqbin, noise=None, logavg=True):
        """Bin power spectrum.

        Parameters:
            freq: array of frequencies
            rpsd: array of raw powers
            fqbin: binning dict to be passed to @misc.group_array
                to bin the frequency axis
            noise: array of noise.
            logavg: do averaging in log-space, and correct for
                bias. Otherwise it is simple averaging

        Returns:
            fq, psd, psde, desc; with desc having some useful info

        """

        # ensure the arrays are compatible #
        if len(freq) != len(rpsd):
            raise ValueError('freq and rpsd are not compatible')

        if noise is None: noise = np.zeros_like(freq) + 1e-10

        # group the freq array #
        nfq = len(freq)
        idx  = misc.group_array(freq, do_unique=True, **fqbin)
        fqm  = [len(i) for i in idx]
        fqL = [freq[i].min() for i in idx] + [freq[idx[-1].max()]]

        # do the actual binning #
        # the noise is logavg if needed, without bias correction
        if logavg:
            f  = [10**np.mean(np.log10(freq[i])) for i in idx]
            p  = [10**np.mean(np.log10(rpsd[i])) for i in idx]
            n  = [10**np.mean(np.log10(noise[i])) for i in idx]
            pe = [np.log(10)*p[i]*(0.310/fqm[i])**0.5 for i in range(len(p))]
        else:
            f  = [np.mean(freq[i]) for i in idx]
            p  = [np.mean(rpsd[i]) for i in idx]
            n  = [np.mean(noise[i]) for i in idx]
            pe = [p[i]*(1./fqm[i])**0.5 for i in range(len(p))]

        fq, psd, psde, n = np.array(f), np.array(p), np.array(pe), np.array(n)

        # bias correction #
        #####################################
        # From the simulations in test_lcurve.py:
        # 1- Whenever logavg=True is used, bias correciton needs
        #    to be applied. Logavg=True does better, most of the
        #    times, particularly when averaging neighboring frequencies
        # bias function: bias_f(2) ~ 0.253 in Papadakis93
        # bias_f = lambda k: -sp.digamma(k/2.)/np.log(10)
        #####################################
        import scipy.special as sp
        bias_f = lambda k: -sp.digamma(k/2.)/np.log(10)
        bias = np.zeros_like(psd) + bias_f(2)
        if logavg: psd *= 10**bias

        # return #
        desc = {'fqL': fqL, 'fqm':fqm, 'noise':n, 'bias':bias}
        return fq, psd, psde, desc

    @staticmethod
    def calculate_lag(rate, Rate, dt, fqbin=None, **kwargs):
        """Calculate and bin lags from two lists of light curves.

        Parameters:
            rate: array or list of arrays of lcurve rates
            Rate: array or list of arrays of Reference lcurve rates
            dt: time bin width of the light curve
            fqbin: binning dict to be passed to @misc.group_array
                to bin the frequency axis. If None, return raw lag

        Keywords:
            rerr: array or list of errors on rate. If not give,
                assume, poisson noise.
            bgd: array or list of background rates. In this case,
                rate above is assumed background subtracted.
            Rerr: array or list of errors on Rate. If not give,
                assume, poisson noise.
            Bgd: array or list of background rates for the reference.
                In this case, Rate above is assumed background subtracted.
            phase: return phase lag instead of time lag
            taper: apply Hanning tapering before calculating the fft
                see p388 Bendat & Piersol; the fft need to be multiplied
                by sqrt(8/3) to componsate for the reduced variance. Default: False
            norm: how to normalize the fft during the calculations. None|rms|leahy|var.
                Default is None, so the calculations is done with raw numpy fft

        Return:
            freq, lag, lage, desc;
            desc = {'fqm', 'fqL', 'limit', 'Limit'}
        """

        phase = kwargs.get('phase', False)

        # check input #
        if not isinstance(rate[0], (np.ndarray, list)): rate = [rate]
        if not isinstance(Rate[0], (np.ndarray, list)): Rate = [Rate]

        # check that lc and reference are compatible #
        for r1,r2 in zip(rate, Rate):
            if len(r1) != len(r2):
                raise ValueError('rate and Rate are incompatible')

        # rerr and bgd; for estimating noise level #
        bgd  = kwargs.get('bgd', 0.0)
        Bgd  = kwargs.get('Bgd', 0.0)
        if not isinstance(bgd, (np.ndarray, list)):
            bgd = [bgd for r in rate]
        if not isinstance(Bgd, (np.ndarray, list)):
            Bgd = [Bgd for r in Rate]

        rerr = kwargs.get('rerr', None)
        Rerr = kwargs.get('Rerr', None)
        if rerr is None:
            # TODO: this is not always correct!
            rerr = [np.sqrt((r+b)/dt) for r,b in zip(rate, bgd)]
        if Rerr is None:
            # TODO: this is not always correct!
            Rerr = [np.sqrt((r+b)/dt) for r,b in zip(Rate, Bgd)]

        # make sure error arrays are also ready
        if not isinstance(rerr[0], (np.ndarray, list)): rerr = [rerr]
        if not isinstance(Rerr[0], (np.ndarray, list)): Rerr = [Rerr]

        # tapering ? #
        taper = kwargs.get('taper', True)
        taper_factor = 1.0
        if taper:
            rate = [(r-r.mean()) * np.hanning(len(r)) + r.mean() for r in rate]
            Rate = [(r-r.mean()) * np.hanning(len(r)) + r.mean() for r in Rate]
            taper_factor = np.sqrt(8/3)

        # normalization ? #
        norm = kwargs.get('norm', None)
        if not norm in [None, 'rms', 'leahy', 'var']:
            raise ValueError('Unknown norm value. It should be None|rms|leahy|var')
        expo = {'var':0, 'leahy':1, 'rms':2}
        if norm is None:
            normf = lambda r: 1.0
        else:
            normf = lambda r: (2.*dt / (len(r) * np.mean(r)**expo[norm]))**0.5

        # fft; remove the 0-freq and the nyquist #
        rfft = [np.fft.rfft(r)[1:-1]*taper_factor*normf(r) for r in rate]
        Rfft = [np.fft.rfft(r)[1:-1]*taper_factor*normf(r) for r in Rate]
        freq = [np.fft.rfftfreq(len(r), dt)[1:-1] for r in rate]
        crss = [R*np.conj(r) for r,R in zip(rfft, Rfft)]
        rpsd = [np.abs(r)**2 for r in rfft]
        Rpsd = [np.abs(r)**2 for r in Rfft]

        # noise level in psd. See comments in @calculate_psd #
        # noise level is: <e^2>/(mu^2 fq_nyq) for rms norm; then renorm accordingly
        fnyq = 0.5/dt
        nois = [ff*0+(np.mean(re**2)*len(re)*normf(r)**2)/(fnyq*2*dt)
                    for ff,re,r in zip(freq, rerr, rate)]
        Nois = [ff*0+(np.mean(re**2)*len(re)*normf(r)**2)/(fnyq*2*dt)
                    for ff,re,r in zip(freq, Rerr, Rate)]

        # flattern lists #
        _c = np.concatenate
        freq = _c(freq)
        isort = np.argsort(freq)
        freq = freq[isort]
        crss = _c(crss)[isort]
        rpsd = _c(rpsd)[isort]
        Rpsd = _c(Rpsd)[isort]
        nois = _c(nois)[isort]
        Nois = _c(Nois)[isort]

        # do we need just raw lags? #
        if fqbin is None:
            lag = np.angle(crss) / (1. if phase else 2*np.pi*freq)
            return freq, lag

        # bin the lag #
        _a = np.array
        idx = misc.group_array(freq, do_unique=True, **fqbin)
        fqm = _a([len(i) for i in idx])
        fqL = _a([freq[i].min() for i in idx] + [freq[idx[-1].max()]])

        meanf  = lambda a: np.mean(a)
        lmeanf = lambda a: 10**(np.mean(np.log10(a)))

        f  = _a([lmeanf(freq[i]) for i in idx])
        p  = _a([meanf(rpsd[i])  for i in idx])
        P  = _a([meanf(Rpsd[i])  for i in idx])
        n  = _a([meanf(nois[i])  for i in idx])
        N  = _a([meanf(Nois[i])  for i in idx])
        c  = _a([meanf(crss[i])  for i in idx])

        # phase lag and its error #
        # g2 is caluclated without noise subtraciton
        # see paragraph after eq. 17 in Nowak+99
        # see eq. 11, 12 in Uttley+14. Nowak (and Uttley too) clearly
        # states that the noise shouldn't be subtracted)
        lag = np.angle(c)
        n2  = ((p - n)*N + (P - N)*n + n*N) / fqm
        g2  = (np.abs(c)**2) / (p * P)

        # mask out points where coherence is undefined #
        g2  = np.clip(g2, 1e-5, 1.0)
        lag_e = np.clip(np.sqrt((1 - g2) / (2*g2*fqm)), 0, np.pi)

        # coherence gamma_2 #
        # here we subtract the noise; see eq. 8
        # in Vaughan+97 and related definitions
        coh   = (np.abs(c)**2 - n2) / ((p-n) * (P-N))
        coh = np.clip(coh, 1e-5, 1-1e-5)
        dcoh  = (2/fqm)**0.5 * (1 - coh)/np.sqrt(coh)
        coh_e = coh * (fqm**-0.5) * ((2*n2*n2*fqm)/(np.abs(c)**2 - n2)**2 +
                (n**2/(p-n)**2) + (N**2/(P-N)**2) + (fqm*dcoh/coh**2))**0.5

        # rms spectrum from psd; error from eq. 14 in Uttley+14 #
        # the rms here is in absolute not fractional units
        dfq  = fqL[1:]-fqL[:-1]
        mu   = np.mean([np.mean(r) for r in rate])
        Mu   = np.mean([np.mean(r) for r in Rate])
        rms  = mu * (dfq * np.abs(p - n))**0.5
        sigx2  = rms**2
        sigxn2 = dfq * n * mu**2
        rmse = ((2*sigx2*sigxn2 + sigxn2**2) / (2*fqm*sigx2) ) **0.5
        ibad = p<n
        rms[ibad] = 0.0
        rmse[ibad] = np.max(np.concatenate((rms, rmse)))

        # covariance: eq. 13, 15 in Uttley+14 #
        # again in absolute not fractional units #
        cov  = ( (np.abs(c)**2 - n2) * mu * mu * dfq / (P-N) )
        ibad = (cov < 0) | (P<=N)
        cov  = np.abs(cov)**0.5
        sigy2 = dfq * np.abs(P-N) * Mu**2
        sigyn2 = dfq * N * Mu**2
        cove = ((sigyn2*cov**2 + sigy2*sigxn2 + sigxn2*sigyn2) / (2*fqm*sigy2))**0.5
        cov[ibad] = 0.0
        cove[ibad] = np.max(np.concatenate((cov, cove)))

        # limits on lag measurements due to poisson noise #
        # equation 30 in Vaughan+2003 #
        limit = np.sqrt(np.abs(n/(fqm * g2 * (p-n))))
        Limit = np.sqrt(np.abs(N/(fqm * g2 * (P-N))))
        limit = np.clip(limit, -np.pi, np.pi)
        Limit = np.clip(Limit, -np.pi, np.pi)

        # do we need time lag instead of phase lag? #
        if not phase:
            lag   /= (2*np.pi*f)
            lag_e /= (2*np.pi*f)
            limit /= (2*np.pi*f)
            Limit /= (2*np.pi*f)

        # return #
        desc = {'fqL': fqL, 'fqm':fqm, 'limit':limit, 'Limit':Limit,
                'limit_avg':(limit+Limit)/2, 'coh': _a([coh, coh_e]),
                'psd': p, 'nois': n, 'Psd': P, 'Nois': N, 'cxd': c, 'n2': n2,
                'g2': g2, 'idx': idx, 'crss': crss, 'freq': f,
                'rfreq':freq, 'rpsd':rpsd, 'Rpsd':Rpsd, 'rnois':nois, 'RNois': Nois,
                'rms': _a([rms, rmse]), 'cov': _a([cov, cove]),}

        return f, lag, lag_e, desc

    @staticmethod
    def create_segments(Lc, seglen, strict=False, **kwargs):
        """Split an LCurve or a list of them to segments.
        Useful to be used with calculate_psd|lag etc.

        Parameters:
            Lc: an LCurve or a list of them
            seglen: segment length in seconds.
            strict: force all segments to have length length. Some data
                may be discarded

        Keywords:
            uneven: The light curves are uneven, so the splitting produces
                segments that have the same number of points. Default: False
            **other arguments to be passed to az.misc.split_array

        Returns:
            rate, rerr, time, seg_idx
            seg_idx is the indices used to create the segments.

        """
        # Keyworkds
        uneven = kwargs.get('uneven', False)

        if not type(Lc) == list:
            Lc = [Lc]

        # assert the same sampling #
        dt = Lc[0].dt
        for l in Lc:
            if dt != l.dt:
                raise ValueError('There is a difference in the time sampling between light curves')

        # segments details #
        iseglen = np.int(seglen/dt)

        # make sure the LCurve objects are evenly sampled #
        if not uneven:
            Lc = [l.make_even() for l in Lc]

        # split the rate arrays #
        splt = [misc.split_array(l.rate, iseglen, strict, l.rerr, l.time, **kwargs)
                        for l in Lc]

        # flatten the segments into on large list #
        rate = [i for s in splt for i in s[0]]
        rerr = [i for s in splt for i in s[2]]
        time = [i for s in splt for i in s[3]]
        seg_idx = [s[1] for s in splt]
        return rate, rerr, time, seg_idx

    @staticmethod
    def prepare_en_segments(rate, rerr, ibin, iref=None, **kwargs):
        """Create a light curve array at some energy bin(s), and
            a corresponding reference band if needed.

        Parameters:
            rate: a list or array of rate values with shape: nen, nseg, ...
            rerr: the errors array corresponding to rate
            ibin: the bin number of a interest. int or a list of int giving
                the indices of interest
            iref: the bin number or a list of bin numbers to create a secondary
                reference band. The ibin value(s) will be removed from the
                reference light curve if ibin_exclude is true (default).
                -1 means use all available bins (excluding ibin)

        Keywords:
            ibin_exclude: exclude ibin from iref. Default True

        Returns:
            rate, rerr, Rate, Rerr
            each has dims: (nseg, ...). The first 2 corresponds to summing over ibin
            and and the last two are for summing over iref (or None if iref=None). Errors
            are propagated quadratically from the input rerr

        """

        # keywords
        ibin_exclude = kwargs.get('ibin_exclude', True)

        nen = len(rate)

        # make sure we are dealing with lists #
        if not isinstance(ibin, list): ibin = [ibin]

        # the rate and error at the bins of interest #
        r  = np.sum(np.array(rate, dtype=object)[ibin], 0)
        re = np.sum(np.square(np.array(rerr, dtype=object))[ibin], 0)**0.5

        # reference #
        R, Re = [], []
        if not iref is None:
            if not isinstance(iref, list):
                iref = list(range(nen)) if iref == -1 else [iref]
            if ibin_exclude:
                iref = [i for i in iref if not i in ibin]
            R  = np.sum(np.array(rate, dtype=object)[iref], 0)
            Re = np.sum(np.square(np.array(rerr, dtype=object))[iref], 0)**0.5

        return r, re, R, Re

max_stars_repo_name: abduzoghbi/aztools

max_issues_repo_name: abduzoghbi/aztools

max_forks_repo_name: abduzoghbi/aztools

reason: import numpy,import scipy,import astropy


text:
# Copyright 2021 The Distla Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
from jax import lax
import numpy as np

from distla_core.analysis.errors import errors_polar
from distla_core.utils import pops

if __name__ == "__main__":
  local_rows = np.array([128, 256, 512, 1024, 2048])
  OUT_ROWS = tuple(pops.NROWS * local_rows)
  OUT_COLS = (None, )
  DTYPES = (np.float32,)  # dtype of the matrices
  P_SZS = (256,)  # panel size of the SUMMA multiplications
  PRECISIONS = (lax.Precision.DEFAULT, lax.Precision.HIGH,
                lax.Precision.HIGHEST,) # ASIC matmul precision
  SEEDS = (None,)  # Random seed to initialize input; system clock if None.
  SERIAL = (False,) # Whether to run in serial or distributed mode.
  EPS = (None, ) # Convergence threshold.
  MAXITER = (50,) # When to terminate if convergence stagnates.
  S_MIN_EST = (None, -1) # Estimated lowest singular value;
                         # None means machine epsilon; -1 means the true value.
  S_THRESH = (0., 0.1) # When to switch to Newton-Schulz from `rogue` iteration.
  S_MIN = (1E-5, 1E-4, 1E-3, 1E-2, 0.1) # Smallest nonzero singular value.
  S_MAX = (1.0,) # Largest singular value of the input matrix.
  N_ZERO_SVS = (0, 10) # The number of zero singular values in the input matrix.
  SV_DISTRIBUTION = ("linear",) # `linear` or `geometric` distribution of
                                # singular values in the input matrix.
  COMPUTE_RHO = (True,) # Whether to compute the positive-semidefinite factor
                         # in addition to the unitary one.
  BATCH_SIZE = 1 # How many runs to assemblage
  REDUCTION_MODE = "min"  # how to assemblage results
  OUTPUT_DIR_PATH = None  # directory of output; CWD if None
  OUTPUT_NAME = "errors_polar" # output saved to OUTPUT_NAME.csv
  _ = errors_polar.errors_polar(
    OUT_ROWS, out_cols=OUT_COLS, dtypes=DTYPES, p_szs=P_SZS,
    precisions=PRECISIONS, seeds=SEEDS, serial=SERIAL, eps=EPS, maxiter=MAXITER,
    s_min_est=S_MIN_EST, s_thresh=S_THRESH, s_min=S_MIN, s_max=S_MAX,
    n_zero_svs=N_ZERO_SVS, sv_distribution=SV_DISTRIBUTION,
    compute_rho=COMPUTE_RHO, batch_size=BATCH_SIZE,
    reduction_mode=REDUCTION_MODE, output_dir_path=OUTPUT_DIR_PATH,
    output_name=OUTPUT_NAME)

reason: import numpy,from jax