Fix logging of road traffic injury statistics (#939)

Author: matt-graham

src/tlo/methods/rti.py

@@ -5387,35 +5387,26 @@ def apply(self, population):
         # Make some summary statistics
         # Get the dataframe and isolate the important information
         df = population.props
-        # dump dataframe each month if population size is large (used to find the minimum viable population size)
-        thoseininjuries = df.loc[df.rt_road_traffic_inc]
+        population_with_injuries = df.loc[df.rt_road_traffic_inc]
         # ================================= Injury severity ===========================================================
-        sev = thoseininjuries['rt_inj_severity']
-        rural_injuries = df.loc[df.rt_road_traffic_inc & ~df.li_urban]
-        if len(rural_injuries) > 0:
-            percent_sev_rural = \
-                len(rural_injuries.loc[rural_injuries['rt_inj_severity'] == 'severe']) / len(rural_injuries)
-        else:
-            percent_sev_rural = 'none_injured'
-        urban_injuries = df.loc[df.rt_road_traffic_inc & df.li_urban]
-        if len(urban_injuries) > 0:
-            percent_sev_urban = \
-                len(urban_injuries.loc[urban_injuries['rt_inj_severity'] == 'severe']) / len(urban_injuries)
-        else:
-            percent_sev_urban = 'none_injured'
-        severity, severitycount = np.unique(sev, return_counts=True)
-        if 'mild' in severity:
-            idx = np.where(severity == 'mild')
-            self.totmild += len(idx)
-        if 'severe' in severity:
-            idx = np.where(severity == 'severe')
-            self.totsevere += len(idx)
+        population_subsets_with_injuries = {
+            "rural": population_with_injuries.loc[~population_with_injuries.li_urban],
+            "urban": population_with_injuries.loc[population_with_injuries.li_urban],
+        }
+        proportion_severely_injured = {
+            label: (
+                len(pop_subset.loc[pop_subset['rt_inj_severity'] == 'severe'])
+                / len(pop_subset)
+            ) if len(pop_subset) > 0 else "none_injured"
+            for label, pop_subset in population_subsets_with_injuries.items()
+        }
+        self.totmild += (population_with_injuries.rt_inj_severity == "mild").sum()
+        self.totsevere += (population_with_injuries.rt_inj_severity == "severe").sum()
         dict_to_output = {
             'total_mild_injuries': self.totmild,
-            ''
-            '_severe_injuries': self.totsevere,
-            'Percent_severe_rural': percent_sev_rural,
-            'Percent_severe_urban': percent_sev_urban
+            'total_severe_injuries': self.totsevere,
+            'proportion_severe_rural': proportion_severely_injured["rural"],
+            'proportion_severe_urban': proportion_severely_injured["urban"],
         }
         logger.info(key='injury_severity',
                     data=dict_to_output,