State serialisers

State Serialisation Components for JUNE Checkpointing

This module contains specialised serialisers for different aspects of the simulation state. Each serialiser handles the conversion of live simulation objects into checkpoint-compatible data structures.

BaseStateSerialiser

Bases: ABC

Base class for all state serialisers

Source code in june/checkpointing/state_serialisers.py

class BaseStateSerialiser(ABC):
    """Base class for all state serialisers"""

    def __init__(self, simulator):
        self.simulator = simulator

    @abstractmethod
    def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
        """Serialise state to checkpoint-compatible format

        Args:
            checkpoint_type (str, optional): (Default value = "full")

        """
        pass

serialise(checkpoint_type='full') abstractmethod

Serialise state to checkpoint-compatible format

Parameters:

Name	Type	Description	Default
checkpoint_type	str	(Default value = "full")	'full'

Source code in june/checkpointing/state_serialisers.py

@abstractmethod
def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
    """Serialise state to checkpoint-compatible format

    Args:
        checkpoint_type (str, optional): (Default value = "full")

    """
    pass

InteractionStateSerialiser

Bases: BaseStateSerialiser

Serialises interaction transmission tracking state including initial infected IDs and all transmission statistics needed for recreating print_transmission_statistics.

Source code in june/checkpointing/state_serialisers.py

class InteractionStateSerialiser(BaseStateSerialiser):
    """Serialises interaction transmission tracking state including initial infected IDs
    and all transmission statistics needed for recreating print_transmission_statistics.

    """

    def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
        """Serialise interaction transmission tracking state.

        Args:
            checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") (Default value = "full")

        Returns:
            Dict[str, Any]: Serialised interaction state data

        """
        logger.debug("Serialising interaction transmission tracking state")

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

        # Ensure initial infected IDs are loaded
        interaction._load_initial_infected_ids()

        interaction_data = {
            # Core transmission tracking - convert numpy types
            'initial_infected_ids': [int(id) for id in interaction._initial_infected_ids],
            'initial_infected_transmission_counts': {
                int(person_id): int(count) for person_id, count in interaction.initial_infected_transmission_counts.items()
            },
            'previous_transmission_counts': {
                int(person_id): int(count) for person_id, count in interaction.previous_transmission_counts.items()
            },
            'timestep_count': int(interaction.timestep_count),

            # Configuration flags
            'initial_infected_ids_loaded': interaction._initial_infected_ids_loaded,

            # Debug tracking data (optional but useful) - convert numpy types
            'debug_cross_rank_infectors': {
                int(person_id): list(ranks) for person_id, ranks in interaction.debug_cross_rank_infectors.items()
            },
            'debug_timestep_infectors': {
                int(timestep): {int(person_id): int(count) for person_id, count in counts.items()}
                for timestep, counts in interaction.debug_timestep_infectors.items()
            },
            'debug_infector_venues': {
                int(person_id): venues for person_id, venues in interaction.debug_infector_venues.items()
            },

            # Interaction configuration (for validation)
            'alpha_physical': interaction.alpha_physical,
            'betas': dict(interaction.betas) if interaction.betas else {},
            'beta_reductions': dict(interaction.beta_reductions) if interaction.beta_reductions else {}
        }

        # Statistics summary - ensure all values are Python types
        stats_summary = {
            'total_initial_infected': int(len(interaction._initial_infected_ids)),
            'total_transmitting_initial_infected': int(len(interaction.initial_infected_transmission_counts)),
            'total_secondary_infections': int(sum(interaction.initial_infected_transmission_counts.values())),
            'current_timestep': int(interaction.timestep_count)
        }

        interaction_data['statistics_summary'] = stats_summary

        logger.debug(f"Serialised interaction state - {stats_summary['total_initial_infected']} initial infected, "
                    f"{stats_summary['total_secondary_infections']} secondary infections, timestep {stats_summary['current_timestep']}")

        return interaction_data

serialise(checkpoint_type='full')

Serialise interaction transmission tracking state.

Parameters:

Name	Type	Description	Default
checkpoint_type	str	Type of checkpoint ("full" or "delta") (Default value = "full")	'full'

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: Serialised interaction state data

Source code in june/checkpointing/state_serialisers.py

def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
    """Serialise interaction transmission tracking state.

    Args:
        checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") (Default value = "full")

    Returns:
        Dict[str, Any]: Serialised interaction state data

    """
    logger.debug("Serialising interaction transmission tracking state")

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

    # Ensure initial infected IDs are loaded
    interaction._load_initial_infected_ids()

    interaction_data = {
        # Core transmission tracking - convert numpy types
        'initial_infected_ids': [int(id) for id in interaction._initial_infected_ids],
        'initial_infected_transmission_counts': {
            int(person_id): int(count) for person_id, count in interaction.initial_infected_transmission_counts.items()
        },
        'previous_transmission_counts': {
            int(person_id): int(count) for person_id, count in interaction.previous_transmission_counts.items()
        },
        'timestep_count': int(interaction.timestep_count),

        # Configuration flags
        'initial_infected_ids_loaded': interaction._initial_infected_ids_loaded,

        # Debug tracking data (optional but useful) - convert numpy types
        'debug_cross_rank_infectors': {
            int(person_id): list(ranks) for person_id, ranks in interaction.debug_cross_rank_infectors.items()
        },
        'debug_timestep_infectors': {
            int(timestep): {int(person_id): int(count) for person_id, count in counts.items()}
            for timestep, counts in interaction.debug_timestep_infectors.items()
        },
        'debug_infector_venues': {
            int(person_id): venues for person_id, venues in interaction.debug_infector_venues.items()
        },

        # Interaction configuration (for validation)
        'alpha_physical': interaction.alpha_physical,
        'betas': dict(interaction.betas) if interaction.betas else {},
        'beta_reductions': dict(interaction.beta_reductions) if interaction.beta_reductions else {}
    }

    # Statistics summary - ensure all values are Python types
    stats_summary = {
        'total_initial_infected': int(len(interaction._initial_infected_ids)),
        'total_transmitting_initial_infected': int(len(interaction.initial_infected_transmission_counts)),
        'total_secondary_infections': int(sum(interaction.initial_infected_transmission_counts.values())),
        'current_timestep': int(interaction.timestep_count)
    }

    interaction_data['statistics_summary'] = stats_summary

    logger.debug(f"Serialised interaction state - {stats_summary['total_initial_infected']} initial infected, "
                f"{stats_summary['total_secondary_infections']} secondary infections, timestep {stats_summary['current_timestep']}")

    return interaction_data

PopulationHealthSerialiser

Bases: BaseStateSerialiser

Serialises complete population health state including infections, symptoms, and all health-related attributes.

Source code in june/checkpointing/state_serialisers.py

