From ad9df230a542ff271138a9aca2b1d57510938810 Mon Sep 17 00:00:00 2001
From: rmanach <manach.r@msn.com>
Date: Wed, 22 Dec 2021 00:06:18 +0100
Subject: [PATCH] group plotted data by fields

---
 drees.py | 51 +++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/drees.py b/drees.py
index e1be538..37ec0e3 100644
--- a/drees.py
+++ b/drees.py
@@ -179,6 +179,22 @@ def get_values_by_age_vac_field(
     return dates, fields
 
 
+def get_values_by_age_vac(
+    dic_data_grouped: Dict[dt, Any], age: AgeGroup, vac_status: VacStatus
+) -> Tuple[List[dt], List[Dict[str, Any]]]:
+    """
+    get deep fields data by age and vaccine status
+    """
+    dates: List[dt] = list()
+    fields: List[Dict[str, Any]] = list()
+    for date, dic_age_grouped in dic_data_grouped.items():
+        if (dic_vac_status := dic_age_grouped.get(age.value)) is not None:
+            if (dic_field := dic_vac_status.get(vac_status.value)) is not None:
+                fields.append(dic_field)
+                dates.append(date)
+    return dates, fields
+
+
 def plot_cumulative_field(dic_data_grouped: Dict[dt, Any], field: Field) -> None:
     fig = get_plot_fig()
 
@@ -197,37 +213,44 @@ def plot_cumulative_field(dic_data_grouped: Dict[dt, Any], field: Field) -> None
     )
 
 
-def plot_data(
-    dic_data_grouped: Dict[dt, Any], age: AgeGroup, vac_status: VacStatus, field: Field
+def extract_field_values(fields: List[Dict[str, Any]], field: Field) -> np.ndarray:
+    field_values: List[float] = list()
+    for item in fields:
+        if (value := item.get(field)) is not None:
+            field_values.append(value)
+    return np.asarray(field_values)
+
+
+def plot_data_by_field(
+    dic_data_grouped: Dict[dt, Any], age: AgeGroup, vac_status: VacStatus
 ) -> None:
     """
     Plot data by vaccine status, age and field
     """
     fig = get_plot_fig()
 
-    dates, fields = get_values_by_age_vac_field(
-        dic_data_grouped, age, vac_status, field
-    )
+    dates, fields = get_values_by_age_vac(dic_data_grouped, age, vac_status)
 
-    plt.plot(dates, fields, label=f"{field.value}")
-    plt.xlabel("date")
-    plt.ylabel("nombre")
-    plt.title(f"{age}ans - {vac_status}")
+    for field in Field:
+        field_values = extract_field_values(fields, field)
+        plt.plot(dates, field_values, label=f"{field.value}")
+        plt.xlabel("date")
+        plt.ylabel("nombre")
+        plt.title(f"{age}ans - {vac_status}")
 
     save_and_close_fig(
         fig, os.path.join(OUTPUT_REPOSITORY, f"{age}_{vac_status}_{field}.pdf")
     )
 
 
-def build_data_pool_args() -> List[Tuple[AgeGroup, VacStatus, Field]]:
+def build_data_pool_args() -> List[Tuple[AgeGroup, VacStatus]]:
     """
     build tuple arguments to plot all data on multiprocess
     """
-    pool_args: List[Tuple[AgeGroup, VacStatus, Field]] = list()
+    pool_args: List[Tuple[AgeGroup, VacStatus]] = list()
     for age_group in AgeGroup:
         for vac_status in VacStatus:
-            for field in Field:
-                pool_args.append((age_group, vac_status, field))
+            pool_args.append((age_group, vac_status))
     return pool_args
 
 
@@ -264,7 +287,7 @@ if __name__ == "__main__":
     )
 
     plot_data_pool_args = build_data_pool_args()
-    f = partial(plot_data, dic_data_grouped)
+    f = partial(plot_data_by_field, dic_data_grouped)
     with Pool() as pool:
         pool.starmap(f, plot_data_pool_args)