Simulator

Module containing the Simulator class, which handles much of managing how the world changes as infection spreads.

Simulator

Source code in june/simulator.py

class Simulator:
    """ """
    ActivityManager = ActivityManager

    def __init__(
        self,
        world: World,
        interaction: Interaction,
        timer: Timer,
        activity_manager: ActivityManager,
        epidemiology: Epidemiology,
        events: Optional[Events] = None,
        record: Optional[Record] = None,
        feature_flags: Optional[dict] = None
    ):
        """
        Class to run an epidemic spread simulation on the world.

        Parameters:
          world (World):
            Instance of the world class.
          interaction (Interaction):
            Instance of the Interaction class
          timer (Timer):
            The timer class keeps track of what stage of the day it is,
            and which types of activities are allowed
          activity_manager (ActivityManager):
            Helps manage allocating which activities people choose to do in a
            given timestep.
          epidemiology (Epidemiology):
            Contains information about how the disease propogates.
          events (Optional[Events]):
            Events object, adding in special events.
          record (Optional[Record]):
            Record object to keep track of the simulation as it goes.
          feature_flags (Optional[dict]):
            Dictionary that says which features are enabled. 
        """
        # Process feature flags with defaults
        if feature_flags is None:
            feature_flags = {}

        self.friend_hangouts_enabled = feature_flags.get("friend_hangouts_enabled", False)
        self.sexual_encounters_enabled = feature_flags.get("sexual_encounters_enabled", False)
        self.test_and_trace_enabled = feature_flags.get("test_and_trace_enabled", False)
        self.ratty_dynamics_enabled = feature_flags.get("ratty_dynamics_enabled", False)
        self.rat_data_saving_enabled = feature_flags.get("rat_data_saving_enabled", True)  # Default to True

        # Original initialisation code
        self.activity_manager = activity_manager
        self.world = world
        self.interaction = interaction
        self.events = events
        self.timer = timer
        self.epidemiology = epidemiology

        if self.epidemiology:
            self.epidemiology.set_medical_care(
                world=world, activity_manager=activity_manager
            )
            self.epidemiology.set_immunity(self.world)
            self.epidemiology.set_past_vaccinations(
                people=self.world.people, date=self.timer.date, record=record
            )

        if self.events is not None:
            self.events.init_events(world=world)

        self.record = record
        if self.record is not None and self.record.record_static_data:
            self.record.static_data(world=world)

        # Only initialise rat manager if feature is enabled
        self.rat_manager = None
        self.save_rat_data = False

        if self.ratty_dynamics_enabled:
            output_logger.info("Ratty dynamics enabled, initialising rat manager")
            from june.zoonosis.rat_manager import RatManager

            # Use record path for rat density files if available
            if self.record and hasattr(self.record, 'record_path'):
                rat_density_path = str(self.record.record_path / "rat_density_map.npy")
            else:
                rat_density_path = "rat_density_map.npy"  # Fallback

            # Normal initialisation - will be overwritten if checkpoint is restored
            self.rat_manager = RatManager(world=world, precomputed_density_path=rat_density_path)
            self.save_rat_data = self.rat_data_saving_enabled
            output_logger.info(f"Rat data saving: {self.save_rat_data}")
            output_logger.info(f"Rat density path: {rat_density_path}")

        # Initialise contact manager if test and trace is enabled or sexual encounters are enabled
        self.contact_manager = None
        if self.test_and_trace_enabled or self.sexual_encounters_enabled:
            output_logger.info("Test and Trace or Sexual Encounters enabled, initialising contact manager")
            self.contact_manager = ContactManager(self)

            # Configure the contact retention based on policy settings            
            disease_config = None

            disease_config = GlobalContext.get_disease_config()
            output_logger.info("Using disease_config from GlobalContext")

            # Use disease_config if available
            if disease_config is not None:
                tracing_data = disease_config.policy_manager.get_policy_data("tracing")
                self.contact_retention_days = tracing_data.get("contact_retention_days", 14) if tracing_data else 14
                output_logger.info(f"Using contact retention days from config: {self.contact_retention_days}")

            # Connect contact manager to leisure system if it exists
            if self.activity_manager.leisure is not None:
                output_logger.info("Connecting contact manager to leisure system")
                self.activity_manager.leisure.set_contact_manager(self.contact_manager)

            # Connect contact manager to sexual encounter system if it exists
            if self.activity_manager.sexual_encounter is not None:
                output_logger.info("Connecting contact manager to sexual encounter system")
                self.activity_manager.sexual_encounter.contact_manager = self.contact_manager
        else:
            output_logger.info("Test and Trace and Sexual Encounters disabled, skipping contact manager initialisation")

        # Initialise zoonotic transmission if rat_manager is properly initialised and feature is enabled
        self.zoonotic_transmission = None
        if self.ratty_dynamics_enabled and self.rat_manager is not None and ZoonoticTransmission is not None:
            output_logger.info("Initialising zoonotic transmission module")
            self.zoonotic_transmission = ZoonoticTransmission(rat_manager=self.rat_manager)
            output_logger.info("Zoonotic transmission module initialised successfully")

        # Initialise checkpointing system placeholder
        # Actual checkpointing setup happens in from_file() method with config
        self.checkpointing = None


    @classmethod
    def from_file(
        cls,
        world: World,
        interaction: Interaction,
        policies: Optional[Policies] = None,
        events: Optional[Events] = None,
        epidemiology: Optional[Epidemiology] = None,
        leisure: Optional[Leisure] = None,
        travel: Optional[Travel] = None,
        sexual_encounter: Optional[SexualEncounter] = None,
        config_filename: str = default_config_filename,
        record: Optional[Record] = None,
    ) -> "Simulator":
        """Load config for simulator from world.yaml

        Args:
            world (World): 
            interaction (Interaction): 
            policies (Optional[Policies], optional): (Default value = None)
            events (Optional[Events], optional): (Default value = None)
            epidemiology (Optional[Epidemiology], optional): (Default value = None)
            leisure (Optional[Leisure], optional): (Default value = None)
            travel (Optional[Travel], optional): (Default value = None)
            sexual_encounter (Optional[SexualEncounter], optional): (Default value = None)
            config_filename (str, optional): (Default value = default_config_filename)
            record (Optional[Record], optional): (Default value = None)
        """
        # Load the configuration file to get feature flags
        with open(config_filename) as f:
            config = yaml.load(f, Loader=yaml.FullLoader)

        # Extract feature flags with defaults in case they're not in the config
        features = config.get("features", {})

        #Friendship networks
        friend_hangouts_config = features.get("friend_hangouts", {"enabled": False})
        friend_hangouts_enabled = friend_hangouts_config.get("enabled", False)

        # Sexual encounters settings
        sexual_encounters_config = features.get("sexual_encounters", {"enabled": False})
        sexual_encounters_enabled = sexual_encounters_config.get("enabled", False)

        # Test and trace settings
        test_and_trace_config = features.get("test_and_trace", {"enabled": False})
        test_and_trace_enabled = test_and_trace_config.get("enabled", False)

        # Ratty dynamics settings
        ratty_dynamics_config = features.get("ratty_dynamics", {"enabled": False})
        ratty_dynamics_enabled = ratty_dynamics_config.get("enabled", False)
        rat_data_saving_enabled = ratty_dynamics_config.get("save_data", True)  # Default to True


        output_logger.info(f"Feature flags from config: Friend hangouts: {friend_hangouts_enabled}, "
                        f"Sexual encounters: {sexual_encounters_enabled}, "
                        f"Test and Trace: {test_and_trace_enabled}, "
                        f"Ratty Dynamics: {ratty_dynamics_enabled}, "
                        f"Rat Data Saving: {rat_data_saving_enabled}, ")

        # Continue with original method
        timer = Timer.from_file(config_filename=config_filename)
        activity_manager = cls.ActivityManager.from_file(
            config_filename=config_filename,
            world=world,
            leisure=leisure,
            travel=travel,
            sexual_encounter=sexual_encounter,
            policies=policies,
            timer=timer,
            record=record,
        )

        # Store feature flags to pass to the constructor
        feature_flags = {
            "friend_hangouts_enabled": friend_hangouts_enabled,
            "sexual_encounters_enabled": sexual_encounters_enabled,
            "test_and_trace_enabled": test_and_trace_enabled,
            "ratty_dynamics_enabled": ratty_dynamics_enabled,
            "rat_data_saving_enabled": rat_data_saving_enabled
        }

        simulator = cls(
            world=world,
            interaction=interaction,
            timer=timer,
            events=events,
            activity_manager=activity_manager,
            epidemiology=epidemiology,
            record=record,
            feature_flags=feature_flags
        )

        # Initialise checkpointing system with config file
        if CHECKPOINTING_AVAILABLE:
            # Replace the default checkpointing with config-based one

            # Get the checkpoint directories (separate read/write for child runs)
            checkpoint_read_dir = "checkpoints"
            checkpoint_write_dir = "checkpoints"

            if hasattr(simulator, 'record') and simulator.record and hasattr(simulator.record, 'record_path'):
                run_dir = simulator.record.record_path.parent
                own_checkpoint_dir = run_dir / "checkpoints"

                # Check if this might be a child run by looking for metadata
                metadata_file = run_dir / "metadata.json"
                if metadata_file.exists():
                    try:
                        import json
                        with open(metadata_file, 'r') as f:
                            metadata = json.load(f)

                        parent_run_id = metadata.get("parent_run_id")
                        if parent_run_id:
                            # This is a child run - separate read and write directories
                            parent_checkpoint_dir = run_dir.parent / parent_run_id / "checkpoints"
                            if parent_checkpoint_dir.exists():
                                checkpoint_read_dir = str(parent_checkpoint_dir)
                                checkpoint_write_dir = str(own_checkpoint_dir)
                                output_logger.info(f"Child run detected:")
                                output_logger.info(f"  - Reading checkpoints from parent: {checkpoint_read_dir}")
                                output_logger.info(f"  - Writing checkpoints to child: {checkpoint_write_dir}")
                            else:
                                checkpoint_read_dir = str(own_checkpoint_dir)
                                checkpoint_write_dir = str(own_checkpoint_dir)
                                output_logger.warning(f"Parent checkpoint directory not found, using child directory: {checkpoint_write_dir}")
                        else:
                            # Regular run - use own checkpoint directory for both
                            checkpoint_read_dir = str(own_checkpoint_dir)
                            checkpoint_write_dir = str(own_checkpoint_dir)
                            output_logger.info(f"Regular run using own checkpoint directory: {checkpoint_write_dir}")
                    except Exception as e:
                        checkpoint_read_dir = str(own_checkpoint_dir)
                        checkpoint_write_dir = str(own_checkpoint_dir)
                        output_logger.warning(f"Failed to read metadata, using default directory: {checkpoint_write_dir}")
                else:
                    checkpoint_read_dir = str(own_checkpoint_dir)
                    checkpoint_write_dir = str(own_checkpoint_dir)
                    output_logger.info(f"No metadata file, using default directory: {checkpoint_write_dir}")

            # Replace with config-based checkpointing
            simulator.checkpointing = add_checkpointing_from_config_with_directories(
                simulator, 
                config_path=config_filename,
                checkpoint_read_dir=checkpoint_read_dir,
                checkpoint_write_dir=checkpoint_write_dir
            )
            output_logger.info(f"Checkpointing system initialised from config: {config_filename}")

        return simulator


    def clear_world(self):
        """Removes everyone from all possible groups, sets everyone's busy attribute
        to False, and cleans up skinny persons that have moved back home.

        """
        # Clear all groups first
        for super_group_name in self.activity_manager.all_super_groups:
            if "visits" in super_group_name:
                continue
            try:
                grouptype = getattr(self.world, super_group_name, None)
                if grouptype is not None:
                    for group in grouptype.members:
                        group.clear()
            except AttributeError:
                # If the attribute doesn't exist, just continue
                continue

        # Reset busy flags and leisure subgroups
        for person in self.world.people.members:            
            person.busy = False
            person.subgroups.leisure = None

        # Clean up skinny persons
        from june.demography import Person
        to_remove = []

        # Find all skinny persons (persons not on their home rank)
        for person_id, person in list(Person._persons.items()):
            if hasattr(person, '_current_rank') and person._current_rank != mpi_rank:
                to_remove.append(person_id)

        # Remove all identified skinny persons
        for person_id in to_remove:
            if person_id in Person._persons:
                del Person._persons[person_id]

        # Synchronise removal across ranks
        mpi_comm.Barrier()


    def do_timestep(self):
        """Perform a time step in the simulation. First, ActivityManager is called
        to send people to the corresponding subgroups according to the current daytime.
        Then we iterate over all the groups and create an InteractiveGroup object, which
        extracts the relevant information of each group to carry the interaction in it.
        We then pass the interactive group to the interaction module, which returns the ids
        of the people who got infected. We record the infection locations, update the health
        status of the population, and distribute scores among the infectors to calculate R0.

        """
        output_logger.info("==================== timestep ====================")
        tick_s, tickw_s = perf_counter(), wall_clock()
        tick, tickw = perf_counter(), wall_clock()

        if self.activity_manager.policies is not None:
            self.activity_manager.policies.interaction_policies.apply(
                date=self.timer.date, interaction=self.interaction
            )
            self.activity_manager.policies.regional_compliance.apply(
                date=self.timer.date, regions=self.world.regions
            )

        activities = self.timer.activities
        # apply events
        if self.events is not None:
            self.events.apply(
                date=self.timer.date,
                world=self.world,
                activities=activities,
                day_type=self.timer.day_type,
                simulator=self,
            )

        if not activities or len(activities) == 0:
            return

        (
            people_from_abroad_dict,
            n_people_from_abroad,
            n_people_going_abroad,
            to_send_abroad,  # useful for knowing who's MPI-ing, so can send extra info as needed.
        ) = self.activity_manager.do_timestep(record=self.record)
        tick_interaction = perf_counter()
        if self.interaction:
            self.interaction.current_time = self.timer.now

        # get the supergroup instances that are active in this time step:
        active_super_groups = self.activity_manager.active_super_groups
        super_group_instances = []
        for super_group_name in active_super_groups:
            if "visits" not in super_group_name:
                try:
                    super_group_instance = getattr(self.world, super_group_name, None)
                    if super_group_instance is not None and len(super_group_instance) > 0:
                        super_group_instances.append(super_group_instance)
                except (AttributeError, TypeError) as e:
                    output_logger.warning(f"Could not access supergroup {super_group_name}: {e}")
                    continue

        # Initialise counters for people tracking
        initial_people = len(self.world.people)  # People before movement
        n_cemetery = sum(len(cemetery.people) for cemetery in self.world.cemeteries.members)

        # Track infections
        infected_ids = []  # ids of the newly infected people
        infection_ids = []  # ids of the viruses they got

        output_logger.info(
            f"Info for rank {mpi_rank}, "
            f"Date = {self.timer.date}, "
            f"number of deaths = {n_cemetery}, "
            f"number of infected = {len(self.world.people.infected)}"
        )

       # Process groups for infections only
        for super_group in super_group_instances:
            for group in super_group:
                if group.external:
                    continue

                # Get people from abroad for this group
                people_from_abroad = people_from_abroad_dict.get(
                    group.spec, {}
                ).get(group.id, None)
                # Only track new infections, ignore group size
                new_infected, new_infections, _ = self.interaction.time_step_for_group(
                    group=group,
                    people_from_abroad=people_from_abroad,
                    delta_time=self.timer.duration,
                    record=self.record,
                )

                infected_ids.extend(new_infected)
                infection_ids.extend(new_infections)

        mpi_comm.Barrier()


        # Calculate final people count after movement
        final_people = len(self.world.people)  # People after movement

        tock_interaction = perf_counter()
        rank_logger.info(
            f"Rank {mpi_rank} -- interaction -- {tock_interaction-tick_interaction}"
        )

        self.epidemiology.do_timestep(
            simulator=self,
            world=self.world,
            timer=self.timer,
            record=self.record,
            infected_ids=infected_ids,
            infection_ids=infection_ids,
            people_from_abroad_dict=people_from_abroad_dict
        )

        self.interaction.print_transmission_statistics()

        is_end_of_day = self.timer.is_end_of_day()

        if is_end_of_day:
            # Sync school incident counts across MPI ranks for school avoidance behavior
            self._sync_school_incidents()

        # Clean old contacts at the end of the day
            if self.test_and_trace_enabled and self.contact_manager is not None:

                self.contact_manager.process_test_results(self.timer.now)

                mpi_comm.Barrier()    

                output_logger.info("Cleaning old contacts in the contact manager")
                output_logger.info(f"Current simulation day (fractional): {self.timer.now}")
                self.contact_manager.clean_old_contacts(
                    current_timestamp=self.timer.now,
                    days_to_remember=self.contact_retention_days,
                    force=True
                )

            # Add rat simulation step
            if self.rat_manager is not None:
                print("\n=== Running Rat Disease Simulation Step ===")
                # Pass the current time step duration to the rat_manager
                rat_results = self.rat_manager.time_step(1)
                # Optionally log some results
                print(f"Number of Rats: {self.rat_manager.num_rats}"),
                print(f"Rat infections: {rat_results['infected']}")
                print(f"Rats with high immunity: {rat_results['immunity_08']}")
                print(f"Rats with medium immunity: {rat_results['immunity_05']}")

                if self.epidemiology is not None and self.zoonotic_transmission is not None:
                    print("\n=== Processing Rat-to-Human Transmissions ===")
                    # Use the zoonotic transmission module instead
                    human_infections = self.zoonotic_transmission.process_rat_to_human_infections(
                        world=self.world,
                        timer=self.timer,
                        epidemiology=self.epidemiology,
                        record=self.record,
                        duration=1.0  # Full day duration (24 hours)
                    )
                    print(f"New human infections from rats: {human_infections}")
                    human2ratinfections = self.zoonotic_transmission.process_human_to_rat_infections(
                        world=self.world, 
                        timer=self.timer,
                        epidemiology=self.epidemiology,
                        record=self.record,
                        duration=1.0
                    )
                    print(f"New rat infections from humans: {human2ratinfections}")

                # Save rat simulation data for post-processing (all ranks)
                if self.save_rat_data:
                    # Determine output directory for data
                    if self.record and hasattr(self.record, 'record_path'):
                        rat_data_dir = str(self.record.record_path / "rat_data")
                    else:
                        rat_data_dir = "outputs/rat_data"  # Fallback

                    # Save simulation data (much faster than creating images)
                    self.rat_manager.save_simulation_data(
                        day=int(self.timer.now),
                        date=self.timer.date,
                        output_dir=rat_data_dir
                    )

        tick, tickw = perf_counter(), wall_clock()
        mpi_comm.Barrier()
        tock, tockw = perf_counter(), wall_clock()
        rank_logger.info(f"Rank {mpi_rank} -- interaction_waiting -- {tock-tick}")

        # Ensure all ranks have finished processing
        mpi_comm.Barrier()

        # Gather detailed counts from all ranks
        local_counts = (
            initial_people,  # Initial people count
            final_people,    # Final people count
            n_people_from_abroad,
            n_people_going_abroad
        )
        all_counts = mpi_comm.allgather(local_counts)

        # Calculate global totals
        total_initial = sum(c[0] for c in all_counts)
        total_final = sum(c[1] for c in all_counts)
        total_from_abroad = sum(c[2] for c in all_counts)
        total_going_abroad = sum(c[3] for c in all_counts)

        # Verify global conservation of people
        if total_initial != total_final:
            movement_delta = total_final - total_initial
            abroad_delta = total_from_abroad - total_going_abroad
            raise SimulatorError(
                f"Global people conservation error on rank {mpi_rank}:\n"
                f"Total initial people: {total_initial}\n"
                f"Total final people: {total_final}\n"
                f"Net change in people: {movement_delta}\n"
                f"Expected net change (from_abroad - going_abroad): {abroad_delta}\n"
                f"Discrepancy: {movement_delta - abroad_delta}\n"
                f"Movement details:\n"
                f"- Total from abroad: {total_from_abroad}\n"
                f"- Total going abroad: {total_going_abroad}\n"
                f"Local counts on this rank:\n"
                f"- Initial people: {initial_people}\n"
                f"- Final people: {final_people}\n"
                f"- From abroad: {n_people_from_abroad}\n"
                f"- Going abroad: {n_people_going_abroad}\n"
                f"- Net movement: {final_people - initial_people}"
            )

        if self.test_and_trace_enabled:
            #from june.records.event_recording import are_test_and_trace_policies_active
            #if are_test_and_trace_policies_active():
                #print_tt_simulation_report(days_simulated=self.timer.total_days)
            tt_recorder = GlobalContext.get_tt_event_recorder()
            if tt_recorder:
                tt_recorder.time_step(self.timer.now)


        # remove everyone from their active groups
        self.clear_world()

        tock, tockw = perf_counter(), wall_clock()
        output_logger.info(
            f"CMS: Timestep for rank {mpi_rank}/{mpi_size} - {tock - tick_s},"
            f"{tockw-tickw_s} - {self.timer.date}\n"
        )
        mpi_logger.info(f"{self.timer.date},{mpi_rank},timestep,{tock-tick_s}")

        # Force synchronisation point for random state consistency
        if mpi_available:
            mpi_comm.Barrier()

    def run(self):
        """Run simulation with n_seed initial infections"""

        # Calculate remaining days to run
        current_time = self.timer.now
        remaining_days = self.timer.total_days - current_time

        output_logger.info(
            f"Starting simulation for {self.timer.total_days} days at day {self.timer.date},"
            f"to run for {remaining_days:.2f} days"
        )

        if self.test_and_trace_enabled:
            # Check if TTEventRecorder already exists (from checkpoint restoration)
            existing_recorder = GlobalContext.get_tt_event_recorder()

            if existing_recorder is not None:
                output_logger.info("Test and Trace enabled, using existing TTEventRecorder from checkpoint")
                output_logger.info(f"Existing TTEventRecorder filename: {existing_recorder.filename}")
            else:
                output_logger.info("Test and Trace enabled, no existing TTEventRecorder found")
                output_logger.info("Initialising new TTEventRecorder")

                # Use record path for TTEventRecorder if available
                if self.record and hasattr(self.record, 'record_path'):
                    tt_output_dir = self.record.record_path / "h5_test_trace"
                else:
                    tt_output_dir = "./results/h5_test_trace"  # Fallback

                tt_recorder = TTEventRecorder(output_dir=str(tt_output_dir))
                GlobalContext.set_tt_event_recorder(tt_recorder)
                output_logger.info(f"New TTEventRecorder filename: {tt_recorder.filename}")
                output_logger.info(f"TTEventRecorder output directory: {tt_output_dir}")

                # Apply any pending TTEventRecorder data from checkpoint restoration
                if hasattr(self, '_pending_tt_event_recorder_data'):
                    output_logger.info("Applying pending TTEventRecorder data from checkpoint")
                    self._apply_pending_tt_event_recorder_data(tt_recorder)
                    delattr(self, '_pending_tt_event_recorder_data')
                    output_logger.info("TTEventRecorder data restoration completed")

        else:
            output_logger.info("Test and Trace disabled, skipping TTEventRecorder")
            GlobalContext.set_tt_event_recorder(None)

        # Rat manager is already initialised normally

        GlobalContext.set_simulator(self)

        # Update registered members home ranks after potential checkpoint restoration
        self.update_registered_members_home_ranks()

        # Update friend home ranks after domain splitting
        if self.friend_hangouts_enabled:
            self.update_friends_home_ranks()
            mpi_comm.Barrier()

        self.clear_world()

        if self.record is not None:
            try:
                self.record.parameters(
                    interaction=self.interaction,
                    epidemiology=self.epidemiology,
                    activity_manager=self.activity_manager,
                )
            except Exception as e:
                output_logger.error(f"Error recording parameters: {e}")

        #START OF THE SIMULATION LOOP
        while self.timer.date < self.timer.final_date:            

            if self.epidemiology:
                self.epidemiology.infection_seeds_timestep(
                    self.timer, record=self.record
                )

            mpi_comm.Barrier()
            if mpi_rank == 0:
                rank_logger.info("Next timestep")
            self.do_timestep()

            next(self.timer)

            # Check for checkpoint creation AFTER advancing to next day
            # This way checkpoint represents "ready to start the next day"
            if self.checkpointing is not None:
                integrate_checkpointing_in_simulation_loop(self.checkpointing)

        # Ensure all ranks have finished saving data 
        if self.save_rat_data:
            mpi_comm.Barrier()  # Wait for all ranks to finish saving data

            # Just report data location - no post-processing during simulation
            if mpi_rank == 0 and self.rat_manager is not None:
                if self.record and hasattr(self.record, 'record_path'):
                    rat_data_dir = self.record.record_path / "rat_data"
                else:
                    rat_data_dir = "outputs/rat_data"

                output_logger.info(f"Rat simulation data saved to: {rat_data_dir}")
                output_logger.info("To create visualizations, run:")
                output_logger.info(f"   python plot_maker/post_process_rat_data.py {rat_data_dir}")

        if self.record:
            self.record.combine_outputs()

        if self.test_and_trace_enabled:
            from june.records.test_trace_event_recording import export_simulation_results

            # Use record path for test and trace export if available
            if self.record and hasattr(self.record, 'record_path'):
                tt_export_dir = str(self.record.record_path)
            else:
                tt_export_dir = "./results"  # Fallback

            export_simulation_results(output_dir=tt_export_dir)

        # Print final validation summary for comparison
        #from june.checkpointing.simulator_checkpointing import print_validation_summary
        #print_validation_summary(self, "FINAL")


    def update_registered_members_home_ranks(self):
        """For each group type (companies, care homes, schools, universities),
        update the registered_members_ids to include the rank where each member lives.

        Works in both MPI and non-MPI environments:
        - In MPI mode: Ranks collaborate to find member home ranks
        - In non-MPI mode: All members are assigned to rank 0

        """

        if mpi_available:
            output_logger.info(f"Updating registered members home ranks on rank {mpi_rank}")
        else:
            output_logger.info("Updating registered members home ranks (non-MPI mode)")

        # Group types to check
        group_types = [
            ("companies", "Company"), 
            ("care_homes", "CareHome"), 
            ("schools", "School"),
            ("universities", "University")
        ]

        # Step 1: Collect all member IDs from groups on this rank
        all_member_ids = set()
        groups_with_members = []

        for group_type_name, _ in group_types:
            group_type = getattr(self.world, group_type_name, None)
            if group_type is None or not hasattr(group_type, "members"):
                continue

            for group in group_type.members:
                if not hasattr(group, "registered_members_ids"):
                    continue

                for subgroup_id, member_ids in group.registered_members_ids.items():
                    for i, member_info in enumerate(member_ids):
                        # Extract the member ID (handle both plain IDs and tuples)
                        if isinstance(member_info, tuple) and len(member_info) == 2:
                            member_id = member_info[0]
                        else:
                            member_id = member_info

                        all_member_ids.add(member_id)

                # Track groups for later updating
                groups_with_members.append((group_type_name, group))

        # Convert to a sorted list for consistent ordering
        all_member_ids = sorted(list(all_member_ids))

        from june.demography.person import Person

        # Different logic for MPI vs non-MPI
        if mpi_available:
            # MPI mode: Distributed ID lookup

            # Step 2: Check which IDs exist on this rank
            local_id_to_rank = {}
            unknown_ids = []

            for member_id in all_member_ids:
                person = Person.find_by_id(member_id)
                if person is not None:
                    # Found locally
                    local_id_to_rank[member_id] = mpi_rank
                else:
                    # Not found locally - add to list of IDs to query other ranks
                    unknown_ids.append(member_id)

            # Step 3: Share only the unknown IDs with other ranks
            all_unknown_ids = mpi_comm.allgather(unknown_ids)

            # Step 4: Check if any of the IDs from other ranks exist on this rank
            additional_mappings = {}

            for rank, rank_unknown_ids in enumerate(all_unknown_ids):
                if rank == mpi_rank:
                    # Skip our own unknown IDs
                    continue

                for member_id in rank_unknown_ids:
                    person = Person.find_by_id(member_id)
                    if person is not None:
                        # We found an ID that another rank was looking for
                        additional_mappings[member_id] = mpi_rank

            # Step 5: Share these additional mappings
            all_additional_mappings = mpi_comm.allgather(additional_mappings)

            # Build the global ID-to-rank mapping
            global_id_to_rank = {}
            # Add our local findings first
            global_id_to_rank.update(local_id_to_rank)

            # Add additional mappings from all ranks
            for rank_mappings in all_additional_mappings:
                global_id_to_rank.update(rank_mappings)

        else:
            # Non-MPI mode: All IDs are on rank 0
            global_id_to_rank = {member_id: 0 for member_id in all_member_ids}

        # Step 6: Update all groups with the correct home rank information
        for group_type_name, group in groups_with_members:
            for subgroup_id, member_ids in list(group.registered_members_ids.items()):
                updated_member_ids = []

                for member_info in member_ids:
                    # Extract the member ID (handle both plain IDs and tuples)
                    if isinstance(member_info, tuple) and len(member_info) == 2:
                        member_id = member_info[0]
                    else:
                        member_id = member_info

                    # Get the home rank from the global mapping
                    home_rank = global_id_to_rank.get(member_id)

                    if home_rank is not None:
                        # Add to the updated list
                        updated_member_ids.append((member_id, home_rank))
                    else:
                        # ID not found in any rank, use default of 0
                        # This should not happen if everything is working correctly
                        updated_member_ids.append((member_id, 0))

                # Replace the old list with the updated one
                group.registered_members_ids[subgroup_id] = updated_member_ids

        # Ensure all ranks are synchronised in MPI mode
        mpi_comm.Barrier()

    def update_friends_home_ranks(self):
        """Update the home rank for each person's friends after domain splitting.
        Replaces the default home rank (0) with the actual home rank where each friend resides.

        Works in both MPI and non-MPI environments:
        - In MPI mode: Ranks collaborate to find friend home ranks
        - In non-MPI mode: All friends are assigned to rank 0

        """

        if not mpi_available:
            return

        # Step 1: Collect all friend IDs from people on this rank
        all_friend_ids = set()
        people_with_friends = []

        for person in self.world.people.members:
            if hasattr(person, 'friends') and person.friends:
                # Handle both old format (just home_rank) and new format (dict)
                for friend_id, friend_data in person.friends.items():
                    all_friend_ids.add(friend_id)
                people_with_friends.append(person)

        # Convert to sorted list for consistent ordering
        all_friend_ids = sorted(list(all_friend_ids))

        from june.demography.person import Person

        # Step 2: Check which friend IDs exist on this rank
        local_id_to_rank = {}
        unknown_ids = []

        for friend_id in all_friend_ids:
            person = Person.find_by_id(friend_id)
            if person is not None:
                # Found locally - store their home rank
                local_id_to_rank[friend_id] = getattr(person, '_home_rank', mpi_rank)
            else:
                # Not found locally - add to list of IDs to query other ranks
                unknown_ids.append(friend_id)

        # Step 3: Share unknown IDs with other ranks
        all_unknown_ids = mpi_comm.allgather(unknown_ids)

        # Step 4: Check if any of the IDs from other ranks exist on this rank
        additional_mappings = {}

        for rank, rank_unknown_ids in enumerate(all_unknown_ids):
            if rank == mpi_rank:
                # Skip our own unknown IDs
                continue

            for friend_id in rank_unknown_ids:
                person = Person.find_by_id(friend_id)
                if person is not None:
                    # Found a friend that another rank was looking for
                    additional_mappings[friend_id] = getattr(person, '_home_rank', mpi_rank)

        # Step 5: Share these additional mappings
        all_additional_mappings = mpi_comm.allgather(additional_mappings)

        # Build the global friend_id-to-home_rank mapping
        global_friend_to_home_rank = {}
        # Add our local findings first
        global_friend_to_home_rank.update(local_id_to_rank)

        # Add additional mappings from all ranks
        for rank_mappings in all_additional_mappings:
            global_friend_to_home_rank.update(rank_mappings)

        # Step 6: Update each person's friends dictionary with correct home ranks
        updated_people = 0
        updated_friends = 0

        for person in people_with_friends:
            person_updated = False

            # Create a new friends dictionary with updated home ranks
            updated_friends_dict = {}

            for friend_id, friend_data in person.friends.items():
                # Handle both old format (just home_rank) and new format (dict)
                if isinstance(friend_data, dict):
                    # New format - update home_rank, keep hobbies
                    current_home_rank = friend_data.get("home_rank", 0)
                    actual_home_rank = global_friend_to_home_rank.get(friend_id)

                    if actual_home_rank is not None:
                        updated_friends_dict[friend_id] = {
                            "home_rank": actual_home_rank,
                            "hobbies": friend_data.get("hobbies", [])
                        }

                        # Track if this friend's home rank was updated
                        if current_home_rank != actual_home_rank:
                            updated_friends += 1
                            person_updated = True
                    else:
                        # Friend ID not found in any rank, keep current values
                        updated_friends_dict[friend_id] = friend_data
                        output_logger.warning(f"Friend ID {friend_id} not found in any rank for person {person.id}")

            # Replace the person's friends dictionary
            person.friends = updated_friends_dict

            if person_updated:
                updated_people += 1

        # Log statistics
        # Gather statistics from all ranks
        all_updated_people = mpi_comm.allgather(updated_people)
        all_updated_friends = mpi_comm.allgather(updated_friends)

        if mpi_rank == 0:
            total_updated_people = sum(all_updated_people)
            total_updated_friends = sum(all_updated_friends)
            output_logger.info(f"Friends home ranks updated: {total_updated_people} people, {total_updated_friends} friend relationships")

        # Ensure all ranks are synchronised
        mpi_comm.Barrier()

    def restore_from_checkpoint(self, checkpoint_name: str) -> bool:
        """Restore simulation from a checkpoint.

        Args:
            checkpoint_name (str): Name of the checkpoint to restore from

        Returns:
            bool: True if restoration was successful
        """
        if self.checkpointing is None:
            output_logger.warning("Checkpointing system not available")
            return False

        try:
            return self.checkpointing.restore_from_checkpoint(checkpoint_name)
        except Exception as e:
            output_logger.error(f"Error restoring from checkpoint: {e}")
            return False

    def _apply_pending_tt_event_recorder_data(self, tt_recorder):
        """Apply pending TTEventRecorder data from checkpoint restoration.

        Args:
            tt_recorder (TTEventRecorder): The newly created TTEventRecorder to restore data into
        """
        if not hasattr(self, '_pending_tt_event_recorder_data'):
            return

        tt_data = self._pending_tt_event_recorder_data
        output_logger.info(f"Applying TTEventRecorder data from checkpoint")

        try:
            # Restore core counters
            tt_recorder.total_counters.update(tt_data.get('total_counters', {}))

            # Restore daily counters
            daily_data = tt_data.get('daily_counters', {})
            for day_str, counters in daily_data.items():
                day = int(day_str)
                for event_type, count in counters.items():
                    tt_recorder.daily_counters[day][event_type] = count

            # Restore unique IDs (convert back to sets)
            unique_ids_data = tt_data.get('unique_ids', {})
            for event_type, id_list in unique_ids_data.items():
                tt_recorder.unique_ids[event_type].update(id_list)

            # Restore current status (convert back to sets)
            currently_data = tt_data.get('currently', {})
            for status_type, id_list in currently_data.items():
                tt_recorder.currently[status_type].update(id_list)

            # Restore deltas
            tt_recorder.deltas.update(tt_data.get('deltas', {}))

            # Restore state tracking
            tt_recorder.last_delta_reset = tt_data.get('last_delta_reset', 0)
            if 'buffer_size' in tt_data:
                tt_recorder._buffer_size = tt_data['buffer_size']

            # Note: Skip event buffer restoration to avoid state changes
            # Event buffer should typically be empty after checkpointing anyway

            output_logger.info(f"TTEventRecorder data applied successfully:")
            output_logger.info(f"  - Total tests: {tt_recorder.total_counters.get('tested', 0)}")
            output_logger.info(f"  - Unique people tested: {len(tt_recorder.unique_ids.get('tested', set()))}")
            output_logger.info(f"  - Days with data: {len(tt_recorder.daily_counters)}")
            output_logger.info(f"  - Currently quarantined: {len(tt_recorder.currently.get('quarantined', set()))}")
            output_logger.info(f"  - Currently isolated: {len(tt_recorder.currently.get('isolated', set()))}")

        except Exception as e:
            output_logger.error(f"Failed to apply TTEventRecorder data: {e}")
            # Continue anyway - TTEventRecorder will work with default empty state

    def _sync_school_incidents(self):
        """Synchronize school incident counts (deaths and ICU transfers) across MPI ranks.
        Called once per day at end of timestep to ensure all ranks have consistent data
        for school avoidance behavior calculations.

        """
        if not mpi_available or mpi_size == 1:
            return  # No need to sync in single-rank mode

        # Collect local school incidents AND cross-rank incidents
        local_data = {
            'schools': {},  # school_id -> {'deaths': count, 'icu': count} for local schools
            'external_incidents': []  # [{'school_id': id, 'type': 'death'/'icu'}] for external schools
        }

        # Add local school counts
        if hasattr(self.world, 'schools'):
            for school in self.world.schools:
                if hasattr(school, 'student_deaths') and hasattr(school, 'student_icu_transfers'):
                    local_data['schools'][school.id] = {
                        'deaths': school.student_deaths,
                        'icu': school.student_icu_transfers
                    }

        # Add external incidents (stored temporarily during the day)
        if hasattr(self, '_external_school_incidents'):
            local_data['external_incidents'] = self._external_school_incidents
            # Clear after collecting
            self._external_school_incidents = []

        # Gather all data from all ranks
        all_rank_data = mpi_comm.allgather(local_data)

        # Aggregate incident counts
        global_school_incidents = {}

        # Process local school data
        for rank_data in all_rank_data:
            for school_id, incidents in rank_data['schools'].items():
                if school_id not in global_school_incidents:
                    global_school_incidents[school_id] = {'deaths': 0, 'icu': 0}
                global_school_incidents[school_id]['deaths'] += incidents['deaths'] 
                global_school_incidents[school_id]['icu'] += incidents['icu']

        # Process external incidents (cross-rank)
        for rank_data in all_rank_data:
            for incident in rank_data['external_incidents']:
                school_id = incident['school_id']
                incident_type = incident['type']

                if school_id not in global_school_incidents:
                    global_school_incidents[school_id] = {'deaths': 0, 'icu': 0}

                if incident_type == 'death':
                    global_school_incidents[school_id]['deaths'] += 1
                elif incident_type == 'icu':
                    global_school_incidents[school_id]['icu'] += 1

        # Update all local schools with global incident counts
        if hasattr(self.world, 'schools'):
            for school in self.world.schools:
                if school.id in global_school_incidents:
                    school.student_deaths = global_school_incidents[school.id]['deaths']
                    school.student_icu_transfers = global_school_incidents[school.id]['icu']

        # Store global incident data in world for ExternalGroup access
        if not hasattr(self.world, 'global_school_incidents'):
            self.world.global_school_incidents = {}
        self.world.global_school_incidents.update(global_school_incidents)

        # Debug output for rank 0
        if mpi_rank == 0:
            total_incidents = sum(data['deaths'] + data['icu'] for data in global_school_incidents.values())
            if total_incidents > 0:
                logging.info(f"[SCHOOL SYNC] Synchronized {len(global_school_incidents)} schools with incidents across {mpi_size} ranks")

    def _track_external_school_incident(self, school_id, incident_type):
        """Track incidents that happen to students whose schools are in other MPI ranks.
        These are collected and synchronized at end of day.

        Args:
            school_id: 
            incident_type: 
        """
        if not hasattr(self, '_external_school_incidents'):
            self._external_school_incidents = []

        self._external_school_incidents.append({
            'school_id': school_id,
            'type': incident_type
        })

