Merge branch 'estimation' of https://github.com/covid-19-impact-lab/s…

…id-germany into estimation

src/policies/combine_policies_over_periods.py

@@ -1,6 +1,7 @@
 from functools import partial
 
 import pandas as pd
+from src.config import BLD
 
 import src.policies.full_policy_blocks as fpb
 from src.policies.domain_level_policy_blocks import _get_base_policy
@@ -14,6 +15,7 @@ def get_december_to_feb_policies(
     contact_models,
     contact_tracing_multiplier,
     scenario,
+    path=None,
 ):
     """Get policies from December 2020 to February 2021.
     Args:
@@ -25,6 +27,8 @@ def get_december_to_feb_policies(
             reduction multiplier for recurrent and non-recurrent
             contact models.
         scenario (str): One of "optimistic", "pessimistic"
+        path (str or pathlib.Path): Path to a folder in which information on the
+            contact tracing is stored.
     Returns:
         policies (dict): policies dictionary.
     """
@@ -125,22 +129,23 @@ def get_december_to_feb_policies(
                 "prefix": "private-contact-tracing",
             },
             multiplier=contact_tracing_multiplier,
+            path=path,
         )
         to_combine.append(contact_tracing_policies)
     return combine_dictionaries(to_combine)
 
 
-def get_christmas_contact_tracing_policies(contact_models, block_info, multiplier):
+def get_christmas_contact_tracing_policies(contact_models, block_info, multiplier, path=None):
     """"""
     # households, educ contact models and Christmas models don't get adjustment
     models_with_post_christmas_isolation = [
         cm for cm in contact_models if "work" in cm or "other" in cm
     ]
-    christmas_id_groups = [
+    christmas_id_groups = list(set([
         model["assort_by"][0]
         for name, model in contact_models.items()
         if "christmas" in name
-    ]
+    ]))
     policies = {}
     for mod in models_with_post_christmas_isolation:
         policy = _get_base_policy(mod, block_info)
@@ -149,6 +154,7 @@ def get_christmas_contact_tracing_policies(contact_models, block_info, multiplie
             multiplier=multiplier,
             group_ids=christmas_id_groups,
             is_recurrent=contact_models[mod]["is_recurrent"],
+            path=path,
         )
         policies[f"{block_info['prefix']}_{mod}"] = policy
     return policies

src/policies/single_policy_functions.py

@@ -16,6 +16,7 @@
 import pandas as pd
 from scipy.interpolate import interp1d
 from sid.time import get_date
+from src.config import BLD
 
 
 def shut_down_model(states, contacts, seed):
