Symptom Bookkeeping (#1626)

Co-authored-by: Asif Tamuri <[email protected]>
Co-authored-by: Tim Hallett <[email protected]>

Author: 3 people

src/scripts/profiling/symptom_manager_bookkeeping_profiling.py

@@ -0,0 +1,98 @@
+"""
+A script to register so many symptoms and run the has_what function so many times.
+For use in profiling.
+"""
+
+import time
+from pathlib import Path
+
+import numpy as np
+
+from tlo import Date, Simulation
+from tlo.methods import demography, symptommanager
+
+
+def setup_simulation(pop_size=100_000, num_symptoms=50):
+    """Set up a simulation with population and many symptoms."""
+    start_date = Date(2010, 1, 1)
+
+    resource_dir = Path("./resources")
+
+    sim = Simulation(
+        start_date=start_date,
+        seed=0,
+        #log_config={"filename": "symptom_profiling", "directory": "./outputs"},
+        resourcefilepath=resource_dir
+    )
+
+    sim.register(demography.Demography(), symptommanager.SymptomManager())
+
+    # Register symptoms
+    sm = sim.modules['SymptomManager']
+    for i in range(num_symptoms):
+        sm.register_symptom(symptommanager.Symptom(name=f'symptom_{i}'))
+
+    # Initialize population - this will create the symptom properties
+    sim.make_initial_population(n=pop_size)
+
+    return sim
+
+
+def assign_random_symptoms(sim, symptom_prob=0.1):
+    """Assign random symptoms to the population."""
+    df = sim.population.props
+    sm = sim.modules['SymptomManager']
+
+    # Assign symptoms randomly
+    for symptom in sm.symptom_names:
+        # Random subset of population to have this symptom
+        has_symptom = np.random.random(len(df)) < symptom_prob
+        person_ids = df.index[has_symptom].tolist()
+
+        if person_ids:
+            sm.change_symptom(
+                person_id=person_ids,
+                symptom_string=symptom,
+                add_or_remove='+',
+                disease_module=sm,
+                duration_in_days=None
+            )
+
+
+def profile_has_what(sim, num_tests=1000000):
+    """Profiling has_what function by calling it repeatedly."""
+    df = sim.population.props
+    sm = sim.modules['SymptomManager']
+
+    # Get random sample of person_ids to test
+    test_ids = np.random.choice(df.index[df.is_alive], size=num_tests, replace=True)
+
+    # Time the has_what function
+    start_time = time.time()
+
+    results = []
+    for person_id in test_ids:
+        results.append(sm.has_what(person_id))
+
+    elapsed = time.time() - start_time
+    avg_time = elapsed / num_tests
+
+    print(f"Tested has_what() {num_tests} times")
+    print(f"Total time: {elapsed:.4f} seconds")
+    print(f"Average time per call: {avg_time:.6f} seconds")
+    print(f"First 5 results: {results[:5]}")
+
+    return avg_time
+
+
+print("Setting up simulation...")
+sim = setup_simulation(pop_size=100_000, num_symptoms=50)
+
+print("Assigning random symptoms...")
+assign_random_symptoms(sim, symptom_prob=0.6)
+
+print("\nProfiling has_what...")
+avg_time = profile_has_what(sim, num_tests=1000000)
+
+print("\nProfiling complete!")
+print(f"Average time per has_what call: {avg_time:.6f} seconds")

src/tlo/methods/demography.py

@@ -83,7 +83,7 @@ def __init__(self, name=None, equal_allocation_by_district: bool = False):
         self.gbd_causes_of_death = set()  # will store all the causes of death defined in the GBD data
         self.gbd_causes_of_death_not_represented_in_disease_modules = set()
         #  will store causes of death in GBD not represented in the simulation
-        self.other_death_poll = None    # will hold pointer to the OtherDeathPoll object
+        self.other_death_poll = None  # will hold pointer to the OtherDeathPoll object
         self.districts = None  # will store all the districts in a list
 
     OPTIONAL_INIT_DEPENDENCIES = {'ImprovedHealthSystemAndCareSeekingScenarioSwitcher'}
@@ -175,14 +175,14 @@ def read_parameters(self, resourcefilepath: Optional[Path] = None):
         # Lookup dicts to map from district_num_of_residence (in the df) and District name and Region name
         self.districts = self.parameters['pop_2010']['District'].drop_duplicates().to_list()
         self.parameters['district_num_to_district_name'] = \
-            self.parameters['pop_2010'][['District_Num', 'District']].drop_duplicates()\
-                                                                     .set_index('District_Num')['District']\
-                                                                     .to_dict()
+            self.parameters['pop_2010'][['District_Num', 'District']].drop_duplicates() \
+                .set_index('District_Num')['District'] \
+                .to_dict()
 
         self.parameters['district_num_to_region_name'] = \
-            self.parameters['pop_2010'][['District_Num', 'Region']].drop_duplicates()\
-                                                                   .set_index('District_Num')['Region']\
-                                                                   .to_dict()
+            self.parameters['pop_2010'][['District_Num', 'Region']].drop_duplicates() \
+                .set_index('District_Num')['Region'] \
+                .to_dict()
 
         districts_in_region = defaultdict(set)
         for _district in self.parameters['pop_2010'][['District', 'Region']].drop_duplicates().itertuples():
@@ -563,6 +563,10 @@ def do_death(self, individual_id: int, cause: str, originating_module: Module):
             if person.hs_is_inpatient:
                 self.sim.modules['HealthSystem'].remove_beddays_footprint(person_id=individual_id)
 
+        # Clear symptoms for the deceased person
+        if 'SymptomManager' in self.sim.modules:
+            self.sim.modules['SymptomManager'].clear_symptoms_for_deceased_person(individual_id)
+
     def create_mappers_from_causes_of_death_to_label(self):
         """Use a helper function to create mappers for causes of death to label."""
         return create_mappers_from_causes_to_label(
@@ -683,6 +687,7 @@ class OtherDeathPoll(RegularEvent, PopulationScopeEventMixin):
     It does this by computing the GBD death rates that are implied by all the causes of death other than those that are
     represented in the disease module registered in this simulation.
     """
+
     def __init__(self, module):
         super().__init__(module, frequency=DateOffset(months=1))
         self.causes_to_represent = self.module.gbd_causes_of_death_not_represented_in_disease_modules

src/tlo/methods/symptommanager.py

@@ -15,7 +15,7 @@
 
 from collections import defaultdict
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, List, Optional, Sequence, Set
 
 import numpy as np
 import pandas as pd
@@ -26,11 +26,14 @@
 from tlo.util import BitsetHandler
 
 if TYPE_CHECKING:
+    from typing import Union
+
     from tlo.population import IndividualProperties
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
+
 # ---------------------------------------------------------------------------------------------------------
 #   MODULE DEFINITIONS
 # ---------------------------------------------------------------------------------------------------------
@@ -218,6 +221,7 @@ def __init__(self, name=None, spurious_symptoms=None):
 
         self.recognised_module_names = None
         self.spurious_symptom_resolve_event = None
+        self.symptom_tracker = defaultdict(set)
 
     def get_column_name_for_symptom(self, symptom_name):
         """get the column name that corresponds to the symptom_name"""
@@ -395,6 +399,10 @@ def change_symptom(self, person_id, symptom_string, add_or_remove, disease_modul
             self.bsh.set(person_id, disease_module.name, columns=sy_columns)
             self._persons_with_newly_onset_symptoms = self._persons_with_newly_onset_symptoms.union(person_id)
 
+            # Update symptom tracker
+            for pid in person_id:
+                self.symptom_tracker[pid] |= set(symptom_string)
+
             # If a duration is given, schedule the auto-resolve event to turn off these symptoms after specified time.
             if duration_in_days is not None:
                 auto_resolve_event = SymptomManager_AutoResolveEvent(self,
@@ -417,10 +425,17 @@ def change_symptom(self, person_id, symptom_string, add_or_remove, disease_modul
             # Do the remove:
             self.bsh.unset(person_id, disease_module.name, columns=sy_columns)
 
+            # Update symptom tracker. Remove if no other module is causing this symptom.
+            for pid in person_id:
+                for sym in symptom_string:
+                    symptom_col = self.get_column_name_for_symptom(sym)
+                    if self.bsh.is_empty(pid, columns=symptom_col):
+                        self.symptom_tracker[pid].discard(sym)
+
     def who_has(self, list_of_symptoms):
         """
         This is a helper function to look up who has a particular symptom or set of symptoms.
-        It returns a list of indicies for person that have all of the symptoms specified
+        It returns a list of indices for person that has all of the symptoms specified
 
         :param: list_of_symptoms : string or list of strings for the symptoms of interest
         :return: list of person_ids for those with all of the symptoms in list_of_symptoms who are alive
@@ -462,7 +477,7 @@ def who_not_have(self, symptom_string: str) -> pd.Index:
             & self.bsh.is_empty(
                 slice(None), columns=self.get_column_name_for_symptom(symptom_string)
             )
-        ]
+            ]
 
     def has_what(
         self,
@@ -492,6 +507,10 @@ def has_what(
             else True
         ), "Disease Module Name is not recognised"
 
+        # Faster to get current symptoms using tracker when no disease is specified
+        if disease_module is None and person_id is not None:
+            return list(self._get_current_symptoms_from_tracker(person_id))
+
         if individual_details is not None:
             # We are working in an IndividualDetails context, avoid lookups to the
             # population DataFrame as we have this context stored already.
@@ -503,10 +522,10 @@ def has_what(
                     symptom
                     for symptom in self.symptom_names
                     if individual_details[
-                        self.bsh._get_columns(self.get_column_name_for_symptom(symptom))
-                    ]
-                    & int_repr
-                    != 0
+                           self.bsh._get_columns(self.get_column_name_for_symptom(symptom))
+                       ]
+                       & int_repr
+                       != 0
                 ]
             else:
                 return [
@@ -582,6 +601,17 @@ def clear_symptoms(self, person_id: Union[int, Sequence[int]], disease_module: M
         sy_columns = [self.get_column_name_for_symptom(sym) for sym in self.symptom_names]
         self.bsh.unset(person_id, disease_module.name, columns=sy_columns)
 
+        # Update bookkeeping
+        for pid in person_id:
+            for sym in self.symptom_names:
+                symptom_col = self.get_column_name_for_symptom(sym)
+                if self.bsh.is_empty(pid, columns=symptom_col):
+                    self.symptom_tracker[pid].discard(sym)
+
+            # Remove the person's entry from the tracker is the symptom set is empty
+            if pid in self.symptom_tracker and not self.symptom_tracker[pid]:
+                del self.symptom_tracker[pid]
+
     def caused_by(self, disease_module: Module):
         """Find the persons experiencing symptoms due to a particular module.
         Returns a dict of the form {<<person_id>>, <<list_of_symptoms>>}."""
@@ -600,6 +630,17 @@ def get_persons_with_newly_onset_symptoms(self):
     def reset_persons_with_newly_onset_symptoms(self):
         self._persons_with_newly_onset_symptoms.clear()
 
+    def _get_current_symptoms_from_tracker(self, person_id: int) -> Set[str]:
+        """Get the current symptoms for a person. Works with bookkeeping dictionary"""
+        return self.symptom_tracker.get(person_id, set())
+
+    def clear_symptoms_for_deceased_person(self, person_id: int):
+        """Clears symptoms by deleting the dead person's ID in the tracker"""
+        # Remove person from tracker entirely
+        if person_id in self.symptom_tracker:
+            del self.symptom_tracker[person_id]
+
+
 # ---------------------------------------------------------------------------------------------------------
 #   EVENTS
 # ---------------------------------------------------------------------------------------------------------
@@ -696,22 +737,21 @@ def apply(self, population):
             do_not_have_symptom = self.module.who_not_have(symptom_string=symp)
 
             for group in ['children', 'adults']:
-
                 p = self.generic_symptoms['prob_per_day'][group][symp]
                 dur = self.generic_symptoms['duration_in_days'][group][symp]
                 persons_eligible_to_get_symptom = group_indices[group][
                     group_indices[group].isin(do_not_have_symptom)
                 ]
                 persons_to_onset_with_this_symptom = persons_eligible_to_get_symptom[
                     self.rand(len(persons_eligible_to_get_symptom)) < p
-                ]
+                    ]
 
                 # Do onset
                 self.sim.modules['SymptomManager'].change_symptom(
                     symptom_string=symp,
                     add_or_remove='+',
                     person_id=persons_to_onset_with_this_symptom,
-                    duration_in_days=None,   # <- resolution for these is handled by the SpuriousSymptomsResolve Event
+                    duration_in_days=None,  # <- resolution for these is handled by the SpuriousSymptomsResolve Event
                     disease_module=self.module,
                 )