class PopulationHealthSerialiser(BaseStateSerialiser):
    """Serialises complete population health state including infections,
    symptoms, and all health-related attributes.

    """

    def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
        """Serialise population health state for ALL people.

        Args:
            checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") (Default value = "full")

        Returns:
            Dict[str, Any]: Serialised population health data

        """
        logger.debug("Serialising population health state for ALL people")

        population_data = {
            'total_people': len(self.simulator.world.people),
            'people_states': {},
            'immunities': {},  # Immunity for EVERY person
            'people_ids': []   # Complete list of all people IDs
        }

        infected_count = 0
        hospitalised_count = 0
        intensive_care_count = 0

        # Process EVERY person in the population
        for person in self.simulator.world.people:
            person_id = int(person.id)  # Ensure consistent type

            # Always add person ID to complete list
            population_data['people_ids'].append(person_id)

            # Always save immunity for every person (like original HDF5)
            immunity_obj = getattr(person, 'immunity', None)
            if immunity_obj is not None and hasattr(immunity_obj, 'susceptibility_dict'):
                # Serialise complete Immunity object
                immunity_data = {
                    'susceptibility_dict': immunity_obj.susceptibility_dict,
                    'effective_multiplier_dict': getattr(immunity_obj, 'effective_multiplier_dict', {})
                }
                population_data['immunities'][person_id] = immunity_data
            else:
                # Person has no immunity object or it's a simple value - create default
                population_data['immunities'][person_id] = {
                    'susceptibility_dict': {},
                    'effective_multiplier_dict': {}
                }

            # Get person's health state (always process, even if "default")
            person_state = self._serialise_person_health(person)

            # ALWAYS store every person's state (even if empty)
            population_data['people_states'][person_id] = person_state if person_state else {}

            # Count statistics
            if person_state and person_state.get('infected', False):
                infected_count += 1
            if person_state and person_state.get('hospitalised', False):
                hospitalised_count += 1
            if person_state and person_state.get('intensive_care', False):
                intensive_care_count += 1

        # Add cemetery state (dead people)
        cemetery_state = self._serialise_cemetery_state()
        population_data['cemetery_state'] = cemetery_state

        # Add summary statistics
        population_data['statistics'] = {
            'total_infected': infected_count,
            'total_hospitalised': hospitalised_count,
            'total_people_saved': len(population_data['people_states']),
            'total_immunities_saved': len(population_data['immunities']),
            'total_deaths': cemetery_state['total_deaths']
        }

        logger.debug(f"Serialised complete state for {len(population_data['people_states'])} people (ALL people in population)")
        logger.debug(f"Saved immunity data for {len(population_data['immunities'])} people")
        return population_data

    def _serialise_person_health(self, person) -> Optional[Dict[str, Any]]:
        """Serialise individual person's health state.

        Args:
            person (Person): The person whose health state to serialise

        Returns:
            Optional[Dict[str, Any]]: Person's health state, or None if in default state

        """
        person_state = {}

        # Basic health flags
        if person.infected:
            person_state['infected'] = True

        if person.hospitalised:
            person_state['hospitalised'] = True
            person_state['medical_facility'] = {
                'facility_id': person.medical_facility.group.id
            }

        if person.intensive_care:
            person_state['intensive_care'] = True
            person_state['medical_facility'] = {
                'facility_id': person.medical_facility.group.id
            }

        # Infection details
        if person.infection is not None:
            person_state['infection'] = {
                'infection_class': type(person.infection).__name__,  # Critical for proper reconstruction
                'infection_type': type(person.infection).__name__,   # Keep for compatibility
                'start_time': person.infection.start_time,
                'transmission_multiplier': getattr(person.infection, 'transmission_multiplier', 1.0)
            }

            # Serialise symptoms properly
            if person.infection.symptoms is not None:
                symptoms_data = self._serialise_symptoms(person.infection.symptoms)
                person_state['infection']['symptoms'] = symptoms_data

            # Serialise transmission properly  
            if person.infection.transmission is not None:
                transmission_data = self._serialise_transmission(person.infection.transmission)
                person_state['infection']['transmission'] = transmission_data

            # NOTE: infection.tag is a property that returns symptoms.tag
            # It's already saved in symptoms, so we don't need to save it separately



        # Return None if person is in completely default state
        return person_state if person_state else None

    def _serialise_symptoms(self, symptoms) -> Optional[Dict[str, Any]]:
        """Serialise symptoms object using HDF5 saver approach.

        This matches the HDF5 symptoms_saver.py pattern exactly.

        Args:
            symptoms (Symptoms): The symptoms object to serialise

        Returns:
            Optional[Dict[str, Any]]: Serialised symptoms data matching HDF5 format

        """
        if symptoms is None:
            return None

        symptoms_data = {}

        # Save basic attributes (matching HDF5 saver)
        if hasattr(symptoms, 'max_tag'):
            symptoms_data['max_tag'] = int(symptoms.max_tag) if symptoms.max_tag is not None else None

        if hasattr(symptoms, 'tag'):
            symptoms_data['tag'] = int(symptoms.tag) if symptoms.tag is not None else None

        if hasattr(symptoms, 'max_severity'):
            symptoms_data['max_severity'] = float(symptoms.max_severity) if symptoms.max_severity is not None else None

        if hasattr(symptoms, 'stage'):
            symptoms_data['stage'] = int(symptoms.stage) if symptoms.stage is not None else None

        if hasattr(symptoms, 'time_of_symptoms_onset'):
            symptoms_data['time_of_symptoms_onset'] = float(symptoms.time_of_symptoms_onset) if symptoms.time_of_symptoms_onset is not None else None

        # Save trajectory exactly like HDF5 saver (separate times and symptoms)
        if hasattr(symptoms, 'trajectory') and symptoms.trajectory is not None:
            trajectory_times = []
            trajectory_symptoms = []

            for time, symp in symptoms.trajectory:
                # Save exactly as they are - preserve all types
                trajectory_times.append(time)
                trajectory_symptoms.append(symp)

            if trajectory_times:  # Only add if not empty
                symptoms_data['trajectory_times'] = trajectory_times
                symptoms_data['trajectory_symptoms'] = trajectory_symptoms

        return symptoms_data if symptoms_data else None

    def _serialise_transmission(self, transmission) -> Optional[Dict[str, Any]]:
        """Serialise transmission object by extracting individual attributes.
        Based on the HDF5 transmission_saver.py approach.

        Args:
            transmission (Transmission): The transmission object to serialise

        Returns:
            Optional[Dict[str, Any]]: Serialised transmission data

        """
        if transmission is None:
            return None

        transmission_data = {
            'transmission_type': type(transmission).__name__
        }

        # Determine transmission type and extract appropriate attributes
        transmission_type = type(transmission).__name__

        if transmission_type == 'TransmissionXNExp':
            # From HDF5 saver: ["time_first_infectious", "norm_time", "n", "norm", "alpha"]
            for attr in ['time_first_infectious', 'norm_time', 'n', 'norm', 'alpha']:
                if hasattr(transmission, attr):
                    value = getattr(transmission, attr)
                    transmission_data[attr] = float(value) if value is not None else None

        elif transmission_type == 'TransmissionGamma':
            # From HDF5 saver: ["shape", "shift", "scale", "norm"]
            for attr in ['shape', 'shift', 'scale', 'norm']:
                if hasattr(transmission, attr):
                    value = getattr(transmission, attr)
                    transmission_data[attr] = float(value) if value is not None else None

        elif transmission_type == 'TransmissionConstant':
            # From HDF5 saver: ["probability"]
            if hasattr(transmission, 'probability'):
                value = getattr(transmission, 'probability')
                transmission_data['probability'] = float(value) if value is not None else None

        if hasattr(transmission, 'probability'):
            transmission_data['current_probability'] = float(transmission.probability)

        # Note: Removed convert_numpy_types to preserve exact numerical accuracy

        return transmission_data if len(transmission_data) > 1 else None  # >1 because we always have transmission_type

    def _serialise_cemetery_state(self) -> Dict[str, Any]:
        """Serialise cemetery state by finding all dead people directly from population.


        Returns:
            Dict[str, Any]: Cemetery state data with simple dead_id list

        """
        # Collect dead people directly from population to avoid cemetery tuple issues
        dead_ids = [person.id for person in self.simulator.world.people if person.dead]

        cemetery_data = {
            'total_deaths': len(dead_ids),
            'dead_id': [int(dead_id) for dead_id in dead_ids]  # Simple int conversion instead of convert_numpy_types
        }

        logger.debug(f"Serialised {cemetery_data['total_deaths']} dead people")
        return cemetery_data

