text:
from pade.misc.utility import display_message, defer_to_thread
from pade.core.agent import Agent
from pade.acl.aid import AID
from pade.acl.messages import ACLMessage
from pade.behaviours.protocols import FipaRequestProtocol
from pade.behaviours.protocols import FipaContractNetProtocol
from pade.behaviours.protocols import FipaSubscribeProtocol
from pade.drivers.mosaik_driver import MosaikCon

from calc_methods import demand_curve
from grid_optimization import solve_dso_reschedule
from grid_optimization import run_powerflow_in_pandapower, create_grid_in_pandapower
from grid_optimization import create_branch_dict, create_power_flow_matrix
from grid_optimization import calc_grid_restrictions
from grid_optimization import solve_dso_operation

from mosaik_mygrid import model

import numpy as np
import random
import datetime as dt
import h5py
import json
import pickle
from ast import literal_eval
from pyfiglet import Figlet

MOSAIK_MODELS = {
    'api_version': '2.2',
    'models': {
        'DSOAgent': {
            'public': True,
            'params': [],
            'attrs': ['dso_param'],
        },
    },
}

with open('force.json', 'r') as f:
    GRID_ELEMENTS, ENTITY_MODEL = model.load_case(f)

SCHEDULE_TIME = 5 * 60
TIME_INTERVAL = 15 * 60
OPERATION_TIME_DELAY = 5 * 60

class MosaikSim(MosaikCon):

    def __init__(self, agent):
        super(MosaikSim, self).__init__(MOSAIK_MODELS, agent)
        self.entities = list()

    def init(self, sid, eid_prefix, start, step_size):
        self.eid_prefix = eid_prefix
        self.eid = '{}{}'.format(self.eid_prefix, '0')
        self.start = start
        self.step_size = step_size
        return MOSAIK_MODELS

    def create(self, num, model):
        entities_info = list()
        for i in range(num):
            entities_info.append(
                {'eid': '{}.{}'.format(self.sim_id, i), 'type': model, 'rel': []})
        return entities_info

    def step(self, time, inputs):
        """

        """

        # ENVIO DE COMANDOS PARA O BESS AGENT VIA MOSAIK

        from_ = self.eid
        data = {from_: {}}
        for i in range(2, 5):
            bess_ref = 'BESSAgentSim{}-0.BESSAgent_{}'.format(i, i)
            to_ = bess_ref
            data[from_][to_] = {'status': None}
        # yield self.set_data_async(data)

        # COMPORTAMENTOS DE SCHEDULE

        if time == SCHEDULE_TIME:
            message = ACLMessage(ACLMessage.CFP)
            message.set_protocol(ACLMessage.FIPA_CONTRACT_NET_PROTOCOL)
            for bess in self.agent.bess_list:
                message.add_receiver(AID(name='bess{}'.format(bess)))
            message.set_content('DSO_SCHEDULE')
            self.agent.call_later(0.1, self.launch_contract_net_protocol, message, 'SCHEDULE')
            f = Figlet()
            print(f.renderText('DSO to BESS. Schedule CFP'))
            return

        # COMPORTAMENTOS DE OPERAÇÃO
        if ((time - OPERATION_TIME_DELAY) % TIME_INTERVAL) == 0 and time != OPERATION_TIME_DELAY:
            message = ACLMessage(ACLMessage.CFP)
            message.set_protocol(ACLMessage.FIPA_CONTRACT_NET_PROTOCOL)
            for bess in self.agent.bess_list:
                message.add_receiver(AID(name='bess{}'.format(bess)))
            message.set_content('DSO_AUCTION')
            self.agent.call_later(0.1, self.launch_contract_net_protocol, message, 'AUCTION')
            f = Figlet()
            print(f.renderText('DSO to BESS. Auction CFP'))
            return

        return time + self.step_size

    def get_data(self, outputs):
        data = {}
        for eid, attrs in outputs.items():
            data[eid] = {}
            for attr in attrs:
                if attr not in MOSAIK_MODELS['models']['DSOAgent']['attrs']:
                    raise ValueError('Unknown output attribute: {}'.format(attr))
                data[eid][attr] = getattr(self.agent, 'dso_param')
        return data

    def launch_contract_net_protocol(self, message, role):
        display_message(self.agent.aid.name, 'Launching FIPA-ContractNet Protocol...')
        self.agent.auction_finished = False
        if self.agent.comp_cnp_1 is None:
            self.agent.comp_cnp_1 = DSOInitiatorContractNetProtocol(self.agent, message, role)
            self.agent.behaviours.append(self.agent.comp_cnp_1)
            self.agent.comp_cnp_1.on_start()
        else:
            self.agent.comp_cnp_1.message = message
            self.agent.comp_cnp_1.role = role
            self.agent.comp_cnp_1.on_start()

    def step_done_(self, time):
        message = self._create_message(1, self.msg_id_step, self.time + self.step_size + time)
        self.agent.mosaik_connection.transport.write(message)
        self.agent.mosaik_connection.message = None
        self.agent.mosaik_connection.mosaik_msg_id = None
        self.agent.mosaik_connection.await_gen = None

class DSOInitiatorRequestProtocol(FipaRequestProtocol):
    def __init__(self, agent, message, role=None):
        super(DSOInitiatorRequestProtocol, self).__init__(agent=agent,
                                                          message=message,
                                                          is_initiator=True)
        self.role = role

    def handle_inform(self, message):
        config = pickle.loads(literal_eval(message.content))

        if self.role == 'MARKET_SCHEDULE':
            self.agent.total_storage_detailed_load = {str(i): j * 1e3 for i, j in config['total_storage_detailed_load'].items()}
            self.agent.prosumer_storage_detailed_load = {str(i): j * 1e3 for i, j in config['prosumer_storage_detailed_load'].items()}
            self.agent.bess_storage_detailed_load = {str(i): j * 1e3 for i, j in config['bess_storage_detailed_load'].items()}

            self.agent.old_prosumer_storage_detailed_load = {str(i): j * 1e3 for i, j in config['old_prosumer_storage_detailed_load'].items()}
            self.agent.new_detailed_load_dso = {str(i): j * 1e3 for i, j in config['new_detailed_load_dso'].items()}

            config = dict()
            config['prosumer_storage_detailed_load'] = self.agent.prosumer_storage_detailed_load
            config['bess_storage_detailed_load'] = self.agent.bess_storage_detailed_load
            config['base_detailed_load'] = self.agent.detailed_load
            content = pickle.dumps(config)

            self.agent.comp_cnp_2.answer.set_content(content)
            self.agent.send(self.agent.comp_cnp_2.answer)

        elif self.role == 'MARKET_AUCTION_2':
            self.agent.real_time_total_storage_detailed_load = {str(i): j for i, j in config['total_storage_detailed_load'].items()}
            self.agent.real_time_prosumer_storage_detailed_load = {str(i): j for i, j in config['prosumer_storage_detailed_load'].items()}
            self.agent.real_time_bess_storage_detailed_load = {str(i): j for i, j in config['bess_storage_detailed_load'].items()}

            self.agent.real_time_old_prosumer_storage_detailed_load = {str(i): j for i, j in config['old_prosumer_storage_detailed_load'].items()}
            self.agent.real_time_new_detailed_load_dso = config['new_detailed_load_dso']

            config = dict()
            config['prosumer_storage_detailed_load'] = self.agent.real_time_new_detailed_load_dso
            content = pickle.dumps(config)

            self.agent.comp_cnp_2.answer.set_content(content)
            self.agent.send(self.agent.comp_cnp_2.answer)

