IIabm/model.py

import agentpy as ap
import pandas as pd
import numpy as np
import random
import networkx as nx
from firm import FirmAgent
from product import ProductAgent
from orm import db_session, Result
import platform
import json


class Model(ap.Model):
    def setup(self):
        self.sample = self.p.sample
        self.int_stop_times, self.int_stop_t = 0, None
        self.random = random.Random(self.p.seed)
        self.nprandom = np.random.default_rng(self.p.seed)
        self.int_n_iter = int(self.p.n_iter)

        self.dct_lst_remove_firm_prod = self.p.dct_lst_init_remove_firm_prod

        self.int_n_max_trial = int(self.p.n_max_trial)
        self.flt_netw_pref_cust_n = float(self.p.netw_pref_cust_n)
        self.flt_netw_pref_cust_size = float(self.p.netw_pref_cust_size)
        self.flt_cap_limit = int(self.p.cap_limit)
        self.flt_diff_remove = float(self.p.diff_remove)

        # init graph bom
        G_bom = nx.adjacency_graph(json.loads(self.p.g_bom))
        self.product_network = ap.Network(self, G_bom)

        # init graph firm
        Firm = pd.read_csv("Firm_amended.csv")
        Firm['Code'] = Firm['Code'].astype('string')
        Firm.fillna(0, inplace=True)
        Firm_attr = Firm.loc[:, ["Code", "Name", "Type_Region", "Revenue_Log"]]
        firm_product = []
        for _, row in Firm.loc[:, '1':].iterrows():
            firm_product.append(row[row == 1].index.to_list())
        Firm_attr.loc[:, 'Product_Code'] = firm_product
        Firm_attr.set_index('Code', inplace=True)
        G_Firm = nx.MultiDiGraph()
        G_Firm.add_nodes_from(Firm["Code"])

        firm_labels_dict = {}
        for code in G_Firm.nodes:
            firm_labels_dict[code] = Firm_attr.loc[code].to_dict()
        nx.set_node_attributes(G_Firm, firm_labels_dict)

        # add edge to G_firm according to G_bom
        for node in nx.nodes(G_Firm):
            for product_code in G_Firm.nodes[node]['Product_Code']:
                for succ_product_code in list(G_bom.successors(product_code)):
                    # for each product of a certain firm
                    # get each successor (finished product) of this product
                    # get a list of firm producing this successor
                    lst_succ_firm = Firm['Code'][Firm[succ_product_code] ==
                                                 1].to_list()
                    lst_succ_firm_size_damp = \
                        [G_Firm.nodes[succ_firm]['Revenue_Log'] **
                         self.flt_netw_pref_cust_size
                         for succ_firm in lst_succ_firm
                         ]
                    lst_prob = \
                        [size_damp / sum(lst_succ_firm_size_damp)
                         for size_damp in lst_succ_firm_size_damp
                         ]
                    # select multiple customer
                    # based on relative size of this firm
                    lst_same_prod_firm = Firm['Code'][Firm[product_code] ==
                                                      1].to_list()
                    lst_same_prod_firm_size = [
                        G_Firm.nodes[f]['Revenue_Log']
                        for f in lst_same_prod_firm
                    ]
                    share = G_Firm.nodes[node]['Revenue_Log'] / sum(
                        lst_same_prod_firm_size)
                    # damp share
                    share = share ** self.flt_netw_pref_cust_n
                    n_succ_firm = round(share * len(lst_succ_firm)) if round(
                        share * len(lst_succ_firm)) > 0 else 1  # at least one
                    lst_choose_firm = self.nprandom.choice(lst_succ_firm,
                                                           n_succ_firm,
                                                           p=lst_prob)
                    lst_choose_firm = list(
                        set(lst_choose_firm
                            ))  # nprandom.choice may have duplicates
                    lst_add_edge = [(node, succ_firm, {
                        'Product': product_code
                    }) for succ_firm in lst_choose_firm]
                    G_Firm.add_edges_from(lst_add_edge)

        self.sample.g_firm = json.dumps(nx.adjacency_data(G_Firm))
        self.firm_network = ap.Network(self, G_Firm)

        # init product
        for ag_node, attr in self.product_network.graph.nodes(data=True):
            product = ProductAgent(self, code=ag_node.label, name=attr['Name'])
            self.product_network.add_agents([product], [ag_node])
        self.a_lst_total_products = ap.AgentList(self,
                                                 self.product_network.agents)

        # init firm
        for ag_node, attr in self.firm_network.graph.nodes(data=True):
            firm_agent = FirmAgent(
                self,
                code=ag_node.label,
                name=attr['Name'],
                type_region=attr['Type_Region'],
                revenue_log=attr['Revenue_Log'],
                a_lst_product=self.a_lst_total_products.select([
                    code in attr['Product_Code']
                    for code in self.a_lst_total_products.code
                ]))
            # init extra capacity based on discrete uniform distribution
            for product in firm_agent.a_lst_product:
                firm_agent.dct_prod_capacity[product] = \
                    self.nprandom.integers(firm_agent.revenue_log / 5,
                                           firm_agent.revenue_log / 5 +
                                           self.flt_cap_limit)

            self.firm_network.add_agents([firm_agent], [ag_node])
        self.a_lst_total_firms = ap.AgentList(self, self.firm_network.agents)

        # init dct_list_remove_firm_prod (from string to agent)
        t_dct = {}
        for firm_code, lst_product in self.dct_lst_remove_firm_prod.items():
            firm = self.a_lst_total_firms.select(
                self.a_lst_total_firms.code == firm_code)[0]
            t_dct[firm] = self.a_lst_total_products.select([
                code in lst_product for code in self.a_lst_total_products.code
            ])
        self.dct_lst_remove_firm_prod = t_dct
        self.dct_lst_disrupt_firm_prod = t_dct

        # init output
        self.lst_dct_lst_remove_firm_prod = []
        self.lst_dct_lst_disrupt_firm_prod = []

        # set the initial firm product that are removed
        for firm, a_lst_product in self.dct_lst_remove_firm_prod.items():
            for product in a_lst_product:
                assert product in firm.a_lst_product, \
                    f"product {product.code} not in firm {firm.code}"
                firm.a_lst_product_removed.append(product)

        # draw network
        self.draw_network()

        print(self.a_lst_total_firms)
        print(self.a_lst_total_firms[133])
        print(self.a_lst_total_firms[133].code)
        print(self.firm_network.positions[self.a_lst_total_firms[133]])
        print(
            self.firm_network.graph[self.firm_network.positions[self.a_lst_total_firms[133]]])
        iter_g = iter(
            self.firm_network.graph[self.firm_network.positions[self.a_lst_total_firms[133]]])
        f_t = next(iter_g)
        print(f_t)
        iter_g = iter(self.firm_network.graph[f_t])
        print(next(iter_g))

        print(self.a_lst_total_firms[133].get_firm_network_node())
        print(list(self.dfs_edges_prod(self.firm_network.graph,
              self.a_lst_total_firms[133],
              self.a_lst_total_firms[133].a_lst_product[0])))
        print(self.betweenness_centrality_firm_prod(self.firm_network.graph))
        print()

    def update(self):
        self.a_lst_total_firms.clean_before_time_step()
        # output
        self.lst_dct_lst_remove_firm_prod.append(
            (self.t, self.dct_lst_remove_firm_prod))
        self.lst_dct_lst_disrupt_firm_prod.append(
            (self.t, self.dct_lst_disrupt_firm_prod))

        # stop simulation if reached terminal number of iteration
        if self.t == self.int_n_iter or len(
                self.dct_lst_remove_firm_prod) == 0:
            self.int_stop_times = self.t
            self.stop()

    def step(self):
        # print('\n', '=' * 20, 'step', self.t, '=' * 20)
        # print(
        #     'dct_list_remove_firm_prod', {
        #         key.name: value.code
        #         for key, value in self.dct_lst_remove_firm_prod.items()
        #     })

        # remove_edge_to_cus_and_cus_up_prod
        for firm, a_lst_product in self.dct_lst_remove_firm_prod.items():
            for product in a_lst_product:
                firm.remove_edge_to_cus_remove_cus_up_prod(product)

        for n_trial in range(self.int_n_max_trial):
            # print('=' * 10, 'trial', n_trial, '=' * 10)
            # seek_alt_supply
            # shuffle self.a_lst_total_firms
            self.a_lst_total_firms = self.a_lst_total_firms.shuffle()
            for firm in self.a_lst_total_firms:
                if len(firm.a_lst_up_product_removed) > 0:
                    firm.seek_alt_supply()

            # handle_request
            # shuffle self.a_lst_total_firms
            self.a_lst_total_firms = self.a_lst_total_firms.shuffle()
            for firm in self.a_lst_total_firms:
                if len(firm.dct_request_prod_from_firm) > 0:
                    firm.handle_request()

            # reset dct_request_prod_from_firm
            self.a_lst_total_firms.clean_before_trial()
            # do not use:
            # self.a_lst_total_firms.dct_request_prod_from_firm = {} why?

        # based on a_lst_up_product_removed
        # update a_lst_product_disrupted / a_lst_product_removed
        # update dct_lst_disrupt_firm_prod / dct_lst_remove_firm_prod
        self.dct_lst_remove_firm_prod = {}
        self.dct_lst_disrupt_firm_prod = {}
        for firm in self.a_lst_total_firms:
            if len(firm.a_lst_up_product_removed) > 0:
                # print(firm.name, 'a_lst_up_product_removed', [
                #     product.code for product in firm.a_lst_up_product_removed
                # ])
                for product in firm.a_lst_product:
                    n_up_product_removed = 0
                    for up_product_removed in firm.a_lst_up_product_removed:
                        if product in up_product_removed.a_successors():
                            n_up_product_removed += 1
                    if n_up_product_removed == 0:
                        continue
                    else:
                        # update a_lst_product_disrupted
                        # update dct_lst_disrupt_firm_prod
                        if product not in firm.a_lst_product_disrupted:
                            firm.a_lst_product_disrupted.append(product)
                            if firm in self.dct_lst_disrupt_firm_prod.keys():
                                self.dct_lst_disrupt_firm_prod[firm].append(
                                    product)
                            else:
                                self.dct_lst_disrupt_firm_prod[
                                    firm] = ap.AgentList(
                                        self.model, [product])
                        # update a_lst_product_removed
                        # update dct_list_remove_firm_prod
                        # mark disrupted firm as removed based conditionally
                        lost_percent = n_up_product_removed / len(
                            product.a_predecessors())
                        lst_size = self.a_lst_total_firms.revenue_log
                        lst_size = [firm.revenue_log for firm
                                    in self.a_lst_total_firms
                                    if product in firm.a_lst_product
                                    and product
                                    not in firm.a_lst_product_removed
                                    ]
                        std_size = (firm.revenue_log - min(lst_size) +
                                    1) / (max(lst_size) - min(lst_size) + 1)
                        prod_remove = 1 - std_size * (1 - lost_percent)
                        # damp prod
                        prod_remove = prod_remove ** self.flt_diff_remove
                        if self.nprandom.choice([True, False],
                                                p=[prod_remove,
                                                   1 - prod_remove]):
                            firm.a_lst_product_removed.append(product)
                            if firm in self.dct_lst_remove_firm_prod.keys():
                                self.dct_lst_remove_firm_prod[firm].append(
                                    product)
                            else:
                                self.dct_lst_remove_firm_prod[
                                    firm] = ap.AgentList(
                                        self.model, [product])

        # print(
        #     'dct_list_remove_firm_prod', {
        #         key.name: value.code
        #         for key, value in self.dct_lst_remove_firm_prod.items()
        #     })

    def end(self):
        # print('/' * 20, 'output', '/' * 20)
        # print('dct_list_remove_firm_prod')
        # for t, dct in self.lst_dct_lst_remove_firm_prod:
        #     for firm, a_lst_product in dct.items():
        #         for product in a_lst_product:
        #             print(t, firm.name, product.code)
        # print('dct_lst_disrupt_firm_prod')
        # for t, dct in self.lst_dct_lst_disrupt_firm_prod:
        #     for firm, a_lst_product in dct.items():
        #         for product in a_lst_product:
        #             print(t, firm.name, product.code)

        qry_result = db_session.query(Result).filter_by(s_id=self.sample.id)
        if qry_result.count() == 0:
            lst_result_info = []
            for t, dct in self.lst_dct_lst_disrupt_firm_prod:
                for firm, a_lst_product in dct.items():
                    for product in a_lst_product:
                        db_r = Result(s_id=self.sample.id,
                                      id_firm=firm.code,
                                      id_product=product.code,
                                      ts=t,
                                      is_disrupted=True)
                        lst_result_info.append(db_r)
            db_session.bulk_save_objects(lst_result_info)
            db_session.commit()
            for t, dct in self.lst_dct_lst_remove_firm_prod:
                for firm, a_lst_product in dct.items():
                    for product in a_lst_product:
                        # only firm disrupted can be removed theoretically
                        qry_f_p = db_session.query(Result).filter(
                            Result.s_id == self.sample.id,
                            Result.id_firm == firm.code,
                            Result.id_product == product.code)
                        if qry_f_p.count() == 1:
                            qry_f_p.update({"is_removed": True})
                            db_session.commit()
        self.sample.is_done_flag = 1
        self.sample.computer_name = platform.node()
        self.sample.stop_t = self.int_stop_times
        db_session.commit()

    def draw_network(self):
        import matplotlib.pyplot as plt
        plt.rcParams['font.sans-serif'] = 'SimHei'
        pos = nx.nx_agraph.graphviz_layout(self.firm_network.graph,
                                           prog="twopi",
                                           args="")
        node_label = nx.get_node_attributes(self.firm_network.graph, 'Name')
        node_degree = dict(self.firm_network.graph.out_degree())
        node_label = {
            key: f"{key} {node_label[key]} {node_degree[key]}"
            for key in node_label.keys()
        }
        node_size = list(
            nx.get_node_attributes(self.firm_network.graph,
                                   'Revenue_Log').values())
        node_size = list(map(lambda x: x**2, node_size))
        edge_label = nx.get_edge_attributes(self.firm_network.graph, "Product")
        # multi(di)graphs, the keys are 3-tuples
        edge_label = {(n1, n2): label
                      for (n1, n2, _), label in edge_label.items()}
        plt.figure(figsize=(12, 12), dpi=300)
        nx.draw(self.firm_network.graph,
                pos,
                node_size=node_size,
                labels=node_label,
                font_size=6)
        nx.draw_networkx_edge_labels(self.firm_network.graph,
                                     pos,
                                     edge_label,
                                     font_size=4)
        plt.savefig("network.png")

    @ staticmethod
    def dfs_edges_prod(G, start_firm, start_prod, depth_limit=None):
        assert start_prod in start_firm.a_lst_product, \
            f"product {start_prod.code} not in firm {start_firm.code}"

        visited = set()
        if depth_limit is None:
            depth_limit = len(G)

        start_node = start_firm.get_firm_network_node()
        visited.add(start_firm)
        stack = [(start_firm, start_prod, depth_limit, iter(G[start_node]))]
        while stack:
            parent_firm, parent_prod, depth_now, children_node = stack[-1]
            try:
                child_node = next(children_node)
                child_firm = \
                    ap.AgentIter(start_firm.model, child_node).to_list()[0]
                if child_firm not in visited:
                    for child_prod in child_firm.a_lst_product:
                        if child_prod not in parent_prod.a_successors():
                            continue
                        yield parent_firm.code, parent_prod.code, child_firm.code, child_prod.code
                        visited.add(child_firm)
                        if depth_now > 1:
                            stack.append((child_firm,
                                          child_prod,
                                          depth_now - 1,
                                          iter(G[child_node])))
            except StopIteration:
                stack.pop()

    @ staticmethod
    def betweenness_centrality_firm_prod(G):
        def _single_source_shortest_path_basic(G, s):
            S = []
            P = {}
            for v in G:
                P[v] = []
            sigma = dict.fromkeys(G, 0.0)  # sigma[v]=0 for v in G
            D = {}
            sigma[s] = 1.0
            D[s] = 0
            Q = [s]
            while Q:  # use BFS to find shortest paths
                v = Q.pop(0)
                S.append(v)
                Dv = D[v]
                sigmav = sigma[v]
                for w in G[v]:
                    if w not in D:
                        Q.append(w)
                        D[w] = Dv + 1
                    if D[w] == Dv + 1:  # this is a shortest path, count paths
                        sigma[w] += sigmav
                        P[w].append(v)  # predecessors
            return S, P, sigma

        def _accumulate_basic(betweenness, S, P, sigma, s):
            delta = dict.fromkeys(S, 0)
            while S:
                w = S.pop()
                coeff = (1 + delta[w]) / sigma[w]
                for v in P[w]:
                    delta[v] += sigma[v] * coeff
                if w != s:
                    betweenness[w] += delta[w]
            return betweenness

        betweenness = dict.fromkeys(G, 0.0)  # b[v]=0 for v in G
        nodes = G
        for s in nodes:
            # single source shortest paths
            # use BFS for unweighted g
            S, P, sigma = _single_source_shortest_path_basic(G, s)
            # accumulation without endpoint
            betweenness = _accumulate_basic(betweenness, S, P, sigma, s)

        return betweenness