Creating Custom Compartmental Models: Building an SIR Model

This tutorial demonstrates how to create custom compartmental models in simcronomicon by implementing the classic SIR (Susceptible-Infectious-Recovered) model from scratch. You’ll learn the essential components needed to build any compartmental model and understand the inheritance structure that makes simcronomicon extensible.

Understanding the SIR Model

The SIR model is the simplest epidemic model, dividing the population into three compartments:

S - Susceptible: Healthy individuals who can become infected.
I - Infectious: Individuals who can transmit the disease to susceptible people.
R - Recovered: Individuals who have recovered and gained immunity.

Disease Flow: S → I → R (unidirectional)

Why Build SIR When SEIR Exists?

While simcronomicon includes a more complex SEIR model, building SIR teaches you:

Model Architecture: Understanding the inheritance hierarchy
Parameter Design: How to define model-specific parameters
Agent Behavior: Implementing custom disease progression logic
Validation: Ensuring your model works correctly

This knowledge enables you to create any compartmental model (SIRD, SEIQR, etc.) for your specific research needs.

Step 1: Define Model Parameters

Every compartmental model needs a parameters class that inherits from AbstractModelParameters:

# File: sir_model.py
from simcronomicon.compartmental_models.abstract_model import AbstractModelParameters

class SIRModelParameters(AbstractModelParameters):
    """
    Parameters for the SIR compartmental model.

    Parameters
    ----------
    max_energy : int
        Maximum daily energy for agents (affects movement patterns)
    beta : float
        Transmission rate - probability of infection per contact
    gamma : float
        Recovery rate - days until recovery from infectious status
    """

    def __init__(self, max_energy, beta, gamma):
        super().__init__(max_energy)

        # Validate parameters
        assert isinstance(beta, (int, float)) and 0 <= beta <= 1, \
            "beta must be a number between 0 and 1"
        assert isinstance(gamma, (int, float)) and gamma > 0, \
            "gamma must be a positive number"

        self.beta = beta
        self.gamma = gamma

    def to_metadata_dict(self):
        """Convert parameters to dictionary for saving/loading simulations."""
        return {
            'max_energy': self.max_energy,
            'beta': self.beta,
            'gamma': self.gamma,
            'model_type': 'SIR'
        }

Key Design Principles:

Validation: Always validate parameter types and ranges
Documentation: Clear docstrings explain each parameter’s meaning
Metadata: The to_metadata_dict() method enables simulation persistence

Step 2: Create the Agent Class

The agent class defines how individuals behave and transition between disease states:

from simcronomicon.compartmental_models.abstract_model import AbstractFolk

class FolkSIR(AbstractFolk):
    """
    SIR agent with disease-specific behaviors.

    Inherits all basic agent functionality from AbstractFolk and adds
    SIR-specific interaction and status transition logic.
    """

    def interact(self, folks_here, current_place_type, status_dict_t, model_params, dice):
        """
        Handle agent interactions at a location, including disease transmission.

        Parameters
        ----------
        folks_here : list of FolkSIR
            All agents at this location
        current_place_type : str
            Type of location (e.g., 'commercial', 'accommodation')
        status_dict_t : dict
            Current status counts for the simulation
        model_params : SIRModelParameters
            Model configuration parameters
        dice : float
            Random number between 0 and 1 for stochastic decisions
        """
        # Only susceptible agents can become infected
        if self.status == 'S':
            # Check if any infectious people are present
            num_contact = len(
             [folk for folk in folks_here if folk != self and folk.status in ['I']])

             # conversion_prob * I / N is the non-linear term that defines conversion
             # This inverse bernoulli function is an interpretation of the term
             # in agent-based modeling
             transmission_prob = self.inverse_bernoulli(
                    num_contact, conversion_prob / len(folks_here))

             # Attempt infection
             if transmission_prob > dice:
                    self.convert('I', status_dict_t)

        # Reduce energy from social interaction
        if self.energy > 0:
            self.energy -= 1

    def sleep(self, status_dict_t, model_params):
        """
        End-of-day processing: energy reset and disease progression.

        Parameters
        ----------
        status_dict_t : dict
            Current status counts for the simulation
        model_params : SIRModelParameters
            Model configuration parameters
        """
        # Call parent sleep method (resets energy, increments status_step_streak)
        super().sleep()

        # Handle disease progression based on time in status
        if self.status == 'I' and self.status_step_streak >= model_params.gamma:
            # Infectious agents recover after gamma days
            self.convert('R', status_dict_t)

Key Agent Concepts:

Interaction Logic: Only susceptible agents can become infected
Inverse Bernoulli: Handles multiple infectious contacts realistically
Time-Based Transitions: Status changes based on status_step_streak
Energy Management: Social interactions consume energy

Step 3: Implement the SIR Model

Custom models need to override the initialize_sim_population method to handle their specific status assignments and initial conditions:

