Checkpoint restorer

Checkpoint Restoration System for JUNE Simulations

This module handles the restoration of simulation state from checkpoint files, reconstructing the complete simulation environment including all critical states.

CheckpointRestorer

Handles restoration of simulation state from checkpoint files.

Coordinates the reconstruction of all simulation components from checkpointed data, ensuring consistency across MPI ranks.

Source code in june/checkpointing/checkpoint_restorer.py

class CheckpointRestorer:
    """Handles restoration of simulation state from checkpoint files.

    Coordinates the reconstruction of all simulation components from
    checkpointed data, ensuring consistency across MPI ranks.

    """

    def __init__(self, simulator):
        """
        Initialise the checkpoint restorer.

        Parameters
        ----------
        simulator : Simulator
            The JUNE simulator instance to restore state into
        """
        self.simulator = simulator
        self.restoration_stats = {}
        self.random_state_manager = RandomStateManager()

    def restore_from_checkpoint(self, checkpoint_path: Path) -> bool:
        """Restore simulation from checkpoint files.

        Args:
            checkpoint_path (Path): Directory containing checkpoint files

        Returns:
            bool: True if restoration was successful

        """

        print(f"Rank {mpi_rank}: Starting checkpoint restoration from {checkpoint_path}")

        # Validate checkpoint directory
        if not self._validate_checkpoint_directory(checkpoint_path):
            logger.error(f"Rank {mpi_rank}: Invalid checkpoint directory: {checkpoint_path}")
            return False

        # Load checkpoint metadata
        metadata = self._load_checkpoint_metadata(checkpoint_path)
        if not metadata:
            logger.error(f"Rank {mpi_rank}: Could not load checkpoint metadata")
            return False

        # Load rank-specific checkpoint data
        rank_file = checkpoint_path / f"checkpoint_rank_{mpi_rank}.h5"
        if not rank_file.exists():
            logger.error(f"Rank {mpi_rank}: Checkpoint file not found: {rank_file}")
            return False

        checkpoint_data = self._load_checkpoint_data(rank_file)
        if not checkpoint_data:
            logger.error(f"Rank {mpi_rank}: Could not load checkpoint data")
            return False

        # EARLY VALIDATION: Check for disease model mismatch before attempting restoration
        if not self._validate_disease_model_compatibility(metadata):
            logger.error(f"Rank {mpi_rank}: Checkpoint restoration aborted due to disease model mismatch")
            return False

        # Coordinate across MPI ranks before restoration
        if mpi_available:
            mpi_comm.Barrier()

        # Restore simulation state components
        success = self._restore_simulation_state(checkpoint_data, metadata, checkpoint_path)

        if success:
            # Mark this simulator as resumed from checkpoint
            self.simulator._is_resumed_from_checkpoint = True
            self.simulator._is_resumed_and_first_round = True

            # Final coordination
            if mpi_available:
                mpi_comm.Barrier()

            print(f"Rank {mpi_rank}: Checkpoint restoration completed successfully")
            self._log_restoration_summary()
            return True
        else:
            logger.error(f"Rank {mpi_rank}: Checkpoint restoration failed")
            return False

    def _validate_checkpoint_directory(self, checkpoint_path: Path) -> bool:
        """Validate that checkpoint directory contains required files

        Args:
            checkpoint_path (Path): 

        """
        if not checkpoint_path.exists() or not checkpoint_path.is_dir():
            return False

        # Check for metadata file
        metadata_file = checkpoint_path / "checkpoint_metadata.json"
        if not metadata_file.exists():
            logger.warning(f"Checkpoint metadata file not found: {metadata_file}")
            return False

        # Check for rank-specific file
        rank_file = checkpoint_path / f"checkpoint_rank_{mpi_rank}.h5"
        if not rank_file.exists():
            logger.warning(f"Rank-specific checkpoint file not found: {rank_file}")
            return False

        return True

    def _load_checkpoint_metadata(self, checkpoint_path: Path) -> Optional[Dict[str, Any]]:
        """Load checkpoint metadata from JSON file

        Args:
            checkpoint_path (Path): 

        """
        metadata_file = checkpoint_path / "checkpoint_metadata.json"

        with open(metadata_file, 'r') as f:
            metadata = json.load(f)

        logger.debug(f"Rank {mpi_rank}: Loaded checkpoint metadata: {metadata.get('checkpoint_type', 'unknown')} "
                    f"from simulation time {metadata.get('simulation_time', 'unknown')}")
        return metadata


    def _load_checkpoint_data(self, checkpoint_file: Path) -> Optional[Dict[str, Any]]:
        """Load checkpoint data from HDF5 file

        Args:
            checkpoint_file (Path): 

        """

        checkpoint_data = {}

        with h5py.File(checkpoint_file, 'r') as f:
            # Load metadata
            if 'metadata' in f:
                metadata_group = f['metadata']
                checkpoint_data['_metadata'] = {
                    attr: metadata_group.attrs[attr] for attr in metadata_group.attrs.keys()
                }

            # Load each component's data
            for component_name in f.keys():
                if component_name != 'metadata':
                    checkpoint_data[component_name] = self._load_component_data(f[component_name])

        logger.debug(f"Rank {mpi_rank}: Loaded checkpoint data for components: {list(checkpoint_data.keys())}")
        return checkpoint_data

    def _load_component_data(self, hdf5_group) -> Dict[str, Any]:
        """Load a component's data from HDF5 group

        Args:
            hdf5_group: 

        """
        component_data = {}

        # Load attributes with enhanced type preservation
        for attr_name in hdf5_group.attrs.keys():
            value = hdf5_group.attrs[attr_name]

            if isinstance(value, bytes):
                value = value.decode('utf-8')
            if value == "None":
                value = None
            elif isinstance(value, str) and value.startswith('{') and value.endswith('}'):
                value = json.loads(value)

            component_data[attr_name] = value

        # Load datasets
        for dataset_name in hdf5_group.keys():
            dataset = hdf5_group[dataset_name]

            if dataset.dtype.kind == 'S':  # String dataset
                # Handle string arrays (might be JSON-encoded objects)
                string_data = [item.decode('utf-8') if isinstance(item, bytes) else item for item in dataset[:]]

                # Try to decode as JSON objects
                component_data[dataset_name] = [json.loads(item) for item in string_data]
            else:
                # Numeric dataset
                component_data[dataset_name] = dataset[:]

        return component_data

    def _restore_simulation_state(self, checkpoint_data: Dict[str, Any], metadata: Dict[str, Any], checkpoint_path: Path) -> bool:
        """Restore all simulation state components.

        Args:
            checkpoint_data (Dict[str, Any]): Loaded checkpoint data
            metadata (Dict[str, Any]): Checkpoint metadata
            checkpoint_path (Path): 

        Returns:
            bool: True if restoration was successful

        """
        print(f"Rank {mpi_rank}: Restoring simulation state components")

        restoration_success = True
        self.restoration_stats = {}

        # Check for feature activations (newly enabled features)
        checkpoint_features = metadata.get('feature_flags', {})
        current_features = self._get_current_feature_flags()
        feature_activations = self._detect_feature_activations(checkpoint_features, current_features)

        # Activate newly enabled features before restoring state
        if feature_activations:
            print(f"Rank {mpi_rank}: Activating newly enabled features: {list(feature_activations.keys())}")

            # Validate feature activations are safe
            validation_result = self._validate_feature_activations(feature_activations, metadata)
            if not validation_result['valid']:
                logger.error(f"Rank {mpi_rank}: Feature activation validation failed: {validation_result['reason']}")
                return False

            success = self._activate_new_features(feature_activations)
            if not success:
                logger.error(f"Rank {mpi_rank}: Failed to activate new features")
                return False

        # Restore timer state first (other components depend on it)
        if 'timer' in checkpoint_data:
            success = self._restore_timer_state(checkpoint_data['timer'])
            restoration_success = restoration_success and success
            logger.debug(f"Rank {mpi_rank}: Timer restoration: {'success' if success else 'failed'}")

        # Restore population health state
        if 'population_health' in checkpoint_data:
            success = self._restore_population_health(checkpoint_data['population_health'])
            restoration_success = restoration_success and success
            logger.debug(f"Rank {mpi_rank}: Population health restoration: {'success' if success else 'failed'}")

        # Restore interaction transmission tracking state
        if 'interaction' in checkpoint_data:
            success = self._restore_interaction_state(checkpoint_data['interaction'])
            restoration_success = restoration_success and success
            logger.debug(f"Rank {mpi_rank}: Interaction state restoration: {'success' if success else 'failed'}")

        # Restore random number generator states
        random_states = self.random_state_manager.load_states(checkpoint_path)
        if random_states is None:
            logger.error(f"Rank {mpi_rank}: Failed to load random states")
            return False

        if not self.random_state_manager.restore_states(random_states):
            logger.error(f"Rank {mpi_rank}: Failed to restore random states")
            return False

        self.restoration_stats['random_states'] = {
            'restored_successfully': True,
            'manager_used': True
        }
        logger.debug(f"Rank {mpi_rank}: Random state restoration: success")

        # Restore test and trace state (or skip if newly activated)
        if 'test_and_trace' in checkpoint_data:
            success = self._restore_test_and_trace_state(checkpoint_data['test_and_trace'])
            restoration_success = restoration_success and success
            logger.debug(f"Rank {mpi_rank}: Test and trace state restoration: {'success' if success else 'failed'}")
        elif feature_activations.get('test_and_trace_enabled'):
            # Test and trace was newly activated - initialize fresh state
            print(f"Rank {mpi_rank}: Test and trace newly activated - starting with fresh state")
            self.restoration_stats['test_and_trace'] = {
                'enabled': True,
                'newly_activated': True,
                'restoration_successful': True
            }

        # Restore rat dynamics state (AFTER random state restoration) or skip if newly activated
        if 'rat_dynamics' in checkpoint_data:
            success = self._restore_rat_dynamics_state(checkpoint_data['rat_dynamics'])
            restoration_success = restoration_success and success
            logger.debug(f"Rank {mpi_rank}: Rat dynamics state restoration: {'success' if success else 'failed'}")
        elif feature_activations.get('ratty_dynamics_enabled'):
            # Rat dynamics was newly activated - initialize fresh state
            print(f"Rank {mpi_rank}: Rat dynamics newly activated - starting with fresh state")
            self.restoration_stats['rat_dynamics'] = {
                'enabled': True,
                'newly_activated': True,
                'restoration_successful': True
            }

        # Restore TTEventRecorder state (daily/cumulative test and trace data)
        if 'tt_event_recorder' in checkpoint_data:
            success = self._restore_tt_event_recorder_state(checkpoint_data['tt_event_recorder'])
            restoration_success = restoration_success and success
            logger.debug(f"Rank {mpi_rank}: TTEventRecorder state restoration: {'success' if success else 'failed'}")

        # Restore school incident tracking state (for NotSendingKidsToSchool policy)
        if 'school_incidents' in checkpoint_data:
            success = self._restore_school_incident_state(checkpoint_data['school_incidents'])
            restoration_success = restoration_success and success
            logger.debug(f"Rank {mpi_rank}: School incident state restoration: {'success' if success else 'failed'}")

        return restoration_success

    def _restore_timer_state(self, timer_data: Dict[str, Any]) -> bool:
        """Restore simulation timer state.

        Args:
            timer_data (Dict[str, Any]): Timer state data

        Returns:
            bool: True if restoration was successful

        """

        timer = self.simulator.timer

        # Store original timer configuration
        original_total_days = timer.total_days
        original_initial_date = timer.initial_date

        # Restore timer state by setting the date directly
        # The 'now' property is calculated from the date, so we need to set the date
        if 'current_date' in timer_data:
            restored_date = datetime.datetime.fromisoformat(timer_data['current_date'])
            timer.date = restored_date
            logger.debug(f"Rank {mpi_rank}: Timer date restored to {restored_date}")

        # Set the shift to 0 to start from the beginning of the day
        # This ensures we're at the correct position within the day
        timer.shift = 0
        timer.delta_time = datetime.timedelta(hours=timer.shift_duration)

        # Restore total_days if available
        if 'total_days' in timer_data:
            timer.total_days = timer_data['total_days']

        current_simulation_time = timer.now
        remaining_days = original_total_days - current_simulation_time

        if remaining_days > 0:
            # Recalculate final_date from current date + remaining time
            timer.final_date = timer.date + datetime.timedelta(days=remaining_days)
            print(f"Rank {mpi_rank}: Final date recalculated to {timer.final_date} (remaining: {remaining_days:.2f} days)")
        else:
            # If we've already completed the simulation, set final_date to current date
            timer.final_date = timer.date
            logger.warning(f"Rank {mpi_rank}: Simulation appears to be complete (current time: {current_simulation_time}, total: {original_total_days})")

        # Restore other timer attributes if they exist
        for attr in ['time_step_size', 'activities_start_time', 'activities_end_time']:
            if attr in timer_data and hasattr(timer, attr):
                setattr(timer, attr, timer_data[attr])

        # Verify the restoration worked correctly
        restored_simulation_time = timer.now
        expected_time = timer_data.get('current_time', 0)

        self.restoration_stats['timer'] = {
            'current_time': restored_simulation_time,
            'current_date': str(timer.date),
            'final_date': str(timer.final_date),
            'expected_time': expected_time,
            'time_difference': abs(restored_simulation_time - expected_time),
            'remaining_days': remaining_days,
            'original_total_days': original_total_days
        }

        print(f"Rank {mpi_rank}: Timer restored to simulation time {restored_simulation_time}, date {timer.date}")
        print(f"Rank {mpi_rank}: Simulation will continue until {timer.final_date} ({remaining_days:.2f} days remaining)")

        # Check if the restoration was accurate
        time_diff = abs(restored_simulation_time - expected_time)
        if time_diff > 0.1:  # Allow small floating point differences
            logger.warning(f"Rank {mpi_rank}: Timer restoration time mismatch. Expected: {expected_time}, Got: {restored_simulation_time}")

        return True


    def _restore_population_health(self, population_data: Dict[str, Any]) -> bool:
        """Restore population health states.

        Args:
            population_data (Dict[str, Any]): Population health data

        Returns:
            bool: True if restoration was successful

        """

        people_states = population_data.get('people_states', {})
        immunities = population_data.get('immunities', {})
        people_ids = population_data.get('people_ids', [])

        restored_count = 0
        infected_count = 0
        hospitalised_count = 0
        immunity_count = 0
        intensive_care_count = 0

        # Create a lookup dictionary for people by ID
        people_by_id = {person.id: person for person in self.simulator.world.people}

        print(f"Rank {mpi_rank}: Restoring complete population state - {len(people_states)} people with health states, {len(immunities)} immunity records")

        for person_id_str, person_state in people_states.items():
            person_id = int(person_id_str)

            if person_id not in people_by_id:
                logger.warning(f"Rank {mpi_rank}: Person {person_id} not found in current world")
                continue

            person = people_by_id[person_id]
            success = self._restore_person_health(person, person_state)

            if success:
                restored_count += 1
                if person_state.get('infected', False):
                    infected_count += 1
                if person_state.get('hospitalised', False):
                    hospitalised_count += 1
                if person_state.get('intensive_care', False):
                    intensive_care_count += 1

        # Restore immunity for ALL people
        for person_id_str, immunity_data in immunities.items():
            person_id = int(person_id_str)

            if person_id not in people_by_id:
                logger.warning(f"Rank {mpi_rank}: Person {person_id} not found for immunity restoration")
                continue

            person = people_by_id[person_id]

            # Reconstruct Immunity object from saved data
            from june.epidemiology.infection import Immunity

            # Complete immunity object data with FIXED key types
            susceptibility_dict = immunity_data.get('susceptibility_dict', {})
            effective_multiplier_dict = immunity_data.get('effective_multiplier_dict', {})

            fixed_susceptibility_dict = {}
            for key, value in susceptibility_dict.items():
                try:
                    # Convert string keys back to integers
                    int_key = int(key) if isinstance(key, str) else key
                    fixed_susceptibility_dict[int_key] = value
                except (ValueError, TypeError):
                    # Keep original key if conversion fails
                    fixed_susceptibility_dict[key] = value

            fixed_effective_multiplier_dict = {}
            for key, value in effective_multiplier_dict.items():
                try:
                    # Convert string keys back to integers
                    int_key = int(key) if isinstance(key, str) else key
                    fixed_effective_multiplier_dict[int_key] = value
                except (ValueError, TypeError):
                    # Keep original key if conversion fails
                    fixed_effective_multiplier_dict[key] = value

            # Create proper Immunity with fixed dictionaries
            person.immunity = Immunity(
                susceptibility_dict=fixed_susceptibility_dict,
                effective_multiplier_dict=fixed_effective_multiplier_dict
            )
            immunity_count += 1
            logger.debug(f"Rank {mpi_rank}: Restored Immunity object for person {person_id} with {len(susceptibility_dict)} susceptibilities")

        self.restoration_stats['population_health'] = {
            'people_restored': restored_count,
            'infected_restored': infected_count,
            'hospitalised_restored': hospitalised_count,
            'immunity_restored': immunity_count,
            'total_people': len(self.simulator.world.people)
        }

        # Restore cemetery state (dead people)
        if 'cemetery_state' in population_data:
            cemetery_success = self._restore_cemetery_state(population_data['cemetery_state'])
            if cemetery_success:
                # Get dead people count from the restored data
                dead_count = population_data['cemetery_state'].get('total_deaths', 
                            len(population_data['cemetery_state'].get('dead_id', [])))
                self.restoration_stats['population_health']['dead_people_restored'] = dead_count

        print(f"Rank {mpi_rank}: Restored complete population state:")
        print(f"  - Health states: {restored_count} people")
        print(f"  - Infections: {infected_count} people")
        print(f"  - Hospitalised: {hospitalised_count} people")
        print(f"  - Immunity: {immunity_count} people")
        return True


    def _restore_person_health(self, person, person_state: Dict[str, Any]) -> bool:
        """Restore individual person's health state.

        Args:
            person (Person): The person to restore state for
            person_state (Dict[str, Any]): The person's health state data

        Returns:
            bool: True if restoration was successful

        """
        # NOTE: hospitalised and intensive_care are properties computed from medical_facility
        # They cannot be set directly - they're computed based on medical facility assignment

        # Restore infection details
        if 'infection' in person_state:
            infection_data = person_state['infection']
            success = self._restore_person_infection(person, infection_data)
            if not success:
                logger.warning(f"Rank {mpi_rank}: Failed to restore infection for person {person.id}")
                return False

        # Restore health object details
        if 'health' in person_state and hasattr(person, 'health'):
            health_data = person_state['health']

            if 'immunity' in health_data and hasattr(person.health, 'immunity'):
                person.health.immunity = health_data['immunity']

            if 'vaccination_history' in health_data and hasattr(person.health, 'vaccination_history'):
                person.health.vaccination_history = health_data['vaccination_history']

        # Restore medical facility assignment (AFTER other state is restored)
        if 'medical_facility' in person_state:
            facility_data = person_state['medical_facility']
            self._restore_hospitalisation(person, facility_data, person_state)

        return True

    def _restore_hospitalisation(self, person, facility_data: Dict[str, Any], person_state: Dict[str, Any]):
        """Restore a person's hospitalization state.

        This needs to respect the hospitalization policy logic:
        - People can be in hospital due to severe symptoms OR
        - People can be kept in hospital while waiting for test results

        Args:
            person: 
            facility_data (Dict[str, Any]): 
            person_state (Dict[str, Any]): 

        """
        medical_facility_id = np.int64(facility_data["facility_id"])
        medical_facility = self.simulator.world.hospitals.get_from_id(medical_facility_id)

        # Check if person should be in hospital based on current symptoms
        from june.global_context import GlobalContext
        disease_config = GlobalContext.get_disease_config()

        if disease_config:
            hospitalised_tags = set(disease_config.symptom_manager._resolve_tags("hospitalised_stage"))
            intensive_care_tags = set(disease_config.symptom_manager._resolve_tags("intensive_care_stage"))
            person_symptoms_tag = person.infection.tag if person.infection else None

            # Case 1: Person has symptoms requiring hospitalization
            if person_symptoms_tag in hospitalised_tags or person_symptoms_tag in intensive_care_tags:
                try:
                    result = medical_facility.allocate_patient(person)
                except Exception as e:
                    # If allocation fails, try manual assignment
                    if person.hospitalised:
                        medical_facility.add_to_ward(person)
                    elif person.intensive_care:
                        medical_facility.add_to_icu(person)

            # Case 2: Person doesn't have hospitalising symptoms but might be waiting for test results
            else:
                # Get the saved hospital state from the checkpoint data
                from june.checkpointing.state_serialisers import PopulationHealthSerialiser

                # We need to determine where this person was in the hospital based on checkpoint data
                # Since we can't rely on person.hospitalised/intensive_care (they're properties)
                was_hospitalised = person_state.get('hospitalised', False) 
                was_intensive_care = person_state.get('intensive_care', False)

                # Check if they have test_and_trace and are waiting for results
                if (person.test_and_trace is not None and 
                    person.test_and_trace.test_result is None):

                    # Manually assign to hospital without going through allocate_patient
                    # since they're being kept for test results, not symptoms
                    if was_intensive_care:
                        medical_facility.add_to_icu(person)
                    elif was_hospitalised:
                        medical_facility.add_to_ward(person)
                else:
                    # This person was hospitalised but doesn't meet current criteria
                    # We should still restore them to maintain checkpoint consistency

                    if was_intensive_care:
                        medical_facility.add_to_icu(person)
                    elif was_hospitalised:
                        medical_facility.add_to_ward(person)

    def _restore_person_infection(self, person, infection_data: Dict[str, Any]) -> bool:
        """Restore a person's infection by reconstructing the infection object.

        Args:
            person (Person): The person to restore infection for
            infection_data (Dict[str, Any]): Infection data from checkpoint

        Returns:
            bool: True if restoration was successful

        """
        # Import infection classes
        from june.epidemiology.infection import infection as infection_module

        # Get infection class - prefer infection_class over infection_type for accuracy
        infection_type = infection_data.get('infection_class') or infection_data.get('infection_type')
        if not infection_type:
            logger.error(f"Rank {mpi_rank}: No infection class/type specified for person {person.id}")
            return False


        # Get the infection class from the module
        if not hasattr(infection_module, infection_type):
            logger.error(f"Rank {mpi_rank}: Unknown infection type: {infection_type}")
            return False

        infection_class = getattr(infection_module, infection_type)

        # Restore symptoms and transmission first
        symptoms = None
        transmission = None

        if 'symptoms' in infection_data:
            symptoms = self._restore_symptoms(infection_data['symptoms'])
            if symptoms is None:
                logger.error(f"Rank {mpi_rank}: Failed to restore symptoms for person {person.id}")
                return False

        if 'transmission' in infection_data:
            transmission = self._restore_transmission(infection_data['transmission'])
            if transmission is None:
                logger.error(f"Rank {mpi_rank}: Failed to restore transmission for person {person.id}")
                return False

        # Create the infection object
        # Most infections take transmission and symptoms as constructor parameters
        infection = infection_class(
            transmission=transmission,
            symptoms=symptoms
        )

        # Restore other infection attributes
        if 'start_time' in infection_data:
            infection.start_time = infection_data['start_time']

        if 'transmission_multiplier' in infection_data:
            infection.transmission_multiplier = infection_data['transmission_multiplier']

        # NOTE: infection.tag is a read-only property that returns symptoms.tag
        # We don't need to set it since we already restored the symptoms

        # Restore additional attributes if they exist on the infection class
        for attr in ['infectiousness', 'severity_multiplier', 'time_of_recovery']:
            if attr in infection_data and hasattr(infection, attr):
                setattr(infection, attr, infection_data[attr])

        # Assign the infection to the person
        person.infection = infection

        logger.debug(f"Rank {mpi_rank}: Successfully restored {infection_type} infection for person {person.id}")
        return True

    def _restore_symptoms(self, symptoms_data: Dict[str, Any]):
        """Restore symptoms object from serialised data by avoiding constructor trajectory generation.

        Args:
            symptoms_data (Dict[str, Any]): Serialised symptoms data

        Returns:
            Symptoms or None: Reconstructed symptoms object or None if failed

        """
        from june.epidemiology.infection import Symptoms, SymptomTag
        from june.global_context import GlobalContext

        # Get disease_config from GlobalContext
        disease_config = GlobalContext.get_disease_config()
        if disease_config is None:
            logger.error(f"Rank {mpi_rank}: Cannot restore symptoms - no disease_config available in GlobalContext")
            return None

        # Create symptoms object WITHOUT triggering trajectory generation
        # We'll manually set all attributes to avoid constructor side effects
        symptoms = object.__new__(Symptoms)  # Create instance without calling __init__

        # Set the disease_config attribute that would normally be set by constructor
        symptoms.disease_config = disease_config

        # Restore all attributes directly from saved data
        # Note: SymptomTag is dynamically loaded, so we work with raw integer values
        if 'max_tag' in symptoms_data and symptoms_data['max_tag'] is not None:
            # Store as raw integer value - the enum structure may not be available
            symptoms.max_tag = int(symptoms_data['max_tag'])
            logger.debug(f"Rank {mpi_rank}: Restored max_tag as integer: {symptoms.max_tag}")
        else:
            symptoms.max_tag = None

        if 'tag' in symptoms_data and symptoms_data['tag'] is not None:
            # Store as raw integer value - the enum structure may not be available
            symptoms.tag = int(symptoms_data['tag'])
            logger.debug(f"Rank {mpi_rank}: Restored tag as integer: {symptoms.tag}")
        else:
            symptoms.tag = None

        if 'max_severity' in symptoms_data and symptoms_data['max_severity'] is not None:
            symptoms.max_severity = symptoms_data['max_severity']
        else:
            symptoms.max_severity = None

        if 'stage' in symptoms_data and symptoms_data['stage'] is not None:
            symptoms.stage = symptoms_data['stage']
        else:
            symptoms.stage = None

        if 'time_of_symptoms_onset' in symptoms_data and symptoms_data['time_of_symptoms_onset'] is not None:
            symptoms.time_of_symptoms_onset = symptoms_data['time_of_symptoms_onset']
        else:
            symptoms.time_of_symptoms_onset = None

        # Simple trajectory restoration (match new serialization format)
        if 'trajectory_times' in symptoms_data and 'trajectory_symptoms' in symptoms_data:
            times = symptoms_data['trajectory_times']
            symptom_tags = symptoms_data['trajectory_symptoms']

            logger.debug(f"Rank {mpi_rank}: Restoring trajectory with {len(times)} points")

            # Reconstruct trajectory preserving all types
            trajectory = []
            # Convert all times to numpy float64 before the loop
            times_np = np.array(times, dtype=np.float64)

            # Then use the converted array in your loop
            for i, (time, symp_int) in enumerate(zip(times_np, symptom_tags)):
                try:
                    # time is already numpy.float64 now
                    trajectory.append((time, symp_int))
                except (ValueError, TypeError) as e:
                    logger.warning(f"Rank {mpi_rank}: Failed to restore trajectory point {i}: {e}. Skipping.")
                    continue

            symptoms.trajectory = tuple(trajectory)
            logger.debug(f"Rank {mpi_rank}: Successfully restored trajectory with {len(trajectory)} valid points")
        else:
            symptoms.trajectory = tuple()  # Empty trajectory

        logger.debug(f"Rank {mpi_rank}: Successfully restored symptoms with {len(symptoms.trajectory) if symptoms.trajectory else 0} trajectory points")
        return symptoms

    def _restore_transmission(self, transmission_data: Dict[str, Any]):
        """Restore transmission object from serialised data.

        Args:
            transmission_data (Dict[str, Any]): Serialised transmission data

        Returns:
            Transmission or None: Reconstructed transmission object or None if failed

        """
        from june.epidemiology.infection import (
            TransmissionGamma,
            TransmissionConstant,
            TransmissionXNExp
        )

        # Map transmission type names to classes
        transmission_classes = {
            'TransmissionXNExp': TransmissionXNExp,
            'TransmissionGamma': TransmissionGamma,
            'TransmissionConstant': TransmissionConstant
        }

        transmission_type = transmission_data.get('transmission_type')
        if not transmission_type:
            logger.error(f"Rank {mpi_rank}: No transmission type specified in transmission data")
            return None

        if transmission_type not in transmission_classes:
            logger.error(f"Rank {mpi_rank}: Unknown transmission type: {transmission_type}")
            return None

        # Create transmission object
        transmission_class = transmission_classes[transmission_type]
        transmission = transmission_class()

        # Restore attributes based on transmission type
        if transmission_type == 'TransmissionXNExp':
            for attr in ['time_first_infectious', 'norm_time', 'n', 'norm', 'alpha']:
                if attr in transmission_data and transmission_data[attr] is not None:
                    setattr(transmission, attr, transmission_data[attr])

        elif transmission_type == 'TransmissionGamma':
            for attr in ['shape', 'shift', 'scale', 'norm']:
                if attr in transmission_data and transmission_data[attr] is not None:
                    setattr(transmission, attr, transmission_data[attr])

        elif transmission_type == 'TransmissionConstant':
            if 'probability' in transmission_data and transmission_data['probability'] is not None:
                transmission.probability = transmission_data['probability']

        if 'current_probability' in transmission_data:
            transmission.probability = transmission_data['current_probability']
            logger.debug(f"Rank {mpi_rank}: Restored transmission probability directly: {transmission.probability}")

        return transmission

    def _restore_cemetery_state(self, cemetery_data: Dict[str, Any]) -> bool:
        """Restore cemetery state by reconstructing dead people assignments.

        Args:
            cemetery_data (Dict[str, Any]): Cemetery state data including dead people IDs

        Returns:
            bool: True if restoration was successful

        """
        # Check if we have dead_id data in the expected format
        if 'dead_id' not in cemetery_data:
            logger.debug(f"Rank {mpi_rank}: No dead_id data to restore")
            return True

        dead_ids = cemetery_data['dead_id']
        people_ids = {person.id for person in self.simulator.world.people}

        print(f"Rank {mpi_rank}: Restoring {len(dead_ids)} dead people to cemeteries")

        restored_deaths = 0
        for dead_id in dead_ids:
            if dead_id not in people_ids:
                continue

            person = self.simulator.world.people.get_from_id(dead_id)
            person.dead = True
            cemetery = self.simulator.world.cemeteries.get_nearest(person)
            cemetery.add(person)
            person.subgroups = Activities(None, None, None, None, None, None, None)
            restored_deaths += 1

        print(f"Rank {mpi_rank}: Successfully restored {restored_deaths} dead people to cemeteries")
        return True

    def _restore_interaction_state(self, interaction_data: Dict[str, Any]) -> bool:
        """Restore interaction transmission tracking state.

        Args:
            interaction_data (Dict[str, Any]): Interaction state data from checkpoint

        Returns:
            bool: True if restoration was successful

        """

        # Get the interaction object from the simulator
        interaction = getattr(self.simulator, 'interaction', None)
        if interaction is None:
            logger.error(f"Rank {mpi_rank}: No interaction object found in simulator")
            return False

        # Restore configuration flags - set this BEFORE restoring IDs
        if 'initial_infected_ids_loaded' in interaction_data:
            interaction._initial_infected_ids_loaded = interaction_data['initial_infected_ids_loaded']

        # Restore core transmission tracking
        if 'initial_infected_ids' in interaction_data:
            interaction._initial_infected_ids = set(interaction_data['initial_infected_ids'])
            interaction._initial_infected_ids_loaded = True  # Force flag to True after restoration
            logger.debug(f"Rank {mpi_rank}: Restored {len(interaction._initial_infected_ids)} initial infected IDs")

            # Debug: verify the restored data
            print(f"Rank {mpi_rank}: DEBUG - Restored initial infected IDs: {len(interaction._initial_infected_ids)} IDs")
            if len(interaction._initial_infected_ids) > 0:
                sample_ids = list(interaction._initial_infected_ids)[:5]
                print(f"Rank {mpi_rank}: DEBUG - Sample restored IDs: {sample_ids}")
        else:
            logger.warning(f"Rank {mpi_rank}: No initial_infected_ids found in checkpoint data")

        if 'initial_infected_transmission_counts' in interaction_data:
            from collections import defaultdict
            interaction.initial_infected_transmission_counts = defaultdict(int)

            for key, value in interaction_data['initial_infected_transmission_counts'].items():
                try:
                    int_key = int(key)
                    interaction.initial_infected_transmission_counts[int_key] = value
                except (ValueError, TypeError):
                    logger.warning(f"Rank {mpi_rank}: Could not convert transmission count key '{key}' to int")

            logger.debug(f"Rank {mpi_rank}: Restored transmission counts for {len(interaction.initial_infected_transmission_counts)} transmitters")

            # DEBUG: Print detailed transmission counts for debugging
            print(f"[Rank {mpi_rank}] Restored transmission counts:")
            print(f"  Total transmitters restored: {len(interaction.initial_infected_transmission_counts)}")
            print(f"  Total transmissions restored: {sum(interaction.initial_infected_transmission_counts.values())}")
            if len(interaction.initial_infected_transmission_counts) > 0:
                sample_transmitters = list(interaction.initial_infected_transmission_counts.items())[:5]
                print(f"  Sample transmitters: {sample_transmitters}")
                print(f"  Sample transmitter key types: {[type(k) for k, v in sample_transmitters]}")

        if 'previous_transmission_counts' in interaction_data:
            from collections import defaultdict
            interaction.previous_transmission_counts = defaultdict(int)

            for key, value in interaction_data['previous_transmission_counts'].items():
                try:
                    int_key = int(key)
                    interaction.previous_transmission_counts[int_key] = value
                except (ValueError, TypeError):
                    logger.warning(f"Rank {mpi_rank}: Could not convert previous transmission count key '{key}' to int")

            logger.debug(f"Rank {mpi_rank}: Restored {len(interaction.previous_transmission_counts)} previous transmission counts")

            # Debug: verify the restoration preserves the difference for timestep calculation
            current_total = sum(interaction.initial_infected_transmission_counts.values())
            previous_total = sum(interaction.previous_transmission_counts.values())
            logger.debug(f"Rank {mpi_rank}: Current total transmissions: {current_total}, Previous total: {previous_total}")

        # NOTE: previous_transmission_counts is already restored from checkpoint data above.
        # DO NOT overwrite it with current counts, as that would make ALL transmissions 
        # appear as "new" in the next timestep calculation.
        # The checkpoint data contains the correct previous counts for proper timestep diff calculation.

        if 'timestep_count' in interaction_data:
            interaction.timestep_count = interaction_data['timestep_count']
            logger.debug(f"Rank {mpi_rank}: Restored timestep count: {interaction.timestep_count}")

        # Restore debug tracking data
        if 'debug_cross_rank_infectors' in interaction_data:
            from collections import defaultdict
            interaction.debug_cross_rank_infectors = defaultdict(set)
            for person_id, ranks in interaction_data['debug_cross_rank_infectors'].items():
                interaction.debug_cross_rank_infectors[int(person_id)] = set(ranks)
            logger.debug(f"Rank {mpi_rank}: Restored debug cross-rank infector data")

        if 'debug_timestep_infectors' in interaction_data:
            from collections import defaultdict
            interaction.debug_timestep_infectors = defaultdict(lambda: defaultdict(int))
            for timestep, counts in interaction_data['debug_timestep_infectors'].items():
                interaction.debug_timestep_infectors[int(timestep)].update(counts)
            logger.debug(f"Rank {mpi_rank}: Restored debug timestep infector data")

        if 'debug_infector_venues' in interaction_data:
            from collections import defaultdict
            interaction.debug_infector_venues = defaultdict(list)
            for person_id, venues in interaction_data['debug_infector_venues'].items():
                interaction.debug_infector_venues[int(person_id)] = venues
            logger.debug(f"Rank {mpi_rank}: Restored debug venue data")

        # Restore interaction configuration (for validation)
        if 'alpha_physical' in interaction_data:
            if abs(interaction.alpha_physical - interaction_data['alpha_physical']) > 1e-6:
                logger.warning(f"Rank {mpi_rank}: Alpha physical mismatch. Current: {interaction.alpha_physical}, Restored: {interaction_data['alpha_physical']}")

        # Mark interaction as restored from checkpoint for verification
        interaction._is_resumed_from_checkpoint = True

        # Set flag to indicate this interaction was restored from checkpoint
        interaction._restored_from_checkpoint = True

        # Log restoration summary
        stats = interaction_data.get('statistics_summary', {})
        print(f"Rank {mpi_rank}: Restored interaction transmission state:")
        print(f"  - Initial infected: {stats.get('total_initial_infected', len(interaction._initial_infected_ids))}")
        print(f"  - Transmitting initial infected: {stats.get('total_transmitting_initial_infected', len(interaction.initial_infected_transmission_counts))}")
        print(f"  - Total secondary infections: {stats.get('total_secondary_infections', sum(interaction.initial_infected_transmission_counts.values()))}")
        print(f"  - Current timestep: {stats.get('current_timestep', interaction.timestep_count)}")
        print(f"  - IDs loaded flag: {interaction._initial_infected_ids_loaded}")
        print(f"  - Current initial infected IDs count: {len(interaction._initial_infected_ids)}")

        self.restoration_stats['interaction'] = {
            'initial_infected_count': len(interaction._initial_infected_ids),
            'transmitters_count': len(interaction.initial_infected_transmission_counts),
            'total_secondary_infections': sum(interaction.initial_infected_transmission_counts.values()),
            'timestep_count': interaction.timestep_count,
            'restoration_successful': True
        }

        return True

    def _restore_test_and_trace_state(self, test_trace_data: Dict[str, Any]) -> bool:
        """Restore test and trace system state.

        Args:
            test_trace_data (Dict[str, Any]): Test and trace state data from checkpoint

        Returns:
            bool: True if restoration was successful

        """
        print(f"Rank {mpi_rank}: Restoring test and trace system state")

        # Check if test and trace was enabled in the checkpoint
        if not test_trace_data.get('enabled', False):
            print(f"Rank {mpi_rank}: Test and trace was disabled in checkpoint - skipping restoration")
            self.restoration_stats['test_and_trace'] = {
                'enabled': False,
                'restoration_successful': True
            }
            return True

        # Check if test and trace is currently enabled
        current_enabled = getattr(self.simulator, 'test_and_trace_enabled', False)
        if not current_enabled:
            print(f"Rank {mpi_rank}: Test and trace was enabled in checkpoint but is disabled in current simulation")
            print(f"Rank {mpi_rank}: Skipping test and trace state restoration (feature disabled)")
            self.restoration_stats['test_and_trace'] = {
                'enabled': False,
                'checkpoint_had_enabled': True,
                'restoration_successful': True,
                'note': 'feature_disabled_in_current_simulation'
            }
            return True

        restoration_success = True
        stats = {
            'enabled': True,
            'people_restored': 0,
            'contact_manager_restored': False,
            'policy_configs_restored': False
        }

        # Restore individual TestAndTrace objects
        if 'person_test_trace_states' in test_trace_data:
            success = self._restore_person_test_trace_states(test_trace_data['person_test_trace_states'])
            restoration_success = restoration_success and success
            stats['people_restored'] = len(test_trace_data['person_test_trace_states'])
            logger.debug(f"Rank {mpi_rank}: Person test and trace states restoration: {'success' if success else 'failed'}")

        # Restore contact manager state
        if 'contact_manager_state' in test_trace_data and hasattr(self.simulator, 'contact_manager'):
            success = self._restore_contact_manager_state(test_trace_data['contact_manager_state'])
            restoration_success = restoration_success and success
            stats['contact_manager_restored'] = success
            logger.debug(f"Rank {mpi_rank}: Contact manager state restoration: {'success' if success else 'failed'}")

        # Policy configurations are informational - restoration is optional
        if 'policy_configurations' in test_trace_data:
            try:
                self._validate_policy_configurations(test_trace_data['policy_configurations'])
                stats['policy_configs_restored'] = True
            except Exception as e:
                logger.warning(f"Rank {mpi_rank}: Policy configuration validation failed: {e}")
                stats['policy_configs_restored'] = False

        self.restoration_stats['test_and_trace'] = stats
        self.restoration_stats['test_and_trace']['restoration_successful'] = restoration_success

        print(f"Rank {mpi_rank}: Test and trace restoration completed - {stats['people_restored']} people restored")
        return restoration_success

    def _restore_person_test_trace_states(self, person_states: Dict[str, Any]) -> bool:
        """Restore TestAndTrace objects for all people who had them.

        Args:
            person_states (Dict[str, Any]): Person test and trace states from checkpoint

        Returns:
            bool: True if restoration was successful

        """
        people_by_id = {person.id: person for person in self.simulator.world.people}
        restored_count = 0

        for person_id_str, tt_data in person_states.items():
            person_id = int(person_id_str)

            if person_id not in people_by_id:
                logger.warning(f"Rank {mpi_rank}: Person {person_id} not found in current world for test and trace restoration")
                continue

            person = people_by_id[person_id]
            success = self._restore_single_test_trace(person, tt_data)

            if success:
                restored_count += 1
            else:
                logger.warning(f"Rank {mpi_rank}: Failed to restore test and trace for person {person_id}")

        print(f"Rank {mpi_rank}: Restored test and trace state for {restored_count} people")
        return True

    def _restore_single_test_trace(self, person, tt_data: Dict[str, Any]) -> bool:
        """Restore a single TestAndTrace object for a person.

        Args:
            person (Person): The person to restore test and trace for
            tt_data (Dict[str, Any]): Test and trace data from checkpoint

        Returns:
            bool: True if restoration was successful

        """
        try:
            from june.epidemiology.test_and_trace import TestAndTrace

            # Create new TestAndTrace object
            person.test_and_trace = TestAndTrace()
            tt = person.test_and_trace

            # Restore all attributes
            if 'notification_time' in tt_data:
                tt.notification_time = tt_data['notification_time']

            if 'scheduled_test_time' in tt_data:
                tt.scheduled_test_time = tt_data['scheduled_test_time']

            if 'contacts_traced' in tt_data:
                tt.contacts_traced = tt_data['contacts_traced']

            if 'time_of_testing' in tt_data:
                tt.time_of_testing = tt_data['time_of_testing']

            if 'time_of_result' in tt_data:
                tt.time_of_result = tt_data['time_of_result']

            if 'pending_test_result' in tt_data:
                tt.pending_test_result = tt_data['pending_test_result']

            if 'test_result' in tt_data:
                tt.test_result = tt_data['test_result']

            if 'emited_quarantine_start_event' in tt_data:
                tt.emited_quarantine_start_event = tt_data['emited_quarantine_start_event']

            if 'emited_quarantine_end_event' in tt_data:
                tt.emited_quarantine_end_event = tt_data['emited_quarantine_end_event']

            if 'isolation_start_time' in tt_data:
                tt.isolation_start_time = tt_data['isolation_start_time']

            if 'isolation_end_time' in tt_data:
                tt.isolation_end_time = tt_data['isolation_end_time']

            if 'tracer_id' in tt_data:
                tt.tracer_id = tt_data['tracer_id']

            if 'contact_reason' in tt_data:
                tt.contact_reason = tt_data['contact_reason']

            return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error restoring test and trace for person {person.id}: {e}")
            return False

    def _restore_contact_manager_state(self, cm_data: Dict[str, Any]) -> bool:
        """Restore contact manager state.

        Args:
            cm_data (Dict[str, Any]): Contact manager state data from checkpoint

        Returns:
            bool: True if restoration was successful

        """
        contact_manager = self.simulator.contact_manager

        try:
            # Restore leisure companions
            if 'leisure_companions' in cm_data:
                contact_manager.leisure_companions.clear()

                for person_id_str, companions in cm_data['leisure_companions'].items():
                    person_id = int(person_id_str)
                    contact_manager.leisure_companions[person_id] = {}

                    for companion_id_str, companion_info in companions.items():
                        companion_id = int(companion_id_str)
                        contact_manager.leisure_companions[person_id][companion_id] = {
                            'timestamp': companion_info['timestamp'],
                            'activity': companion_info['activity'],
                            'home_rank': companion_info.get('home_rank', 0)
                        }

            # Restore pending tests
            if 'tests_ids_pending' in cm_data:
                contact_manager.tests_ids_pending = []
                for test_info in cm_data['tests_ids_pending']:
                    # Convert string keys back to appropriate types
                    restored_test = {}
                    for key, value in test_info.items():
                        if key in ['person_id', 'residence_id', 'primary_activity_group_id', 'primary_activity_subgroup_type', 'pa_domain_id']:
                            restored_test[key] = int(value) if value != -1 else -1
                        elif key in ['result_time']:
                            restored_test[key] = float(value)
                        elif key in ['is_pa_external']:
                            restored_test[key] = bool(value)
                        else:
                            restored_test[key] = value
                    contact_manager.tests_ids_pending.append(restored_test)

            # Restore cleanup timestamp
            if 'last_cleanup' in cm_data:
                contact_manager.last_cleanup = cm_data['last_cleanup']

            print(f"Rank {mpi_rank}: Contact manager state restored - "
                  f"{len(contact_manager.leisure_companions)} people with leisure companions, "
                  f"{len(contact_manager.tests_ids_pending)} pending tests")

            return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error restoring contact manager state: {e}")
            return False

    def _validate_policy_configurations(self, policy_configs: Dict[str, Any]) -> bool:
        """Validate that current policy configurations match checkpoint configurations.

        Args:
            policy_configs (Dict[str, Any]): Policy configurations from checkpoint

        Returns:
            bool: True if configurations are compatible

        """
        # This is primarily for validation/warning purposes
        # Policy configurations are set at simulation start and don't need restoration

        try:
            from june.global_context import GlobalContext
            disease_config = GlobalContext.get_disease_config()

            if disease_config and hasattr(disease_config, 'policy_manager'):
                # Check key policy parameters for compatibility
                for policy_name in ['testing', 'tracing']:
                    if policy_name in policy_configs:
                        checkpoint_config = policy_configs[policy_name]
                        current_config = disease_config.policy_manager.get_policy_data(policy_name)

                        if current_config:
                            # Check critical parameters
                            if policy_name == 'testing':
                                checkpoint_accuracy = checkpoint_config.get('test_accuracy')
                                current_accuracy = current_config.get('test_accuracy')
                                if checkpoint_accuracy != current_accuracy:
                                    logger.warning(f"Rank {mpi_rank}: Test accuracy mismatch - checkpoint: {checkpoint_accuracy}, current: {current_accuracy}")

                            elif policy_name == 'tracing':
                                checkpoint_contacts = checkpoint_config.get('max_contacts_to_trace')
                                current_contacts = current_config.get('max_contacts_to_trace')
                                if checkpoint_contacts != current_contacts:
                                    logger.warning(f"Rank {mpi_rank}: Max contacts to trace mismatch - checkpoint: {checkpoint_contacts}, current: {current_contacts}")

            return True

        except Exception as e:
            logger.warning(f"Rank {mpi_rank}: Policy configuration validation error: {e}")
            return False

    def _restore_rat_dynamics_state(self, rat_data: Dict[str, Any]) -> bool:
        """Restore rat dynamics state including rat populations and disease models.

        Args:
            rat_data (Dict[str, Any]): Rat dynamics data from checkpoint

        Returns:
            bool: True if restoration was successful

        """
        logger.debug(f"Rank {mpi_rank}: Restoring rat dynamics state")

        try:
            # Check if rat dynamics was enabled
            if not rat_data.get('enabled', False):
                logger.debug(f"Rank {mpi_rank}: Rat dynamics was disabled in checkpoint")
                self.restoration_stats['rat_dynamics'] = {'enabled': False}
                return True

            # Check if rat dynamics was properly initialised
            if not rat_data.get('initialised', False):
                logger.warning(f"Rank {mpi_rank}: Rat dynamics was not properly initialised in checkpoint")
                self.restoration_stats['rat_dynamics'] = {'enabled': True, 'initialised': False}
                return True

            # Get the rat manager from simulator
            rat_manager = getattr(self.simulator, 'rat_manager', None)
            if rat_manager is None:
                logger.error(f"Rank {mpi_rank}: No rat manager found in simulator for restoration")
                return False

            print(f"Rank {mpi_rank}: Restoring rat dynamics - {rat_data['statistics']['total_rats']} rats")

            # Restore rat manager configuration
            if not self._restore_rat_manager_config(rat_manager, rat_data['rat_manager_config']):
                return False

            # Restore rat population state
            if not self._restore_rat_population_state(rat_manager, rat_data['rat_population_state']):
                return False

            # Restore density calculator state
            if not self._restore_density_calculator_state(rat_manager, rat_data['density_calculator_state']):
                return False

            # Restore disease model state
            if not self._restore_disease_model_state(rat_manager, rat_data['disease_model_state']):
                return False

            # Restore spatial grid state
            if not self._restore_spatial_grid_state(rat_manager, rat_data['spatial_grid_state']):
                return False

            # Restore area mapper state
            if not self._restore_area_mapper_state(rat_manager, rat_data['area_mapper_state']):
                return False

            # Ensure spatial grid is properly initialised after restoration
            if rat_manager.spatial_grid.spatial_grid_array is None and rat_manager.num_rats > 0:
                rat_manager.spatial_grid.initialise_spatial_grid()

            # Update restoration statistics
            self.restoration_stats['rat_dynamics'] = {
                'enabled': True,
                'initialised': True,
                'total_rats_restored': rat_data['statistics']['total_rats'],
                'infected_rats_restored': rat_data['statistics']['infected_rats'],
                'restored_successfully': True
            }

            print(f"Rank {mpi_rank}: Rat dynamics restoration completed successfully")
            return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error restoring rat dynamics state: {e}")
            import traceback
            traceback.print_exc()
            return False

    def _restore_rat_manager_config(self, rat_manager, config_data: Dict[str, Any]) -> bool:
        """Restore RatManager configuration with detailed debugging

        Args:
            rat_manager: 
            config_data (Dict[str, Any]): 

        """
        logger.debug(f"=== RAT MANAGER CONFIG DEBUG - AFTER CHECKPOINT RESTORATION ===")
        logger.debug(f"BEFORE RESTORATION:")
        logger.debug(f"  num_rats: {rat_manager.num_rats}")
        logger.debug(f"  dt: {rat_manager.dt}")
        logger.debug(f"  rat_to_human_factor: {rat_manager.rat_to_human_factor}")
        logger.debug(f"  human_to_rat_factor: {rat_manager.human_to_rat_factor}")
        logger.debug(f"  initial_infections: {rat_manager.initial_infections}")

        rat_manager.num_rats = config_data['num_rats']
        rat_manager.dt = config_data['dt']
        rat_manager.rat_to_human_factor = config_data['rat_to_human_factor']
        rat_manager.human_to_rat_factor = config_data['human_to_rat_factor']
        rat_manager.initial_infections = config_data['initial_infections']

        # Restore grid dimensions and bounds
        rat_manager.minx = config_data.get('minx')
        rat_manager.miny = config_data.get('miny')
        rat_manager.maxx = config_data.get('maxx')
        rat_manager.maxy = config_data.get('maxy')
        rat_manager.sim_rows = config_data.get('sim_rows')
        rat_manager.sim_cols = config_data.get('sim_cols')

        logger.debug(f"AFTER RESTORATION:")
        logger.debug(f"  num_rats: {rat_manager.num_rats}")
        logger.debug(f"  dt: {rat_manager.dt}")
        logger.debug(f"  rat_to_human_factor: {rat_manager.rat_to_human_factor}")
        logger.debug(f"  human_to_rat_factor: {rat_manager.human_to_rat_factor}")
        logger.debug(f"  initial_infections: {rat_manager.initial_infections}")
        logger.debug(f"  Grid bounds - minx: {rat_manager.minx}, miny: {rat_manager.miny}, maxx: {rat_manager.maxx}, maxy: {rat_manager.maxy}")
        logger.debug(f"  Grid dimensions - rows: {rat_manager.sim_rows}, cols: {rat_manager.sim_cols}")

        return True

    def _restore_rat_population_state(self, rat_manager, population_data: Dict[str, Any]) -> bool:
        """Restore rat population arrays with comprehensive debugging

        Args:
            rat_manager: 
            population_data (Dict[str, Any]): 

        """
        if population_data['num_rats'] == 0:
            logger.debug(f"=== RAT POPULATION STATE DEBUG - NO RATS TO RESTORE ===")
            return True

        print(f"=== RAT POPULATION STATE DEBUG - AFTER CHECKPOINT RESTORATION ===")
        print(f"BEFORE RESTORATION:")
        print(f"  Total rats: {rat_manager.num_rats}")

        # Debug BEFORE restoration
        if rat_manager.positions is not None:
            print(f"  Existing positions shape: {rat_manager.positions.shape}")
            print(f"  Existing positions sample: {rat_manager.positions[:3].tolist()}")
        else:
            print(f"  Existing positions: None")

        if rat_manager.states is not None:
            state_counts = np.bincount(rat_manager.states, minlength=3)
            print(f"  Existing state distribution - S: {state_counts[0]}, I: {state_counts[1]}, R: {state_counts[2]}")
        else:
            print(f"  Existing states: None")

        # Restore population arrays with explicit type conversion
        rat_manager.positions = np.array(population_data['positions'], dtype=float) if population_data['positions'] is not None else None
        rat_manager.states = np.array(population_data['states'], dtype=int) if population_data['states'] is not None else None
        rat_manager.infection_age = np.array(population_data['infection_age'], dtype=float) if population_data['infection_age'] is not None else None
        rat_manager.immunity = np.array(population_data['immunity'], dtype=float) if population_data['immunity'] is not None else None
        rat_manager.personal_delta = np.array(population_data['personal_delta'], dtype=float) if population_data['personal_delta'] is not None else None
        rat_manager.grid_indices = np.array(population_data['grid_indices'], dtype=int) if population_data['grid_indices'] is not None else None

        print(f"AFTER RESTORATION:")
        # Debug positions after restoration
        if rat_manager.positions is not None:
            print(f"  Restored positions shape: {rat_manager.positions.shape}")
            print(f"  Restored positions dtype: {rat_manager.positions.dtype}")
            print(f"  First 5 restored positions: {rat_manager.positions[:5].tolist()}")
            print(f"  Restored position bounds - X: [{np.min(rat_manager.positions[:, 0]):.6f}, {np.max(rat_manager.positions[:, 0]):.6f}]")
            print(f"  Restored position bounds - Y: [{np.min(rat_manager.positions[:, 1]):.6f}, {np.max(rat_manager.positions[:, 1]):.6f}]")
            print(f"  Restored position mean - X: {np.mean(rat_manager.positions[:, 0]):.6f}, Y: {np.mean(rat_manager.positions[:, 1]):.6f}")

        # Debug states after restoration
        if rat_manager.states is not None:
            state_counts = np.bincount(rat_manager.states, minlength=3)
            print(f"  Restored states shape: {rat_manager.states.shape}, dtype: {rat_manager.states.dtype}")
            print(f"  Restored state distribution - Susceptible(0): {state_counts[0]}, Infected(1): {state_counts[1]}, Recovered(2): {state_counts[2]}")
            print(f"  First 10 restored states: {rat_manager.states[:10].tolist()}")

        # Debug infection ages after restoration
        if rat_manager.infection_age is not None:
            infected_mask = rat_manager.states == 1
            print(f"  Restored infection ages shape: {rat_manager.infection_age.shape}, dtype: {rat_manager.infection_age.dtype}")
            print(f"  Restored infection age range: [{np.min(rat_manager.infection_age):.6f}, {np.max(rat_manager.infection_age):.6f}]")
            if np.any(infected_mask):
                print(f"  Restored infected rats infection ages: {rat_manager.infection_age[infected_mask][:10].tolist()}")

        # Debug immunity after restoration
        if rat_manager.immunity is not None:
            print(f"  Restored immunity shape: {rat_manager.immunity.shape}, dtype: {rat_manager.immunity.dtype}")
            print(f"  Restored immunity range: [{np.min(rat_manager.immunity):.6f}, {np.max(rat_manager.immunity):.6f}]")
            print(f"  Restored immunity mean: {np.mean(rat_manager.immunity):.6f}")
            print(f"  First 10 restored immunity values: {rat_manager.immunity[:10].tolist()}")

        # Debug personal delta after restoration
        if rat_manager.personal_delta is not None:
            print(f"  Restored personal delta shape: {rat_manager.personal_delta.shape}, dtype: {rat_manager.personal_delta.dtype}")
            print(f"  Restored personal delta range: [{np.min(rat_manager.personal_delta):.6f}, {np.max(rat_manager.personal_delta):.6f}]")
            print(f"  Restored personal delta mean: {np.mean(rat_manager.personal_delta):.6f}")

        # Debug grid indices after restoration
        if rat_manager.grid_indices is not None:
            print(f"  Restored grid indices shape: {rat_manager.grid_indices.shape}, dtype: {rat_manager.grid_indices.dtype}")
            print(f"  Restored grid indices range - Row: [{np.min(rat_manager.grid_indices[:, 0])}, {np.max(rat_manager.grid_indices[:, 0])}]")
            print(f"  Restored grid indices range - Col: [{np.min(rat_manager.grid_indices[:, 1])}, {np.max(rat_manager.grid_indices[:, 1])}]")
            print(f"  First 5 restored grid indices: {rat_manager.grid_indices[:5].tolist()}")

        print(f"=== RAT POPULATION RESTORATION COMPLETE ===")
        return True

    def _restore_density_calculator_state(self, rat_manager, density_data: Dict[str, Any]) -> bool:
        """Restore density calculator state

        Args:
            rat_manager: 
            density_data (Dict[str, Any]): 

        """
        density_calc = rat_manager.density_calculator

        # Restore configuration
        density_calc.rat_ratio = density_data['rat_ratio']
        density_calc.cell_size = density_data['cell_size']
        density_calc.gaussian_sigma = density_data['gaussian_sigma']
        density_calc.precomputed_density_path = density_data['precomputed_density_path']
        density_calc.total_population = density_data['total_population']
        density_calc.coordinate_system = density_data.get('coordinate_system')

        # Restore grid bounds
        density_calc.minx = density_data.get('minx')
        density_calc.miny = density_data.get('miny')
        density_calc.maxx = density_data.get('maxx')
        density_calc.maxy = density_data.get('maxy')
        density_calc.sim_rows = density_data.get('sim_rows')
        density_calc.sim_cols = density_data.get('sim_cols')

        # Restore rat density array
        if density_data['rat_density'] is not None:
            density_calc.rat_density = np.array(density_data['rat_density'], dtype=float)

        return True

    def _restore_disease_model_state(self, rat_manager, disease_data: Dict[str, Any]) -> bool:
        """Restore disease model state with detailed debugging

        Args:
            rat_manager: 
            disease_data (Dict[str, Any]): 

        """
        disease_model = rat_manager.disease_model

        logger.debug(f"=== RAT DISEASE MODEL STATE DEBUG - AFTER CHECKPOINT RESTORATION ===")
        logger.debug(f"BEFORE RESTORATION:")
        logger.debug(f"  Disease parameters - beta: {disease_model.beta}, alpha: {disease_model.alpha}, gamma: {disease_model.gamma}")
        logger.debug(f"  History lengths - infected: {len(disease_model.infected_history)}, immunity_08: {len(disease_model.immunity_08_history)}")
        logger.debug(f"  Total global seeds: {disease_model.total_global_seeds}")

        # Restore disease parameters
        disease_model.beta = disease_data['beta']
        disease_model.alpha = disease_data['alpha']
        disease_model.gamma = disease_data['gamma']
        disease_model.delta_mean = disease_data['delta_mean']
        disease_model.delta_std = disease_data['delta_std']
        disease_model.max_trans_distance = disease_data['max_trans_distance']
        disease_model.p_global_seed = disease_data['p_global_seed']
        disease_model.global_seed_min = disease_data['global_seed_min']
        disease_model.global_seed_max = disease_data['global_seed_max']
        disease_model.global_seed_immunity_threshold = disease_data['global_seed_immunity_threshold']
        disease_model.infectiousness_threshold = disease_data['infectiousness_threshold']

        # Restore disease history tracking
        disease_model.infected_history = disease_data['infected_history']
        disease_model.immunity_08_history = disease_data['immunity_08_history']
        disease_model.immunity_05_history = disease_data['immunity_05_history']
        disease_model.global_seeds_history = disease_data['global_seeds_history']
        disease_model.total_global_seeds = disease_data['total_global_seeds']

        # Restore transmission kernel
        if disease_data['transmission_kernel'] is not None:
            disease_model.transmission_kernel = np.array(disease_data['transmission_kernel'], dtype=float)

        logger.debug(f"AFTER RESTORATION:")
        logger.debug(f"  Disease parameters - beta: {disease_model.beta}, alpha: {disease_model.alpha}, gamma: {disease_model.gamma}")
        logger.debug(f"  Immunity decay - delta_mean: {disease_model.delta_mean}, delta_std: {disease_model.delta_std}")
        logger.debug(f"  Transmission - max_trans_distance: {disease_model.max_trans_distance}, infectiousness_threshold: {disease_model.infectiousness_threshold}")
        logger.debug(f"  Global seeding - p_global_seed: {disease_model.p_global_seed}, range: [{disease_model.global_seed_min}, {disease_model.global_seed_max}]")
        logger.debug(f"  Global seed immunity threshold: {disease_model.global_seed_immunity_threshold}")

        # Debug history tracking after restoration
        logger.debug(f"  Restored history lengths - infected: {len(disease_model.infected_history)}, immunity_08: {len(disease_model.immunity_08_history)}")
        logger.debug(f"  Restored history lengths - immunity_05: {len(disease_model.immunity_05_history)}, global_seeds: {len(disease_model.global_seeds_history)}")
        logger.debug(f"  Restored total global seeds: {disease_model.total_global_seeds}")

        if disease_model.infected_history:
            logger.debug(f"  Restored recent infected history (last 5): {disease_model.infected_history[-5:]}")
        if disease_model.global_seeds_history:
            logger.debug(f"  Restored recent global seeds history (last 5): {disease_model.global_seeds_history[-5:]}")

        # Debug transmission kernel after restoration
        if disease_model.transmission_kernel is not None:
            logger.debug(f"  Restored transmission kernel shape: {disease_model.transmission_kernel.shape}")
            logger.debug(f"  Restored transmission kernel sum: {np.sum(disease_model.transmission_kernel):.6f}")
            logger.debug(f"  Restored transmission kernel max: {np.max(disease_model.transmission_kernel):.6f}")
        else:
            logger.debug(f"  Restored transmission kernel: None")

        return True

    def _restore_spatial_grid_state(self, rat_manager, spatial_data: Dict[str, Any]) -> bool:
        """Restore spatial grid configuration with debugging

        Args:
            rat_manager: 
            spatial_data (Dict[str, Any]): 

        """
        spatial_grid = rat_manager.spatial_grid

        logger.debug(f"=== RAT SPATIAL GRID STATE DEBUG - AFTER CHECKPOINT RESTORATION ===")
        logger.debug(f"BEFORE RESTORATION:")
        logger.debug(f"  Cell size: {spatial_grid.cell_size}, max_trans_distance: {spatial_grid.max_trans_distance}")
        logger.debug(f"  Movement enabled: {spatial_grid.enable_movement}, p_move: {spatial_grid.p_move}")

        # Restore configuration
        spatial_grid.cell_size = spatial_data['cell_size']
        spatial_grid.max_trans_distance = spatial_data['max_trans_distance']
        spatial_grid.enable_movement = spatial_data['enable_movement']
        spatial_grid.p_move = spatial_data['p_move']
        spatial_grid.lambda_move = spatial_data['lambda_move']
        spatial_grid.Rmax_move = spatial_data['Rmax_move']
        spatial_grid.grid_cell_size = spatial_data['grid_cell_size']
        spatial_grid.grid_cell_rows = spatial_data.get('grid_cell_rows')
        spatial_grid.grid_cell_cols = spatial_data.get('grid_cell_cols')

        # Restore movement offsets
        if spatial_data['move_offsets'] is not None:
            spatial_grid.move_offsets = np.array(spatial_data['move_offsets'], dtype=int)
        if spatial_data['move_offset_dists'] is not None:
            spatial_grid.move_offset_dists = np.array(spatial_data['move_offset_dists'], dtype=int)

        logger.debug(f"AFTER RESTORATION:")
        logger.debug(f"  Restored cell size: {spatial_grid.cell_size}, max_trans_distance: {spatial_grid.max_trans_distance}")
        logger.debug(f"  Restored movement enabled: {spatial_grid.enable_movement}, p_move: {spatial_grid.p_move}")
        logger.debug(f"  Restored movement parameters - lambda_move: {spatial_grid.lambda_move}, Rmax_move: {spatial_grid.Rmax_move}")
        logger.debug(f"  Restored grid cell size: {spatial_grid.grid_cell_size}")
        logger.debug(f"  Restored grid cell dimensions - rows: {spatial_grid.grid_cell_rows}, cols: {spatial_grid.grid_cell_cols}")

        # Debug movement offsets after restoration
        if spatial_grid.move_offsets is not None:
            logger.debug(f"  Restored move offsets shape: {spatial_grid.move_offsets.shape}")
            logger.debug(f"  Restored move offset distances shape: {spatial_grid.move_offset_dists.shape if spatial_grid.move_offset_dists is not None else 'None'}")
            logger.debug(f"  First 5 restored move offsets: {spatial_grid.move_offsets[:5].tolist()}")
            logger.debug(f"  Restored move offset distance range: [{np.min(spatial_grid.move_offset_dists)}, {np.max(spatial_grid.move_offset_dists)}]")

        spatial_grid.spatial_grid_array = None
        logger.debug(f"  Spatial grid array cleared to avoid stale data")

        return True

    def _restore_area_mapper_state(self, rat_manager, mapper_data: Dict[str, Any]) -> bool:
        """Restore area mapper state

        Args:
            rat_manager: 
            mapper_data (Dict[str, Any]): 

        """
        area_mapper = rat_manager.area_mapper

        # Restore basic mappings
        area_mapper.area_to_msoa = mapper_data['area_to_msoa']
        area_mapper._msoa_mappings_initialised = mapper_data['_msoa_mappings_initialised']

        # Restore msoa_to_areas (convert area names back to Area objects)
        area_mapper.msoa_to_areas = {}
        for msoa_id, area_names in mapper_data['msoa_to_areas'].items():
            area_objects = []
            for area_name in area_names:
                # Find the Area object by name
                for super_area in self.simulator.world.super_areas:
                    for area in super_area.areas:
                        if hasattr(area, 'name') and area.name == area_name:
                            area_objects.append(area)
                            break
            if area_objects:
                area_mapper.msoa_to_areas[msoa_id] = area_objects

        # Restore msoa_to_cells
        area_mapper.msoa_to_cells = {}
        for msoa_id, cells_list in mapper_data['msoa_to_cells'].items():
            if cells_list:
                # Convert back to list of tuples with explicit int conversion
                area_mapper.msoa_to_cells[msoa_id] = [tuple(int(coord) for coord in cell) for cell in cells_list]

        return True

    def _restore_tt_event_recorder_state(self, tt_data: Dict[str, Any]) -> bool:
        """Restore TTEventRecorder state including daily/cumulative counters,
        unique IDs, current status, and event buffer.

        Args:
            tt_data (Dict[str, Any]): TTEventRecorder state data

        Returns:
            bool: True if restoration was successful

        """
        print(f"Rank {mpi_rank}: Restoring TTEventRecorder state")

        # Check if TTEventRecorder was enabled
        if not tt_data.get('enabled', False):
            print(f"Rank {mpi_rank}: TTEventRecorder was disabled in checkpoint - skipping restoration")
            self.restoration_stats['tt_event_recorder'] = {'enabled': False, 'restored': False}
            return True

        # TTEventRecorder restoration strategy:
        # Store the data for later restoration after TTEventRecorder is naturally created
        # This avoids triggering premature TTEventRecorder creation during restoration
        print(f"Rank {mpi_rank}: Storing TTEventRecorder data for deferred restoration")

        # Store the TTEventRecorder data in the simulator for later restoration
        if not hasattr(self.simulator, '_pending_tt_event_recorder_data'):
            self.simulator._pending_tt_event_recorder_data = {}

        self.simulator._pending_tt_event_recorder_data = tt_data
        print(f"Rank {mpi_rank}: TTEventRecorder data stored for restoration after natural creation")

        self.restoration_stats['tt_event_recorder'] = {
            'enabled': True, 
            'restored': 'deferred', 
            'reason': 'stored_for_deferred_restoration',
            'note': 'Data will be applied after TTEventRecorder is naturally created'
        }
        return True

    def _log_restoration_summary(self):
        """Log comprehensive restoration summary"""
        print(f"\n=== CHECKPOINT RESTORATION SUMMARY (Rank {mpi_rank}) ===")

        for component, stats in self.restoration_stats.items():
            print(f"\n{component.upper()}:")
            for key, value in stats.items():
                print(f"  {key}: {value}")

        print(f"\n=== RESTORATION COMPLETE (Rank {mpi_rank}) ===")

    def _get_current_feature_flags(self) -> Dict[str, bool]:
        """Get the current feature flags from the simulator.


        Returns:
            Dict[str, bool]: Current feature flag states

        """
        return {
            'test_and_trace_enabled': getattr(self.simulator, 'test_and_trace_enabled', False),
            'ratty_dynamics_enabled': getattr(self.simulator, 'ratty_dynamics_enabled', False),
            'friend_hangouts_enabled': getattr(self.simulator, 'friend_hangouts_enabled', False)
        }

    def _detect_feature_activations(self, checkpoint_features: Dict[str, bool], current_features: Dict[str, bool]) -> Dict[str, bool]:
        """Detect which features are newly enabled (were disabled in checkpoint but enabled now).

        Args:
            checkpoint_features (Dict[str, bool]): Feature flags from the checkpoint
            current_features (Dict[str, bool]): Current feature flags

        Returns:
            Dict[str, bool]: Features that are newly enabled

        """
        activations = {}

        for feature_name, current_enabled in current_features.items():
            checkpoint_enabled = checkpoint_features.get(feature_name, False)

            # Feature is newly activated if it's enabled now but was disabled in checkpoint
            if current_enabled and not checkpoint_enabled:
                activations[feature_name] = True
                logger.info(f"Rank {mpi_rank}: Feature '{feature_name}' newly activated (checkpoint: {checkpoint_enabled} -> current: {current_enabled})")
            elif current_enabled and checkpoint_enabled:
                logger.debug(f"Rank {mpi_rank}: Feature '{feature_name}' remains enabled")
            elif not current_enabled and checkpoint_enabled:
                logger.info(f"Rank {mpi_rank}: Feature '{feature_name}' disabled (checkpoint: {checkpoint_enabled} -> current: {current_enabled})")
            else:
                logger.debug(f"Rank {mpi_rank}: Feature '{feature_name}' remains disabled")

        return activations

    def _validate_feature_activations(self, feature_activations: Dict[str, bool], metadata: Dict[str, Any]) -> Dict[str, Any]:
        """Validate that feature activations are safe and feasible.

        Args:
            feature_activations (Dict[str, bool]): Features to be activated
            metadata (Dict[str, Any]): Checkpoint metadata for context

        Returns:
            Dict[str, Any]: Validation result with 'valid' boolean and 'reason' if invalid

        """
        validation_warnings = []

        # Check if simulator has required dependencies for each feature
        for feature_name in feature_activations:
            if feature_name == 'test_and_trace_enabled':
                # Check if disease config supports test and trace
                try:
                    from june.global_context import GlobalContext
                    disease_config = GlobalContext.get_disease_config()
                    if disease_config is None:
                        return {
                            'valid': False,
                            'reason': 'No disease configuration available for test and trace initialization'
                        }

                    # Check if required policies exist
                    tracing_policy = disease_config.policy_manager.get_policy_data("tracing")
                    if not tracing_policy:
                        validation_warnings.append("No tracing policy found - using defaults")

                except Exception as e:
                    return {
                        'valid': False,
                        'reason': f'Error checking test and trace dependencies: {e}'
                    }

            elif feature_name == 'ratty_dynamics_enabled':
                # Check if rat dynamics modules are available
                try:
                    from june.zoonosis.rat_manager import RatManager
                    from june.zoonosis.zoonotic_transmission import ZoonoticTransmission

                    # Check if world has required attributes for rat dynamics
                    if not hasattr(self.simulator.world, 'super_areas'):
                        return {
                            'valid': False,
                            'reason': 'World does not have required super_areas for rat dynamics'
                        }

                except ImportError as e:
                    return {
                        'valid': False,
                        'reason': f'Rat dynamics modules not available: {e}'
                    }

            elif feature_name == 'friend_hangouts_enabled':
                # Check if leisure system is available
                if not hasattr(self.simulator.activity_manager, 'leisure') or self.simulator.activity_manager.leisure is None:
                    validation_warnings.append("No leisure system available - friend hangouts may have limited functionality")

        # Log any warnings
        for warning in validation_warnings:
            logger.warning(f"Rank {mpi_rank}: Feature activation warning: {warning}")

        return {
            'valid': True,
            'warnings': validation_warnings
        }

    def _activate_new_features(self, feature_activations: Dict[str, bool]) -> bool:
        """Activate newly enabled features by initializing their components.

        Args:
            feature_activations (Dict[str, bool]): Features that need to be activated

        Returns:
            bool: True if all activations were successful

        """
        activation_success = True

        # Activate test and trace if newly enabled
        if feature_activations.get('test_and_trace_enabled'):
            print(f"Rank {mpi_rank}: Activating test and trace system...")
            success = self._activate_test_and_trace()
            activation_success = activation_success and success
            if success:
                print(f"Rank {mpi_rank}: Test and trace system activated successfully")
            else:
                logger.error(f"Rank {mpi_rank}: Failed to activate test and trace system")

        # Activate rat dynamics if newly enabled
        if feature_activations.get('ratty_dynamics_enabled'):
            print(f"Rank {mpi_rank}: Activating rat dynamics system...")
            success = self._activate_rat_dynamics()
            activation_success = activation_success and success
            if success:
                print(f"Rank {mpi_rank}: Rat dynamics system activated successfully")
            else:
                logger.error(f"Rank {mpi_rank}: Failed to activate rat dynamics system")

        # Activate friend hangouts if newly enabled
        if feature_activations.get('friend_hangouts_enabled'):
            print(f"Rank {mpi_rank}: Activating friend hangouts system...")
            success = self._activate_friend_hangouts()
            activation_success = activation_success and success
            if success:
                print(f"Rank {mpi_rank}: Friend hangouts system activated successfully")
            else:
                logger.error(f"Rank {mpi_rank}: Failed to activate friend hangouts system")

        # Activate sexual encounters if newly enabled
        if feature_activations.get('sexual_encounter_enabled'):
            print(f"Rank {mpi_rank}: Activating sexual encounter system...")
            success = self._activate_sexual_encounters()
            activation_success = activation_success and success
            if success:
                print(f"Rank {mpi_rank}: Sexual encounter system activated successfully")
            else:
                logger.error(f"Rank {mpi_rank}: Failed to activate sexual encounter system")

        return activation_success

    def _activate_test_and_trace(self) -> bool:
        """Activate the test and trace system from restoration point.


        Returns:
            bool: True if activation was successful

        """
        try:
            # Import required modules
            from june.groups.contact import ContactManager
            from june.global_context import GlobalContext

            # Initialize contact manager
            if self.simulator.contact_manager is None:
                print(f"Rank {mpi_rank}: Initializing contact manager for test and trace")
                self.simulator.contact_manager = ContactManager(self.simulator)

                # Configure contact retention based on policy settings
                disease_config = GlobalContext.get_disease_config()
                if disease_config is not None:
                    tracing_data = disease_config.policy_manager.get_policy_data("tracing")
                    self.simulator.contact_retention_days = tracing_data.get("contact_retention_days", 14) if tracing_data else 14
                    print(f"Rank {mpi_rank}: Contact retention days set to: {self.simulator.contact_retention_days}")
                else:
                    self.simulator.contact_retention_days = 14  # Default
                    print(f"Rank {mpi_rank}: Using default contact retention days: 14")

                # Connect contact manager to leisure system if it exists
                if self.simulator.activity_manager.leisure is not None:
                    print(f"Rank {mpi_rank}: Connecting contact manager to leisure system")
                    self.simulator.activity_manager.leisure.set_contact_manager(self.simulator.contact_manager)

                print(f"Rank {mpi_rank}: Test and trace system activated with fresh contact manager")
                return True
            else:
                print(f"Rank {mpi_rank}: Contact manager already exists")
                return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error activating test and trace: {e}")
            import traceback
            traceback.print_exc()
            return False

    def _activate_rat_dynamics(self) -> bool:
        """Activate the rat dynamics system from restoration point.


        Returns:
            bool: True if activation was successful

        """
        try:
            # Import required modules
            from june.zoonosis.rat_manager import RatManager
            from june.zoonosis.zoonotic_transmission import ZoonoticTransmission

            # Initialize rat manager if not already present
            if self.simulator.rat_manager is None:
                print(f"Rank {mpi_rank}: Initializing rat manager for rat dynamics")

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

                # Initialize rat manager with fresh state
                self.simulator.rat_manager = RatManager(
                    world=self.simulator.world, 
                    precomputed_density_path=rat_density_path
                )
                print(f"Rank {mpi_rank}: Rat manager initialized with density path: {rat_density_path}")

                # Set animation flag if enabled
                # Animation handling removed - now done by post-processing script
                print(f"Rank {mpi_rank}: Rat data saving enabled: {getattr(self.simulator, 'save_rat_data', False)}")
            else:
                print(f"Rank {mpi_rank}: Rat manager already exists")

            # Initialize zoonotic transmission if not already present
            if self.simulator.zoonotic_transmission is None and ZoonoticTransmission is not None:
                print(f"Rank {mpi_rank}: Initializing zoonotic transmission module")
                self.simulator.zoonotic_transmission = ZoonoticTransmission(rat_manager=self.simulator.rat_manager)
                print(f"Rank {mpi_rank}: Zoonotic transmission module initialized successfully")

            print(f"Rank {mpi_rank}: Rat dynamics system activated with fresh state")
            return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error activating rat dynamics: {e}")
            import traceback
            traceback.print_exc()
            return False

    def _activate_friend_hangouts(self) -> bool:
        """Activate the friend hangouts system from restoration point.


        Returns:
            bool: True if activation was successful

        """
        try:
            # Friend hangouts activation is primarily handled through the activity manager
            # and leisure system. The main requirement is that the feature flag is set,
            # which should already be done in the simulator initialization.

            print(f"Rank {mpi_rank}: Friend hangouts feature activated")
            print(f"Rank {mpi_rank}: Social networks will be built from this point forward")

            # The actual friend network building happens dynamically during leisure activities
            # No additional initialization is required beyond the feature flag

            return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error activating friend hangouts: {e}")
            return False

    def _activate_sexual_encounters(self) -> bool:
        """Activate sexual encounter system for newly enabled feature.


        Returns:
            bool: True if activation was successful

        """
        try:
            # Sexual encounters are stateless - just ensure the feature is enabled
            self.simulator.sexual_encounter_enabled = True

            # Initialize sexual encounter system if it doesn't exist
            if not hasattr(self.simulator, 'sexual_encounter') or self.simulator.sexual_encounter is None:
                from june.sexual_encounter.sexual_encounter import SexualEncounter
                self.simulator.sexual_encounter = SexualEncounter(
                    private_rooms=getattr(self.simulator, 'private_rooms', None),
                    contact_manager=getattr(self.simulator, 'contact_manager', None)
                )

            print(f"Rank {mpi_rank}: Sexual encounter feature activated")
            print(f"Rank {mpi_rank}: Sexual encounter system will process invitations from this point forward")

            # No additional state needs to be initialized here, as the sexual encounter system
            # will handle the creation of private rooms and invitations dynamically.

            return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error activating sexual encounters: {e}")
            return False

    def _validate_disease_model_compatibility(self, metadata: Dict[str, Any]) -> bool:
        """Validate that checkpoint disease model matches current disease configuration.

        This method checks the disease name stored in checkpoint metadata against the currently
        configured disease model to prevent KeyError during simulation.

        Args:
            metadata (Dict[str, Any]): The loaded checkpoint metadata

        Returns:
            bool: True if checkpoint disease model is compatible with current config

        """
        try:
            # Get current disease name from simulator
            current_disease_name = None
            if (hasattr(self.simulator, 'epidemiology') and 
                self.simulator.epidemiology and 
                hasattr(self.simulator.epidemiology, 'infection_selectors') and
                self.simulator.epidemiology.infection_selectors._infection_selectors):
                current_disease_name = self.simulator.epidemiology.infection_selectors._infection_selectors[0].disease_name

            if not current_disease_name:
                logger.warning(f"Rank {mpi_rank}: Cannot validate disease compatibility - no disease name found in current config")
                return True  # Allow restoration to proceed if we can't validate

            # Get checkpoint disease name from metadata
            checkpoint_disease_name = metadata.get('disease_name')

            if not checkpoint_disease_name:
                logger.warning(f"Rank {mpi_rank}: No disease name found in checkpoint metadata - checkpoint may be from older version")
                return True  # Allow restoration for older checkpoints without disease metadata

            # Check for mismatch
            if checkpoint_disease_name.lower() != current_disease_name.lower():
                logger.critical(
                    f"Rank {mpi_rank}: DISEASE MODEL MISMATCH - Cannot restore checkpoint!\n"
                    f"Checkpoint disease model: '{checkpoint_disease_name}'\n"
                    f"Current configuration disease model: '{current_disease_name}'\n"
                    f"This mismatch would cause KeyError during simulation.\n"
                    f"Please either:\n"
                    f"  1) Update your disease model in config from '{current_disease_name}' to '{checkpoint_disease_name}', or\n"
                    f"  2) Start a fresh simulation without loading checkpoints."
                )
                return False

            logger.info(f"Rank {mpi_rank}: Disease compatibility validated - both checkpoint and config use '{current_disease_name}'")
            return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error during disease compatibility validation: {e}")
            return False

    def _restore_school_incident_state(self, school_data: Dict[str, Any]) -> bool:
        """Restore school incident tracking state for NotSendingKidsToSchool policy.

        Args:
            school_data (Dict[str, Any]): School incident tracking data from checkpoint

        Returns:
            bool: True if restoration was successful

        """
        try:
            logger.debug("Starting school incident state restoration")

            # Restore local school incident data
            if 'local_schools' in school_data:
                restored_schools = 0
                for school_id_str, incident_data in school_data['local_schools'].items():
                    # Find the school object by ID (convert string back to original type)
                    school = None
                    for s in self.simulator.world.schools:
                        if str(s.id) == school_id_str:
                            school = s
                            break

                    if school is not None:
                        school.student_deaths = incident_data.get('student_deaths', 0)
                        school.student_icu_transfers = incident_data.get('student_icu_transfers', 0)
                        school.households_avoiding_school = set(incident_data.get('households_avoiding_school', []))
                        # Convert string keys back to original types for household decision tracking
                        school.households_last_decision_at_death_count = {int(k): v for k, v in incident_data.get('households_last_decision_at_death_count', {}).items()}
                        school.households_last_decision_at_icu_count = {int(k): v for k, v in incident_data.get('households_last_decision_at_icu_count', {}).items()}
                        restored_schools += 1

                logger.debug(f"Restored incident data for {restored_schools} local schools")

            # Restore global school incidents registry
            if 'global_school_incidents' in school_data:
                if not hasattr(self.simulator.world, 'global_school_incidents'):
                    self.simulator.world.global_school_incidents = {}
                # Convert string keys back to original types
                for school_id_str, incidents in school_data['global_school_incidents'].items():
                    self.simulator.world.global_school_incidents[int(school_id_str)] = incidents
                logger.debug(f"Restored {len(school_data['global_school_incidents'])} global school incidents")

            # Restore global household decisions registry  
            if 'global_household_decisions' in school_data:
                if not hasattr(self.simulator.world, 'global_household_decisions'):
                    self.simulator.world.global_household_decisions = {}

                # Convert lists back to sets and string keys back to original types
                for school_key_str, decisions in school_data['global_household_decisions'].items():
                    self.simulator.world.global_household_decisions[school_key_str] = {
                        'avoiding_households': set(decisions.get('avoiding_households', [])),
                        'last_death_decision': {int(k): v for k, v in decisions.get('last_death_decision', {}).items()},
                        'last_icu_decision': {int(k): v for k, v in decisions.get('last_icu_decision', {}).items()}
                    }

                logger.debug(f"Restored {len(school_data['global_household_decisions'])} global household decisions")

            self.restoration_stats['school_incidents'] = {
                'enabled': True,
                'restoration_successful': True
            }

            logger.debug("School incident state restoration completed successfully")
            return True

        except Exception as e:
            logger.error(f"Rank {mpi_rank}: Error restoring school incident state: {e}")
            return False