class DSOInitiatorContractNetProtocol(FipaContractNetProtocol):
    """AuctionClear

    Initial FIPA-ContractNet Behaviour that sends CFP messages
    to ProsumerAgents connected to BESSAgents asking for proposals.
    This behaviour also aggregate the proposals to find the
    clear value for the network state."""

    def __init__(self, agent, message, role):
        super(DSOInitiatorContractNetProtocol, self).__init__(
            agent=agent, message=message, is_initiator=True)
        self.cfp = message
        self.role = role
        self.restriction = dict()

    def handle_all_proposes(self, proposes):
        """
        """

        super(DSOInitiatorContractNetProtocol, self).handle_all_proposes(proposes)

        if self.role == 'AUCTION':
            self._handle_auction_proposes(proposes)
        elif self.role == 'SCHEDULE':
            self._handle_schedule_proposes(proposes)

    def _handle_schedule_proposes(self, proposes):
        display_message(self.agent.aid.name, 'Analyzing proposals...')

        agg_demand = np.zeros(96)
        # dict com chaves sendo nomes dos nós e valores lista
        # com 96 pos. com schedule de carga
        self.agent.detailed_load = dict()
        self.agent.init_prosumer_storage_detailed_load = dict()
        self.agent.init_bess_storage_detailed_load = dict()
        self.agent.prosumer_storage_nodes = dict()
        self.agent.bess_storage_nodes = dict()
        self.agent.config_bess_dso_dict = dict()

        for propose in proposes:
            config = pickle.loads(literal_eval(propose.content))
            agg_demand += config['aggregate_demand']
            self.agent.detailed_load.update(config['detailed_load'])
            self.agent.init_prosumer_storage_detailed_load.update(config['prosumer_storage_detailed_load'])
            self.agent.init_bess_storage_detailed_load.update(config['bess_storage_detailed_load'])
            self.agent.config_bess_dso_dict.update(config['bess_storage_config'])

            bess_id = propose.sender.localname.split('bess')[1]

            self.agent.prosumer_storage_nodes[bess_id] = config['prosumer_storage_nodes']
            self.agent.bess_storage_nodes[bess_id] = config['bess_storage_nodes']

        # chama o metodo de analise do DSO em thread separada para não causar erro no loop do PADE
        defer_to_thread(self.analyse_schedule_grid_restrictions,
                        self.after_schedule_analyse_actions)

    def analyse_schedule_grid_restrictions(self):

        # verify if all the detailed_load and storage_detailed_load values have the same lenght
        aux1 = [len(v) for k, v in self.agent.detailed_load.items()]
        aux2 = [len(v) for k, v in self.agent.init_prosumer_storage_detailed_load.items()]
        if aux1[1:] == aux1[:-1] and aux2[1:] == aux2[:-1] and aux1[0] == aux2[0]:
            time_qtd = aux1[0]
        else:
            raise(Exception('The length of data in detailed_load is not the same!'))

        # adequação dos dicionarios de dados de carga
        # os pontos que não contém informação dos dicionários
        # detailed_load, storage_detailed_load e control_storage_detailed_load
        #  são preenchidos com (time_qtd) valores zero.
        for node in self.agent.grid_data['nodes']:
            self.agent.detailed_load.setdefault(str(node['name']), np.zeros(time_qtd))
            self.agent.init_prosumer_storage_detailed_load.setdefault(str(node['name']), np.zeros(time_qtd))
            self.agent.init_bess_storage_detailed_load.setdefault(str(node['name']), np.zeros(time_qtd))
        # retira o nó '0' pois o Jacobiano não contem o nó referencia do sistema
        self.agent.init_prosumer_storage_detailed_load.pop('0')
        self.agent.init_bess_storage_detailed_load.pop('0')

        # define o dicionário com os lagrangianos para a otimização no caso de
        # transactive control
        self.agent.lagrangian_multiplier = {int(n): np.zeros(time_qtd) for n in self.agent.init_prosumer_storage_detailed_load.keys()}

        # cria o objecto net com as informções necessárias para execução de fluxo
        # de carga no pandapower
        self.agent.pp_net_schedule = create_grid_in_pandapower(self.agent.detailed_load)

        # executa fluxo de carga para cada um dos períodos estabelecidos nos dados
        # contidos nas informações de detailed_load
        (self.agent.v0_pu_list,
         self.agent.jac_inv_21_list) = run_powerflow_in_pandapower(self.agent.detailed_load,
                                                                   self.agent.pp_net_schedule,
                                                                   save=True)
        self.agent.branch_dict = create_branch_dict(self.agent.pp_net_schedule)
        self.agent.power_flow_matrix = create_power_flow_matrix(self.agent.branch_dict, self.agent.pp_net_schedule)

        return None

    def after_schedule_analyse_actions(self, result):
        # TODO: Esse metodo precisa ser refatorado por completo
        # pois nao representa mais o fluxo de informação proposto
        # storage_detailed_load = results[0]
        # new_scheduled_storage_load = results[1]
        # bess_storage_nodes = results[2]

        # self.agent.bess_nodes_schedule = new_scheduled_storage_load

        # answer = ACLMessage(ACLMessage.REJECT_PROPOSAL)
        # answer.set_protocol(ACLMessage.FIPA_CONTRACT_NET_PROTOCOL)

        # for bess_name, bess_nodes in bess_storage_nodes.items():
        #     new_schedules = dict()
        #     for bess_node in bess_nodes:
        #         new_schedules[str(bess_node)] = new_scheduled_storage_load[str(bess_node)]

        #     answer.set_content(pickle.dumps(new_schedules))
        #     answer.add_receiver(AID(name='bess{}'.format(bess_name)))
        #     self.agent.send(answer)

        self.agent.mosaik_sim.step_done_(time=60)
        display_message(self.agent.aid.name, '>>> STEP DONE')

    def _handle_auction_proposes(self, proposes):
        display_message(self.agent.aid.name, 'Analyzing proposals...')

        accepted_proposes_aids = list()
        not_accepted_proposes_aids = list()
        total_load = 0.0
        total_storage = 0.0
        detailed_load = dict()
        detailed_storage = dict()

        for message in proposes:
            display_message(self.agent.aid.name, 'PROPOSE message from {}'.format(message.sender.name))
            accepted_proposes_aids.append(message.sender.name)
            config = pickle.loads(literal_eval(message.content))

            for agent_id, load_value in config['real_time_load_value'].items():
                total_load += load_value
                detailed_load[str(agent_id)] = load_value

            for agent_id, storage_value in config['real_time_storage_value'].items():
                total_storage += storage_value
                detailed_storage[str(agent_id)] = storage_value

        #  tempo referente ao instante de tempo t = t0 + 2c
        time_index = int((self.agent.mosaik_sim.time -
                          OPERATION_TIME_DELAY) / TIME_INTERVAL) + 2
        # variavel com os valores de programação dos dispositivos
        # de armazenamento controláveis, para o instante de tempo
        # time_index
        self.agent.init_real_time_bess_storage_detailed_load = {i: np.array([self.agent.bess_storage_detailed_load[i][time_index]])
            for i in self.agent.bess_storage_detailed_load.keys()}

        # adequação dos dicionarios de dados de carga
        # dos pontos que não contém informação dos dicionários
        # detailed_load, storage_detailed_load e control_storage_detailed_load
        # são preenchidos com (time_qtd) valores zero.
        for node in self.agent.grid_data['nodes']:
            detailed_load.setdefault(str(node['name']), np.array([0.0]))
            detailed_storage.setdefault(str(node['name']), np.array([0.0]))
            self.agent.init_real_time_bess_storage_detailed_load.setdefault(str(node['name']), np.array([0.0]))

        detailed_load.pop('0')
        detailed_storage.pop('0')
        self.agent.init_real_time_bess_storage_detailed_load.pop('0')

        # avalia a diferença o que foi programado e o que será executado
        diff = dict()
        for k in self.agent.detailed_load.keys():
            if k in detailed_load.keys() and k in detailed_storage.keys():
                try:
                    a = self.agent.detailed_load[k][time_index] + \
                        self.agent.prosumer_storage_detailed_load[k][time_index]
                    b = detailed_load[k][0] + \
                        detailed_storage[k][0]
                    diff[k] = (a - b)
                    display_message(self.agent.aid.name,
                                    'PROSUMER {} --> DIFF: {:5.3f}'.format(k, diff[k]))
                except KeyError:
                    print('KeyError, chave nao encontrada no dicionario')

        total_detailed_load_kw = {i: [detailed_load[i][0] +
                                      detailed_storage[i][0] +
                                      self.agent.init_real_time_bess_storage_detailed_load[i][0]]
                                  for i in self.agent.bess_storage_detailed_load.keys()}

        if not self.agent.pp_net_operation:
            self.agent.pp_net_operation = create_grid_in_pandapower(self.agent.detailed_load)

        defer_to_thread(self.analyse_auction_grid_restrictions,
                        self.after_auction_analyse_actions,
                        total_detailed_load_kw,
                        detailed_load,
                        detailed_storage)

    def analyse_auction_grid_restrictions(self,
                                          total_detailed_load_kw,
                                          detailed_load,
                                          detailed_storage):

        (max_current_restriction,
         max_voltage_restriction,
         min_voltage_restriction) = calc_grid_restrictions(self.agent.pp_net_operation,
                                                           total_detailed_load_kw,
                                                           self.agent.power_flow_matrix,
                                                           self.agent.branch_dict)

        self.agent.max_current_restriction = max_current_restriction
        self.agent.max_voltage_restriction = max_voltage_restriction
        self.agent.min_voltage_restriction = min_voltage_restriction

        if (self.agent.max_current_restriction.any() or
            self.agent.max_voltage_restriction.any() or
            self.agent.min_voltage_restriction.any()):

            display_message(self.agent.aid.name, 'RESTRICAO VIOLADA NA OPERACAO')

            self.agent.grid_restriction = True

            detailed_load_plus_storage = {i: detailed_load[i] +
                                          detailed_storage[i]
                                          for i in detailed_load.keys()}

            v0_pu_list, jac_inv_21_list = run_powerflow_in_pandapower(detailed_load_plus_storage,
                                                                      self.agent.pp_net_operation,
                                                                      save=False)

            operation_model = solve_dso_operation(detailed_load_plus_storage,
                                                  self.agent.init_real_time_bess_storage_detailed_load,
                                                  self.agent.bess_storage_nodes,
                                                  v0_pu_list,
                                                  jac_inv_21_list,
                                                  self.agent.power_flow_matrix,
                                                  self.agent.branch_dict,
                                                  self.agent.config_bess_dso_dict)

            if operation_model:
                display_message(self.agent.aid.name, 'ALTERACAO CALCULADA E POSSIVEL')
                self.agent.bess_storage_alt = True

                # envia as novas programações para os dispositivos de
                # armazenamento controláveis
                self.agent.real_time_bess_storage_detailed_load = dict()
                for n in operation_model.set_node:
                    aux1 = list()
                    for t in operation_model.set_time:
                        aux1.append(operation_model.var_control_power[n, t].value * 1e3)
                    self.agent.real_time_bess_storage_detailed_load[str(n)] = np.array(aux1)

                config = dict()
                config['message_type'] = 'STORAGE_SCHEDULE_2'
                config['bess_storage_detailed_load'] = self.agent.real_time_bess_storage_detailed_load
                content = pickle.dumps(config)

                message = ACLMessage(ACLMessage.INFORM)
                message.set_protocol(ACLMessage.FIPA_SUBSCRIBE_PROTOCOL)
                message.set_content(content)

                self.agent.comp_pubsub_1.accept_confirmation_count = 0
                self.agent.comp_pubsub_1.schedule_accept_message = None
                self.agent.comp_pubsub_1.notify(message)

            else:
                display_message(self.agent.aid.name, 'ALTERACAO NAO E POSSIVEL')

                self.agent.real_time_detailed_load = detailed_load
                self.agent.init_real_time_prosumer_storage_detailed_load = detailed_storage

                (self.agent.v0_pu_list,
                 self.agent.jac_inv_21_list) = run_powerflow_in_pandapower(self.agent.real_time_detailed_load,
                                                                          self.agent.pp_net_operation,
                                                                          save=False)

                self.agent.lagrangian_multiplier = {int(n): np.zeros(1)
                                                    for n in self.agent.init_real_time_prosumer_storage_detailed_load.keys()}

                self.agent.bess_storage_alt = False
        else:
            display_message(self.agent.aid.name, 'NAO HA VIOLACAO DE RESTRICOES NA OPERACAO')
            self.agent.grid_restriction = False
            self.agent.bess_storage_alt = False

        return None

    def after_auction_analyse_actions(self, results):

        # answer = ACLMessage(ACLMessage.REJECT_PROPOSAL)
        # answer.set_protocol(ACLMessage.FIPA_CONTRACT_NET_PROTOCOL)
        # answer.add_receiver(AID(name='bess{}'.format(bess_name)))
        # config = dict()
        # content = pickle.dumps(config)
        # answer.set_content(content)
        # self.agent.send(answer)
        if not self.agent.bess_storage_alt:
            self.agent.auction_finished = True
            self.agent.mosaik_sim.step_done()
            display_message(self.agent.aid.name, '>>> STEP DONE')

    def handle_inform(self, message):
        """
        """
        super(DSOInitiatorContractNetProtocol, self).handle_inform(message)

        display_message(self.agent.aid.name, 'INFORM message received')

    def handle_refuse(self, message):
        """
        """
        super(DSOInitiatorContractNetProtocol, self).handle_refuse(message)

        display_message(self.agent.aid.name, 'REFUSE message received')

    def handle_propose(self, message):
        """
        """
        super(DSOInitiatorContractNetProtocol, self).handle_propose(message)

        display_message(self.agent.aid.name, 'PROPOSE message received')