serialise(checkpoint_type='full')

Serialise population health state for ALL people.

Parameters:

Name	Type	Description	Default
checkpoint_type	str	Type of checkpoint ("full" or "delta") (Default value = "full")	'full'

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: Serialised population health data

Source code in june/checkpointing/state_serialisers.py

def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
    """Serialise population health state for ALL people.

    Args:
        checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") (Default value = "full")

    Returns:
        Dict[str, Any]: Serialised population health data

    """
    logger.debug("Serialising population health state for ALL people")

    population_data = {
        'total_people': len(self.simulator.world.people),
        'people_states': {},
        'immunities': {},  # Immunity for EVERY person
        'people_ids': []   # Complete list of all people IDs
    }

    infected_count = 0
    hospitalised_count = 0
    intensive_care_count = 0

    # Process EVERY person in the population
    for person in self.simulator.world.people:
        person_id = int(person.id)  # Ensure consistent type

        # Always add person ID to complete list
        population_data['people_ids'].append(person_id)

        # Always save immunity for every person (like original HDF5)
        immunity_obj = getattr(person, 'immunity', None)
        if immunity_obj is not None and hasattr(immunity_obj, 'susceptibility_dict'):
            # Serialise complete Immunity object
            immunity_data = {
                'susceptibility_dict': immunity_obj.susceptibility_dict,
                'effective_multiplier_dict': getattr(immunity_obj, 'effective_multiplier_dict', {})
            }
            population_data['immunities'][person_id] = immunity_data
        else:
            # Person has no immunity object or it's a simple value - create default
            population_data['immunities'][person_id] = {
                'susceptibility_dict': {},
                'effective_multiplier_dict': {}
            }

        # Get person's health state (always process, even if "default")
        person_state = self._serialise_person_health(person)

        # ALWAYS store every person's state (even if empty)
        population_data['people_states'][person_id] = person_state if person_state else {}

        # Count statistics
        if person_state and person_state.get('infected', False):
            infected_count += 1
        if person_state and person_state.get('hospitalised', False):
            hospitalised_count += 1
        if person_state and person_state.get('intensive_care', False):
            intensive_care_count += 1

    # Add cemetery state (dead people)
    cemetery_state = self._serialise_cemetery_state()
    population_data['cemetery_state'] = cemetery_state

    # Add summary statistics
    population_data['statistics'] = {
        'total_infected': infected_count,
        'total_hospitalised': hospitalised_count,
        'total_people_saved': len(population_data['people_states']),
        'total_immunities_saved': len(population_data['immunities']),
        'total_deaths': cemetery_state['total_deaths']
    }

    logger.debug(f"Serialised complete state for {len(population_data['people_states'])} people (ALL people in population)")
    logger.debug(f"Saved immunity data for {len(population_data['immunities'])} people")
    return population_data

RatDynamicsSerialiser

Bases: BaseStateSerialiser

Serialises rat dynamics state including rat populations, disease states, and spatial mappings needed for zoonotic transmission.

Source code in june/checkpointing/state_serialisers.py

