From e9bc7f6eb43396c48f056e2db43916d57784f4ff Mon Sep 17 00:00:00 2001
From: Matt Graham <matthew.m.graham@gmail.com>
Date: Tue, 9 Aug 2022 13:43:17 +0100
Subject: [PATCH] Avoid use of apply in computing infection status

---
 src/tlo/methods/schisto.py | 33 ++++++++++++++++-----------------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/src/tlo/methods/schisto.py b/src/tlo/methods/schisto.py
index 0fd9633050..e8aabf2cf0 100644
--- a/src/tlo/methods/schisto.py
+++ b/src/tlo/methods/schisto.py
@@ -518,18 +518,21 @@ def update_infectious_status_and_symptoms(self, idx: pd.Index) -> None:
         if not len(idx) > 0:
             return
 
-        def _inf_status(age: int, agg_wb: int) -> str:
-            if age < 5:
-                if agg_wb >= params['high_intensity_threshold_PSAC']:
-                    return 'High-infection'
-
-            if agg_wb >= params['high_intensity_threshold']:
-                return 'High-infection'
-
-            if agg_wb >= params['low_intensity_threshold']:
-                return 'Low-infection'
-
-            return 'Non-infected'
+        def _get_infection_status(population: pd.DataFrame) -> pd.Series:
+            age = population["age_years"]
+            agg_wb = population[prop("aggregate_worm_burden")]
+            status = pd.Series(
+                "Non-infected",
+                index=population.index,
+                dtype=population[prop("infection_status")].dtype
+            )
+            high_group = (
+                (age < 5) & (agg_wb >= params["high_intensity_threshold_PSAC"])
+            ) | (agg_wb >= params["high_intensity_threshold"])
+            low_group = ~high_group & (agg_wb >= params["low_intensity_threshold"])
+            status[high_group] = "High-infection"
+            status[low_group] = "Low-infection"
+            return status
 
         def _impose_symptoms_of_high_intensity_infection(idx: pd.Index) -> None:
             """Assign symptoms to the person with high intensity infection.
@@ -547,11 +550,7 @@ def _impose_symptoms_of_high_intensity_infection(idx: pd.Index) -> None:
                         disease_module=schisto_module
                     )
 
-        correct_status = df.loc[idx].apply(
-            lambda x: _inf_status(x['age_years'], x[prop('aggregate_worm_burden')]),
-            axis=1
-        )
-
+        correct_status = _get_infection_status(df.loc[idx])
         original_status = df.loc[idx, prop('infection_status')]
 
         # Impose symptoms for those newly having 'High-infection' status