class DSOParticipantContractNetProtocol(FipaContractNetProtocol):
    def __init__(self, agent):
        super(DSOParticipantContractNetProtocol, self).__init__(agent=agent,
                                                                message=None,
                                                                is_initiator=False)
        self.role = None
        self.schedule_accept_message = None

    def handle_cfp(self, message):
        """
        """
        self.agent.call_later(random.uniform(0.1, 0.2), self._handle_cfp, message)

    def _handle_cfp(self, message):
        """
        """
        super(DSOParticipantContractNetProtocol, self).handle_cfp(message)
        self.message = message

        display_message(self.agent.aid.name, 'CFP message received')

        if message.content == 'MARKET_SCHEDULE':
            self._handle_schedule_request(message)
            self.role = 'MARKET_SCHEDULE'
        elif message.content == 'MARKET_AUCTION':
            self._handle_auction_request(message)
            self.role = 'MARKET_AUCTION'
        elif message.content == 'MARKET_AUCTION_2':
            self._handle_auction_2_request(message)
            self.role = 'MARKET_AUCTION_2'

    def _handle_schedule_request(self, message):

        config = dict()
        config['detailed_load'] = self.agent.detailed_load  # ok
        config['storage_detailed_load'] = self.agent.init_prosumer_storage_detailed_load  # ok
        config['control_storage_detailed_load'] = self.agent.init_bess_storage_detailed_load  # ok
        config['bess_storage_nodes'] = self.agent.prosumer_storage_nodes  # ok
        config['control_bess_storage_nodes'] = self.agent.bess_storage_nodes  # ok
        config['v0_pu_list'] = self.agent.v0_pu_list  # ok
        config['jac_inv_21_list'] = self.agent.jac_inv_21_list  # ok
        config['lagrangian_multiplier'] = self.agent.lagrangian_multiplier  # ok
        config['power_flow_matrix'] = self.agent.power_flow_matrix  # ok
        config['branch_dict'] = self.agent.branch_dict  # ok
        config['config_bess_dso_dict'] = self.agent.config_bess_dso_dict  # ok

        message = ACLMessage(ACLMessage.REQUEST)
        message.set_protocol(ACLMessage.FIPA_REQUEST_PROTOCOL)
        message.add_receiver(AID(name='solver'))
        content = pickle.dumps(config)
        message.set_content(content)

        if self.agent.comp_request_1 is None:
            self.agent.comp_request_1 = DSOInitiatorRequestProtocol(self.agent,
                                                                    message,
                                                                    'MARKET_SCHEDULE')
            self.agent.behaviours.append(self.agent.comp_request_1)
            self.agent.comp_request_1.on_start()
        else:
            self.agent.comp_request_1.message = message
            self.agent.comp_request_1.role = 'MARKET_SCHEDULE'
            self.agent.comp_request_1.on_start()

        self.answer = self.message.create_reply()
        self.answer.set_performative(ACLMessage.PROPOSE)

    def _handle_auction_request(self, message):
        answer = self.message.create_reply()
        answer.set_performative(ACLMessage.PROPOSE)

        config = dict()
        config['grid_restriction'] = self.agent.grid_restriction
        config['bess_storage_alt'] = self.agent.bess_storage_alt
        content = pickle.dumps(config)

        answer.set_content(content)
        self.agent.send(answer)

    def _handle_auction_2_request(self, message):
        # self.agent.init_real_time_prosumer_storage_detailed_load.setdefault('0', np.array([0.0]))
        # self.agent.init_real_time_bess_storage_detailed_load.setdefault('0', np.array([0.0]))

        config = dict()
        config['detailed_load'] = self.agent.real_time_detailed_load  # ok
        config['storage_detailed_load'] = self.agent.init_real_time_prosumer_storage_detailed_load  # ok
        config['control_storage_detailed_load'] = self.agent.init_real_time_bess_storage_detailed_load # ok
        config['bess_storage_nodes'] = self.agent.prosumer_storage_nodes  # ok
        config['control_bess_storage_nodes'] = self.agent.bess_storage_nodes  # ok
        config['v0_pu_list'] = self.agent.v0_pu_list  # ok
        config['jac_inv_21_list'] = self.agent.jac_inv_21_list  # ok
        config['lagrangian_multiplier'] = self.agent.lagrangian_multiplier  # ok
        config['power_flow_matrix'] = self.agent.power_flow_matrix  # ok
        config['branch_dict'] = self.agent.branch_dict  # ok
        config['config_bess_dso_dict'] = self.agent.config_bess_dso_dict  # ok

        message = ACLMessage(ACLMessage.REQUEST)
        message.set_protocol(ACLMessage.FIPA_REQUEST_PROTOCOL)
        message.add_receiver(AID(name='solver'))
        content = pickle.dumps(config)
        message.set_content(content)

        if self.agent.comp_request_1 is None:
            self.agent.comp_request_1 = DSOInitiatorRequestProtocol(self.agent,
                                                                    message,
                                                                    'MARKET_AUCTION_2')
            self.agent.behaviours.append(self.agent.comp_request_1)
            self.agent.comp_request_1.on_start()
        else:
            self.agent.comp_request_1.message = message
            self.agent.comp_request_1.role = 'MARKET_AUCTION_2'
            self.agent.comp_request_1.on_start()

        self.answer = self.message.create_reply()
        self.answer.set_performative(ACLMessage.PROPOSE)

    def handle_accept_propose(self, message):
        super(DSOParticipantContractNetProtocol, self).handle_accept_propose(message)

        if self.role == 'MARKET_SCHEDULE':
            self.schedule_accept_message = message

            config = pickle.loads(literal_eval(message.content))
            self.agent.lagrangian_multiplier = config['lagrangian_multiplier']

            config = dict()
            config['message_type'] = 'STORAGE_SCHEDULE'
            config['bess_storage_detailed_load'] = self.agent.bess_storage_detailed_load
            content = pickle.dumps(config)

            message_ = ACLMessage(ACLMessage.INFORM)
            message_.set_protocol(ACLMessage.FIPA_SUBSCRIBE_PROTOCOL)
            message_.set_content(content)

            # configura o comportamento do protocolo FIPA Subscriber
            # para contar as confirmações de mensagens recebidas
            # e enviar a confirmação para o Market Agent
            self.agent.comp_pubsub_1.accept_confirmation_count = 0
            self.agent.comp_pubsub_1.schedule_accept_message = message
            self.agent.comp_pubsub_1.notify(message_)

            defer_to_thread(self.analyse_new_schedule,
                            self.after_analyse_new_schedule)
        elif self.role == 'MARKET_AUCTION_2':
            config = pickle.loads(literal_eval(message.content))
            self.agent.lagrangian_multiplier = config['lagrangian_multiplier']

            config = dict()
            config['message_type'] = 'STORAGE_SCHEDULE_2'
            config['bess_storage_detailed_load'] = self.agent.real_time_bess_storage_detailed_load
            content = pickle.dumps(config)

            message_ = ACLMessage(ACLMessage.INFORM)
            message_.set_protocol(ACLMessage.FIPA_SUBSCRIBE_PROTOCOL)
            message_.set_content(content)

            # configura o comportamento do protocolo FIPA Subscriber
            # para contar as confirmações de mensagens recebidas
            # e enviar a confirmação para o Market Agent
            self.agent.comp_pubsub_1.accept_confirmation_count = 0
            self.agent.comp_pubsub_1.schedule_accept_message = message
            self.agent.comp_pubsub_1.notify(message_)

    def analyse_new_schedule(self):
        # executa fluxo de carga com novo schedule de carga
        new_detailed_load = dict()
        for node in self.agent.detailed_load.keys():
            if node == '0':
                new_detailed_load[node] = self.agent.detailed_load[node]
            else:
                new_detailed_load[node] = (self.agent.detailed_load[node] +
                                           np.array(self.agent.total_storage_detailed_load[node]))

        v0_pu_list, jac_inv_21_list = run_powerflow_in_pandapower(new_detailed_load,
                                                                  self.agent.pp_net_schedule,
                                                                  save=True,
                                                                  file_name='new_power_flow_data')

    def after_analyse_new_schedule(self, results):
        pass

    def handle_reject_propose(self, message):
        super(DSOParticipantContractNetProtocol, self).handle_reject_propose(message)
        if self.role == 'MARKET_SCHEDULE':
            config = pickle.loads(literal_eval(message.content))
            self.agent.lagrangian_multiplier = config['lagrangian_multiplier']
        elif self.role == 'MARKET_AUCTION_2':
            config = pickle.loads(literal_eval(message.content))
            self.agent.lagrangian_multiplier = config['lagrangian_multiplier']

    def handle_inform(self, message):
        pass

class DSOPublisherProtocol(FipaSubscribeProtocol):

    def __init__(self, agent):
        super(DSOPublisherProtocol, self).__init__(agent,
                                                    message=None,
                                                    is_initiator=False)
        self.schedule_accept_message = None
        self.accept_confirmation_qtd = 0
        self.accept_confirmation_count = 0

    def handle_subscribe(self, message):
        self.register(message.sender)
        self.accept_confirmation_qtd = len(self.subscribers)
        display_message(self.agent.aid.name, message.content)
        answear = message.create_reply()
        answear.set_performative(ACLMessage.AGREE)
        answear.set_content('Subscribe message accepted')
        self.agent.send(answear)

    def handle_cancel(self, message):
        self.deregister(self, message.sender)
        display_message(self.agent.aid.name, message.content)

    def notify(self, message):
        super(DSOPublisherProtocol, self).notify(message)

    def handle_inform(self, message):

        config = pickle.loads(literal_eval(message.content))
        if config['message_type'] == 'SUBSCRIBER CONFIRM':
            self.accept_confirmation_count += 1
            display_message(self.agent.aid.name,
                            'S: {}/R: {}'.format(self.accept_confirmation_qtd,
                                                 self.accept_confirmation_count
                                                 )
                            )
            if self.accept_confirmation_count == self.accept_confirmation_qtd:
                if self.schedule_accept_message is not None:
                    reply = self.schedule_accept_message.create_reply()
                    reply.set_performative(ACLMessage.INFORM)
                    config = dict()
                    config['message_type'] = 'DSO CONFIRM'
                    content = pickle.dumps(config)
                    reply.set_content(content)
                    self.agent.send(reply)
                else:
                    self.agent.auction_finished = True
                    self.agent.mosaik_sim.step_done()
                    display_message(self.agent.aid.name, '>>> STEP DONE')

class DSOAgent(Agent):
    def __init__(self, aid, bess_list, mode):
        super(DSOAgent, self).__init__(aid=aid, debug=False, mode=mode)
        self.mosaik_sim = MosaikSim(self)
        self.bess_list = bess_list
        self.dso_param = None
        self.pp_net_schedule = None
        self.pp_net_operation = None
        self.runnning_power_flow = False

        self.grid_restriction = None
        self.bess_storage_alt = None
        self.max_current_restriction = None
        self.max_voltage_restriction = None
        self.min_voltage_restriction = None

        # variaveis para a fase de planejamento
        self.detailed_load = None
        self.total_storage_detailed_load = None
        self.init_prosumer_storage_detailed_load = None
        self.prosumer_storage_detailed_load = None
        self.init_bess_storage_detailed_load = None
        self.bess_storage_detailed_load = None
        self.prosumer_storage_nodes = None
        self.bess_storage_nodes = None
        self.v0_pu_list = None
        self.jac_inv_21_list = None
        self.lagrangian_multiplier = None
        self.power_flow_matrix = None
        self.branch_dict = None
        self.config_bess_dso_dict = None

        # variaveis para a fase de operação

        self.real_time_detailed_load = None
        self.init_real_time_prosumer_storage_detailed_load = None
        self.init_real_time_bess_storage_detailed_load = None

        self.real_time_total_storage_detailed_load = None
        self.real_time_prosumer_storage_detailed_load = None
        self.real_time_bess_storage_detailed_load = None

        self.real_time_old_prosumer_storage_detailed_load = None
        self.real_time_new_detailed_load_dso = None

        # argumentos retornados pelo solver
        # como resltados do processo de otimização
        self.old_prosumer_storage_detailed_load = None
        self.new_detailed_load_dso = None

        with open('force.json', 'r') as f:
            self.grid_data = json.load(f)

        self.comp_cnp_1 = None
        self.comp_cnp_2 = DSOParticipantContractNetProtocol(self)
        self.behaviours.append(self.comp_cnp_2)

        self.comp_pubsub_1 = DSOPublisherProtocol(self)
        self.behaviours.append(self.comp_pubsub_1)

        self.comp_request_1 = None

max_stars_repo_name: yanaspaula/smad-cim

max_issues_repo_name: yanaspaula/smad-cim

max_forks_repo_name: yanaspaula/smad-cim


text:
from collections import deque
import numpy as np

class MovingAverage(object):
    """An implementation of moving average. It has an internal queue of the values.

    Args:
        size (int): The length of the value vector.
        window_size (int): The window size for calculation of the moving average.
    """
    def __init__(self, size=1, window_size=10):
        self.size = size
        self.window_size = window_size

        self.state = {}
        self.state["current_size"] = 0
        self.state["memory"] = np.empty((window_size, size))

    def append(self, item):
        # Because of 'maxlen' the first item will be dropped from the deque after reaching 'maxlen'.
        assert len(item)==self.size, "Item should have the same size as was initialized = (" + str(self.size) + ")."
        self.state["current_size"] += 1
        if self.state["current_size"] > self.window_size:
            self.state["memory"] = np.roll(self.state["memory"], shift=-1, axis=0)
            self.state["current_size"] -= 1
        self.state["memory"][self.state["current_size"]-1,:] = item

    @property
    def data(self):
        return self.state["memory"][0:self.state["current_size"],:]

    @property
    def mean(self):
        return np.mean(self.state["memory"][0:self.state["current_size"],:], axis=0)

    @property
    def std(self):
        return np.std(self.state["memory"][0:self.state["current_size"],:], axis=0)

    @property
    def median(self):
        return np.median(self.state["memory"][0:self.state["current_size"],:], axis=0)

    @property
    def min(self):
        return np.min(self.state["memory"][0:self.state["current_size"],:], axis=0)

    @property
    def max(self):
        return np.max(self.state["memory"][0:self.state["current_size"],:], axis=0)