class RatDynamicsSerialiser(BaseStateSerialiser):
    """Serialises rat dynamics state including rat populations, disease states,
    and spatial mappings needed for zoonotic transmission.

    """

    def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
        """Serialise rat dynamics state from the RatManager.

        Args:
            checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") (Default value = "full")

        Returns:
            Dict[str, Any]: Serialised rat dynamics data

        """
        logger.debug("Serialising rat dynamics state")

        # Check if rat dynamics is enabled
        ratty_dynamics_enabled = getattr(self.simulator, 'ratty_dynamics_enabled', False)
        if not ratty_dynamics_enabled:
            logger.debug("Rat dynamics disabled - skipping serialisation")
            return {'enabled': False}

        # Check if rat manager exists
        rat_manager = getattr(self.simulator, 'rat_manager', None)
        if rat_manager is None:
            logger.warning("Rat manager not found - rat dynamics may not be properly initialised")
            return {'enabled': True, 'initialised': False}

        rat_data = {
            'enabled': True,
            'initialised': True,
            'rat_manager_config': {},
            'rat_population_state': {},
            'density_calculator_state': {},
            'disease_model_state': {},
            'spatial_grid_state': {},
            'area_mapper_state': {},
            'statistics': {}
        }

        # Serialise RatManager configuration
        rat_data['rat_manager_config'] = self._serialise_rat_manager_config(rat_manager)

        # Serialise rat population arrays
        rat_data['rat_population_state'] = self._serialise_rat_population_state(rat_manager)

        # Serialise density calculator state
        rat_data['density_calculator_state'] = self._serialise_density_calculator_state(rat_manager)

        # Serialise disease model state
        rat_data['disease_model_state'] = self._serialise_disease_model_state(rat_manager)

        # Serialise spatial grid configuration
        rat_data['spatial_grid_state'] = self._serialise_spatial_grid_state(rat_manager)

        # Serialise area mapping state
        rat_data['area_mapper_state'] = self._serialise_area_mapper_state(rat_manager)

        # Generate statistics
        rat_data['statistics'] = self._generate_rat_statistics(rat_manager)

        logger.debug(f"Serialised rat dynamics state - {rat_data['statistics']['total_rats']} rats, "
                    f"{rat_data['statistics']['infected_rats']} infected")

        return rat_data

    def _serialise_rat_manager_config(self, rat_manager) -> Dict[str, Any]:
        """Serialise RatManager configuration parameters with detailed debugging

        Args:
            rat_manager: 

        """
        print(f"=== RAT MANAGER CONFIG DEBUG - BEFORE CHECKPOINT ===")
        print(f"num_rats: {rat_manager.num_rats}")
        print(f"dt: {rat_manager.dt}")
        print(f"rat_to_human_factor: {rat_manager.rat_to_human_factor}")
        print(f"human_to_rat_factor: {rat_manager.human_to_rat_factor}")
        print(f"initial_infections: {rat_manager.initial_infections}")
        print(f"Grid bounds - minx: {rat_manager.minx}, miny: {rat_manager.miny}, maxx: {rat_manager.maxx}, maxy: {rat_manager.maxy}")
        print(f"Grid dimensions - rows: {rat_manager.sim_rows}, cols: {rat_manager.sim_cols}")

        config_data = {
            'num_rats': int(rat_manager.num_rats),
            'dt': float(rat_manager.dt),
            'rat_to_human_factor': float(rat_manager.rat_to_human_factor),
            'human_to_rat_factor': float(rat_manager.human_to_rat_factor),
            'initial_infections': int(rat_manager.initial_infections),

            # Grid dimensions and bounds
            'minx': float(rat_manager.minx) if rat_manager.minx is not None else None,
            'miny': float(rat_manager.miny) if rat_manager.miny is not None else None,
            'maxx': float(rat_manager.maxx) if rat_manager.maxx is not None else None,
            'maxy': float(rat_manager.maxy) if rat_manager.maxy is not None else None,
            'sim_rows': int(rat_manager.sim_rows) if rat_manager.sim_rows is not None else None,
            'sim_cols': int(rat_manager.sim_cols) if rat_manager.sim_cols is not None else None
        }

        return config_data

    def _serialise_rat_population_state(self, rat_manager) -> Dict[str, Any]:
        """Serialise rat population arrays and individual rat states with comprehensive debugging

        Args:
            rat_manager: 

        """
        if rat_manager.num_rats == 0:
            logger.debug("=== RAT POPULATION STATE DEBUG - NO RATS ===")
            return {'num_rats': 0}

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

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

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

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

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

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

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

        population_data = {
            'num_rats': int(rat_manager.num_rats),
            'positions': rat_manager.positions.tolist() if rat_manager.positions is not None else None,
            'states': rat_manager.states.tolist() if rat_manager.states is not None else None,
            'infection_age': rat_manager.infection_age.tolist() if rat_manager.infection_age is not None else None,
            'immunity': rat_manager.immunity.tolist() if rat_manager.immunity is not None else None,
            'personal_delta': rat_manager.personal_delta.tolist() if rat_manager.personal_delta is not None else None,
            'grid_indices': rat_manager.grid_indices.tolist() if rat_manager.grid_indices is not None else None
        }

        print(f"=== RAT POPULATION SERIALIZATION COMPLETE ===")
        return population_data

    def _serialise_density_calculator_state(self, rat_manager) -> Dict[str, Any]:
        """Serialise rat density calculator state

        Args:
            rat_manager: 

        """
        density_calc = rat_manager.density_calculator

        density_data = {
            'rat_ratio': int(density_calc.rat_ratio),
            'cell_size': float(density_calc.cell_size),
            'gaussian_sigma': float(density_calc.gaussian_sigma),
            'precomputed_density_path': str(density_calc.precomputed_density_path),
            'total_population': int(density_calc.total_population),
            'coordinate_system': str(density_calc.coordinate_system) if density_calc.coordinate_system is not None else None,

            # Grid bounds
            'minx': float(density_calc.minx) if density_calc.minx is not None else None,
            'miny': float(density_calc.miny) if density_calc.miny is not None else None,
            'maxx': float(density_calc.maxx) if density_calc.maxx is not None else None,
            'maxy': float(density_calc.maxy) if density_calc.maxy is not None else None,
            'sim_rows': int(density_calc.sim_rows) if density_calc.sim_rows is not None else None,
            'sim_cols': int(density_calc.sim_cols) if density_calc.sim_cols is not None else None,

            # Rat density array (this might be large, but necessary for exact restoration)
            'rat_density': density_calc.rat_density.tolist() if density_calc.rat_density is not None else None
        }

        return density_data

    def _serialise_disease_model_state(self, rat_manager) -> Dict[str, Any]:
        """Serialise rat disease model state with detailed debugging

        Args:
            rat_manager: 

        """
        disease_model = rat_manager.disease_model

        logger.debug(f"=== RAT DISEASE MODEL STATE DEBUG - BEFORE CHECKPOINT ===")
        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
        logger.debug(f"History lengths - infected: {len(disease_model.infected_history)}, immunity_08: {len(disease_model.immunity_08_history)}")
        logger.debug(f"History lengths - immunity_05: {len(disease_model.immunity_05_history)}, global_seeds: {len(disease_model.global_seeds_history)}")
        logger.debug(f"Total global seeds: {disease_model.total_global_seeds}")

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

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

        disease_data = {
            # Disease parameters
            'beta': float(disease_model.beta),
            'alpha': float(disease_model.alpha),
            'gamma': float(disease_model.gamma),
            'delta_mean': float(disease_model.delta_mean),
            'delta_std': float(disease_model.delta_std),
            'max_trans_distance': int(disease_model.max_trans_distance),
            'p_global_seed': float(disease_model.p_global_seed),
            'global_seed_min': int(disease_model.global_seed_min),
            'global_seed_max': int(disease_model.global_seed_max),
            'global_seed_immunity_threshold': float(disease_model.global_seed_immunity_threshold),
            'infectiousness_threshold': float(disease_model.infectiousness_threshold),

            # Disease history tracking
            'infected_history': [int(x) for x in disease_model.infected_history],
            'immunity_08_history': [int(x) for x in disease_model.immunity_08_history],
            'immunity_05_history': [int(x) for x in disease_model.immunity_05_history],
            'global_seeds_history': [int(x) for x in disease_model.global_seeds_history],
            'total_global_seeds': int(disease_model.total_global_seeds),

            # Transmission kernel (if built)
            'transmission_kernel': disease_model.transmission_kernel.tolist() if disease_model.transmission_kernel is not None else None
        }

        return disease_data

    def _serialise_spatial_grid_state(self, rat_manager) -> Dict[str, Any]:
        """Serialise spatial grid configuration with debugging

        Args:
            rat_manager: 

        """
        spatial_grid = rat_manager.spatial_grid

        logger.debug(f"=== RAT SPATIAL GRID STATE DEBUG - BEFORE CHECKPOINT ===")
        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}")
        logger.debug(f"Movement parameters - lambda_move: {spatial_grid.lambda_move}, Rmax_move: {spatial_grid.Rmax_move}")
        logger.debug(f"Grid cell size: {spatial_grid.grid_cell_size}")
        logger.debug(f"Grid cell dimensions - rows: {spatial_grid.grid_cell_rows}, cols: {spatial_grid.grid_cell_cols}")

        # Debug movement offsets
        if spatial_grid.move_offsets is not None:
            logger.debug(f"Move offsets shape: {spatial_grid.move_offsets.shape}")
            logger.debug(f"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 move offsets: {spatial_grid.move_offsets[:5].tolist()}")
            logger.debug(f"Move offset distance range: [{np.min(spatial_grid.move_offset_dists)}, {np.max(spatial_grid.move_offset_dists)}]")

        # Debug spatial grid array if available
        if hasattr(spatial_grid, 'spatial_grid_array') and spatial_grid.spatial_grid_array is not None:
            total_rats_in_grid = sum(len(cell) for row in spatial_grid.spatial_grid_array for cell in row)
            non_empty_cells = sum(1 for row in spatial_grid.spatial_grid_array for cell in row if len(cell) > 0)
            logger.debug(f"Spatial grid populated - total rats: {total_rats_in_grid}, non-empty cells: {non_empty_cells}")

        spatial_data = {
            'cell_size': float(spatial_grid.cell_size),
            'max_trans_distance': int(spatial_grid.max_trans_distance),
            'enable_movement': bool(spatial_grid.enable_movement),
            'p_move': float(spatial_grid.p_move),
            'lambda_move': float(spatial_grid.lambda_move),
            'Rmax_move': int(spatial_grid.Rmax_move),
            'grid_cell_size': float(spatial_grid.grid_cell_size),
            'grid_cell_rows': int(spatial_grid.grid_cell_rows) if spatial_grid.grid_cell_rows is not None else None,
            'grid_cell_cols': int(spatial_grid.grid_cell_cols) if spatial_grid.grid_cell_cols is not None else None,

            # Movement offsets (precomputed)
            'move_offsets': spatial_grid.move_offsets.tolist() if spatial_grid.move_offsets is not None else None,
            'move_offset_dists': spatial_grid.move_offset_dists.tolist() if spatial_grid.move_offset_dists is not None else None
        }

        return spatial_data

    def _serialise_area_mapper_state(self, rat_manager) -> Dict[str, Any]:
        """Serialise area mapper state including MSOA mappings

        Args:
            rat_manager: 

        """
        area_mapper = rat_manager.area_mapper

        mapper_data = {
            # Basic mapping dictionaries
            'area_to_msoa': dict(area_mapper.area_to_msoa) if area_mapper.area_to_msoa else {},
            'msoa_to_areas': {},  # Will be reconstructed from area_to_msoa
            'msoa_to_cells': {},
            '_msoa_mappings_initialised': bool(area_mapper._msoa_mappings_initialised)
        }

        # Serialise msoa_to_areas (convert Area objects to names)
        for msoa_id, areas_list in area_mapper.msoa_to_areas.items():
            if areas_list:
                mapper_data['msoa_to_areas'][msoa_id] = [area.name for area in areas_list if hasattr(area, 'name')]

        # Serialise msoa_to_cells
        for msoa_id, cells_list in area_mapper.msoa_to_cells.items():
            if cells_list:
                mapper_data['msoa_to_cells'][msoa_id] = cells_list  # List of (row, col) tuples

        return mapper_data

    def _generate_rat_statistics(self, rat_manager) -> Dict[str, Any]:
        """Generate summary statistics for the rat population

        Args:
            rat_manager: 

        """
        stats = {
            'total_rats': int(rat_manager.num_rats),
            'infected_rats': 0,
            'recovered_rats': 0,
            'susceptible_rats': 0,
            'high_immunity_rats': 0,
            'medium_immunity_rats': 0
        }

        if rat_manager.num_rats > 0 and rat_manager.states is not None:
            stats['infected_rats'] = int(np.sum(rat_manager.states == 1))
            stats['recovered_rats'] = int(np.sum(rat_manager.states == 2))
            stats['susceptible_rats'] = int(np.sum(rat_manager.states == 0))

            if rat_manager.immunity is not None:
                stats['high_immunity_rats'] = int(np.sum(rat_manager.immunity > 0.8))
                stats['medium_immunity_rats'] = int(np.sum(rat_manager.immunity > 0.5))

        return stats