__init__(world, interaction, timer, activity_manager, epidemiology, events=None, record=None, feature_flags=None)

Class to run an epidemic spread simulation on the world.

Parameters:

Name	Type	Description	Default
world	World	Instance of the world class.	required
interaction	Interaction	Instance of the Interaction class	required
timer	Timer	The timer class keeps track of what stage of the day it is, and which types of activities are allowed	required
activity_manager	ActivityManager	Helps manage allocating which activities people choose to do in a given timestep.	required
epidemiology	Epidemiology	Contains information about how the disease propogates.	required
events	Optional[Events]	Events object, adding in special events.	None
record	Optional[Record]	Record object to keep track of the simulation as it goes.	None
feature_flags	Optional[dict]	Dictionary that says which features are enabled.	None

Source code in june/simulator.py

def __init__(
    self,
    world: World,
    interaction: Interaction,
    timer: Timer,
    activity_manager: ActivityManager,
    epidemiology: Epidemiology,
    events: Optional[Events] = None,
    record: Optional[Record] = None,
    feature_flags: Optional[dict] = None
):
    """
    Class to run an epidemic spread simulation on the world.

    Parameters:
      world (World):
        Instance of the world class.
      interaction (Interaction):
        Instance of the Interaction class
      timer (Timer):
        The timer class keeps track of what stage of the day it is,
        and which types of activities are allowed
      activity_manager (ActivityManager):
        Helps manage allocating which activities people choose to do in a
        given timestep.
      epidemiology (Epidemiology):
        Contains information about how the disease propogates.
      events (Optional[Events]):
        Events object, adding in special events.
      record (Optional[Record]):
        Record object to keep track of the simulation as it goes.
      feature_flags (Optional[dict]):
        Dictionary that says which features are enabled. 
    """
    # Process feature flags with defaults
    if feature_flags is None:
        feature_flags = {}

    self.friend_hangouts_enabled = feature_flags.get("friend_hangouts_enabled", False)
    self.sexual_encounters_enabled = feature_flags.get("sexual_encounters_enabled", False)
    self.test_and_trace_enabled = feature_flags.get("test_and_trace_enabled", False)
    self.ratty_dynamics_enabled = feature_flags.get("ratty_dynamics_enabled", False)
    self.rat_data_saving_enabled = feature_flags.get("rat_data_saving_enabled", True)  # Default to True

    # Original initialisation code
    self.activity_manager = activity_manager
    self.world = world
    self.interaction = interaction
    self.events = events
    self.timer = timer
    self.epidemiology = epidemiology

    if self.epidemiology:
        self.epidemiology.set_medical_care(
            world=world, activity_manager=activity_manager
        )
        self.epidemiology.set_immunity(self.world)
        self.epidemiology.set_past_vaccinations(
            people=self.world.people, date=self.timer.date, record=record
        )

    if self.events is not None:
        self.events.init_events(world=world)

    self.record = record
    if self.record is not None and self.record.record_static_data:
        self.record.static_data(world=world)

    # Only initialise rat manager if feature is enabled
    self.rat_manager = None
    self.save_rat_data = False

    if self.ratty_dynamics_enabled:
        output_logger.info("Ratty dynamics enabled, initialising rat manager")
        from june.zoonosis.rat_manager import RatManager

        # Use record path for rat density files if available
        if self.record and hasattr(self.record, 'record_path'):
            rat_density_path = str(self.record.record_path / "rat_density_map.npy")
        else:
            rat_density_path = "rat_density_map.npy"  # Fallback

        # Normal initialisation - will be overwritten if checkpoint is restored
        self.rat_manager = RatManager(world=world, precomputed_density_path=rat_density_path)
        self.save_rat_data = self.rat_data_saving_enabled
        output_logger.info(f"Rat data saving: {self.save_rat_data}")
        output_logger.info(f"Rat density path: {rat_density_path}")

    # Initialise contact manager if test and trace is enabled or sexual encounters are enabled
    self.contact_manager = None
    if self.test_and_trace_enabled or self.sexual_encounters_enabled:
        output_logger.info("Test and Trace or Sexual Encounters enabled, initialising contact manager")
        self.contact_manager = ContactManager(self)

        # Configure the contact retention based on policy settings            
        disease_config = None

        disease_config = GlobalContext.get_disease_config()
        output_logger.info("Using disease_config from GlobalContext")

        # Use disease_config if available
        if disease_config is not None:
            tracing_data = disease_config.policy_manager.get_policy_data("tracing")
            self.contact_retention_days = tracing_data.get("contact_retention_days", 14) if tracing_data else 14
            output_logger.info(f"Using contact retention days from config: {self.contact_retention_days}")

        # Connect contact manager to leisure system if it exists
        if self.activity_manager.leisure is not None:
            output_logger.info("Connecting contact manager to leisure system")
            self.activity_manager.leisure.set_contact_manager(self.contact_manager)

        # Connect contact manager to sexual encounter system if it exists
        if self.activity_manager.sexual_encounter is not None:
            output_logger.info("Connecting contact manager to sexual encounter system")
            self.activity_manager.sexual_encounter.contact_manager = self.contact_manager
    else:
        output_logger.info("Test and Trace and Sexual Encounters disabled, skipping contact manager initialisation")

    # Initialise zoonotic transmission if rat_manager is properly initialised and feature is enabled
    self.zoonotic_transmission = None
    if self.ratty_dynamics_enabled and self.rat_manager is not None and ZoonoticTransmission is not None:
        output_logger.info("Initialising zoonotic transmission module")
        self.zoonotic_transmission = ZoonoticTransmission(rat_manager=self.rat_manager)
        output_logger.info("Zoonotic transmission module initialised successfully")

    # Initialise checkpointing system placeholder
    # Actual checkpointing setup happens in from_file() method with config
    self.checkpointing = None