ext: py

reason: import numpy


text:
import sys
import state
import numpy as np

# Class that represents the sliding puzzle board
class Board:
    # Initializer method
    def __init__(self, size):
        self.size = size
        self.board = self.random_board()
        # Test board below
        # self.board = np.array([[7,1,2], [4,6,5], [3,8,0]])
        self.initial_state = state.State(self.size, self.board)

    # Generate victory state
    def winning_board(self):
        board = np.arange(1, self.size*self.size+1).reshape(self.size, self.size)
        board[-1][-1] = 0
        return board

    # Generate random board
    def random_board(self):
        while True:
            board = self.winning_board().flatten()
            np.random.shuffle(board)
            board = board.reshape(self.size, self.size)
            if self.solvable(board):
                break
        return board

    # Check if board is solvable
    def solvable(self, board):
        flat_board = board.flatten()
        # Case where the size of the board is odd
        if (self.size % 2) != 0:
            if (self.invertion_count(flat_board) % 2) == 0:
                return True
        # Case where the size of the board is even
        else:
            if ((self.invertion_count(flat_board) % 2) == 0) and (not self.even_blank_position(board)):
                return True
            if ((self.invertion_count(flat_board) % 2) != 0) and (self.even_blank_position(board)):
                return True
        return False

    # Get the number of inversions of the current board
    def invertion_count(self, board):
        n = len(board)
        invertion_count = 0
        for i in range(n):
            for j in range(i + 1, n):
                if (board[i] > board[j]):
                    invertion_count += 1
        return invertion_count

    # Check if the blank is in and even row, counting from the bottom
    def even_blank_position(self, board):
        position = self.size - np.where(board == 0)[0]
        if (position % 2) == 0:
            return True
        else:
            return False

    # Get board size
    def get_size(self):
        return self.size

    # Get board
    def get_board(self):
        return self.board

    # Set board
    def set_board(self, board):
        self.board = board

    # Get initial state
    def get_initial_state(self):
        return self.initial_state


text:
import pandas as pd
import numpy as np
import pickle as pkl
import os

# dataset = "music:itunes-amazon"
# dataname = "music"
# train_rate_range = [0.6, 0.1]

# dataset = "citation:dblp-scholar"
# dataname = "citation"
# train_rate_range = [0.6, 0.01]

# dataset = "citeacm:dblp-acm"
# dataname = "citeacm"
# train_rate_range = [0.6, 0.01]

# dataset = "dmusic:itunes-amazon"
# dataname = "dmusic"
# train_rate_range = [0.6, 0.1]

# dataset = "dcitation:dblp-scholar"
# dataname = "dcitation"
# train_rate_range = [0.6, 0.01]

dataset = "dciteacm:dblp-acm"
dataname = "dciteacm"
train_rate_range = [0.6, 0.01]

index = pkl.load(open("{}/{}.pair.pkl".format(dataset, dataname), "rb"))

for drop_rate in range(0, 1):
    big_table_data_frame = pd.read_csv("{}/{}.{}.csv".format(dataset, drop_rate, dataname))

    print(train_rate_range)
    for train_rate in train_rate_range:

        val_rate = 0.2
        test_rate = 1 - train_rate - val_rate

        columns = ["label",] + ["left_" + x for x in big_table_data_frame.columns] + ["right_" + x for x in big_table_data_frame.columns]
        idxs = dict()

        idxs["train"] = index[ : int(train_rate * len(index))]
        idxs["test"] = index[int(train_rate * len(index)) : int((train_rate + test_rate) * len(index))]
        idxs["val"] = index[int((train_rate + test_rate) * len(index)) : ]

        idx_left = dict()
        idx_right = dict()
        labels = dict()

        for split in ["train", "val", "test"]:
            idx_left[split] = [x[0] for x in idxs[split]]
            idx_right[split] = [x[1] for x in idxs[split]]
            labels[split] = [x[2] for x in idxs[split]]

        data_values = big_table_data_frame.values
        left_data_values = dict()
        right_data_values = dict()
        label_values = dict()

        for split in ["train", "val", "test"]:
            left_data_values[split] = data_values[idx_left[split]]
            right_data_values[split] = data_values[idx_right[split]]
            label_values[split] = np.array(labels[split]).reshape((-1, 1)).astype(np.int)

            pair_data = np.concatenate((label_values[split], left_data_values[split], right_data_values[split]), axis=1)
            pair_data_frame = pd.DataFrame(pair_data, columns = columns)

            pair_data_frame.to_csv("{}/{},{},{},{}.csv".format(dataset, drop_rate, train_rate, dataname, split), index = True, index_label = "id")

max_stars_repo_name: iamlockelightning/HIF-KAT

max_issues_repo_name: iamlockelightning/HIF-KAT

max_forks_repo_name: iamlockelightning/HIF-KAT


text:
import numpy as np

def euler2quaternion(phi1, Phi, phi2, P = 1):
    # Input - Euler Angles in Radians, Permutation operator (+- 1)
    # Output - Tuple containing quaternion form
    SIGMA = 0.5*(phi1 + phi2)
    DELTA = 0.5*(phi1 - phi2)
    C = np.cos(Phi/2)
    S = np.sin(Phi/2)

    q0 = C*np.cos(SIGMA)
    q1 = -P*S*np.cos(DELTA)
    q2 = -P*S*np.sin(DELTA)
    q3 = -P*C*np.sin(SIGMA)

    if (q0 < 0):
        return (-q0, -q1, -q2, -q3)
    else:
        return (q0, q1, q2, q3)

def quaternion2euler(q0, q1, q2, q3, P = 1):
    # Input - Quaternion, Permutation operator (+- 1)
    # Output - Tuple of Euler Angles in Radians in Bunge Convention
    q03 = np.square(q0) + np.square(q3)
    q12 = np.square(q1) + np.square(q2)
    CHI = np.sqrt(q03*q12)

    if (CHI == 0 and q12 == 0):
        THETA = (np.arctan2(-2*P*q0*q3, q0^2 - q3^2), 0, 0)
    elif (CHI == 0 and q03 == 0):
        THETA = (np.arctan2(2*q1*q2, q1^2 - q2^2), np.pi, 0)
    elif (CHI != 0):
        THETA01 = (q1*q3 - P*q0*q2) / CHI
        THETA02 = (-P*q0*q1 - q2*q3) / CHI
        THETA11 = 2*CHI
        THETA12 = q03 - q12
        THETA21 = (P*q0*q2 + q1*q3) / CHI
        THETA22 = (q2*q3 - P*q0*q1) / CHI

        THETA = []
        for angle in [(THETA01, THETA02), (THETA11, THETA12), (THETA21, THETA22)]:
            val = np.arctan2(angle[0], angle[1])
            if (val < 0):
                val += 2*np.pi
            THETA.append(val)

        THETA = tuple(THETA)
    return THETA

def euler2orimatrix(phi1, Phi, phi2, P = 1):
    # Input - Euler Angle in Radians, Permutation operator (+- 1)
    # Output - Numpy orientation matrix
    C1 = np.cos(phi1)
    C2 = np.cos(phi2)
    S1 = np.sin(phi1)
    S2 = np.sin(phi2)
    C = np.cos(Phi)
    S = np.sin(Phi)

    E1 = C1*C2 - S1*C*S2
    E2 = S1*C2 - C1*C*S2
    E3 = S*S2
    E4 = -C1*S2 - S1*C*C2
    E5 = -S1*S2 + C1*C*C2
    E6 = S*C2
    E7 = S1*S
    E8 = -C1*S
    E9 = C

    return np.matrix([[E1, E2, E3], [E4, E5, E6] , [E7, E8, E9]])

def euler2axisangle(phi1, Phi, phi2, P = 1):
    # Input - Euler Angles in Radians, Permutation operator (+- 1)
    # Output - Tuple containing (axis1, axis2, axis3, angle)
    T = np.tan(Phi / 2)
    SIGMA = (1/2)*(phi1 + phi2)
    DELTA = (1/2)*(phi1 - phi2)
    TAU = np.sqrt(np.square(T) + np.square(np.sin(SIGMA)))
    OMEGA = 2*np.arctan(TAU / np.cos(SIGMA))
    if (OMEGA > np.pi):
        OMEGA = 2*np.pi - OMEGA

    axis1 = (P/TAU)*T*np.cos(DELTA)
    axis2 = (P/TAU)*T*np.sin(DELTA)
    axis3 = (P/TAU)*T*np.sin(SIGMA)

    return (axis1, axis2, axis3, OMEGA)

def orimatrix2euler(mat, P = 1):
    # Input - Numpy 3x3 Orientation Matrix
    # Output - Tuple of Euler Angles in Radians
    ZETA = (1 / np.sqrt(1 - np.square(mat[2,2])))
    if (mat[2,2] == 1):
        THETA1 = np.arctan2(mat[0,1], mat[0,0])
        THETA2 = (np.pi/2)*(1 - mat[2,2])
        THETA3 = 0
    else:
        THETA1 = np.arctan2(mat[2,0]*ZETA, -(mat[2,1]*ZETA))
        THETA2 = np.arccos(mat[2,2])
        THETA3 = np.arctan2(mat[0,2]*ZETA, mat[1,2]*ZETA)

    return (THETA1, THETA2, THETA3)

def orimatrix2quaternion(mat, P = 1):
    # Input - Numpy 3x3 Orientation Matrix
    # Output - Tuple containing (q0, q1, q2, q3)
    q0 = (1/2)*np.sqrt(1 + mat[0,0] + mat[1,1] + mat[2,2])
    q1 = (P/2)*np.sqrt(1 + mat[0,0] - mat[1,1] - mat[2,2])
    q2 = (P/2)*np.sqrt(1 - mat[0,0] + mat[1,1] - mat[2,2])
    q3 = (P/2)*np.sqrt(1 - mat[0,0] - mat[1,1] + mat[2,2])

    if (mat[2,1] < mat[1,2]):
        q1 = -q1

    if (mat[0,2] < mat[2,0]):
        q2 = -q2

    if (mat[1,0] > mat[0,1]):
        q3 = -q3

    MAGNITUDE = np.sqrt(np.square(q0) + np.square(q1) + np.square(q2) + np.square(q3))

    return (q0/MAGNITUDE, q1/MAGNITUDE, q2/MAGNITUDE, q3/MAGNITUDE)

#TODO: Quaternion misorientation
#TODO: GAM
#TODO: Find Symmetry operator quatonion form

def quaternion2axisangle(q0, q1, q2, q3, P = 1):
    # Input - Four quaternion values
    # Output - Tuple containing a list of three axis values and an angle in radians
    OMEGA = 2*np.arccos(q0)

    if OMEGA == 0:
        return ([q1, q2, q3], np.pi)
    else:
        s = np.sign(q0) / np.sqrt(np.square(q1) + np.square(q2) + np.square(q3))
        return ([s*q1, s*q2, s*q3], OMEGA)


text:
#!/usr/bin/env python

from itertools import combinations

import multiprocessing

import scanpy.api as sc
import matplotlib.pyplot as plt
import numpy as np