class SIRModel(AbstractCompartmentalModel):
    """SIR compartmental model implementation."""

    def __init__(self, model_params, step_events=None):
        # Define model-specific attributes BEFORE calling super().__init__
        self.infected_statuses = ['I']  # Only infectious status
        self.all_statuses = ['S', 'I', 'R']  # All possible statuses
        self.folk_class = FolkSIR  # Agent class to use
        self.step_events = step_events  # Custom or default events

        # Initialize parent class (validates our definitions)
        super().__init__(model_params)

    def initialize_sim_population(self, town):
        import random as rd

        # Get basic population parameters from parent
        num_pop, num_init_spreader, num_init_spreader_rd, folks, household_node_indices, assignments = super().initialize_sim_population(town)

        # Randomly assign initial spreaders (not on specified nodes)
        for i in range(num_init_spreader_rd):
            node = rd.choice(town.accommodation_node_ids)
            assignments.append((node, 'I'))  # Infectious status

        # Assign the rest as susceptible
        for i in range(num_pop - num_init_spreader):
            node = rd.choice(town.accommodation_node_ids)
            assignments.append((node, 'S'))  # Susceptible status

        # Assign initial spreaders to specified nodes (if any)
        for node in town.town_params.spreader_initial_nodes:
            assignments.append((node, 'I'))

        # Create folks and update graph/node info
        for i, (node, status) in enumerate(assignments):
            folk = self.create_folk(i, node, self.model_params.max_energy, status)
            folks.append(folk)
            town.town_graph.nodes[node]["folks"].append(folk)

            # Track household nodes (nodes with 2+ people)
            if len(town.town_graph.nodes[node]["folks"]) == 2:
                household_node_indices.add(node)

        # Create initial status dictionary for timestep 0
        status_dict_t0 = {
            'current_event': None,
            'timestep': 0,
            'S': num_pop - num_init_spreader,  # Susceptible count
            'I': num_init_spreader,            # Infectious count
            'R': 0                             # Recovered count (starts at 0)
        }

        return folks, household_node_indices, status_dict_t0

Key Implementation Points:

Call Parent Method: Use super().initialize_sim_population(town) to get base setup
Status Assignment: Assign initial statuses based on your model’s compartments
Node Assignment: Distribute agents across accommodation nodes in the town
Status Dictionary: Initialize counts for all possible statuses at timestep 0
Household Tracking: Update household indices for nodes with multiple residents

Step 4: Test Your SIR Model

Create a simple test script to validate your implementation:

# File: test_sir_model.py
import simcronomicon as scon
from sir_model import SIRModel, SIRModelParameters

def test_sir_model():
    """Test basic SIR model functionality."""

    # Create a simple town
    point = (50.7753, 6.0839)  # Aachen coordinates
    town_params = scon.TownParameters(100, 5)  # 100 people, 5 initial infected
    town = scon.Town.from_point(point, 500, "test_sir", town_params)

    # Configure SIR model
    sir_params = SIRModelParameters(
        max_energy=5,
        beta=0.3,      # 30% transmission probability
        gamma=7        # Recovery after 7 days
    )

    # Create model and simulation
    sir_model = SIRModel(sir_params)
    simulation = scon.Simulation(town, sir_model, timesteps=50)

    # Run simulation
    print("Running SIR simulation...")
    simulation.run()

    # Analyze results
    print("SIR simulation completed successfully!")

    # Visualize results
    scon.plot_status_summary_from_hdf5("simulation_output.h5")

    return True

if __name__ == "__main__":
    test_sir_model()

Step 5: Compare with Mathematical SIR

Validate your agent-based model against the classic SIR differential equations:

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

def compare_sir_models(sir_params, population_size=100):
    """Compare agent-based SIR with ODE solution."""

    # ODE system for SIR
    def sir_ode(t, y):
        S, I, R = y
        N = S + I + R
        dS_dt = -sir_params.beta * S * I / N
        dI_dt = sir_params.beta * S * I / N - I / sir_params.gamma
        dR_dt = I / sir_params.gamma
        return [dS_dt, dI_dt, dR_dt]

    # Initial conditions (normalized)
    S0 = 0.95  # 95% susceptible
    I0 = 0.05  # 5% infectious
    R0 = 0.00  # 0% recovered

    # Solve ODE
    t_span = (0, 50)
    t_eval = np.arange(0, 51)
    solution = solve_ivp(sir_ode, t_span, [S0, I0, R0], t_eval=t_eval)

    # Plot comparison
    plt.figure(figsize=(12, 5))

    # ODE solution
    plt.subplot(1, 2, 1)
    plt.plot(solution.t, solution.y[0], 'b-', label='Susceptible')
    plt.plot(solution.t, solution.y[1], 'r-', label='Infectious')
    plt.plot(solution.t, solution.y[2], 'g-', label='Recovered')
    plt.xlabel('Time (days)')
    plt.ylabel('Proportion')
    plt.title('SIR ODE Model')
    plt.legend()
    plt.grid(True)

    # Agent-based results (load from your simulation)
    plt.subplot(1, 2, 2)
    # Add code to plot your simulation results here
    plt.title('Agent-Based SIR Model')

    plt.tight_layout()
    plt.show()

    print("📊 Model comparison complete!")
    print(f"   Beta (transmission rate): {sir_params.beta}")
    print(f"   Gamma (recovery time): {sir_params.gamma} days")
    print(f"   Basic reproduction number R₀: {sir_params.beta * sir_params.gamma:.2f}")