serialise(checkpoint_type='full')

Serialise rat dynamics state from the RatManager.

Parameters:

Name	Type	Description	Default
checkpoint_type	str	Type of checkpoint ("full" or "delta") (Default value = "full")	'full'

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: Serialised rat dynamics data

Source code in june/checkpointing/state_serialisers.py

def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
    """Serialise rat dynamics state from the RatManager.

    Args:
        checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") (Default value = "full")

    Returns:
        Dict[str, Any]: Serialised rat dynamics data

    """
    logger.debug("Serialising rat dynamics state")

    # Check if rat dynamics is enabled
    ratty_dynamics_enabled = getattr(self.simulator, 'ratty_dynamics_enabled', False)
    if not ratty_dynamics_enabled:
        logger.debug("Rat dynamics disabled - skipping serialisation")
        return {'enabled': False}

    # Check if rat manager exists
    rat_manager = getattr(self.simulator, 'rat_manager', None)
    if rat_manager is None:
        logger.warning("Rat manager not found - rat dynamics may not be properly initialised")
        return {'enabled': True, 'initialised': False}

    rat_data = {
        'enabled': True,
        'initialised': True,
        'rat_manager_config': {},
        'rat_population_state': {},
        'density_calculator_state': {},
        'disease_model_state': {},
        'spatial_grid_state': {},
        'area_mapper_state': {},
        'statistics': {}
    }

    # Serialise RatManager configuration
    rat_data['rat_manager_config'] = self._serialise_rat_manager_config(rat_manager)

    # Serialise rat population arrays
    rat_data['rat_population_state'] = self._serialise_rat_population_state(rat_manager)

    # Serialise density calculator state
    rat_data['density_calculator_state'] = self._serialise_density_calculator_state(rat_manager)

    # Serialise disease model state
    rat_data['disease_model_state'] = self._serialise_disease_model_state(rat_manager)

    # Serialise spatial grid configuration
    rat_data['spatial_grid_state'] = self._serialise_spatial_grid_state(rat_manager)

    # Serialise area mapping state
    rat_data['area_mapper_state'] = self._serialise_area_mapper_state(rat_manager)

    # Generate statistics
    rat_data['statistics'] = self._generate_rat_statistics(rat_manager)

    logger.debug(f"Serialised rat dynamics state - {rat_data['statistics']['total_rats']} rats, "
                f"{rat_data['statistics']['infected_rats']} infected")

    return rat_data

SchoolIncidentSerialiser

Bases: BaseStateSerialiser

Serialises school incident tracking data for NotSendingKidsToSchool policy.

Saves student death/ICU counts and household avoidance decisions to preserve parental decision-making state across checkpoint resumes.

Source code in june/checkpointing/state_serialisers.py

class SchoolIncidentSerialiser(BaseStateSerialiser):
    """Serialises school incident tracking data for NotSendingKidsToSchool policy.

    Saves student death/ICU counts and household avoidance decisions to preserve
    parental decision-making state across checkpoint resumes.

    """

    def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
        """Serialise school incident tracking data.

        Args:
            checkpoint_type (str, optional): (Default value = "full")

        Returns:
            Dict[str, Any]: School incident and household decision data

        """
        logger.debug("Serialising school incident tracking data")

        school_data = {
            'local_schools': {},
            'global_school_incidents': {},
            'global_household_decisions': {}
        }

        # Serialize local school incident data
        if hasattr(self.simulator.world, 'schools'):
            for school in self.simulator.world.schools:
                school_id = str(school.id)  # Convert numpy.int64 to string for JSON compatibility
                school_data['local_schools'][school_id] = {
                    'student_deaths': int(getattr(school, 'student_deaths', 0)),
                    'student_icu_transfers': int(getattr(school, 'student_icu_transfers', 0)),
                    'households_avoiding_school': [int(hh_id) for hh_id in getattr(school, 'households_avoiding_school', set())],
                    'households_last_decision_at_death_count': {str(k): int(v) for k, v in getattr(school, 'households_last_decision_at_death_count', {}).items()},
                    'households_last_decision_at_icu_count': {str(k): int(v) for k, v in getattr(school, 'households_last_decision_at_icu_count', {}).items()}
                }

        # Serialize global registries for external schools
        if hasattr(self.simulator.world, 'global_school_incidents'):
            global_incidents = {}
            for school_id, incidents in self.simulator.world.global_school_incidents.items():
                global_incidents[str(school_id)] = {
                    'deaths': int(incidents.get('deaths', 0)),
                    'icu': int(incidents.get('icu', 0))
                }
            school_data['global_school_incidents'] = global_incidents

        if hasattr(self.simulator.world, 'global_household_decisions'):
            # Convert sets to lists and ensure proper types for JSON serialization
            global_decisions = {}
            for school_key, decisions in self.simulator.world.global_household_decisions.items():
                global_decisions[str(school_key)] = {
                    'avoiding_households': [int(hh_id) for hh_id in decisions.get('avoiding_households', set())],
                    'last_death_decision': {str(k): int(v) for k, v in decisions.get('last_death_decision', {}).items()},
                    'last_icu_decision': {str(k): int(v) for k, v in decisions.get('last_icu_decision', {}).items()}
                }
            school_data['global_household_decisions'] = global_decisions

        logger.debug(f"School incident data serialised: {len(school_data['local_schools'])} local schools, "
                    f"{len(school_data['global_school_incidents'])} global incidents, "
                    f"{len(school_data['global_household_decisions'])} global decisions")

        return school_data