from granatum_sdk import Granatum

# import pandas as pd
# import seaborn as sns

def to_percentage(x):
    return "%{:3.2f}".format(x * 100)

def main():
    gn = Granatum()

    adata = gn.ann_data_from_assay(gn.get_import("assay"))
    num_top_comps = gn.get_arg("num_top_comps")

    sc.pp.pca(adata, 20)

    variance_ratios = adata.uns["pca"]["variance_ratio"]
    pc_labels = ["PC{}".format(x + 1) for x in range(len(variance_ratios))]

    plt.figure()
    plt.bar(pc_labels, variance_ratios)
    plt.tight_layout()
    gn.add_current_figure_to_results("Explained variance (ratio) by each Principal Component (PC)", height=350, dpi=75)

    X_pca = adata.obsm["X_pca"]

    for i, j in combinations(range(num_top_comps), 2):
        xlabel = "PC{}".format(i + 1)
        ylabel = "PC{}".format(j + 1)

        plt.figure()
        plt.scatter(X_pca[:, i], X_pca[:, j], s=5000 / adata.shape[0])
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.tight_layout()
        gn.add_current_figure_to_results("PC{} vs. PC{}".format(i + 1, j + 1), dpi=75)

        pca_export = {
            "dimNames": [xlabel, ylabel],
            "coords": {sample_id: X_pca[k, [i, j]].tolist() for k, sample_id in enumerate(adata.obs_names)},
        }
        gn.export(pca_export, "PC{} vs. PC{}".format(i + 1, j + 1), kind="sampleCoords", meta={})

    gn.commit()

if __name__ == "__main__":
    main()

max_stars_repo_name: granatumx/gbox-py

max_issues_repo_name: lanagarmire/granatumx

max_forks_repo_name: lanagarmire/granatumx


text:
#!/usr/bin/env python
from __future__ import division, with_statement
import warnings
import numpy
from contextlib import contextmanager
from .setting import Var, ConstVal
from .calculation import Calculator
from cogent.util import parallel
from cogent.maths.stats.distribution import chdtri
from cogent.maths.optimisers import MaximumEvaluationsReached

__author__ = "Peter Maxwell"
__copyright__ = "Copyright 2007-2012, The Cogent Project"
__credits__ = ["Peter Maxwell", "Gavin Huttley"]
__license__ = "GPL"
__version__ = "1.5.3"
__maintainer__ = "Peter Maxwell"
__email__ = "pm67nz@gmail.com"
__status__ = "Production"

class ScopeError(KeyError):
    pass

class InvalidScopeError(ScopeError):
    """for scopes including an unknown value for a known dimension"""
    pass

class InvalidDimensionError(ScopeError):
    """for scopes including an unknown dimension"""
    pass

class IncompleteScopeError(ScopeError):
    """For underspecified scope when retrieving values"""
    pass

# Can be passed to _LeafDefn.interpretScopes()
class _ExistentialQualifier(object):
    def __init__(self, cats=None):
        self.cats = cats
    def __repr__(self):
        if self.cats is None:
            return self.__class__.__name__
        else:
            return '%s(%s)' % (self.__class__.__name__, self.cats)

class EACH(_ExistentialQualifier):
    independent = True

class ALL(_ExistentialQualifier):
    independent = False

def theOneItemIn(items):
    assert len(items) == 1, items
    return iter(items).next()

def _indexed(values):
    # This is the core of the redundancy elimination, used to group
    # identical calculations.
    # >>> _indexed({'a':1.0, 'b':2.0, 'c':3.0, 'd':1.0, 'e':1.0})
    # ([1.0, 2.0, 3.0], {'a':0, 'b':1, 'c':2, 'd':0, 'e':0})
    uniq = []
    index = {}
    values = values.items()
    values.sort()
    for (key, value) in values:
        if value in uniq:
            u = uniq.index(value)
        else:
            u = len(uniq)
            uniq.append(value)
        index[key] = u
    return uniq, index

def _fmtrow(width, values, maxwidth):
    if (len(dict([(id(v),1) for v in values])) == 1 and
            len(str(values[0])) > width):
        s = str(values[0]).replace('\n', ' ')
        if len(s) > maxwidth:
            s = s[:maxwidth-4] + '...'
    else:
        template = '%%%ss' % width
        s = ''.join([(template % (v,)).replace('\n', ' ')[:width]
                for v in values])
    return s

class Undefined(object):
    # Placeholder for a value that can't be calculated
    # because input 'name' has not been provided.
    def __init__(self, name):
        self.name = name
    def __repr__(self):
        return 'Undef(%s)' % self.name

def nullor(name, f, recycled=False):
    # If None, record as undefined.
    # If undefined, propagate error up.
    # Otherwise, do the calculation.
    def g(*args):
        undef = [x for x in args if isinstance(x, Undefined)]
        if undef:
            return undef[0]
        elif None in args:
            return Undefined(name)
        else:
            if recycled:
                args = (None,) + args
            return f(*args)
    return g

#  Level1:  D E F I N I T I O N S

#  Each ParamDefn supplied a .calc(args) method. Used to define the
#  calculation as a DAG of ParamDefns.

# A _Defn has two phases in its life: pre activation it just has .args,
# post activation (once it becomes part of a parameter controller) it
# holds a dynamic list of scope assignments.
# This means defn.makeParamController() can only be called once.

class _Defn(object):
    name = '?'
    default = None
    user_param = False

    def __init__(self):
        self.clients = []
        self.selection = {}
        self.assignments = {}
        self.activated = False

    def makeName(self, name, extra_label=None):
        if name is None:
            name = self.name
        if extra_label is not None:
            name += extra_label
        return name

    def getDefaultSetting(self):
        return None

    def addClient(self, client):
        assert not self.activated, self.name
        assert not self.assignments, self.assignments
        self.clients.append(client)

    def acrossDimension(self, dimension, cats):
        return [self.selectFromDimension(dimension, cat) for cat in cats]

    def selectFromDimension(self, dimension, cat):
        return SelectFromDimension(self, **{dimension:cat})

    def getRequiredScopes(self, arg_dimensions):
        # A list of scope dictionaries: [{dimension:value},] that this
        # Defn needs from an input Defn with `arg_dimensions`
        if not self.activated:
            assert not self.clients, self.clients
            raise RuntimeError('Value at "%s" step never used' % self.name)
        if self.assignments:
            result = []
            for scope_t in self.assignments:
                sel = {}
                sel.update(self.selection)
                for (d, c) in zip(self.valid_dimensions, scope_t):
                    if d in arg_dimensions:
                        sel[d] = c
                result.append(sel)
        else:
            result = [self.selection]
        return result

    def addScopes(self, scopes):
        assert not self.activated
        for scope in scopes:
            scope_t = [scope.get(d, 'all') for d in self.valid_dimensions]
            scope_t = tuple(scope_t)
            if scope_t not in self.assignments:
                self.assignments[scope_t] = self.getDefaultSetting()

    def outputOrdinalFor(self, scope):
        scope_t = tuple([scope[d] for d in self.valid_dimensions])
        return self.index[scope_t]

    def usedDimensions(self):
        used = []
        for (d, dim) in enumerate(self.valid_dimensions):
            seen = {}
            for (scope_t, i) in self.index.items():
                rest_of_scope = scope_t[:d]+scope_t[d+1:]
                if rest_of_scope in seen:
                    if i != seen[rest_of_scope]:
                        used.append(dim)
                        break
                else:
                    seen[rest_of_scope] = i
        return tuple(used) + self.internal_dimensions

    def _getPosnForScope(self, *args, **scope):
        scope = self.interpretPositionalScopeArgs(*args, **scope)
        posns = set()
        for scope_t in self.interpretScope(**scope):
            posns.add(self.index[scope_t])
        if len(posns) == 0:
            raise InvalidScopeError("no value for %s at %s" % (self.name, scope))
        if len(posns) > 1:
            raise IncompleteScopeError("%s distinct values of %s within %s" %
                    (len(posns), self.name, scope))
        return theOneItemIn(posns)

    def wrapValue(self, value):
        if isinstance(value, Undefined):
            raise ValueError('Input "%s" is not defined' % value.name)
        if getattr(self, 'array_template', None) is not None:
            value = self.array_template.wrap(value)
        return value

    def unwrapValue(self, value):
        if getattr(self, 'array_template', None) is not None:
            value = self.array_template.unwrap(value)
        return value

    def getCurrentValueForScope(self, *args, **scope):
        posn = self._getPosnForScope(*args, **scope)
        return self.wrapValue(self.values[posn])

    def getCurrentSettingForScope(self, *args, **scope):
        posn = self._getPosnForScope(*args, **scope)
        return self.uniq[posn]

    def interpretPositionalScopeArgs(self, *args, **scope):
        # Carefully turn scope args into scope kwargs
        assert len(args) <= len(self.valid_dimensions), args
        for (dimension, arg) in zip(self.valid_dimensions, args):
            assert dimension not in scope, dimension
            scope[dimension] = arg
        return scope

    def interpretScopes(self, independent=None, **kw):
        """A list of the scopes defined by the selecting keyword arguments.

        Keyword arguments should be of the form dimension=settings,
        where settings are a list of categories from that
        dimension, or an instance of EACH or ALL wrapping such a list.

        A missing list, None, or an uninstantiated ALL / EACH class
        is taken to mean the entire dimension.

        If 'independent' (which defaults to self.independent_by_default)
        is true then category lists not wrapped as an EACH or an ALL will
        be treated as an EACH, otherwise as an ALL.

        There will only be one scope in the resulting list unless at least
        one dimension is set to EACH."""

        if independent is None:
            independent = self.independent_by_default

        # interpretScopes is used for assigning, so should specify
        # the scope exactly
        for d in kw:
            if d not in self.valid_dimensions:
                raise InvalidDimensionError(d)

        # Initially ignore EACH, just get a full ungrouped set
        kw2 = {}
        independent_dimensions = []
        for (i,dimension) in enumerate(self.valid_dimensions):
            selection = kw.get(dimension, None)
            if selection in [EACH, ALL]:
                dimension_independent = selection.independent
                selection = None
            elif isinstance(selection, (EACH, ALL)):
                dimension_independent = selection.independent
                selection = selection.cats
            else:
                dimension_independent = independent
            if dimension_independent:
                independent_dimensions.append(i)
            if selection is not None:
                kw2[dimension] = selection
        all = self.interpretScope(**kw2)

        # Group independent scopes
        result = {}
        for scope_t in all:
            key = tuple([scope_t[i] for i in independent_dimensions])
            if key not in result:
                result[key] = set()
            result[key].add(scope_t)
        return result.values()

    def interpretScope(self, **kw):
        """A set of the scope-tuples that match the input dict like
        {dimension:[categories]}"""
        selector = []
        unused = {}
        valid_dimensions = list(self.valid_dimensions)
        for d in kw:
            if d not in valid_dimensions:
                continue
            if kw[d] is None: #i.e.: ALL
                continue
            if isinstance(kw[d], str):
                kw[d] = [kw[d]]
            assert type(kw[d]) in [tuple, list], (d, kw[d])
            assert len(kw[d]), kw[d]
            selector.append((valid_dimensions.index(d), d, kw[d]))
            unused[d] = kw[d][:]

        result = set()
        for scope_t in self.assignments:
            for (i, d, cs) in selector:
                if scope_t[i] not in cs:
                    break
            else:
                result.add(scope_t)
                for (i, d, cs) in selector:
                    if d in unused:
                        if scope_t[i] in unused[d]:
                            unused[d].remove(scope_t[i])
                        if not unused[d]:
                            del unused[d]

        if unused:
            # print unused, self.assignments.keys()
            raise InvalidScopeError(unused)

        return result

    def fillParValueDict(self, result, dimensions, cell_value_lookup):
        """Low level method for extracting values.  Pushes values of this
        particular parameter/defn into the dict tree 'result',
        eg: length_defn.fillParValueDict(['edge']) populates 'result' like
        {'length':{'human':1.0, 'mouse':1.0}}"""

        assert self.name not in result, self.name

        posns = [
                list(self.valid_dimensions).index(d)
                for d in dimensions
                if d in self.valid_dimensions]
        for (scope_t, i) in self.index.items():
            value = cell_value_lookup(self, i)
            value = self.wrapValue(value)
            scope = tuple([scope_t[i] for i in posns])

            (d,key) = (result, self.name)
            for key2 in scope:
                if key not in d: d[key] = {}
                (d, key) = (d[key], key2)

            if key in d and  value != d[key]:
                msg = 'Multiple values for %s' % self.name
                if scope:
                    msg += ' within scope %s' % '/'.join(scope)
                raise IncompleteScopeError(msg)
            d[key] = value

    def _update_from_assignments(self):
        (self.uniq, self.index) = _indexed(self.assignments)

    def _local_repr(self, col_width, max_width):
        body = []
        for (i, arg) in enumerate(self.args):
            row = []
            if isinstance(arg, SelectFromDimension):
                argname = arg.arg.name
                for nums in self.uniq:
                    num = arg.uniq[nums[i]]
                    row.append(theOneItemIn(num))
            else:
                argname = arg.name
                for nums in self.uniq:
                    row.append(nums[i])
            body.append((['', self.name][i==0], argname, row))

        return '\n'.join(
            ['%-10s%-10s%s' % (label1[:9], label2[:9],
                    _fmtrow(col_width+1, settings, max_width))
            for (label1, label2, settings) in body])

    def __repr__(self):
        return '%s(%s x %s)' % (self.__class__.__name__, self.name,
                len(getattr(self, 'cells', [])))