clear_world()

Removes everyone from all possible groups, sets everyone's busy attribute to False, and cleans up skinny persons that have moved back home.

Source code in june/simulator.py

def clear_world(self):
    """Removes everyone from all possible groups, sets everyone's busy attribute
    to False, and cleans up skinny persons that have moved back home.

    """
    # Clear all groups first
    for super_group_name in self.activity_manager.all_super_groups:
        if "visits" in super_group_name:
            continue
        try:
            grouptype = getattr(self.world, super_group_name, None)
            if grouptype is not None:
                for group in grouptype.members:
                    group.clear()
        except AttributeError:
            # If the attribute doesn't exist, just continue
            continue

    # Reset busy flags and leisure subgroups
    for person in self.world.people.members:            
        person.busy = False
        person.subgroups.leisure = None

    # Clean up skinny persons
    from june.demography import Person
    to_remove = []

    # Find all skinny persons (persons not on their home rank)
    for person_id, person in list(Person._persons.items()):
        if hasattr(person, '_current_rank') and person._current_rank != mpi_rank:
            to_remove.append(person_id)

    # Remove all identified skinny persons
    for person_id in to_remove:
        if person_id in Person._persons:
            del Person._persons[person_id]

    # Synchronise removal across ranks
    mpi_comm.Barrier()

