Use broadcaster to collect events. Need to expand to include HSI events

Author: marghe-molaro

src/tlo/events.py

@@ -11,6 +11,7 @@
 
 import pandas as pd
 
+from tlo.notify import notifier
 from tlo.util import convert_chain_links_into_EAV
 
 import copy
@@ -296,7 +297,12 @@ def run(self):
         """Make the event happen."""
 
         # Collect relevant information before event takes place
+        # If statement outside or inside dispatch notification?
         if self.sim.generate_event_chains:
+        
+            # Dispatch notification that event is about to run
+            notifier.dispatch("event_about_to_run", data={"target": self.target, "EventName": type(self).__name__})
+            
             print_chains, row_before, df_before, mni_row_before, entire_mni_before, mni_instances_before = self.store_chains_to_do_before_event()
 
         self.apply(self.target)
@@ -305,6 +311,11 @@ def run(self):
         # Collect event info + meaningful property changes of individuals. Combined, these will constitute a 'link'
         # in the individual's event chain.
         if self.sim.generate_event_chains and print_chains:
+       
+            print("About to pass")
+            # Dispatch notification that event is about to run
+            notifier.dispatch("event_has_just_ran", data={"target": self.target, "EventName": type(self).__name__})
+            
             chain_links = self.store_chains_to_do_after_event(row_before, df_before, mni_row_before, entire_mni_before, mni_instances_before)
 
             if chain_links:

src/tlo/methods/collect_event_chains.py