class SelectFromDimension(_Defn):
    """A special kind of Defn used to bridge from Defns where a particular
    dimension is just part of the scope rules to later Defns where each
    value has its own Defn, eg: edges of a tree"""

    name = 'select'
    #params = {}

    def __init__(self, arg, **kw):
        assert not arg.activated, arg.name
        _Defn.__init__(self)
        self.args = (arg,)
        self.arg = arg
        self.valid_dimensions = tuple([
            d for d in arg.valid_dimensions if d not in kw])
        self.selection = kw
        arg.addClient(self)

    def update(self):
        for scope_t in self.assignments:
            scope = dict(zip(self.valid_dimensions, scope_t))
            scope.update(self.selection)
            input_num = self.arg.outputOrdinalFor(scope)
            self.assignments[scope_t] = (input_num,)
        self._update_from_assignments()
        self.values = [self.arg.values[i] for (i,) in self.uniq]

    def makeCells(self, input_soup, variable=None):
        arg = input_soup[id(self.arg)]
        outputs = [arg[input_num] for (input_num,) in self.uniq]
        return ([], outputs)

class _NonLeafDefn(_Defn):
    def __init__(self, *args, **kw):
        _Defn.__init__(self)
        valid_dimensions = []
        for arg in args:
            assert isinstance(arg, _Defn), type(arg)
            assert not arg.activated, arg.name
            for dimension in arg.valid_dimensions:
                if dimension not in valid_dimensions:
                    valid_dimensions.append(dimension)
            #print >>sys.stderr, arg.name, '>', valid_dimensions, '>', self.name
            arg.addClient(self)
        valid_dimensions.sort()
        self.valid_dimensions = tuple(valid_dimensions)
        self.args = args
        if 'name' in kw:
            self.name = kw.pop('name')
        self.setup(**kw)

    def setup(self):
        pass

    def update(self):
        for scope_t in self.assignments:
            scope = dict(zip(self.valid_dimensions, scope_t))
            input_nums = [arg.outputOrdinalFor(scope) for arg in self.args]
            self.assignments[scope_t] = tuple(input_nums)
        self._update_from_assignments()
        calc = self.makeCalcFunction()
        self.values = [nullor(self.name, calc, self.recycling)(*[a.values[i] for (i,a) in zip(u, self.args)]) for u in self.uniq]