serialise(checkpoint_type='full')

Serialise school incident tracking data.

Parameters:

Name	Type	Description	Default
checkpoint_type	str	(Default value = "full")	'full'

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: School incident and household decision data

Source code in june/checkpointing/state_serialisers.py

def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
    """Serialise school incident tracking data.

    Args:
        checkpoint_type (str, optional): (Default value = "full")

    Returns:
        Dict[str, Any]: School incident and household decision data

    """
    logger.debug("Serialising school incident tracking data")

    school_data = {
        'local_schools': {},
        'global_school_incidents': {},
        'global_household_decisions': {}
    }

    # Serialize local school incident data
    if hasattr(self.simulator.world, 'schools'):
        for school in self.simulator.world.schools:
            school_id = str(school.id)  # Convert numpy.int64 to string for JSON compatibility
            school_data['local_schools'][school_id] = {
                'student_deaths': int(getattr(school, 'student_deaths', 0)),
                'student_icu_transfers': int(getattr(school, 'student_icu_transfers', 0)),
                'households_avoiding_school': [int(hh_id) for hh_id in getattr(school, 'households_avoiding_school', set())],
                'households_last_decision_at_death_count': {str(k): int(v) for k, v in getattr(school, 'households_last_decision_at_death_count', {}).items()},
                'households_last_decision_at_icu_count': {str(k): int(v) for k, v in getattr(school, 'households_last_decision_at_icu_count', {}).items()}
            }

    # Serialize global registries for external schools
    if hasattr(self.simulator.world, 'global_school_incidents'):
        global_incidents = {}
        for school_id, incidents in self.simulator.world.global_school_incidents.items():
            global_incidents[str(school_id)] = {
                'deaths': int(incidents.get('deaths', 0)),
                'icu': int(incidents.get('icu', 0))
            }
        school_data['global_school_incidents'] = global_incidents

    if hasattr(self.simulator.world, 'global_household_decisions'):
        # Convert sets to lists and ensure proper types for JSON serialization
        global_decisions = {}
        for school_key, decisions in self.simulator.world.global_household_decisions.items():
            global_decisions[str(school_key)] = {
                'avoiding_households': [int(hh_id) for hh_id in decisions.get('avoiding_households', set())],
                'last_death_decision': {str(k): int(v) for k, v in decisions.get('last_death_decision', {}).items()},
                'last_icu_decision': {str(k): int(v) for k, v in decisions.get('last_icu_decision', {}).items()}
            }
        school_data['global_household_decisions'] = global_decisions

    logger.debug(f"School incident data serialised: {len(school_data['local_schools'])} local schools, "
                f"{len(school_data['global_school_incidents'])} global incidents, "
                f"{len(school_data['global_household_decisions'])} global decisions")

    return school_data

TTEventRecorderSerialiser

Bases: BaseStateSerialiser

Serialises Test and Trace event recorder state including daily/cumulative counters, unique IDs, current status, and event buffer for exact restoration.

Source code in june/checkpointing/state_serialisers.py

class TTEventRecorderSerialiser(BaseStateSerialiser):
    """Serialises Test and Trace event recorder state including daily/cumulative counters,
    unique IDs, current status, and event buffer for exact restoration.

    """

    def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
        """Serialise TTEventRecorder state including all counters and data.

        Args:
            checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") - currently only full checkpoints supported (Default value = "full")

        Returns:
            Dict[str, Any]: Serialised TTEventRecorder data

        """
        logger.debug("Serialising TTEventRecorder state")

        # Get the TTEventRecorder from GlobalContext
        from june.global_context import GlobalContext
        tt_recorder = GlobalContext.get_tt_event_recorder()

        if tt_recorder is None:
            logger.debug("No TTEventRecorder found - test and trace may be disabled")
            return {'enabled': False}

        # NOTE: Do NOT process pending events during serialization to avoid changing state
        # The buffer will be serialised as-is and processed after restoration

        tt_data = {
            'enabled': True,
            'filename': str(tt_recorder.filename),
            'output_dir': str(tt_recorder.output_dir),

            # Core counters and statistics
            'total_counters': dict(tt_recorder.total_counters),
            'daily_counters': self._serialise_daily_counters(tt_recorder.daily_counters),
            'unique_ids': self._serialise_unique_ids(tt_recorder.unique_ids),
            'currently': self._serialise_current_status(tt_recorder.currently),
            'deltas': dict(tt_recorder.deltas),

            # State tracking
            'last_delta_reset': int(tt_recorder.last_delta_reset),
            'buffer_size': int(tt_recorder._buffer_size),

            # Event buffer (should be empty after processing, but save for safety)
            'event_buffer': self._serialise_event_buffer(tt_recorder._event_buffer)
        }

        logger.debug(f"Serialised TTEventRecorder - {len(tt_data['daily_counters'])} days of data, "
                    f"{tt_data['total_counters']['tested']} total tests, "
                    f"{len(tt_data['unique_ids']['tested'])} unique tested people")

        return tt_data

    def _serialise_daily_counters(self, daily_counters) -> Dict[str, Dict[str, int]]:
        """Serialise daily counters defaultdict to regular dict.

        Args:
            daily_counters (defaultdict): Daily counters from TTEventRecorder

        Returns:
            Dict[str, Dict[str, int]]: Serialised daily counters

        """
        serialised = {}
        for day, counters in daily_counters.items():
            serialised[str(day)] = dict(counters)  # Convert to regular dict
        return serialised

    def _serialise_unique_ids(self, unique_ids) -> Dict[str, list]:
        """Serialise unique IDs sets to lists.

        Args:
            unique_ids (Dict[str, set]): Unique IDs from TTEventRecorder

        Returns:
            Dict[str, list]: Serialised unique IDs as lists

        """
        serialised = {}
        for event_type, id_set in unique_ids.items():
            serialised[event_type] = list(id_set)  # Convert sets to lists
        return serialised

    def _serialise_current_status(self, currently) -> Dict[str, list]:
        """Serialise current status sets to lists.

        Args:
            currently (Dict[str, set]): Current status from TTEventRecorder

        Returns:
            Dict[str, list]: Serialised current status as lists

        """
        serialised = {}
        for status_type, id_set in currently.items():
            serialised[status_type] = list(id_set)  # Convert sets to lists
        return serialised

    def _serialise_event_buffer(self, event_buffer) -> list:
        """Serialise event buffer to list of dictionaries.

        Args:
            event_buffer (list): Event buffer from TTEventRecorder

        Returns:
            list: Serialised event buffer

        """
        serialised_events = []
        for event in event_buffer:
            event_data = {
                'event_type': event.event_type,
                'person_id': int(event.person_id),
                'timestamp': float(event.timestamp),
                'metadata': dict(event.metadata) if event.metadata else {},
                'creation_time': event.creation_time.isoformat()
            }
            serialised_events.append(event_data)
        return serialised_events

serialise(checkpoint_type='full')