__init__(simulator)

Initialise the checkpoint restorer.

Parameters

simulator : Simulator The JUNE simulator instance to restore state into

Source code in june/checkpointing/checkpoint_restorer.py

def __init__(self, simulator):
    """
    Initialise the checkpoint restorer.

    Parameters
    ----------
    simulator : Simulator
        The JUNE simulator instance to restore state into
    """
    self.simulator = simulator
    self.restoration_stats = {}
    self.random_state_manager = RandomStateManager()

restore_from_checkpoint(checkpoint_path)

Restore simulation from checkpoint files.

Parameters:

Name	Type	Description	Default
checkpoint_path	Path	Directory containing checkpoint files	required

Returns:

Name	Type	Description
bool	bool	True if restoration was successful

Source code in june/checkpointing/checkpoint_restorer.py

def restore_from_checkpoint(self, checkpoint_path: Path) -> bool:
    """Restore simulation from checkpoint files.

    Args:
        checkpoint_path (Path): Directory containing checkpoint files

    Returns:
        bool: True if restoration was successful

    """

    print(f"Rank {mpi_rank}: Starting checkpoint restoration from {checkpoint_path}")

    # Validate checkpoint directory
    if not self._validate_checkpoint_directory(checkpoint_path):
        logger.error(f"Rank {mpi_rank}: Invalid checkpoint directory: {checkpoint_path}")
        return False

    # Load checkpoint metadata
    metadata = self._load_checkpoint_metadata(checkpoint_path)
    if not metadata:
        logger.error(f"Rank {mpi_rank}: Could not load checkpoint metadata")
        return False

    # Load rank-specific checkpoint data
    rank_file = checkpoint_path / f"checkpoint_rank_{mpi_rank}.h5"
    if not rank_file.exists():
        logger.error(f"Rank {mpi_rank}: Checkpoint file not found: {rank_file}")
        return False

    checkpoint_data = self._load_checkpoint_data(rank_file)
    if not checkpoint_data:
        logger.error(f"Rank {mpi_rank}: Could not load checkpoint data")
        return False

    # EARLY VALIDATION: Check for disease model mismatch before attempting restoration
    if not self._validate_disease_model_compatibility(metadata):
        logger.error(f"Rank {mpi_rank}: Checkpoint restoration aborted due to disease model mismatch")
        return False

    # Coordinate across MPI ranks before restoration
    if mpi_available:
        mpi_comm.Barrier()

    # Restore simulation state components
    success = self._restore_simulation_state(checkpoint_data, metadata, checkpoint_path)

    if success:
        # Mark this simulator as resumed from checkpoint
        self.simulator._is_resumed_from_checkpoint = True
        self.simulator._is_resumed_and_first_round = True

        # Final coordination
        if mpi_available:
            mpi_comm.Barrier()

        print(f"Rank {mpi_rank}: Checkpoint restoration completed successfully")
        self._log_restoration_summary()
        return True
    else:
        logger.error(f"Rank {mpi_rank}: Checkpoint restoration failed")
        return False