class _LeafDefn(_Defn):
    """An input to the calculator, ie: a Defn with no inputs itself.

    This class is incomplete - subclasses provide:
        makeDefaultSetting()
        adaptSetting(setting)
        makeCells(input_soup)"""

    args = ()
    name = None
    name_required = True

    # These can be overriden in a subclass or the constuctor.
    valid_dimensions = ()
    numeric = False

    array_template = None
    internal_dimensions = ()

    def __init__(self, name=None, extra_label=None,
            dimensions=None, independent_by_default=None):
        _Defn.__init__(self)
        if dimensions is not None:
            assert type(dimensions) in [list, tuple], type(dimensions)
            self.valid_dimensions = tuple(dimensions)
        if independent_by_default is not None:
            self.independent_by_default = independent_by_default
        if name is not None:
            self.name = name
        if self.name_required:
            assert isinstance(self.name, basestring), self.name
        if extra_label is not None:
            self.name = self.name + extra_label

    def getDefaultSetting(self):
        if (getattr(self, '_default_setting', None) is None or
                self.independent_by_default):
            self._default_setting = self.makeDefaultSetting()
        return self._default_setting

    def update(self):
        self._update_from_assignments()
        gdv = lambda x:x.getDefaultValue()
        self.values = [nullor(self.name, gdv)(u) for u in self.uniq]

    def assignAll(self, scope_spec=None, value=None,
            lower=None, upper=None, const=None, independent=None):
        settings = []
        if const is None:
            const = self.const_by_default

        for scope in self.interpretScopes(
                independent=independent, **(scope_spec or {})):
            if value is None:
                s_value = self.getMeanCurrentValue(scope)
            else:
                s_value = self.unwrapValue(value)

            if const:
                setting = ConstVal(s_value)
            elif not self.numeric:
                if lower is not None or upper is not None:
                    raise ValueError(
                            "Non-scalar input '%s' doesn't support bounds"
                            % self.name)
                setting = Var((None, s_value, None))
            else:
                (s_lower, s_upper) = self.getCurrentBounds(scope)
                if lower is not None:
                    s_lower = lower
                if upper is not None:
                    s_upper = upper

                if s_lower > s_upper:
                    raise ValueError("Bounds: upper < lower")
                elif (s_lower is not None) and s_value < s_lower:
                    s_value = s_lower
                    warnings.warn("Value of %s increased to keep within bounds"
                                % self.name, stacklevel=3)
                elif (s_upper is not None) and s_value > s_upper:
                    s_value = s_upper
                    warnings.warn("Value of %s decreased to keep within bounds"
                                % self.name, stacklevel=3)
                setting = Var((s_lower, s_value, s_upper))
            self.checkSettingIsValid(setting)
            settings.append((scope, setting))

        for (scope, setting) in settings:
            for scope_t in scope:
                assert scope_t in self.assignments, scope_t
                self.assignments[scope_t] = setting

    def getMeanCurrentValue(self, scope):
        values = [self.assignments[s].getDefaultValue() for s in scope]
        if len(values) == 1:
            s_value = values[0]
        else:
            s_value = sum(values) / len(values)
            for value in values:
                if not numpy.all(value==s_value):
                    warnings.warn("Used mean of %s %s values" %
                            (len(values), self.name), stacklevel=4)
                    break
        return s_value

    def getCurrentBounds(self, scope):
        lowest = highest = None
        for s in scope:
            (lower, init, upper)  = self.assignments[s].getBounds()
            if upper == lower: continue
            if lowest is None or lower < lowest: lowest = lower
            if highest is None or upper > highest: highest = upper

        if lowest is None or highest is None:
            # All current settings are consts so use the class defaults
            (lowest, default, highest) = self.getDefaultSetting().getBounds()
        return (lowest, highest)

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__,
                self._local_repr(col_width=6, max_width=60))

    def _local_repr(self, col_width, max_width):
        template = "%%%s.%sf" % (col_width, (col_width-1)//2)
        assignments = []
        for (i,a) in self.assignments.items():
            if a is None:
                assignments.append('None')
            elif a.is_constant:
                if isinstance(a.value, float):
                    assignments.append(template % a.value)
                else:
                    assignments.append(a.value)
            else:
                assignments.append('Var') # %s' % str(i))
        return '%-20s%s' % (self.name[:19],
                _fmtrow(col_width+1, assignments, max_width))

class ParameterController(object):
    """Holds a set of activated CalculationDefns, including their parameter
    scopes.  Makes calculators on demand."""

    def __init__(self, top_defn):
        # topological sort
        indegree = {id(top_defn):0}
        Q = [top_defn]
        while Q:
            pd = Q.pop(0)
            for arg in pd.args:
                arg_id = id(arg)
                if arg_id in indegree:
                    indegree[arg_id] += 1
                else:
                    indegree[arg_id] = 1
                    Q.append(arg)
        topdown = []
        Q = [top_defn]
        while Q:
            pd = Q.pop(0)
            topdown.append(pd)
            for arg in pd.args:
                arg_id = id(arg)
                indegree[arg_id] -= 1
                if indegree[arg_id] == 0:
                    Q.append(arg)

        # propagate categories downwards
        top_defn.assignments = {}
        for pd in topdown:
            pd.assignments = {}
            for client in pd.clients:
                scopes = client.getRequiredScopes(pd.valid_dimensions)
                # print pd.valid_dimensions, pd.name, '<', scopes, '<', client.name, client.valid_dimensions
                pd.addScopes(scopes)
            if not pd.assignments:
                pd.addScopes([{}])
            pd.activated = True

        self.defns = topdown[::-1]

        self.defn_for = {}
        for defn in self.defns:
            #if not defn.args:
            #assert defn.name not in self.defn_for, defn.name
            if defn.name in self.defn_for:
                self.defn_for[defn.name] = None
                # duplicate
            else:
                self.defn_for[defn.name] = defn

        self._changed = set()
        self._update_suspended = False
        self.updateIntermediateValues(self.defns)
        self.setupParallelContext()

    def getParamNames(self, scalar_only=False):
        """The names of the numerical inputs to the calculation."""
        return [defn.name for defn in self.defns if defn.user_param and
            (defn.numeric or not scalar_only)]

    def getUsedDimensions(self, par_name):
        return self.defn_for[par_name].usedDimensions()

    def getParamValue(self, par_name, *args, **kw):
        """The value for 'par_name'.  Additional arguments specify the scope.
        Despite the name intermediate values can also be retrieved this way."""
        callback = self._makeValueCallback(None, None)
        defn = self.defn_for[par_name]
        posn = defn._getPosnForScope(*args, **kw)
        return callback(defn, posn)

    def getParamInterval(self, par_name, *args, **kw):
        """Confidence interval for 'par_name' found by adjusting the
        single parameter until the final result falls by 'dropoff', which
        can be specified directly or via 'p' as chdtri(1, p).  Additional
        arguments are taken to specify the scope."""
        dropoff = kw.pop('dropoff', None)
        p = kw.pop('p', None)
        if dropoff is None and p is None:
            p = 0.05
        callback = self._makeValueCallback(dropoff, p, kw.pop('xtol', None))
        defn = self.defn_for[par_name]
        posn = defn._getPosnForScope(*args, **kw)
        return callback(defn, posn)

    def getFinalResult(self):
        return self.defns[-1].getCurrentValueForScope()

    def getParamValueDict(self, dimensions, p=None, dropoff=None,
            params=None, xtol=None):
        """A dict tree of parameter values, with parameter names as the
        top level keys, and the various dimensions ('edge', 'bin', etc.)
        supplying lower level keys: edge names, bin names etc.
        If 'p' or 'dropoff' is specified returns chi-square intervals instead
        of simple values."""
        callback = self._makeValueCallback(dropoff, p, xtol)
        if params is None:
            params = self.getParamNames(scalar_only=True)
        result = {}
        for param_name in params:
            ev = self.defn_for[param_name]
            ev.fillParValueDict(result, dimensions, callback)
        return result

    def _makeValueCallback(self, dropoff, p, xtol=None):
        """Make a setting -> value function"""
        if p is not None:
            assert dropoff is None, (p, dropoff)
            dropoff = chdtri(1, p) / 2.0
        if dropoff is None:
            def callback(defn, posn):
                return defn.values[posn]
        else:
            assert dropoff > 0, dropoff
            def callback(defn, posn):
                lc = self.makeCalculator(variable=defn.uniq[posn])
                assert len(lc.opt_pars) == 1, lc.opt_pars
                opt_par = lc.opt_pars[0]
                return lc._getCurrentCellInterval(opt_par, dropoff, xtol)
        return callback

    @contextmanager
    def updatesPostponed(self):
        "Temporarily turn off calculation for faster input setting"
        (old, self._update_suspended) = (self._update_suspended, True)
        yield
        self._update_suspended = old
        self._updateIntermediateValues()

    def updateIntermediateValues(self, changed=None):
        if changed is None:
            changed = self.defns # all
        self._changed.update(id(defn) for defn in changed)
        self._updateIntermediateValues()

    def _updateIntermediateValues(self):
        if self._update_suspended:
            return
        # use topological sort order
        # xxx parallel context check?
        for defn in self.defns:
            if id(defn) in self._changed:
                defn.update()
                for c in defn.clients:
                    self._changed.add(id(c))
        self._changed.clear()

    def assignAll(self, par_name, *args, **kw):
        defn = self.defn_for[par_name]
        if not isinstance(defn, _LeafDefn):
            args = ' and '.join(['"%s"' % a.name for a in defn.args])
            msg = '"%s" is not settable as it is derived from %s.' % (
                    par_name, args)
            raise ValueError(msg)
        defn.assignAll(*args, **kw)
        self.updateIntermediateValues([defn])

    def measureEvalsPerSecond(self, *args, **kw):
        return self.makeCalculator().measureEvalsPerSecond(*args, **kw)

    def setupParallelContext(self, parallel_split=None):
        self.overall_parallel_context = parallel.getContext()
        with parallel.split(parallel_split) as parallel_context:
            parallel_context = parallel_context.getCommunicator()
            self.remaining_parallel_context = parallel.getContext()
            if 'parallel_context' in self.defn_for:
                self.assignAll(
                    'parallel_context', value=parallel_context, const=True)

    def makeCalculator(self, calculatorClass=None, variable=None, **kw):
        cells = []
        input_soup = {}
        for defn in self.defns:
            defn.update()
            (newcells, outputs) = defn.makeCells(input_soup, variable)
            cells.extend(newcells)
            input_soup[id(defn)] = outputs
        if calculatorClass is None:
            calculatorClass = Calculator
        kw['overall_parallel_context'] = self.overall_parallel_context
        kw['remaining_parallel_context'] = self.remaining_parallel_context
        return calculatorClass(cells, input_soup, **kw)

    def updateFromCalculator(self, calc):
        changed = []
        for defn in self.defn_for.values():
            if isinstance(defn, _LeafDefn):
                defn.updateFromCalculator(calc)
                changed.append(defn)
        self.updateIntermediateValues(changed)

    def getNumFreeParams(self):
        return sum(defn.getNumFreeParams() for defn in self.defns if isinstance(defn, _LeafDefn))

    def optimise(self, local=None,
            filename=None, interval=None,
            limit_action='warn',  max_evaluations=None,
            tolerance=1e-6, global_tolerance=1e-1, **kw):
        """Find input values that optimise this function.
        'local' controls the choice of optimiser, the default being to run
        both the global and local optimisers. 'filename' and 'interval'
        control checkpointing.  Unknown keyword arguments get passed on to
        the optimiser(s)."""
        return_calculator = kw.pop('return_calculator', False) # only for debug
        for n in ['local', 'filename', 'interval', 'max_evaluations',
                'tolerance', 'global_tolerance']:
            kw[n] = locals()[n]
        lc = self.makeCalculator()
        try:
            lc.optimise(**kw)
        except MaximumEvaluationsReached, detail:
            evals = detail[0]
            err_msg = 'FORCED EXIT from optimiser after %s evaluations' % evals
            if limit_action == 'ignore':
                pass
            elif limit_action == 'warn':
                warnings.warn(err_msg, stacklevel=2)
            else:
                raise ArithmeticError(err_msg)
        finally:
            self.updateFromCalculator(lc)
        if return_calculator:
            return lc

    def graphviz(self, **kw):
        lc = self.makeCalculator()
        return lc.graphviz(**kw)

max_stars_repo_name: sauloal/cnidaria

max_issues_repo_name: sauloal/cnidaria

max_forks_repo_name: sauloal/cnidaria


text:
import cv2
import os
import  glob
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.utils as utils
import  time #tcw20182159tcw
from torch.autograd import Variable
from torch.utils.data import DataLoader
#from tensorboardX import SummaryWriter
from torch.nn.modules.loss import _Loss #TCW20180913TCW
from models import DudeNet
from dataset_r import ImageDataset #tcw201812041630
from utils import *

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

parser = argparse.ArgumentParser(description="DudeNet")
parser.add_argument("--preprocess", type=bool, default=False, help='run prepare_data or not')
parser.add_argument("--batchSize", type=int, default=128, help="Training batch size")
parser.add_argument("--num_of_layers", type=int, default=17, help="Number of total layers")
parser.add_argument("--epochs", type=int, default=70, help="Number of training epochs")
parser.add_argument("--milestone", type=int, default=30, help="When to decay learning rate; should be less than epochs")
parser.add_argument("--lr", type=float, default=1e-3, help="Initial learning rate")
parser.add_argument("--outf", type=str, default="logs", help='path of log files')
parser.add_argument("--mode", type=str, default="S", help='with known noise level (S) or blind training (B)')
parser.add_argument("--noiseL", type=float, default=15, help='noise level; ignored when mode=B')
parser.add_argument("--val_noiseL", type=float, default=15, help='noise level used on validation set')
parser.add_argument("--test_data",type=str,default='cc',help='test on cc')
'''
parser.add_argument("--clip",type=float,default=0.005,help='Clipping Gradients. Default=0.4') #tcw201809131446tcw
parser.add_argument("--momentum",default=0.9,type='float',help = 'Momentum, Default:0.9') #tcw201809131447tcw
parser.add_argument("--weight-decay","-wd",default=1e-3,type=float,help='Weight decay, Default:1e-4') #tcw20180913347tcw
'''
opt = parser.parse_args()
class sum_squared_error(_Loss):  # PyTorch 0.4.1
    """
    Definition: sum_squared_error = 1/2 * nn.MSELoss(reduction = 'sum')
    The backward is defined as: input-target
    """
    def __init__(self, size_average=None, reduce=None, reduction='sum'):
        super(sum_squared_error, self).__init__(size_average, reduce, reduction)

    def forward(self, input, target):
        # return torch.sum(torch.pow(input-target,2), (0,1,2,3)).div_(2)
        return torch.nn.functional.mse_loss(input, target, size_average=None, reduce=None, reduction='sum').div_(2)

def normalize(data):
   return data/255.

def main():
    # Load dataset
    t1 = time.clock()
    save_dir = opt.outf + 'sigma' + str(opt.noiseL) + '_' + time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) + '/'
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    # Load dataset
    print('Loading dataset ...\n')
    aa = "data"
    train_loader = DataLoader(ImageDataset(aa, noisy_image= "fadnettrain_noisyimages",clean_image="fadnettrain_labelimages"),batch_size =opt.batchSize,shuffle=True,num_workers=4)
    net = DudeNet(channels=3, num_of_layers=opt.num_of_layers)
    criterion = nn.MSELoss(size_average=False)
    device_ids = [0]
    model = nn.DataParallel(net, device_ids=device_ids).cuda()
    criterion.cuda() #tcw201810192202tcw
    optimizer = optim.Adam(model.parameters(), lr=opt.lr)
    psnr_list = [] #201809062254tcw. it is used to save the psnr value of each epoch
    patch_num = len(train_loader)
    for epoch in range(opt.epochs):
        if epoch <= opt.milestone:
            current_lr = opt.lr
        if epoch > opt.milestone and  epoch <=60:
            current_lr  =  opt.lr/10.
        if epoch > 60  and  epoch <=90:
            current_lr = opt.lr/100.
        if epoch > 90:
            current_lr = opt.lr/1000.
        for param_group in optimizer.param_groups:
	    param_group["lr"] = current_lr
	    print('learning rate %f' % current_lr)
        for i, data in enumerate(train_loader):
	    model.train()
	    noisy = data['input'].cuda()
	    clean = data['target'].cuda()
	    noisy = Variable(noisy)
	    clean = Variable(clean)
	    our_predicted_cleanimage = model(noisy)
	    loss =  criterion(our_predicted_cleanimage, clean) / (noisy.size()[0]*2)
	    optimizer.zero_grad() #tcw201809112015tcw
	    loss.backward()
	    optimizer.step()
	    model.eval()
	    our_predicted_cleanimage1 = torch.clamp(model(noisy), 0., 1.)
	    psnr_train = batch_PSNR(our_predicted_cleanimage1, clean, 1.)
	    print("[epoch %d][%d/%d] loss: %.4f PSNR_train: %.4f" % (epoch+1, i+1,patch_num, loss.item(), psnr_train))
        model.eval() #tcw20180915tcw
        model_name = 'model'+ '_' + str(epoch+1) + '.pth' #tcw201809071117tcw
        torch.save(model.state_dict(), os.path.join(save_dir, model_name)) #tcw201809062210tcw
        psnr_val = 0
        files_source = glob.glob(os.path.join('data', opt.test_data, '*_real.png'))
        files_source.sort()
        psnr_test = 0
        for f in files_source:
            Img_test = cv2.imread(f)
            Img_test = torch.tensor(Img_test)
            Img_test = Img_test.permute(2,0,1)
            Img_test = Img_test.numpy()
            a1, a2, a3 = Img_test.shape
            Img_test = np.tile(Img_test,(3,1,1,1))
            Img_test = np.float32(normalize(Img_test))
            ISource = torch.Tensor(Img_test)
            ISource = ISource.cuda()
            Img_test_clean = f[:-9] + '_mean.png'
            Img_test_clean = cv2.imread(Img_test_clean)
            Img_test_clean = torch.tensor(Img_test_clean)
            Img_test_clean = Img_test_clean.permute(2,0,1)
            Img_test_clean = Img_test_clean.numpy()
            a1, a2, a3 = Img_test_clean.shape
            Img_test_clean = np.tile(Img_test_clean,(3,1,1,1))
            Img_test_clean = np.float32(normalize(Img_test_clean))
            Img_test_clean = torch.Tensor(Img_test_clean)
            Img_test_clean  = Img_test_clean .cuda()
            with torch.no_grad(): # this can save much memory
                Out  = torch.clamp(model(ISource), 0., 1.)
            psnr = batch_PSNR(Out, Img_test_clean, 1.)
            psnr_list.append(psnr)
            psnr_test += psnr
            print("%s PSNR %f" % (f, psnr))
        psnr_test /= len(files_source)
        print("\nPSNR on test data %f" % psnr_test)
    filename = save_dir + 'psnr.txt' #tcw201809071117tcw
    f = open(filename,'w') #201809071117tcw
    for line in psnr_list:  #201809071117tcw
        f.write(line+'\n') #2018090711117tcw
    f.close()
    t2 = time.clock()
    print t2-t1