Advanced Features

Once your basic SIR model works, you can extend it with advanced features:

Custom Movement Patterns

# Add custom step events for different behaviors
sir_step_events = [
    scon.StepEvent(
        "morning_commute",
        FolkSIR.interact,
        scon.EventType.DISPERSE,
        10000,  # Travel up to 10km
        ['workplace', 'education'],
        scon.log_normal_mobility  # Distance-based probability
    ),
    scon.StepEvent(
        "evening_social",
        FolkSIR.interact,
        scon.EventType.DISPERSE,
        5000,
        ['commercial', 'religious']
    )
]

sir_model = SIRModel(sir_params, step_events=sir_step_events)

Behavioral Interventions

class FolkSIRWithMasks(FolkSIR):
    """SIR agents with mask-wearing behavior."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.wearing_mask = False
        self.mask_effectiveness = 0.5  # 50% transmission reduction

    def interact(self, folks_here, current_place_type, status_dict_t, model_params, dice):
        # Modify transmission probability based on mask wearing
        if self.status == 'S':
            num_contact = len(
             [folk for folk in folks_here if folk != self and folk.status in ['I']])

             # conversion_prob * I / N is the non-linear term that defines conversion
             # This inverse bernoulli function is an interpretation of the term
             # in agent-based modeling
             base_prob = self.inverse_bernoulli(
                    num_contact, conversion_prob / len(folks_here))

             # Reduce transmission if either person wears a mask
             if self.wearing_mask or any(folk.wearing_mask for folk in infectious_folks):
                base_prob *= (1 - self.mask_effectiveness)

             if base_prob > dice:
                self.convert('I', status_dict_t)

        if self.energy > 0:
            self.energy -= 1

Validation and Testing

Always validate your custom model:

Unit Tests: Test individual components (parameter validation, agent transitions)
Integration Tests: Verify the complete model runs without errors
Mathematical Validation: Compare with known analytical solutions when possible
Sensitivity Analysis: Test how parameter changes affect outcomes
Edge Cases: Test extreme parameter values and unusual scenarios

def validate_sir_model():
    """Comprehensive SIR model validation."""

    # Test 1: Parameter validation
    try:
        SIRModelParameters(max_energy=5, beta=1.5, gamma=7)  # Should fail
        assert False, "Should have caught invalid beta"
    except AssertionError:
        print("Parameter validation working")

    # Test 2: Status transitions
    agent = FolkSIR(id=1, home_address=0, max_energy=5, status='S')
    status_dict = {'S': 100, 'I': 0, 'R': 0}

    agent.convert('I', status_dict)
    assert agent.status == 'I' and agent.status_step_streak == 0
    print("Status transitions working")

    # Test 3: Disease progression
    params = SIRModelParameters(max_energy=5, beta=0.3, gamma=3)
    agent.status_step_streak = 3  # At recovery threshold
    agent.sleep(status_dict, params)

    if agent.status_step_streak >= params.gamma:
        agent.convert('R', status_dict)
    assert agent.status == 'R'
    print("Disease progression working")

    print("All validation tests passed!")

Best Practices

When creating custom compartmental models:

1. Start Simple: Begin with basic functionality, then add complexity incrementally
2. Validate Early: Test each component before building the complete model
3. Document Everything: Clear docstrings and comments make models maintainable
4. Follow Conventions: Use consistent naming and structure with existing simcronomicon models
5. Test Thoroughly: Validate against mathematical models when possible
6. Consider Performance: Profile your model with large populations to identify bottlenecks

Common Pitfalls

Incorrect Inheritance Order: Always define model attributes BEFORE calling super().__init__()
Missing Status Validation: Ensure all_statuses includes every possible agent status
Energy Management: Don’t forget to decrement energy during interactions
Status Transition Logic: Use status_step_streak correctly for time-based transitions
Parameter Validation: Always validate parameter ranges and types

Next Steps

Now that you can create custom compartmental models, explore:

SIRD Model: Add death compartment to your SIR model
Age-Structured Models: Different parameters for age groups
Spatial Models: Location-dependent transmission rates
Multi-Strain Models: Competing variants with different characteristics
Economic Models: Incorporate economic factors and interventions

The simcronomicon framework makes it straightforward to implement any compartmental model structure. Your SIR implementation provides the foundation for understanding how to build increasingly sophisticated epidemic models tailored to your specific research questions.