[feat] 광역의회 다양성 지수 계산, DB에 저장

pingpingy1 · Nov 25, 2023 · e77a1a3 · e77a1a3
1 parent 828fac7
commit e77a1a3
Showing 1 changed file with 130 additions and 6 deletions.
diff --git a/analysis/diversity_db.py b/analysis/diversity_db.py
@@ -5,6 +5,11 @@
 from db.client import client
 
 
+# ====================================
+#     Diversity index calculations
+# ====================================
+
+
 def count(data, stair=0):
     """
     Returns a counter object of the data, while stairing them to appropriate bins if stair > 0
@@ -50,7 +55,12 @@ def shannon(data, stair=0, opts=True):
     return sh_idx
 
 
-def save_to_mongo(localId: int, factor: str, stair=0, opts=False) -> None:
+# ====================================
+#  Local council diversity statistics
+# ====================================
+
+
+def save_to_mongo_local(localId: int, factor: str, stair=0, opts=False) -> None:
     factor_field = {"age": "age", "gender": "gender", "party": "jdName"}
     data = [
         councilor[factor_field[factor]]
@@ -68,6 +78,7 @@ def save_to_mongo(localId: int, factor: str, stair=0, opts=False) -> None:
 def calculate_rank_local(factor: str) -> None:
     result = client["stats"]["diversity_index"].aggregate(
         [
+            {"$match": {"localId": {"$ne": None}}},
             {"$sort": {f"{factor}DiversityIndex": -1}},
             {"$group": {"_id": "", "items": {"$push": "$$ROOT"}}},
             {"$unwind": {"path": "$items", "includeArrayIndex": "items.rank"}},
@@ -82,7 +93,7 @@ def calculate_rank_local(factor: str) -> None:
         )
 
 
-def calculate_age_diversity_rank_history() -> None:
+def calculate_age_diversity_rank_history_local() -> None:
     for councilor_type in ["elected", "candidate"]:
         for localId in range(1, 227):
             docs = client["stats"]["age_hist"].find(
@@ -146,12 +157,125 @@ def calculate_age_diversity_rank_history() -> None:
                 )
 
 
+# ====================================
+#  Metro council diversity statistics
+# ====================================
+
+
+def save_to_mongo_metro(metroId: int, factor: str, stair=0, opts=False) -> None:
+    factor_field = {"age": "age", "gender": "gender", "party": "jdName"}
+    data = [
+        councilor[factor_field[factor]]
+        for councilor in client["council"]["metro_councilor"].find({"metroId": metroId})
+    ]
+    # print(f"{metroId} {factor}")
+    # print(data)
+    client["stats"].get_collection("diversity_index").update_one(
+        {"metroId": metroId},
+        {"$set": {f"{factor}DiversityIndex": gini_simpson(data, stair, opts)}},
+        upsert=True,
+    )
+
+
+def calculate_rank_metro(factor: str) -> None:
+    result = client["stats"]["diversity_index"].aggregate(
+        [
+            {"$match": {"metroId": {"$ne": None}}},
+            {"$sort": {f"{factor}DiversityIndex": -1}},
+            {"$group": {"_id": "", "items": {"$push": "$$ROOT"}}},
+            {"$unwind": {"path": "$items", "includeArrayIndex": "items.rank"}},
+            {"$replaceRoot": {"newRoot": "$items"}},
+            {"$addFields": {"rank": {"$add": ["$rank", 1]}}},
+        ]
+    )
+    for doc in result:
+        client["stats"]["diversity_index"].find_one_and_update(
+            {"metroId": doc["metroId"]},
+            {"$set": {f"{factor}DiversityRank": int(doc["rank"])}},
+        )
+
+
+def calculate_age_diversity_rank_history_metro() -> None:
+    for councilor_type in ["elected", "candidate"]:
+        for metroId in range(1, 18):
+            docs = client["stats"]["age_hist"].find(
+                {
+                    "councilorType": councilor_type,
+                    "method": "equal",
+                    "level": 1,
+                    "metroId": metroId,
+                }
+            )
+            for doc in docs:
+                diversity_index = gini_simpson(
+                    [
+                        group["minAge"]
+                        for group in doc["data"]
+                        for _ in range(group["count"])
+                    ],
+                    stair=10,
+                )
+                client["stats"]["age_hist"].find_one_and_update(
+                    {
+                        "councilorType": councilor_type,
+                        "method": "equal",
+                        "level": 1,
+                        "metroId": metroId,
+                        "year": doc["year"],
+                    },
+                    {"$set": {"diversityIndex": diversity_index}},
+                )
+
+        years = list({doc["year"] for doc in client["stats"]["age_hist"].find()})
+
+        for year in years:
+            result = client["stats"]["age_hist"].aggregate(
+                [
+                    {
+                        "$match": {
+                            "councilorType": councilor_type,
+                            "method": "equal",
+                            "level": 1,
+                            "year": year,
+                        }
+                    },
+                    {"$sort": {"diversityIndex": -1}},
+                    {"$group": {"_id": "", "items": {"$push": "$$ROOT"}}},
+                    {"$unwind": {"path": "$items", "includeArrayIndex": "items.rank"}},
+                    {"$replaceRoot": {"newRoot": "$items"}},
+                    {"$addFields": {"rank": {"$add": ["$rank", 1]}}},
+                ]
+            )
+            for doc in result:
+                client["stats"]["age_hist"].find_one_and_update(
+                    {
+                        "councilorType": councilor_type,
+                        "method": "equal",
+                        "level": 1,
+                        "metroId": doc["metroId"],
+                        "year": year,
+                    },
+                    {"$set": {"diversityRank": int(doc["rank"])}},
+                )
+
+
 if __name__ == "__main__":
     # for localId in range(1, 227):
-    #     save_to_mongo(localId, "age", stair=10)
-    #     save_to_mongo(localId, "gender")
-    #     save_to_mongo(localId, "party")
+    #     save_to_mongo_local(localId, "age", stair=10)
+    #     save_to_mongo_local(localId, "gender")
+    #     save_to_mongo_local(localId, "party")
     # calculate_rank_local("age")
     # calculate_rank_local("gender")
     # calculate_rank_local("party")
-    calculate_age_diversity_rank_history()
+    calculate_age_diversity_rank_history_local()
+
+    # for metroId in range(1, 18):
+    #     if metroId in [8, 17]:
+    #         continue
+    #     save_to_mongo_metro(metroId, "age", stair=10)
+    #     save_to_mongo_metro(metroId, "gender")
+    #     save_to_mongo_metro(metroId, "party")
+    # calculate_rank_metro("age")
+    # calculate_rank_metro("gender")
+    # calculate_rank_metro("party")
+    # calculate_age_diversity_rank_history_metro()