Serialise TTEventRecorder state including all counters and data.

Parameters:

Name	Type	Description	Default
checkpoint_type	str	Type of checkpoint ("full" or "delta") - currently only full checkpoints supported (Default value = "full")	'full'

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: Serialised TTEventRecorder data

Source code in june/checkpointing/state_serialisers.py

def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
    """Serialise TTEventRecorder state including all counters and data.

    Args:
        checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") - currently only full checkpoints supported (Default value = "full")

    Returns:
        Dict[str, Any]: Serialised TTEventRecorder data

    """
    logger.debug("Serialising TTEventRecorder state")

    # Get the TTEventRecorder from GlobalContext
    from june.global_context import GlobalContext
    tt_recorder = GlobalContext.get_tt_event_recorder()

    if tt_recorder is None:
        logger.debug("No TTEventRecorder found - test and trace may be disabled")
        return {'enabled': False}

    # NOTE: Do NOT process pending events during serialization to avoid changing state
    # The buffer will be serialised as-is and processed after restoration

    tt_data = {
        'enabled': True,
        'filename': str(tt_recorder.filename),
        'output_dir': str(tt_recorder.output_dir),

        # Core counters and statistics
        'total_counters': dict(tt_recorder.total_counters),
        'daily_counters': self._serialise_daily_counters(tt_recorder.daily_counters),
        'unique_ids': self._serialise_unique_ids(tt_recorder.unique_ids),
        'currently': self._serialise_current_status(tt_recorder.currently),
        'deltas': dict(tt_recorder.deltas),

        # State tracking
        'last_delta_reset': int(tt_recorder.last_delta_reset),
        'buffer_size': int(tt_recorder._buffer_size),

        # Event buffer (should be empty after processing, but save for safety)
        'event_buffer': self._serialise_event_buffer(tt_recorder._event_buffer)
    }

    logger.debug(f"Serialised TTEventRecorder - {len(tt_data['daily_counters'])} days of data, "
                f"{tt_data['total_counters']['tested']} total tests, "
                f"{len(tt_data['unique_ids']['tested'])} unique tested people")

    return tt_data

TestAndTraceSerialiser

Bases: BaseStateSerialiser

Serialises test and trace system state including individual TestAndTrace objects, contact manager state, and policy configurations.

Source code in june/checkpointing/state_serialisers.py

class TestAndTraceSerialiser(BaseStateSerialiser):
    """Serialises test and trace system state including individual TestAndTrace objects,
    contact manager state, and policy configurations.

    """

    def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
        """Serialise test and trace system state.

        Args:
            checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") (Default value = "full")

        Returns:
            Dict[str, Any]: Serialised test and trace data

        """
        logger.debug("Serialising test and trace system state")

        # Check if test and trace is enabled
        test_and_trace_enabled = getattr(self.simulator, 'test_and_trace_enabled', False)
        if not test_and_trace_enabled:
            logger.debug("Test and trace disabled - skipping serialisation")
            return {'enabled': False}

        test_trace_data = {
            'enabled': True,
            'person_test_trace_states': {},
            'contact_manager_state': {},
            'policy_configurations': {},
            'statistics': {}
        }

        # Serialise individual TestAndTrace objects for all people
        person_states, stats = self._serialise_person_test_trace_states()
        test_trace_data['person_test_trace_states'] = person_states
        test_trace_data['statistics'].update(stats)

        # Serialise contact manager state
        if hasattr(self.simulator, 'contact_manager'):
            contact_manager_state = self._serialise_contact_manager_state()
            test_trace_data['contact_manager_state'] = contact_manager_state

        # Serialise policy configurations
        policy_configs = self._serialise_policy_configurations()
        test_trace_data['policy_configurations'] = policy_configs

        logger.debug(f"Serialised test and trace state - {len(test_trace_data['person_test_trace_states'])} people with T&T data")
        return test_trace_data

    def _serialise_person_test_trace_states(self) -> tuple:
        """Serialise TestAndTrace objects for all people who have them.


        Returns:
            tuple: (person_states_dict, statistics_dict)

        """
        person_states = {}
        stats = {
            'total_with_test_trace': 0,
            'pending_tests': 0,
            'completed_tests': 0,
            'positive_results': 0,
            'negative_results': 0,
            'notified_contacts': 0,
            'isolated_people': 0
        }

        for person in self.simulator.world.people:
            if person.test_and_trace is not None:
                person_id = int(person.id)
                tt_data = self._serialise_single_test_trace(person.test_and_trace)
                person_states[person_id] = tt_data
                stats['total_with_test_trace'] += 1

                # Update statistics
                if tt_data.get('pending_test_result') is not None:
                    stats['pending_tests'] += 1
                if tt_data.get('test_result') is not None:
                    stats['completed_tests'] += 1
                    if tt_data['test_result'] == 'Positive':
                        stats['positive_results'] += 1
                    elif tt_data['test_result'] == 'Negative':
                        stats['negative_results'] += 1
                if tt_data.get('notification_time') is not None:
                    stats['notified_contacts'] += 1
                if tt_data.get('isolation_start_time') is not None:
                    stats['isolated_people'] += 1

        return person_states, stats

    def _serialise_single_test_trace(self, test_and_trace) -> Dict[str, Any]:
        """Serialise a single TestAndTrace object.

        Args:
            test_and_trace (TestAndTrace): The TestAndTrace object to serialise

        Returns:
            Dict[str, Any]: Serialised TestAndTrace data

        """
        tt_data = {}

        # Contact tracing information
        if test_and_trace.notification_time is not None:
            tt_data['notification_time'] = float(test_and_trace.notification_time)

        if test_and_trace.scheduled_test_time is not None:
            tt_data['scheduled_test_time'] = float(test_and_trace.scheduled_test_time)

        tt_data['contacts_traced'] = bool(test_and_trace.contacts_traced)

        # Testing information
        if test_and_trace.time_of_testing is not None:
            tt_data['time_of_testing'] = float(test_and_trace.time_of_testing)

        if test_and_trace.time_of_result is not None:
            tt_data['time_of_result'] = float(test_and_trace.time_of_result)

        if test_and_trace.pending_test_result is not None:
            tt_data['pending_test_result'] = str(test_and_trace.pending_test_result)

        if test_and_trace.test_result is not None:
            tt_data['test_result'] = str(test_and_trace.test_result)

        # Event flags
        if test_and_trace.emited_quarantine_start_event is not None:
            tt_data['emited_quarantine_start_event'] = test_and_trace.emited_quarantine_start_event

        if test_and_trace.emited_quarantine_end_event is not None:
            tt_data['emited_quarantine_end_event'] = test_and_trace.emited_quarantine_end_event

        # Isolation information
        if test_and_trace.isolation_start_time is not None:
            tt_data['isolation_start_time'] = float(test_and_trace.isolation_start_time)

        if test_and_trace.isolation_end_time is not None:
            tt_data['isolation_end_time'] = float(test_and_trace.isolation_end_time)

        # Contact tracing source information
        if test_and_trace.tracer_id is not None:
            tt_data['tracer_id'] = int(test_and_trace.tracer_id)

        if test_and_trace.contact_reason is not None:
            tt_data['contact_reason'] = str(test_and_trace.contact_reason)

        return tt_data

    def _serialise_contact_manager_state(self) -> Dict[str, Any]:
        """Serialise contact manager state including leisure companions and pending tests.


        Returns:
            Dict[str, Any]: Serialised contact manager data

        """
        contact_manager = self.simulator.contact_manager

        cm_data = {
            'leisure_companions': {},
            'tests_ids_pending': [],
            'last_cleanup': 0,
            'statistics': {}
        }

        # Serialise leisure companions with careful type conversion
        for person_id, companions in contact_manager.leisure_companions.items():
            person_id_str = str(person_id)
            cm_data['leisure_companions'][person_id_str] = {}

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

        # Serialise pending tests
        cm_data['tests_ids_pending'] = []
        for test_info in contact_manager.tests_ids_pending:
            serialised_test = {}
            for key, value in test_info.items():
                if isinstance(value, (int, float, str, bool)):
                    serialised_test[key] = value
                else:
                    serialised_test[key] = str(value)
            cm_data['tests_ids_pending'].append(serialised_test)

        # Store cleanup timestamp
        cm_data['last_cleanup'] = float(contact_manager.last_cleanup)

        # Generate statistics
        cm_data['statistics'] = {
            'total_people_with_leisure_companions': len(contact_manager.leisure_companions),
            'total_leisure_connections': sum(len(companions) for companions in contact_manager.leisure_companions.values()),
            'pending_tests_count': len(contact_manager.tests_ids_pending)
        }

        logger.debug(f"Serialised contact manager - {cm_data['statistics']['total_people_with_leisure_companions']} people with leisure companions, "
                    f"{cm_data['statistics']['pending_tests_count']} pending tests")

        return cm_data

    def _serialise_policy_configurations(self) -> Dict[str, Any]:
        """Serialise policy configurations for test and trace system.


        Returns:
            Dict[str, Any]: Serialised policy configuration data

        """
        policy_configs = {}

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

        # Serialise testing policy configuration
        testing_config = disease_config.policy_manager.get_policy_data("testing")
        policy_configs['testing'] = dict(testing_config) if testing_config is not None else {}

        # Serialise tracing policy configuration
        tracing_config = disease_config.policy_manager.get_policy_data("tracing")
        policy_configs['tracing'] = dict(tracing_config) if tracing_config is not None else {}

        # Serialise other relevant policy configurations
        for policy_name in ['quarantine', 'isolation', 'hospitalisation']:
            policy_data = disease_config.policy_manager.get_policy_data(policy_name)
            policy_configs[policy_name] = dict(policy_data) if policy_data is not None else {}

        return policy_configs