do_timestep()

Perform a time step in the simulation. First, ActivityManager is called to send people to the corresponding subgroups according to the current daytime. Then we iterate over all the groups and create an InteractiveGroup object, which extracts the relevant information of each group to carry the interaction in it. We then pass the interactive group to the interaction module, which returns the ids of the people who got infected. We record the infection locations, update the health status of the population, and distribute scores among the infectors to calculate R0.

Source code in june/simulator.py

def do_timestep(self):
    """Perform a time step in the simulation. First, ActivityManager is called
    to send people to the corresponding subgroups according to the current daytime.
    Then we iterate over all the groups and create an InteractiveGroup object, which
    extracts the relevant information of each group to carry the interaction in it.
    We then pass the interactive group to the interaction module, which returns the ids
    of the people who got infected. We record the infection locations, update the health
    status of the population, and distribute scores among the infectors to calculate R0.

    """
    output_logger.info("==================== timestep ====================")
    tick_s, tickw_s = perf_counter(), wall_clock()
    tick, tickw = perf_counter(), wall_clock()

    if self.activity_manager.policies is not None:
        self.activity_manager.policies.interaction_policies.apply(
            date=self.timer.date, interaction=self.interaction
        )
        self.activity_manager.policies.regional_compliance.apply(
            date=self.timer.date, regions=self.world.regions
        )

    activities = self.timer.activities
    # apply events
    if self.events is not None:
        self.events.apply(
            date=self.timer.date,
            world=self.world,
            activities=activities,
            day_type=self.timer.day_type,
            simulator=self,
        )

    if not activities or len(activities) == 0:
        return

    (
        people_from_abroad_dict,
        n_people_from_abroad,
        n_people_going_abroad,
        to_send_abroad,  # useful for knowing who's MPI-ing, so can send extra info as needed.
    ) = self.activity_manager.do_timestep(record=self.record)
    tick_interaction = perf_counter()
    if self.interaction:
        self.interaction.current_time = self.timer.now

    # get the supergroup instances that are active in this time step:
    active_super_groups = self.activity_manager.active_super_groups
    super_group_instances = []
    for super_group_name in active_super_groups:
        if "visits" not in super_group_name:
            try:
                super_group_instance = getattr(self.world, super_group_name, None)
                if super_group_instance is not None and len(super_group_instance) > 0:
                    super_group_instances.append(super_group_instance)
            except (AttributeError, TypeError) as e:
                output_logger.warning(f"Could not access supergroup {super_group_name}: {e}")
                continue

    # Initialise counters for people tracking
    initial_people = len(self.world.people)  # People before movement
    n_cemetery = sum(len(cemetery.people) for cemetery in self.world.cemeteries.members)

    # Track infections
    infected_ids = []  # ids of the newly infected people
    infection_ids = []  # ids of the viruses they got

    output_logger.info(
        f"Info for rank {mpi_rank}, "
        f"Date = {self.timer.date}, "
        f"number of deaths = {n_cemetery}, "
        f"number of infected = {len(self.world.people.infected)}"
    )

   # Process groups for infections only
    for super_group in super_group_instances:
        for group in super_group:
            if group.external:
                continue

            # Get people from abroad for this group
            people_from_abroad = people_from_abroad_dict.get(
                group.spec, {}
            ).get(group.id, None)
            # Only track new infections, ignore group size
            new_infected, new_infections, _ = self.interaction.time_step_for_group(
                group=group,
                people_from_abroad=people_from_abroad,
                delta_time=self.timer.duration,
                record=self.record,
            )

            infected_ids.extend(new_infected)
            infection_ids.extend(new_infections)

    mpi_comm.Barrier()


    # Calculate final people count after movement
    final_people = len(self.world.people)  # People after movement

    tock_interaction = perf_counter()
    rank_logger.info(
        f"Rank {mpi_rank} -- interaction -- {tock_interaction-tick_interaction}"
    )

    self.epidemiology.do_timestep(
        simulator=self,
        world=self.world,
        timer=self.timer,
        record=self.record,
        infected_ids=infected_ids,
        infection_ids=infection_ids,
        people_from_abroad_dict=people_from_abroad_dict
    )

    self.interaction.print_transmission_statistics()

    is_end_of_day = self.timer.is_end_of_day()

    if is_end_of_day:
        # Sync school incident counts across MPI ranks for school avoidance behavior
        self._sync_school_incidents()

    # Clean old contacts at the end of the day
        if self.test_and_trace_enabled and self.contact_manager is not None:

            self.contact_manager.process_test_results(self.timer.now)

            mpi_comm.Barrier()    

            output_logger.info("Cleaning old contacts in the contact manager")
            output_logger.info(f"Current simulation day (fractional): {self.timer.now}")
            self.contact_manager.clean_old_contacts(
                current_timestamp=self.timer.now,
                days_to_remember=self.contact_retention_days,
                force=True
            )

        # Add rat simulation step
        if self.rat_manager is not None:
            print("\n=== Running Rat Disease Simulation Step ===")
            # Pass the current time step duration to the rat_manager
            rat_results = self.rat_manager.time_step(1)
            # Optionally log some results
            print(f"Number of Rats: {self.rat_manager.num_rats}"),
            print(f"Rat infections: {rat_results['infected']}")
            print(f"Rats with high immunity: {rat_results['immunity_08']}")
            print(f"Rats with medium immunity: {rat_results['immunity_05']}")

            if self.epidemiology is not None and self.zoonotic_transmission is not None:
                print("\n=== Processing Rat-to-Human Transmissions ===")
                # Use the zoonotic transmission module instead
                human_infections = self.zoonotic_transmission.process_rat_to_human_infections(
                    world=self.world,
                    timer=self.timer,
                    epidemiology=self.epidemiology,
                    record=self.record,
                    duration=1.0  # Full day duration (24 hours)
                )
                print(f"New human infections from rats: {human_infections}")
                human2ratinfections = self.zoonotic_transmission.process_human_to_rat_infections(
                    world=self.world, 
                    timer=self.timer,
                    epidemiology=self.epidemiology,
                    record=self.record,
                    duration=1.0
                )
                print(f"New rat infections from humans: {human2ratinfections}")

            # Save rat simulation data for post-processing (all ranks)
            if self.save_rat_data:
                # Determine output directory for data
                if self.record and hasattr(self.record, 'record_path'):
                    rat_data_dir = str(self.record.record_path / "rat_data")
                else:
                    rat_data_dir = "outputs/rat_data"  # Fallback

                # Save simulation data (much faster than creating images)
                self.rat_manager.save_simulation_data(
                    day=int(self.timer.now),
                    date=self.timer.date,
                    output_dir=rat_data_dir
                )

    tick, tickw = perf_counter(), wall_clock()
    mpi_comm.Barrier()
    tock, tockw = perf_counter(), wall_clock()
    rank_logger.info(f"Rank {mpi_rank} -- interaction_waiting -- {tock-tick}")

    # Ensure all ranks have finished processing
    mpi_comm.Barrier()

    # Gather detailed counts from all ranks
    local_counts = (
        initial_people,  # Initial people count
        final_people,    # Final people count
        n_people_from_abroad,
        n_people_going_abroad
    )
    all_counts = mpi_comm.allgather(local_counts)

    # Calculate global totals
    total_initial = sum(c[0] for c in all_counts)
    total_final = sum(c[1] for c in all_counts)
    total_from_abroad = sum(c[2] for c in all_counts)
    total_going_abroad = sum(c[3] for c in all_counts)

    # Verify global conservation of people
    if total_initial != total_final:
        movement_delta = total_final - total_initial
        abroad_delta = total_from_abroad - total_going_abroad
        raise SimulatorError(
            f"Global people conservation error on rank {mpi_rank}:\n"
            f"Total initial people: {total_initial}\n"
            f"Total final people: {total_final}\n"
            f"Net change in people: {movement_delta}\n"
            f"Expected net change (from_abroad - going_abroad): {abroad_delta}\n"
            f"Discrepancy: {movement_delta - abroad_delta}\n"
            f"Movement details:\n"
            f"- Total from abroad: {total_from_abroad}\n"
            f"- Total going abroad: {total_going_abroad}\n"
            f"Local counts on this rank:\n"
            f"- Initial people: {initial_people}\n"
            f"- Final people: {final_people}\n"
            f"- From abroad: {n_people_from_abroad}\n"
            f"- Going abroad: {n_people_going_abroad}\n"
            f"- Net movement: {final_people - initial_people}"
        )

    if self.test_and_trace_enabled:
        #from june.records.event_recording import are_test_and_trace_policies_active
        #if are_test_and_trace_policies_active():
            #print_tt_simulation_report(days_simulated=self.timer.total_days)
        tt_recorder = GlobalContext.get_tt_event_recorder()
        if tt_recorder:
            tt_recorder.time_step(self.timer.now)


    # remove everyone from their active groups
    self.clear_world()

    tock, tockw = perf_counter(), wall_clock()
    output_logger.info(
        f"CMS: Timestep for rank {mpi_rank}/{mpi_size} - {tock - tick_s},"
        f"{tockw-tickw_s} - {self.timer.date}\n"
    )
    mpi_logger.info(f"{self.timer.date},{mpi_rank},timestep,{tock-tick_s}")

    # Force synchronisation point for random state consistency
    if mpi_available:
        mpi_comm.Barrier()

