model firm

firm.py

@@ -2,14 +2,16 @@ import agentpy as ap
 
 
 class FirmAgent(ap.Agent):
-    def setup(self, code, name, type_region, revenue_log):
+    def setup(self, code, name, type_region, revenue_log, list_product,
+              capacity):
         self.firm_network = self.model.firm_network
 
         self.code = code
         self.name = name
         self.type_region = type_region
         self.revenue_log = revenue_log
+        self.list_product = list_product
+        self.capacity = capacity
 
-        self.is_disrupted = False
-    
-    
+        self.dct_product_is_disrupted = dict.fromkeys(list_product, False)
+        self.dct_product_is_removed = dict.fromkeys(list_product, False)

model.py

@@ -7,12 +7,12 @@ from firm import FirmAgent
 sample = 0
 seed = 0
 n_iter = 3
-list_init_remove_firm_code = [0, 2]
+dct_list_init_remove_firm_prod = {0: ['1.4.4'], 2: ['1.1.3']}
 dct_sample_para = {
     'sample': sample,
     'seed': seed,
     'n_iter': n_iter,
-    'list_init_remove_firm_code': list_init_remove_firm_code
+    'dct_list_init_remove_firm_prod': dct_list_init_remove_firm_prod
 }
 
 
@@ -20,6 +20,7 @@ class Model(ap.Model):
     def setup(self):
         self.sample = self.p.sample
         self.nprandom = np.random.default_rng(self.p.seed)
+        self.dct_list_remove_firm_prod = self.p.dct_list_init_remove_firm_prod
         self.int_n_iter = int(self.p.n_iter)
 
         # init graph bom
@@ -28,7 +29,7 @@ class Model(ap.Model):
         BomCateNet = pd.read_csv('BomCateNet.csv', index_col=0)
         BomCateNet.fillna(0, inplace=True)
 
-        G_bom = nx.from_pandas_adjacency(BomCateNet,
+        G_bom = nx.from_pandas_adjacency(BomCateNet.T,
                                          create_using=nx.MultiDiGraph())
 
         bom_labels_dict = {}
@@ -56,18 +57,15 @@ class Model(ap.Model):
         # add edge to G_firm according to G_bom
         for node in nx.nodes(G_Firm):
             # print(node, '-' * 20)
-            list_pred_product_code = []
             for product_code in G_Firm.nodes[node]['Product_Code']:
-                list_pred_product_code += list(
-                    G_bom.predecessors(product_code))
-            list_pred_product_code = list(set(list_pred_product_code))
-            for pred_product_code in list_pred_product_code:
-                # print(pred_product_code)
-                list_pred_firms = Firm.index[Firm[pred_product_code] ==
+                # print(product_code)
+                for succ_product_code in list(G_bom.successors(product_code)):
+                    # print(succ_product_code)
+                    list_succ_firms = Firm.index[Firm[succ_product_code] ==
                                                  1].to_list()
                     list_revenue_log = [
-                    G_Firm.nodes[pred_firm]['Revenue_Log']
-                    for pred_firm in list_pred_firms
+                        G_Firm.nodes[succ_firm]['Revenue_Log']
+                        for succ_firm in list_succ_firms
                     ]
                     list_prob = [
                         (v - min(list_revenue_log) + 1) /
@@ -78,12 +76,11 @@ class Model(ap.Model):
                         self.nprandom.choice([1, 0], p=[prob, 1 - prob])
                         for prob in list_prob
                     ]
-                # print(list(zip(list_pred_firms,list_flag, list_prob)))
-                list_added_edges = [
-                    (node, pred_firm)
-                    for pred_firm, flag in zip(list_pred_firms, list_flag)
-                    if flag == 1
-                ]
+                    # print(list(zip(list_succ_firms,list_flag,list_prob)))
+                    list_added_edges = [(node, succ_firm, {
+                        'Product': product_code
+                    }) for succ_firm, flag in zip(list_succ_firms, list_flag)
+                                        if flag == 1]
                     G_Firm.add_edges_from(list_added_edges)
                     # print('-' * 20)
 
@@ -95,20 +92,30 @@ class Model(ap.Model):
 
         # init firm
         for ag_node, attr in self.firm_network.graph.nodes(data=True):
-            firm_agent = FirmAgent(self,
+            firm_agent = FirmAgent(
+                self,
                 code=attr['Code'],
                 name=attr['Name'],
                 type_region=attr['Type_Region'],
-                                   revenue_log=attr['Revenue_Log'])
+                revenue_log=attr['Revenue_Log'],
+                list_product=attr['Product_Code'],
+                # init capacity as the degree of out edges
+                capacity=self.firm_network.graph.out_degree(ag_node))
             self.firm_network.add_agents([firm_agent], [ag_node])
-        a_list_total_firms = ap.AgentList(self, self.firm_network.agents)
-        # set the initial firm that is removed
-        list_b_is_removed = list(
-            map(lambda x: x in self.p.list_init_remove_firm_code,
-                self.firm_network.agents.code))
-        self.a_list_firms_removed = a_list_total_firms.select(
-            list_b_is_removed)
-        self.a_list_firms_removed.is_removed = True
+        self.a_list_total_firms = ap.AgentList(self, self.firm_network.agents)
+        # print(
+        #     list(
+        #         zip(self.a_list_total_firms.code, self.a_list_total_firms.name,
+        #             self.a_list_total_firms.capacity)))
+
+        # set the initial firm product that are removed
+        for firm_code, list_product in self.dct_list_remove_firm_prod.items():
+            firm = self.a_list_total_firms.select(
+                self.a_list_total_firms.code == firm_code)[0]
+            for product in list_product:
+                assert product in firm.list_product, \
+                    f"product {product} not in firm {firm_code}"
+                firm.dct_product_is_removed[product] = True
 
     def update(self):
         # Update list of unhappy people
@@ -133,19 +140,33 @@ class Model(ap.Model):
                                            prog="twopi",
                                            args="")
         node_label = nx.get_node_attributes(self.firm_network.graph, 'Name')
+        # print(node_label)
+        node_degree = dict(self.firm_network.graph.out_degree())
+        node_label = {
+            key: f"{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")
 
 
 model = Model(dct_sample_para)
 model.setup()
-# model.draw_network()
+model.draw_network()