From 6b87ca7a63fb863fd1e7b5d71ca6235848bf6e76 Mon Sep 17 00:00:00 2001
From: Cricial <2911646453@qq.com>
Date: Sat, 24 Aug 2024 11:32:09 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BB=A3=E7=90=86=E7=B1=BB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 model.py     | 15 ++++++++-------
 scheduler.py | 27 +++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 7 deletions(-)
 create mode 100644 scheduler.py

diff --git a/model.py b/model.py
index 9756931..e12f4f7 100644
--- a/model.py
+++ b/model.py
@@ -9,13 +9,14 @@ from mesa.datacollection import DataCollector
 
 from firm import FirmAgent
 from product import ProductAgent
+from scheduler import CustomScheduler
 
 
 class MyModel(Model):
     def __init__(self, params):
-        # self.num_agents = params['N']
-        # self.grid = MultiGrid(params['width'], params['height'], True)
-        # self.schedule = RandomActivation(self)
+        self.num_agents = params['N']
+        self.grid = MultiGrid(params['width'], params['height'], True)
+        self.schedule = CustomScheduler(self)
 
         # Initialize parameters from `params`
         self.sample = params['sample']
@@ -56,19 +57,19 @@ class MyModel(Model):
         firm = pd.read_csv("input_data/Firm_amended.csv")
         firm['Code'] = firm['Code'].astype('string')
         firm.fillna(0, inplace=True)
-        Firm_attr = firm[["Code", "Type_Region", "Revenue_Log"]]
+        firm_attr = firm[["Code", "Type_Region", "Revenue_Log"]]
         firm_product = []
         for _, row in firm.loc[:, '1':].iterrows():
             firm_product.append(row[row == 1].index.to_list())
-        Firm_attr['Product_Code'] = firm_product
-        Firm_attr.set_index('Code', inplace=True)
+        firm_attr['Product_Code'] = firm_product
+        firm_attr.set_index('Code', inplace=True)
         G_Firm = nx.MultiDiGraph()
         G_Firm.add_nodes_from(firm["Code"])
 
         # Add node attributes
         firm_labels_dict = {}
         for code in G_Firm.nodes:
-            firm_labels_dict[code] = Firm_attr.loc[code].to_dict()
+            firm_labels_dict[code] = firm_attr.loc[code].to_dict()
         nx.set_node_attributes(G_Firm, firm_labels_dict)
 
         # Add edges based on BOM graph
diff --git a/scheduler.py b/scheduler.py
new file mode 100644
index 0000000..4ee2611
--- /dev/null
+++ b/scheduler.py
@@ -0,0 +1,27 @@
+from mesa.time import BaseScheduler
+
+from firm import FirmAgent
+from product import ProductAgent
+
+
+class CustomScheduler(BaseScheduler):
+    def __init__(self, model):
+        super().__init__(model)
+        self.company_agents = []
+        self.product_agents = []
+
+    def add_agent(self, agent):
+        if isinstance(agent, FirmAgent):
+            self.company_agents.append(agent)
+        elif isinstance(agent, ProductAgent):
+            self.product_agents.append(agent)
+        super().add_agent(agent)
+
+    def step(self):
+        # First, activate all company agents
+        for agent in self.company_agents:
+            agent.step()
+
+        # Then, activate all product agents
+        for agent in self.product_agents:
+            agent.step()