from_file(world, interaction, policies=None, events=None, epidemiology=None, leisure=None, travel=None, sexual_encounter=None, config_filename=default_config_filename, record=None) classmethod

Load config for simulator from world.yaml

Parameters:

Name	Type	Description	Default
world	World		required
interaction	Interaction		required
policies	Optional[Policies]	(Default value = None)	None
events	Optional[Events]	(Default value = None)	None
epidemiology	Optional[Epidemiology]	(Default value = None)	None
leisure	Optional[Leisure]	(Default value = None)	None
travel	Optional[Travel]	(Default value = None)	None
sexual_encounter	Optional[SexualEncounter]	(Default value = None)	None
config_filename	str	(Default value = default_config_filename)	default_config_filename
record	Optional[Record]	(Default value = None)	None

Source code in june/simulator.py

@classmethod
def from_file(
    cls,
    world: World,
    interaction: Interaction,
    policies: Optional[Policies] = None,
    events: Optional[Events] = None,
    epidemiology: Optional[Epidemiology] = None,
    leisure: Optional[Leisure] = None,
    travel: Optional[Travel] = None,
    sexual_encounter: Optional[SexualEncounter] = None,
    config_filename: str = default_config_filename,
    record: Optional[Record] = None,
) -> "Simulator":
    """Load config for simulator from world.yaml

    Args:
        world (World): 
        interaction (Interaction): 
        policies (Optional[Policies], optional): (Default value = None)
        events (Optional[Events], optional): (Default value = None)
        epidemiology (Optional[Epidemiology], optional): (Default value = None)
        leisure (Optional[Leisure], optional): (Default value = None)
        travel (Optional[Travel], optional): (Default value = None)
        sexual_encounter (Optional[SexualEncounter], optional): (Default value = None)
        config_filename (str, optional): (Default value = default_config_filename)
        record (Optional[Record], optional): (Default value = None)
    """
    # Load the configuration file to get feature flags
    with open(config_filename) as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    # Extract feature flags with defaults in case they're not in the config
    features = config.get("features", {})

    #Friendship networks
    friend_hangouts_config = features.get("friend_hangouts", {"enabled": False})
    friend_hangouts_enabled = friend_hangouts_config.get("enabled", False)

    # Sexual encounters settings
    sexual_encounters_config = features.get("sexual_encounters", {"enabled": False})
    sexual_encounters_enabled = sexual_encounters_config.get("enabled", False)

    # Test and trace settings
    test_and_trace_config = features.get("test_and_trace", {"enabled": False})
    test_and_trace_enabled = test_and_trace_config.get("enabled", False)

    # Ratty dynamics settings
    ratty_dynamics_config = features.get("ratty_dynamics", {"enabled": False})
    ratty_dynamics_enabled = ratty_dynamics_config.get("enabled", False)
    rat_data_saving_enabled = ratty_dynamics_config.get("save_data", True)  # Default to True


    output_logger.info(f"Feature flags from config: Friend hangouts: {friend_hangouts_enabled}, "
                    f"Sexual encounters: {sexual_encounters_enabled}, "
                    f"Test and Trace: {test_and_trace_enabled}, "
                    f"Ratty Dynamics: {ratty_dynamics_enabled}, "
                    f"Rat Data Saving: {rat_data_saving_enabled}, ")

    # Continue with original method
    timer = Timer.from_file(config_filename=config_filename)
    activity_manager = cls.ActivityManager.from_file(
        config_filename=config_filename,
        world=world,
        leisure=leisure,
        travel=travel,
        sexual_encounter=sexual_encounter,
        policies=policies,
        timer=timer,
        record=record,
    )

    # Store feature flags to pass to the constructor
    feature_flags = {
        "friend_hangouts_enabled": friend_hangouts_enabled,
        "sexual_encounters_enabled": sexual_encounters_enabled,
        "test_and_trace_enabled": test_and_trace_enabled,
        "ratty_dynamics_enabled": ratty_dynamics_enabled,
        "rat_data_saving_enabled": rat_data_saving_enabled
    }

    simulator = cls(
        world=world,
        interaction=interaction,
        timer=timer,
        events=events,
        activity_manager=activity_manager,
        epidemiology=epidemiology,
        record=record,
        feature_flags=feature_flags
    )

    # Initialise checkpointing system with config file
    if CHECKPOINTING_AVAILABLE:
        # Replace the default checkpointing with config-based one

        # Get the checkpoint directories (separate read/write for child runs)
        checkpoint_read_dir = "checkpoints"
        checkpoint_write_dir = "checkpoints"

        if hasattr(simulator, 'record') and simulator.record and hasattr(simulator.record, 'record_path'):
            run_dir = simulator.record.record_path.parent
            own_checkpoint_dir = run_dir / "checkpoints"

            # Check if this might be a child run by looking for metadata
            metadata_file = run_dir / "metadata.json"
            if metadata_file.exists():
                try:
                    import json
                    with open(metadata_file, 'r') as f:
                        metadata = json.load(f)

                    parent_run_id = metadata.get("parent_run_id")
                    if parent_run_id:
                        # This is a child run - separate read and write directories
                        parent_checkpoint_dir = run_dir.parent / parent_run_id / "checkpoints"
                        if parent_checkpoint_dir.exists():
                            checkpoint_read_dir = str(parent_checkpoint_dir)
                            checkpoint_write_dir = str(own_checkpoint_dir)
                            output_logger.info(f"Child run detected:")
                            output_logger.info(f"  - Reading checkpoints from parent: {checkpoint_read_dir}")
                            output_logger.info(f"  - Writing checkpoints to child: {checkpoint_write_dir}")
                        else:
                            checkpoint_read_dir = str(own_checkpoint_dir)
                            checkpoint_write_dir = str(own_checkpoint_dir)
                            output_logger.warning(f"Parent checkpoint directory not found, using child directory: {checkpoint_write_dir}")
                    else:
                        # Regular run - use own checkpoint directory for both
                        checkpoint_read_dir = str(own_checkpoint_dir)
                        checkpoint_write_dir = str(own_checkpoint_dir)
                        output_logger.info(f"Regular run using own checkpoint directory: {checkpoint_write_dir}")
                except Exception as e:
                    checkpoint_read_dir = str(own_checkpoint_dir)
                    checkpoint_write_dir = str(own_checkpoint_dir)
                    output_logger.warning(f"Failed to read metadata, using default directory: {checkpoint_write_dir}")
            else:
                checkpoint_read_dir = str(own_checkpoint_dir)
                checkpoint_write_dir = str(own_checkpoint_dir)
                output_logger.info(f"No metadata file, using default directory: {checkpoint_write_dir}")

        # Replace with config-based checkpointing
        simulator.checkpointing = add_checkpointing_from_config_with_directories(
            simulator, 
            config_path=config_filename,
            checkpoint_read_dir=checkpoint_read_dir,
            checkpoint_write_dir=checkpoint_write_dir
        )
        output_logger.info(f"Checkpointing system initialised from config: {config_filename}")

    return simulator