@@ -0,0 +1,281 @@
+from tlo.notify import notifier
+
+from pathlib import Path
+from typing import Optional
+from tlo import Module, logging, population
+from tlo.population import Population
+import pandas as pd
+
+from tlo.util import convert_chain_links_into_EAV
+
+import copy
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+class CollectEventChains(Module):
+
+    def __init__(self, name=None):
+        super().__init__(name)
+    
+        # This is how I am passing data from fnc taking place before event to the one after
+        # It doesn't seem very elegant but not sure how else to go about it
+        self.print_chains = False
+        self.df_before = []
+        self.row_before = pd.Series()
+        self.mni_instances_before = False
+        self.mni_row_before = {}
+        self.entire_mni_before = {}
+        
+    def initialise_simulation(self, sim):
+        notifier.add_listener("event_about_to_run", self.on_notification_event_about_to_run)
+        notifier.add_listener("event_has_just_ran", self.on_notification_event_has_just_ran)
+        
+    def read_parameters(self, resourcefilepath: Optional[Path] = None):
+        pass
+        
+    def initialise_population(self, population):
+        pass
+
+    def on_birth(self, mother, child):
+        pass
+        
+    def on_notification_sim_about_to_start(self,data):
+        pass
+        
+    def on_notification_event_about_to_run(self, data):
+        """Do this when notified that an event is about to run. This function checks whether this event should be logged as part of the event chains, and if so stored required information before the event has occurred. """
+        print("This is the data I received ", data)
+
+        # Initialise these variables
+        self.print_chains = False
+        self.df_before = []
+        self.row_before = pd.Series()
+        self.mni_instances_before = False
+        self.mni_row_before = {}
+        self.entire_mni_before = {}
+        
+        print("My Modules")
+        print(self.sim.modules.keys())
+        # Only print event if it belongs to modules of interest and if it is not in the list of events to ignore
+        if all(sub not in str(data['EventName']) for sub in self.sim.generate_event_chains_ignore_events):
+        
+        # Will eventually use this once I can actually GET THE NAME OF THE SELF
+        #if not set(self.sim.generate_event_chains_ignore_events).intersection(str(self)):
+
+            self.print_chains = True
+            
+            # Target is single individual
+            if not isinstance(data["target"], Population):
+
+                # Save row for comparison after event has occurred
+                self.row_before = self.sim.population.props.loc[abs(data['target'])].copy().fillna(-99999)
+                
+                # Check if individual is already in mni dictionary, if so copy her original status
+                if 'PregnancySupervisor' in self.sim.modules:
+                    mni = self.sim.modules['PregnancySupervisor'].mother_and_newborn_info
+                    if data['target'] in mni:
+                        self.mni_instances_before = True
+                        self.mni_row_before = mni[data['target']].copy()
+                else:
+                    self.mni_row_before = None
+                
+            else:
+
+                # This will be a population-wide event. In order to find individuals for which this led to
+                # a meaningful change, make a copy of the while pop dataframe/mni before the event has occurred.
+                self.df_before = self.sim.population.props.copy()
+                if 'PregnancySupervisor' in self.sim.modules:
+                    self.entire_mni_before = copy.deepcopy(self.sim.modules['PregnancySupervisor'].mother_and_newborn_info)
+                else:
+                    self.entire_mni_before = None
+
+        return
+        
+    
+    def on_notification_event_has_just_ran(self, data):
+        """ If print_chains=True, this function logs the event and identifies and logs the any property changes that have occured to one or multiple individuals as a result of the event taking place. """
+        print("This is the data I received ", data)
+        
+        chain_links = {}
+    
+        # Target is single individual
+        if not isinstance(data["target"], Population):
+    
+            # Copy full new status for individual
+            row_after = self.sim.population.props.loc[abs(data['target'])].fillna(-99999)
+            
+            # Check if individual is in mni after the event
+            mni_instances_after = False
+            if 'PregnancySupervisor' in self.sim.modules:
+                mni = self.sim.modules['PregnancySupervisor'].mother_and_newborn_info
+                if data['target'] in mni:
+                    mni_instances_after = True
+            else:
+                mni_instances_after = None
+            
+            # Create and store event for this individual, regardless of whether any property change occurred
+            link_info = {
+                'EventName' : data['EventName'],
+            }
+            
+            # Store (if any) property changes as a result of the event for this individual
+            for key in self.row_before.index:
+                if self.row_before[key] != row_after[key]: # Note: used fillna previously, so this is safe
+                    link_info[key] = row_after[key]
+            
+            if 'PregnancySupervisor' in self.sim.modules:
+                # Now check and store changes in the mni dictionary, accounting for following cases:
+                # Individual is in mni dictionary before and after
+                if self.mni_instances_before and mni_instances_after:
+                    for key in self.mni_row_before:
+                        if self.mni_values_differ(mni_row_before[key], mni[data['target']][key]):
+                            link_info[key] = mni[data['target']][key]
+                # Individual is only in mni dictionary before event
+                elif self.mni_instances_before and not mni_instances_after:
+                    default = self.sim.modules['PregnancySupervisor'].default_all_mni_values
+                    for key in self.mni_row_before:
+                        if self.mni_values_differ(mni_row_before[key], default[key]):
+                            link_info[key] = default[key]
+                # Individual is only in mni dictionary after event
+                elif mni_instances_after and not self.mni_instances_before:
+                    default = self.sim.modules['PregnancySupervisor'].default_all_mni_values
+                    for key in default:
+                        if self.mni_values_differ(default[key], mni[data['target']][key]):
+                            link_info[key] = mni[data['target']][key]
+                # Else, no need to do anything
+                    
+            # Add individual to the chain links
+            chain_links[data['target']] = link_info
+            
+        else:
+            # Target is entire population. Identify individuals for which properties have changed
+            # and store their changes.
+            
+            # Population frame after event
+            df_after = self.sim.population.props
+            if 'PregnancySupervisor' in self.sim.modules:
+                entire_mni_after = copy.deepcopy(self.sim.modules['PregnancySupervisor'].mother_and_newborn_info)
+            else:
+                entire_mni_after = None
+            
+            #  Create and store the event and dictionary of changes for affected individuals
+            chain_links = self.compare_population_dataframe_and_mni(self.df_before, df_after, self.entire_mni_before, entire_mni_after)
+
+            if chain_links:
+                # Convert chain_links into EAV
+                ednav = convert_chain_links_into_EAV(chain_links)
+
+                logger.info(key='event_chains',
+                      data= ednav.to_dict(),
+                      description='Links forming chains of events for simulated individuals')
+                      
+        # Reset variables
+        self.print_chains = False
+        self.df_before = []
+        self.row_before = pd.Series()
+        self.mni_instances_before = False
+        self.mni_row_before = {}
+        self.entire_mni_before = {}
+
+        return
+    
+    def mni_values_differ(self, v1, v2):
+
+        if isinstance(v1, list) and isinstance(v2, list):
+            return v1 != v2  # simple element-wise comparison
+
+        if pd.isna(v1) and pd.isna(v2):
+            return False  # treat both NaT/NaN as equal
+        return v1 != v2
+    
+    def compare_entire_mni_dicts(self,entire_mni_before, entire_mni_after):
+        diffs = {}
+
+        all_individuals = set(entire_mni_before.keys()) | set(entire_mni_after.keys())
+            
+        for person in all_individuals:
+            if person not in entire_mni_before: # but is afterward
+                for key in entire_mni_after[person]:
+                    if self.mni_values_differ(entire_mni_after[person][key],self.sim.modules['PregnancySupervisor'].default_all_mni_values[key]):
+                        if person not in diffs:
+                            diffs[person] = {}
+                        diffs[person][key] = entire_mni_after[person][key]
+                    
+            elif person not in entire_mni_after: # but is beforehand
+                for key in entire_mni_before[person]:
+                    if self.mni_values_differ(entire_mni_before[person][key],self.sim.modules['PregnancySupervisor'].default_all_mni_values[key]):
+                        if person not in diffs:
+                            diffs[person] = {}
+                        diffs[person][key] = self.sim.modules['PregnancySupervisor'].default_all_mni_values[key]
+
+            else: # person is in both
+                # Compare properties
+                for key in entire_mni_before[person]:
+                    if self.mni_values_differ(entire_mni_before[person][key],entire_mni_after[person][key]):
+                        if person not in diffs:
+                            diffs[person] = {}
+                        diffs[person][key] = entire_mni_after[person][key]
+
+        return diffs
+        
+    def compare_population_dataframe_and_mni(self,df_before, df_after, entire_mni_before, entire_mni_after):
+        """ This function compares the population dataframe and mni dictionary before/after a population-wide event has occurred.
+        It allows us to identify the individuals for which this event led to a significant (i.e. property) change, and to store the properties which have changed as a result of it. """
+        
+        # Create a mask of where values are different
+        diff_mask = (df_before != df_after) & ~(df_before.isna() & df_after.isna())
+        if 'PregnancySupervisor' in self.sim.modules:
+            diff_mni = self.compare_entire_mni_dicts(entire_mni_before, entire_mni_after)
+        else:
+            diff_mni = []
+        
+        # Create an empty list to store changes for each of the individuals
+        chain_links = {}
+        len_of_diff = len(diff_mask)
+
+        # Loop through each row of the mask
+        persons_changed = []
+        
+        for idx, row in diff_mask.iterrows():
+            changed_cols = row.index[row].tolist()
+
+            if changed_cols:  # Proceed only if there are changes in the row
+                persons_changed.append(idx)
+                # Create a dictionary for this person
+                # First add event info
+                link_info = {
+                    'EventName': type(self).__name__,
+                }
+                
+                # Store the new values from df_after for the changed columns
+                for col in changed_cols:
+                    link_info[col] = df_after.at[idx, col]
+
+                if idx in diff_mni:
+                    # This person has also undergone changes in the mni dictionary, so add these here
+                    for key in diff_mni[idx]:
+                        link_info[col] = diff_mni[idx][key]
+
+                # Append the event and changes to the individual key
+                chain_links[idx] = link_info
+     
+        if 'PregnancySupervisor' in self.sim.modules:
+            # For individuals which only underwent changes in mni dictionary, save changes here
+            if len(diff_mni)>0:
+                for key in diff_mni:
+                    if key not in persons_changed:
+                        # If individual hadn't been previously added due to changes in pop df, add it here
+                        link_info = {
+                            'EventName': type(self).__name__,
+                        }
+                        
+                        for key_prop in diff_mni[key]:
+                            link_info[key_prop] = diff_mni[key][key_prop]
+                            
+                        chain_links[key] = link_info
+
+        return chain_links
+
+
+        

src/tlo/methods/fullmodel.py

@@ -8,6 +8,7 @@
     cardio_metabolic_disorders,
     care_of_women_during_pregnancy,
     cervical_cancer,
+    collect_event_chains,
     contraception,
     copd,
     demography,
@@ -116,6 +117,7 @@ def fullmodel(
         copd.Copd,
         depression.Depression,
         epilepsy.Epilepsy,
+        collect_event_chains.CollectEventChains,
     ]
     return [
         module_class(

src/tlo/simulation.py

@@ -13,6 +13,10 @@
 import pandas as pd
 import tlo.population
 import numpy as np
+import tlo.methods.collect_event_chains
+
+from tlo.notify import notifier
+from tlo.methods.collect_event_chains import CollectEventChains
 from tlo.util import df_to_EAV, convert_chain_links_into_EAV
 
 try:
@@ -148,6 +152,7 @@ def __init__(
 
         # Whether simulation has been initialised
         self._initialised = False
+        
 
     def _configure_logging(
         self,