@@ -287,7 +288,7 @@ def _interpolate_activity_level(
 
 
 def reduce_contacts_through_private_contact_tracing(
-    contacts, states, seed, multiplier, group_ids, is_recurrent
+    contacts, states, seed, multiplier, group_ids, is_recurrent, path
 ):
     today = get_date(states)
     days_since_christmas = (today - pd.Timestamp("2020-12-26")).days
@@ -303,12 +304,13 @@ def reduce_contacts_through_private_contact_tracing(
         group_ids=group_ids,
         condition=risk_condition,
         is_recurrent=is_recurrent,
+        path=path,
     )
     return reduced
 
 
 def reduce_contacts_when_condition_among_recurrent_contacts(
-    contacts, states, seed, multiplier, group_ids, condition, is_recurrent
+    contacts, states, seed, multiplier, group_ids, condition, is_recurrent, path=None
 ):
     """Reduce contacts when one of your contacts fulfills a condition.
 
@@ -326,13 +328,15 @@ def reduce_contacts_when_condition_among_recurrent_contacts(
             individual fulfills the condition, the individual is marked as having
             had a risk contact (unless (s)he herself fulfills the condition).
         is_recurrent (bool): Whether the contact model is recurrent or not.
+        path (str or pathlib.Path): Path to a folder in which information on the
+            contact tracing is stored.
 
     Returns:
         reduced (pandas.Series): reduced contacts.
 
     """
     with_risk_contacts = _identify_individuals_with_risk_contacts(
-        states, group_ids, condition
+        states, group_ids, condition, path
     )
 
     if is_recurrent:
@@ -345,7 +349,7 @@ def reduce_contacts_when_condition_among_recurrent_contacts(
     return reduced
 
 
-def _identify_individuals_with_risk_contacts(states, group_ids, condition):
+def _identify_individuals_with_risk_contacts(states, group_ids, condition, path=None):
     """Identify those in whose groups someone fulfills the condition.
 
     .. warning::
@@ -363,31 +367,43 @@ def _identify_individuals_with_risk_contacts(states, group_ids, condition):
         condition (str): query string. If any member of any group
             fulfills the condition, an individual is marked as having
             had a risk contact (unless (s)he herself fulfills the condition).
+        path (str or pathlib.Path): Path to a folder in which information on the
+            contact tracing is stored.
 
     Returns:
         risk_in_any_group (pandas.Series): boolean Series with same index
             as states. True for individuals who don't fulfill the condition
             but have a contact in one of their groups who does.
 
     """
-    risk_in_any_group = pd.Series(False, index=states.index)
     today = get_date(states)
-    risk_col = f"is_known_risk_contact_{today}"
+    risk_col = f"has_known_risk_contact_{today.date()}"
     if risk_col not in states.columns:
-        states[risk_col] = states.eval(condition)
-        old_col = f"is_known_risk_contact_{today - pd.Timedelta(days=1)}"
+        old_col = f"has_known_risk_contact_{(today - pd.Timedelta(days=1)).date()}"
         if old_col in states.columns:
             states.drop(
                 columns=[old_col],
                 inplace=True,
             )
-    for col in group_ids:
-        risk_in_this_group = states.groupby(col)[risk_col].transform("any")
-        # those in the -1 group have no contacts
-        risk_in_this_group = risk_in_this_group.where(states[col] != -1, False)
-        risk_in_any_group = risk_in_any_group | risk_in_this_group
-
-    # individuals who are themselves affected (e.g. symptomatic)
-    # reduce behavior through a different function. We don't want to "double" this.
-    risk_in_any_group = risk_in_any_group.where(~states.eval(condition), False)
+
+        helper = states[group_ids].copy()
+        helper["is_risk_contact"] = states.eval(condition)
+
+        risk_in_any_group = pd.Series(False, index=states.index)
+        for col in group_ids:
+            risk_in_this_group = helper.groupby(col)["is_risk_contact"].transform("any")
+            # those in the -1 group have no contacts
+            risk_in_this_group = risk_in_this_group.where(helper[col] != -1, False)
+            risk_in_any_group = risk_in_any_group | risk_in_this_group
+
+        # individuals who are themselves affected (e.g. symptomatic)
+        # reduce behavior through a different function. We don't want to "double" this.
+        risk_in_any_group = risk_in_any_group.where(~states.eval(condition), False)
+
+        states[risk_col] = risk_in_any_group
+        if path is not None:
+            risk_in_any_group.to_pickle(path / f"{risk_col}.pkl")
+    else:
+        risk_in_any_group = states[risk_col]
+
     return risk_in_any_group

src/simulation/task_simulate_christmas_scenarios.py

@@ -98,6 +98,7 @@ def task_simulation_christmas_scenarios(
         contact_models=contact_models,
         contact_tracing_multiplier=contact_tracing_multiplier,
         scenario=scenario,
+        path=path,
     )
 
     initial_conditions = create_initial_conditions(
@@ -121,8 +122,8 @@ def task_simulation_christmas_scenarios(
         saved_columns={
             "time": ["date"],
             "initial_states": False,
-            "disease_states": ["symptomatic"],
-            "other": ["n_has_infected", "newly_infected", "new_known_case"],
+            "disease_states": ["symptomatic", "ever_infected", "newly_infected"],
+            "other": ["n_has_infected", "new_known_case"],
         },
     )