restore_from_checkpoint(checkpoint_name)

Restore simulation from a checkpoint.

Parameters:

Name	Type	Description	Default
checkpoint_name	str	Name of the checkpoint to restore from	required

Returns:

Name	Type	Description
bool	bool	True if restoration was successful

Source code in june/simulator.py

def restore_from_checkpoint(self, checkpoint_name: str) -> bool:
    """Restore simulation from a checkpoint.

    Args:
        checkpoint_name (str): Name of the checkpoint to restore from

    Returns:
        bool: True if restoration was successful
    """
    if self.checkpointing is None:
        output_logger.warning("Checkpointing system not available")
        return False

    try:
        return self.checkpointing.restore_from_checkpoint(checkpoint_name)
    except Exception as e:
        output_logger.error(f"Error restoring from checkpoint: {e}")
        return False

run()

Run simulation with n_seed initial infections

Source code in june/simulator.py

def run(self):
    """Run simulation with n_seed initial infections"""

    # Calculate remaining days to run
    current_time = self.timer.now
    remaining_days = self.timer.total_days - current_time

    output_logger.info(
        f"Starting simulation for {self.timer.total_days} days at day {self.timer.date},"
        f"to run for {remaining_days:.2f} days"
    )

    if self.test_and_trace_enabled:
        # Check if TTEventRecorder already exists (from checkpoint restoration)
        existing_recorder = GlobalContext.get_tt_event_recorder()

        if existing_recorder is not None:
            output_logger.info("Test and Trace enabled, using existing TTEventRecorder from checkpoint")
            output_logger.info(f"Existing TTEventRecorder filename: {existing_recorder.filename}")
        else:
            output_logger.info("Test and Trace enabled, no existing TTEventRecorder found")
            output_logger.info("Initialising new TTEventRecorder")

            # Use record path for TTEventRecorder if available
            if self.record and hasattr(self.record, 'record_path'):
                tt_output_dir = self.record.record_path / "h5_test_trace"
            else:
                tt_output_dir = "./results/h5_test_trace"  # Fallback

            tt_recorder = TTEventRecorder(output_dir=str(tt_output_dir))
            GlobalContext.set_tt_event_recorder(tt_recorder)
            output_logger.info(f"New TTEventRecorder filename: {tt_recorder.filename}")
            output_logger.info(f"TTEventRecorder output directory: {tt_output_dir}")

            # Apply any pending TTEventRecorder data from checkpoint restoration
            if hasattr(self, '_pending_tt_event_recorder_data'):
                output_logger.info("Applying pending TTEventRecorder data from checkpoint")
                self._apply_pending_tt_event_recorder_data(tt_recorder)
                delattr(self, '_pending_tt_event_recorder_data')
                output_logger.info("TTEventRecorder data restoration completed")

    else:
        output_logger.info("Test and Trace disabled, skipping TTEventRecorder")
        GlobalContext.set_tt_event_recorder(None)

    # Rat manager is already initialised normally

    GlobalContext.set_simulator(self)

    # Update registered members home ranks after potential checkpoint restoration
    self.update_registered_members_home_ranks()

    # Update friend home ranks after domain splitting
    if self.friend_hangouts_enabled:
        self.update_friends_home_ranks()
        mpi_comm.Barrier()

    self.clear_world()

    if self.record is not None:
        try:
            self.record.parameters(
                interaction=self.interaction,
                epidemiology=self.epidemiology,
                activity_manager=self.activity_manager,
            )
        except Exception as e:
            output_logger.error(f"Error recording parameters: {e}")

    #START OF THE SIMULATION LOOP
    while self.timer.date < self.timer.final_date:            

        if self.epidemiology:
            self.epidemiology.infection_seeds_timestep(
                self.timer, record=self.record
            )

        mpi_comm.Barrier()
        if mpi_rank == 0:
            rank_logger.info("Next timestep")
        self.do_timestep()

        next(self.timer)

        # Check for checkpoint creation AFTER advancing to next day
        # This way checkpoint represents "ready to start the next day"
        if self.checkpointing is not None:
            integrate_checkpointing_in_simulation_loop(self.checkpointing)

    # Ensure all ranks have finished saving data 
    if self.save_rat_data:
        mpi_comm.Barrier()  # Wait for all ranks to finish saving data

        # Just report data location - no post-processing during simulation
        if mpi_rank == 0 and self.rat_manager is not None:
            if self.record and hasattr(self.record, 'record_path'):
                rat_data_dir = self.record.record_path / "rat_data"
            else:
                rat_data_dir = "outputs/rat_data"

            output_logger.info(f"Rat simulation data saved to: {rat_data_dir}")
            output_logger.info("To create visualizations, run:")
            output_logger.info(f"   python plot_maker/post_process_rat_data.py {rat_data_dir}")

    if self.record:
        self.record.combine_outputs()

    if self.test_and_trace_enabled:
        from june.records.test_trace_event_recording import export_simulation_results

        # Use record path for test and trace export if available
        if self.record and hasattr(self.record, 'record_path'):
            tt_export_dir = str(self.record.record_path)
        else:
            tt_export_dir = "./results"  # Fallback

        export_simulation_results(output_dir=tt_export_dir)