serialise(checkpoint_type='full')

Serialise test and trace system state.

Parameters:

Name	Type	Description	Default
checkpoint_type	str	Type of checkpoint ("full" or "delta") (Default value = "full")	'full'

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: Serialised test and trace data

Source code in june/checkpointing/state_serialisers.py

def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
    """Serialise test and trace system state.

    Args:
        checkpoint_type (str, optional): Type of checkpoint ("full" or "delta") (Default value = "full")

    Returns:
        Dict[str, Any]: Serialised test and trace data

    """
    logger.debug("Serialising test and trace system state")

    # Check if test and trace is enabled
    test_and_trace_enabled = getattr(self.simulator, 'test_and_trace_enabled', False)
    if not test_and_trace_enabled:
        logger.debug("Test and trace disabled - skipping serialisation")
        return {'enabled': False}

    test_trace_data = {
        'enabled': True,
        'person_test_trace_states': {},
        'contact_manager_state': {},
        'policy_configurations': {},
        'statistics': {}
    }

    # Serialise individual TestAndTrace objects for all people
    person_states, stats = self._serialise_person_test_trace_states()
    test_trace_data['person_test_trace_states'] = person_states
    test_trace_data['statistics'].update(stats)

    # Serialise contact manager state
    if hasattr(self.simulator, 'contact_manager'):
        contact_manager_state = self._serialise_contact_manager_state()
        test_trace_data['contact_manager_state'] = contact_manager_state

    # Serialise policy configurations
    policy_configs = self._serialise_policy_configurations()
    test_trace_data['policy_configurations'] = policy_configs

    logger.debug(f"Serialised test and trace state - {len(test_trace_data['person_test_trace_states'])} people with T&T data")
    return test_trace_data

TimerStateSerialiser

Bases: BaseStateSerialiser

Serialises simulation timer state including current time, date progression, and temporal scheduling information.

Source code in june/checkpointing/state_serialisers.py

class TimerStateSerialiser(BaseStateSerialiser):
    """Serialises simulation timer state including current time,
    date progression, and temporal scheduling information.

    """

    def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
        """Serialise timer state with explicit type preservation.

        Args:
            checkpoint_type (str, optional): (Default value = "full")

        Returns:
            Dict[str, Any]: Serialised timer data

        """
        logger.debug("Serialising timer state")

        timer = self.simulator.timer

        timer_data = {
            'current_time': float(timer.now),
            'current_date': timer.date.isoformat(),
            'initial_date': timer.initial_date.isoformat(),
            'total_days': float(timer.total_days),  # Explicitly ensure numeric type
            'time_step_size': float(getattr(timer, 'time_step_size', 1.0))
        }

        # Add any other timer attributes that might be relevant
        for attr in ['activities_start_time', 'activities_end_time']:
            if hasattr(timer, attr):
                value = getattr(timer, attr)
                if value is not None:
                    timer_data[attr] = value

        logger.debug(f"Serialised timer state at simulation time {timer.now}")
        return timer_data

serialise(checkpoint_type='full')

Serialise timer state with explicit type preservation.

Parameters:

Name	Type	Description	Default
checkpoint_type	str	(Default value = "full")	'full'

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: Serialised timer data

Source code in june/checkpointing/state_serialisers.py

def serialise(self, checkpoint_type: str = "full") -> Dict[str, Any]:
    """Serialise timer state with explicit type preservation.

    Args:
        checkpoint_type (str, optional): (Default value = "full")

    Returns:
        Dict[str, Any]: Serialised timer data

    """
    logger.debug("Serialising timer state")

    timer = self.simulator.timer

    timer_data = {
        'current_time': float(timer.now),
        'current_date': timer.date.isoformat(),
        'initial_date': timer.initial_date.isoformat(),
        'total_days': float(timer.total_days),  # Explicitly ensure numeric type
        'time_step_size': float(getattr(timer, 'time_step_size', 1.0))
    }

    # Add any other timer attributes that might be relevant
    for attr in ['activities_start_time', 'activities_end_time']:
        if hasattr(timer, attr):
            value = getattr(timer, attr)
            if value is not None:
                timer_data[attr] = value

    logger.debug(f"Serialised timer state at simulation time {timer.now}")
    return timer_data

convert_numpy_types(obj)

Convert numpy types to Python native types for HDF5 compatibility.

Parameters:

Name	Type	Description	Default
obj	Any	Object that might contain numpy types	required

Returns:

Name	Type	Description
Any		Object with numpy types converted to Python native types

Source code in june/checkpointing/state_serialisers.py

def convert_numpy_types(obj):
    """Convert numpy types to Python native types for HDF5 compatibility.

    Args:
        obj (Any): Object that might contain numpy types

    Returns:
        Any: Object with numpy types converted to Python native types

    """
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {convert_numpy_types(k): convert_numpy_types(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        # Handle lists more carefully to avoid shape issues
        converted_list = []
        for item in obj:
            converted_list.append(convert_numpy_types(item))
        return converted_list
    else:
        return obj