if __name__ == "__main__":
    if opt.preprocess:
        main()

max_stars_repo_name: hellloxiaotian/DudeNet

max_issues_repo_name: goldfishxiao/DudeNet

max_forks_repo_name: goldfishxiao/DudeNet


text:
from __future__ import print_function
import setuptools
from distutils.core import setup
from distutils.extension import Extension
import numpy as np
from distutils.ccompiler import new_compiler
import os
import sys
import tempfile
"""
Check for OpenMP based on
https://github.com/MDAnalysis/mdanalysis/tree/develop/package/setup.py
retrieved 06/15/15
"""
def get_requires():
    reqs = []
    for line in open('requirements.txt', 'r').readlines():
        reqs.append(line)
    return reqs
def hasfunction(cc, funcname, include=None, extra_postargs=None):
	# From http://stackoverflow.com/questions/
	#            7018879/disabling-output-when-compiling-with-distutils
	tmpdir = tempfile.mkdtemp(prefix='hasfunction-')
	devnull = oldstderr = None
	try:
		try:
			fname = os.path.join(tmpdir, 'funcname.c')
			f = open(fname, 'w')
			if include is not None:
				f.write('#include %s\n' % include)
			f.write('int main(void) {\n')
			f.write('    %s;\n' % funcname)
			f.write('}\n')
			f.close()
			# Redirect stderr to /dev/null to hide any error messages
			# from the compiler.
			# This will have to be changed if we ever have to check
			# for a function on Windows.
			devnull = open('/dev/null', 'w')
			oldstderr = os.dup(sys.stderr.fileno())
			os.dup2(devnull.fileno(), sys.stderr.fileno())
			objects = cc.compile([fname], output_dir=tmpdir, extra_postargs=extra_postargs)
			cc.link_executable(objects, os.path.join(tmpdir, "a.out"))
		except Exception as e:
			return False
		return True
	finally:
		if oldstderr is not None:
			os.dup2(oldstderr, sys.stderr.fileno())
		if devnull is not None:
			devnull.close()
"""
Check for OpenMP based on
https://github.com/MDAnalysis/mdanalysis/tree/develop/package/setup.py
"""
def detect_openmp():
	"""Does this compiler support OpenMP parallelization?"""
	compiler = new_compiler()
	print("Checking for OpenMP support... ")
	hasopenmp = hasfunction(compiler, 'omp_get_num_threads()')
	needs_gomp = hasopenmp
	if not hasopenmp:
		compiler.add_library('gomp')
	hasopenmp = hasfunction(compiler, 'omp_get_num_threads()')
	needs_gomp = hasopenmp
	if hasopenmp: print("Compiler supports OpenMP")
	else: print( "Did not detect OpenMP support.")
	return hasopenmp, needs_gomp
has_openmp, needs_gomp = detect_openmp()
parallel_args = ['-fopenmp', '-std=c99'] if has_openmp else ['-std=c99']
parallel_libraries = ['gomp'] if needs_gomp else []
def main():
    setup(name="lofti_gaia",
          version="2.0.3",
          description="Orbit fitting with Gaia astrometry",
          author="Logan Pearce",
		  url='https://github.com/logan-pearce/lofti_gaia',
          download_url='https://github.com/logan-pearce/lofti_gaia/archive/2.0.0.tar.gz',
		  packages =['lofti_gaia'],
		  zip_safe=False,
		  classifiers=[
			  "Development Status :: 5 - Production/Stable",
			  "Intended Audience :: Science/Research",
			  "License :: OSI Approved :: MIT License",
			  "Programming Language :: Python :: 3 :: Only",
			  "Programming Language :: C"
		  ],
          author_email="loganpearce1@email.arizona.edu",
		  licence="MIT",
          include_dirs = [np.get_include()],
          install_requires = get_requires(),
          ext_modules=[Extension("lofti_gaia.cFunctions", ["c_src/cFunctions.c"],extra_compile_args = parallel_args,libraries = parallel_libraries)])

if __name__ == "__main__":
    main()


text:
"""
Helper functions for the tests
"""
import os

import numpy as np

from msl.io import read

def read_sample(filename, **kwargs):
    """Read a file in the 'samples' directory.

    Parameters
    ----------
    filename : str
        The name of the file in the samples/ directory

    Returns
    -------
    A root object
    """
    return read(os.path.join(os.path.dirname(__file__), 'samples', filename), **kwargs)

def metadata_equal(m1, m2):
    """Assert that two Metadata objects are equal."""
    assert len(m1) == len(m2)
    for k1, v1 in m1.items():
        v2 = m2[k1]
        if isinstance(v1, (list, tuple, np.ndarray)):
            assert np.array_equal(v1, v2), '{}\n{}'.format(v1, v2)
        else:
            assert v1 == v2, '{} != {}'.format(v1, v2)
    return True

def datasets_equal(d1, d2):
    """Assert that two Dataset objects are equal."""
    assert d1.name == d2.name, '{} != {}'.format(d1.name, d2.name)
    assert np.array_equal(d1.data, d2.data), '{}\n{}'.format(d1.data, d2.data)
    assert metadata_equal(d1.metadata, d2.metadata)
    return True

def roots_equal(r1, r2):
    """Assert that two Root objects are equal."""
    assert metadata_equal(r1.metadata, r2.metadata)

    groups1 = list(r1.groups())
    groups1.sort(key=lambda x: x.name)
    groups2 = list(r2.groups())
    groups2.sort(key=lambda x: x.name)
    assert len(groups1) == len(groups2)
    for g1, g2 in zip(groups1, groups2):
        assert g1.name == g2.name, '{} != {}'.format(g1.name, g2.name)
        assert metadata_equal(g1.metadata, g2.metadata)

    datasets1 = list(r1.datasets())
    datasets1.sort(key=lambda x: x.name)
    datasets2 = list(r2.datasets())
    datasets2.sort(key=lambda x: x.name)
    assert len(datasets1) == len(datasets2)
    for d1, d2 in zip(datasets1, datasets2):
        assert datasets_equal(d1, d2)

    return True

max_stars_repo_name: MSLNZ/MSL-IO

max_issues_repo_name: MSLNZ/MSL-IO

max_forks_repo_name: MSLNZ/MSL-IO


text:
# ===============================================================================
# Copyright 2016 Jake Ross
#
# 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.
# ===============================================================================

# ============= enthought library imports =======================
from traits.api import Bool, Any, Float, \
    Button, Instance, List
# ============= standard library imports ========================
from threading import Thread, Event
from numpy import hstack, array
import time
import os
# ============= local library imports  ==========================
from pychron.graph.graph import Graph
from pychron.paths import paths
from pychron.persistence_loggable import PersistenceLoggable

class BaseScanner(PersistenceLoggable):
    spectrometer = Any

    step = Float(0.1)

    clear_graph_button = Button
    start_scanner = Button
    stop_scanner = Button
    new_scanner = Button('New Magnet Scan')
    start_scanner_enabled = Bool
    stop_scanner_enabled = Bool
    new_scanner_enabled = Bool(True)

    graph = Instance(Graph)
    plotid = 0

    _cancel_event = None

    def stop(self):
        """
        Cancel the current scan

        :return:
        """
        if self._cancel_event:
            self._cancel_event.set()

    def scan(self):
        """
        Start a scan thread. calls ``Scanner._scan`` on a new thread

        :return:
        """
        # if self.plotid>0:
        # self.graph.new_series()
        self.set_plot_visiblity()

        self._cancel_event = Event()

        t = Thread(target=self._scan)
        t.start()

    def set_plot_visibility(self):
        pass

    def reset(self):
        self._clear_graph_button_fired()
        self._reset_hook()

    # private
    def _scan(self):
        graph = self.graph

        plot = graph.plots[0]
        try:
            line = plot.plots['plot{}'.format(self.plotid)][0]
        except KeyError:
            line, _ = graph.new_series()
        # xs = line.index.get_data()
        # ys = line.index.get_data()

        spec = self.spectrometer
        magnet = spec.magnet

        period = spec.integration_time - magnet.settling_time
        # period = 0.1
        st = time.time()

        limits = self._get_limits()
        graph.set_x_limits(*limits, pad='0.1')
        refdet = self.spectrometer.reference_detector

        for i, si in enumerate(self._calculate_steps(*limits)):
            if self._cancel_event.is_set():
                self.debug('exiting scan. dac={}'.format(si))
                break

            self._do_step(magnet, si)
            time.sleep(period)
            if i == 0:
                time.sleep(3)

            ks, ss, t = spec.get_intensities()

            refsig = refdet.intensity
            refk = '{}y{}'.format(refdet, self.plotid)
            rys = plot.data.get_data(refk)
            if i == 0:
                rys = array([refsig])
                xs = array([si])
            else:
                rys = hstack((rys, refsig))
                xs = hstack((xs, si))

            plot.data.update_data({'x{}'.format(self.plotid): xs})
            plot.data.set_data(refk, rys)

            ref_mi, ref_ma = mi, ma = rys.min(), rys.max()
            ref_r = rys.max() - ref_mi
            for det, sig in zip(ks, ss):
                if det == refdet.name:
                    continue

                oys = None
                k = 'odata{}_{}'.format(i, self.plotid)
                if hasattr(plot, k):
                    oys = getattr(plot, k)

                oys = array([sig]) if oys is None else hstack((oys, sig))
                setattr(plot, k, oys)

                mir = oys.min()
                r = oys.max() - mir
                oys = (oys - mir) * ref_r / r + ref_mi

                plot.data.update_data({'{}y{}'.format(det, self.plotid): oys})
                det = self.spectrometer.get_detector(det)
                if det.active:
                    mi, ma = min(mi, min(oys)), max(ma, max(oys))

            self.graph.set_y_limits(min_=mi, max_=ma, pad='0.05',
                                    pad_style='upper')

        self.plotid += 1
        self.debug('duration={:0.3f}'.format(time.time() - st))
        self.new_scanner_enabled = True
        self.start_scanner_enabled = True
        self.stop_scanner_enabled = False

    def _do_step(self, magnet, step):
        raise NotImplementedError

    def _calculate_steps(self):
        raise NotImplementedError

    def _reset_hook(self):
        pass

    # persistence
    @property
    def persistence_path(self):
        return os.path.join(paths.hidden_dir, self.__class__.__name__.lower())

    def _graph_factory(self):
        g = Graph()
        p = g.new_plot(padding_top=30, padding_right=10)
        self._setup_graph(g, p)
        return g

    def _setup_graph(self, g, p):
        pass

    def _graph_default(self):
        g = self._graph_factory()
        return g

    # handlers
    def _start_scanner_fired(self):
        print('start scanner')
        self.info('starting scanner')
        self.new_scanner_enabled = False
        self.start_scanner_enabled = False
        self.stop_scanner_enabled = True
        self.scan()

    def _stop_scanner_fired(self):
        self.stop()
        self.stop_scanner_enabled = False
        self.new_scanner_enabled = True

    def _new_scanner_fired(self):
        print('new scanner')
        self.info('new scanner')
        self.new_scanner_enabled = False
        self.start_scanner_enabled = True

    def _clear_graph_button_fired(self):
        self.graph.clear_plots()
        self.plotid = 0
        self.graph.redraw()

# ============= EOF =============================================

max_stars_repo_name: WiscAr/pychron

max_issues_repo_name: WiscAr/pychron

max_forks_repo_name: UManPychron/pychron

reason: from numpy