update_friends_home_ranks()

Update the home rank for each person's friends after domain splitting. Replaces the default home rank (0) with the actual home rank where each friend resides.

Works in both MPI and non-MPI environments: - In MPI mode: Ranks collaborate to find friend home ranks - In non-MPI mode: All friends are assigned to rank 0

Source code in june/simulator.py

def update_friends_home_ranks(self):
    """Update the home rank for each person's friends after domain splitting.
    Replaces the default home rank (0) with the actual home rank where each friend resides.

    Works in both MPI and non-MPI environments:
    - In MPI mode: Ranks collaborate to find friend home ranks
    - In non-MPI mode: All friends are assigned to rank 0

    """

    if not mpi_available:
        return

    # Step 1: Collect all friend IDs from people on this rank
    all_friend_ids = set()
    people_with_friends = []

    for person in self.world.people.members:
        if hasattr(person, 'friends') and person.friends:
            # Handle both old format (just home_rank) and new format (dict)
            for friend_id, friend_data in person.friends.items():
                all_friend_ids.add(friend_id)
            people_with_friends.append(person)

    # Convert to sorted list for consistent ordering
    all_friend_ids = sorted(list(all_friend_ids))

    from june.demography.person import Person

    # Step 2: Check which friend IDs exist on this rank
    local_id_to_rank = {}
    unknown_ids = []

    for friend_id in all_friend_ids:
        person = Person.find_by_id(friend_id)
        if person is not None:
            # Found locally - store their home rank
            local_id_to_rank[friend_id] = getattr(person, '_home_rank', mpi_rank)
        else:
            # Not found locally - add to list of IDs to query other ranks
            unknown_ids.append(friend_id)

    # Step 3: Share unknown IDs with other ranks
    all_unknown_ids = mpi_comm.allgather(unknown_ids)

    # Step 4: Check if any of the IDs from other ranks exist on this rank
    additional_mappings = {}

    for rank, rank_unknown_ids in enumerate(all_unknown_ids):
        if rank == mpi_rank:
            # Skip our own unknown IDs
            continue

        for friend_id in rank_unknown_ids:
            person = Person.find_by_id(friend_id)
            if person is not None:
                # Found a friend that another rank was looking for
                additional_mappings[friend_id] = getattr(person, '_home_rank', mpi_rank)

    # Step 5: Share these additional mappings
    all_additional_mappings = mpi_comm.allgather(additional_mappings)

    # Build the global friend_id-to-home_rank mapping
    global_friend_to_home_rank = {}
    # Add our local findings first
    global_friend_to_home_rank.update(local_id_to_rank)

    # Add additional mappings from all ranks
    for rank_mappings in all_additional_mappings:
        global_friend_to_home_rank.update(rank_mappings)

    # Step 6: Update each person's friends dictionary with correct home ranks
    updated_people = 0
    updated_friends = 0

    for person in people_with_friends:
        person_updated = False

        # Create a new friends dictionary with updated home ranks
        updated_friends_dict = {}

        for friend_id, friend_data in person.friends.items():
            # Handle both old format (just home_rank) and new format (dict)
            if isinstance(friend_data, dict):
                # New format - update home_rank, keep hobbies
                current_home_rank = friend_data.get("home_rank", 0)
                actual_home_rank = global_friend_to_home_rank.get(friend_id)

                if actual_home_rank is not None:
                    updated_friends_dict[friend_id] = {
                        "home_rank": actual_home_rank,
                        "hobbies": friend_data.get("hobbies", [])
                    }

                    # Track if this friend's home rank was updated
                    if current_home_rank != actual_home_rank:
                        updated_friends += 1
                        person_updated = True
                else:
                    # Friend ID not found in any rank, keep current values
                    updated_friends_dict[friend_id] = friend_data
                    output_logger.warning(f"Friend ID {friend_id} not found in any rank for person {person.id}")

        # Replace the person's friends dictionary
        person.friends = updated_friends_dict

        if person_updated:
            updated_people += 1

    # Log statistics
    # Gather statistics from all ranks
    all_updated_people = mpi_comm.allgather(updated_people)
    all_updated_friends = mpi_comm.allgather(updated_friends)

    if mpi_rank == 0:
        total_updated_people = sum(all_updated_people)
        total_updated_friends = sum(all_updated_friends)
        output_logger.info(f"Friends home ranks updated: {total_updated_people} people, {total_updated_friends} friend relationships")

    # Ensure all ranks are synchronised
    mpi_comm.Barrier()

update_registered_members_home_ranks()

For each group type (companies, care homes, schools, universities), update the registered_members_ids to include the rank where each member lives.

Works in both MPI and non-MPI environments: - In MPI mode: Ranks collaborate to find member home ranks - In non-MPI mode: All members are assigned to rank 0

Source code in june/simulator.py

def update_registered_members_home_ranks(self):
    """For each group type (companies, care homes, schools, universities),
    update the registered_members_ids to include the rank where each member lives.

    Works in both MPI and non-MPI environments:
    - In MPI mode: Ranks collaborate to find member home ranks
    - In non-MPI mode: All members are assigned to rank 0

    """

    if mpi_available:
        output_logger.info(f"Updating registered members home ranks on rank {mpi_rank}")
    else:
        output_logger.info("Updating registered members home ranks (non-MPI mode)")

    # Group types to check
    group_types = [
        ("companies", "Company"), 
        ("care_homes", "CareHome"), 
        ("schools", "School"),
        ("universities", "University")
    ]

    # Step 1: Collect all member IDs from groups on this rank
    all_member_ids = set()
    groups_with_members = []

    for group_type_name, _ in group_types:
        group_type = getattr(self.world, group_type_name, None)
        if group_type is None or not hasattr(group_type, "members"):
            continue

        for group in group_type.members:
            if not hasattr(group, "registered_members_ids"):
                continue

            for subgroup_id, member_ids in group.registered_members_ids.items():
                for i, member_info in enumerate(member_ids):
                    # Extract the member ID (handle both plain IDs and tuples)
                    if isinstance(member_info, tuple) and len(member_info) == 2:
                        member_id = member_info[0]
                    else:
                        member_id = member_info

                    all_member_ids.add(member_id)

            # Track groups for later updating
            groups_with_members.append((group_type_name, group))

    # Convert to a sorted list for consistent ordering
    all_member_ids = sorted(list(all_member_ids))

    from june.demography.person import Person

    # Different logic for MPI vs non-MPI
    if mpi_available:
        # MPI mode: Distributed ID lookup

        # Step 2: Check which IDs exist on this rank
        local_id_to_rank = {}
        unknown_ids = []

        for member_id in all_member_ids:
            person = Person.find_by_id(member_id)
            if person is not None:
                # Found locally
                local_id_to_rank[member_id] = mpi_rank
            else:
                # Not found locally - add to list of IDs to query other ranks
                unknown_ids.append(member_id)

        # Step 3: Share only the unknown IDs with other ranks
        all_unknown_ids = mpi_comm.allgather(unknown_ids)

        # Step 4: Check if any of the IDs from other ranks exist on this rank
        additional_mappings = {}

        for rank, rank_unknown_ids in enumerate(all_unknown_ids):
            if rank == mpi_rank:
                # Skip our own unknown IDs
                continue

            for member_id in rank_unknown_ids:
                person = Person.find_by_id(member_id)
                if person is not None:
                    # We found an ID that another rank was looking for
                    additional_mappings[member_id] = mpi_rank

        # Step 5: Share these additional mappings
        all_additional_mappings = mpi_comm.allgather(additional_mappings)

        # Build the global ID-to-rank mapping
        global_id_to_rank = {}
        # Add our local findings first
        global_id_to_rank.update(local_id_to_rank)

        # Add additional mappings from all ranks
        for rank_mappings in all_additional_mappings:
            global_id_to_rank.update(rank_mappings)

    else:
        # Non-MPI mode: All IDs are on rank 0
        global_id_to_rank = {member_id: 0 for member_id in all_member_ids}

    # Step 6: Update all groups with the correct home rank information
    for group_type_name, group in groups_with_members:
        for subgroup_id, member_ids in list(group.registered_members_ids.items()):
            updated_member_ids = []

            for member_info in member_ids:
                # Extract the member ID (handle both plain IDs and tuples)
                if isinstance(member_info, tuple) and len(member_info) == 2:
                    member_id = member_info[0]
                else:
                    member_id = member_info

                # Get the home rank from the global mapping
                home_rank = global_id_to_rank.get(member_id)

                if home_rank is not None:
                    # Add to the updated list
                    updated_member_ids.append((member_id, home_rank))
                else:
                    # ID not found in any rank, use default of 0
                    # This should not happen if everything is working correctly
                    updated_member_ids.append((member_id, 0))

            # Replace the old list with the updated one
            group.registered_members_ids[subgroup_id] = updated_member_ids

    # Ensure all ranks are synchronised in MPI mode
    mpi_comm.Barrier()

enable_mpi_debug(results_folder)

Parameters:

Name	Type	Description	Default
results_folder	str	Path to the folder where the results are being held.	required

Source code in june/simulator.py

def enable_mpi_debug(results_folder):
    """

    Args:
        results_folder (str): Path to the folder where the results are being held.
    """
    from june.logging import MPIFileHandler

    logging_file = Path(results_folder) / "mpi.log"
    with open(logging_file, "w"):
        pass
    mh = MPIFileHandler(logging_file)
    rank_logger.addHandler(mh)