Building Models¶

The ModelBuilder class provides a fluent API for constructing epidemiological models.

ModelBuilder Basics¶

Creating a Builder¶

from epimodel import ModelBuilder

builder = ModelBuilder(
    name="My Model",
    version="1.0",
    description="Optional description"
)

Chaining Methods¶

The builder uses method chaining for a clean, readable API:

model = (
    ModelBuilder(name="SIR Model")
    .add_disease_state(id="S", name="Susceptible")
    .add_disease_state(id="I", name="Infected")
    .add_disease_state(id="R", name="Recovered")
    .add_parameter(id="beta", value=0.3)
    .add_transition(id="infection", source=["S"], target=["I"], rate="beta * S * I / N")
    .build(ModelTypes.DIFFERENCE_EQUATIONS)
)

Adding Disease States¶

builder.add_disease_state(
    id="S",                    # Required: Unique identifier
    name="Susceptible",        # Required: Display name
    description="Population susceptible to infection"  # Optional
)

Best Practices¶

Use short, clear IDs (S, I, R, E, etc.)
Provide descriptive names

Adding Stratifications¶

Stratifications create population subgroups:

builder.add_stratification(
    id="age_group",
    categories=["0-17", "18-64", "65+"],
    description="Age-based stratification"
)

Multiple Stratifications¶

builder.add_stratification(id="age", categories=["young", "old"])
builder.add_stratification(id="location", categories=["urban", "rural"])

This creates compartments: S_young_urban, S_young_rural, S_old_urban, S_old_rural, etc.

Adding Parameters¶

Parameters are global constants used in formulas:

builder.add_parameter(
    id="beta",
    value=0.3,
    description="Transmission rate per contact per day"
)

Parameter Guidelines¶

Use meaningful IDs (beta, gamma, R0, etc.)
Document units and meaning
Ensure values are realistic for your model

Adding Transitions¶

Transitions move populations between states.

Understanding Transition Rates¶

The rate parameter accepts mathematical expressions that can include:

Parameters: Reference parameter IDs (e.g., "gamma")
Disease states: Use state populations (e.g., "S", "I")
Special variables: N (total population), step or t (current time step), pi, e
Mathematical operations: +, -, *, /, ** (power)
Functions: sin(), cos(), exp(), log(), sqrt(), max(), min(), etc.

For the complete list of functions and advanced examples, see Mathematical Expressions.

Simple Parameter-Based Rates¶

builder.add_transition(
    id="recovery",
    source=["I"],
    target=["R"],
    rate="gamma"  # References parameter id
)

Formula-Based Rates¶

builder.add_transition(
    id="infection",
    source=["S"],
    target=["I"],
    rate="beta * S * I / N"  # Mathematical expression
)

Constant Rates¶

builder.add_transition(
    id="birth",
    source=[],      # Empty = enters system
    target=["S"],
    rate="0.001"    # Fixed rate
)

Time-Dependent Rates¶

builder.add_transition(
    id="seasonal_infection",
    source=["S"],
    target=["I"],
    rate="beta * (1 + 0.3 * sin(2 * pi * t / 365)) * S * I / N"
)

See Mathematical Expressions for more complex rate formulas.

Multi-State Transitions¶

# Death from any compartment
builder.add_transition(
    id="death",
    source=["S", "I", "R"],
    target=[],  # Empty = leaves system
    rate="mu"
)

Stratified Transitions¶

When a model includes stratifications, you often need different transition rates for different subgroups. The add_transition method supports this via the stratified_rates parameter.

This parameter takes a list of dictionaries, where each dictionary defines a rate for a specific combination of stratification categories.

Single Stratification¶

Let's define different recovery rates for different age groups.

builder.add_stratification(id="age", categories=["child", "adult", "elderly"])
builder.add_parameter(id="gamma_child", value=0.15)
builder.add_parameter(id="gamma_adult", value=0.1)
builder.add_parameter(id="gamma_elderly", value=0.08)

builder.add_transition(
    id="recovery",
    source=["I"],
    target=["R"],
    stratified_rates=[
        {
            "conditions": [{"stratification": "age", "category": "child"}],
            "rate": "gamma_child"
        },
        {
            "conditions": [{"stratification": "age", "category": "adult"}],
            "rate": "gamma_adult"
        },
        {
            "conditions": [{"stratification": "age", "category": "elderly"}],
            "rate": "gamma_elderly"
        },
    ]
)

Multi-Stratification Transitions¶

To define rates for intersections of multiple stratifications, add multiple conditions to a single rate entry.

For example, let's model different infection rates for high-risk adults in urban areas.

builder.add_stratification(id="age", categories=["child", "adult"])
builder.add_stratification(id="risk", categories=["low", "high"])
builder.add_stratification(id="location", categories=["urban", "rural"])

builder.add_parameter(id="beta_urban_adult_high_risk", value=0.8)
builder.add_parameter(id="beta_default", value=0.3)

builder.add_transition(
    id="infection",
    source=["S"],
    target=["I"],
    rate="beta_default * S * I / N",  # Fallback rate
    stratified_rates=[
        {
            "conditions": [
                {"stratification": "age", "category": "adult"},
                {"stratification": "risk", "category": "high"},
                {"stratification": "location", "category": "urban"},
            ],
            "rate": "beta_urban_adult_high_risk * S * I / N"
        }
    ]
)

In this example:

The rate parameter acts as a fallback for any compartment that doesn't match a specific stratified rate.
The stratified_rates entry defines a high infection rate that only applies to compartments matching all three conditions (e.g., S_adult_high_urban).

Setting Initial Conditions¶

Basic Setup¶

builder.set_initial_conditions(
    population_size=1000,
    disease_state_fractions=[
        {"disease_state": "S", "fraction": 0.99},
        {"disease_state": "I", "fraction": 0.01},
        {"disease_state": "R", "fraction": 0.0}
    ]
)

With Stratifications¶

builder.set_initial_conditions(
    population_size=10000,
    disease_state_fractions=[
        {"disease_state": "S", "fraction": 0.99},
        {"disease_state": "I", "fraction": 0.01},
        {"disease_state": "R", "fraction": 0.0}
    ],
    stratification_fractions=[
        {
            "stratification": "age_group",
            "fractions": [
                {"category": "young", "fraction": 0.3},
                {"category": "adult", "fraction": 0.5},
                {"category": "elderly", "fraction": 0.2}
            ]
        },
        {
            "stratification": "risk",
            "fractions": [
                {"category": "low", "fraction": 0.8},
                {"category": "high", "fraction": 0.2}
            ]
        }
    ]
)

Building the Model¶

Once all components are added, build the model:

from epimodel.constants import ModelTypes

model = builder.build(typology=ModelTypes.DIFFERENCE_EQUATIONS)

Validation¶

The build process validates:

All disease state fractions sum to 1.0
All stratification fractions sum to 1.0
Transition sources/targets reference valid states
Mathematical expressions are syntactically correct
No security issues in formulas

If validation fails, a descriptive error is raised.

Advanced: Conditional Transitions¶

Create transitions that only occur under certain conditions:

# Create a condition
condition = builder.create_condition(
    logic="and",
    rules=[
        {"variable": "state:I", "operator": "gt", "value": 100},
        {"variable": "step", "operator": "gt", "value": 30}
    ]
)

# Add conditional transition
builder.add_transition(
    id="intervention",
    source=["S"],
    target=["S"],
    rate="0.5 * beta",  # Reduced transmission
    condition=condition
)

Loading from JSON¶

Load pre-defined models from JSON files:

from epimodel import ModelLoader

model = ModelLoader.from_json("path/to/model.json")

JSON Structure¶

{
  "name": "SIR Model",
  "version": "1.0",
  "population": {
    "disease_states": [
      { "id": "S", "name": "Susceptible" },
      { "id": "I", "name": "Infected" },
      { "id": "R", "name": "Recovered" }
    ]
  },
  "parameters": [
    { "id": "beta", "value": 0.3 },
    { "id": "gamma", "value": 0.1 }
  ],
  "dynamics": {
    "typology": "difference_equations",
    "transitions": [
      {
        "id": "infection",
        "source": ["S"],
        "target": ["I"],
        "rate": "beta * S * I / N"
      }
    ]
  }
}

Complete Example¶

from epimodel import ModelBuilder, Simulation
from epimodel.constants import ModelTypes

# Build SEIR model
model = (
    ModelBuilder(name="SEIR Model", version="1.0")
    .add_disease_state(id="S", name="Susceptible")
    .add_disease_state(id="E", name="Exposed")
    .add_disease_state(id="I", name="Infected")
    .add_disease_state(id="R", name="Recovered")
    .add_parameter(id="beta", value=0.4, description="Transmission rate")
    .add_parameter(id="sigma", value=0.2, description="Incubation rate")
    .add_parameter(id="gamma", value=0.1, description="Recovery rate")
    .add_transition(id="exposure", source=["S"], target=["E"], rate="beta * S * I / N")
    .add_transition(id="infection", source=["E"], target=["I"], rate="sigma")
    .add_transition(id="recovery", source=["I"], target=["R"], rate="gamma")
    .set_initial_conditions(
        population_size=1000,
        disease_state_fractions=[
            {"disease_state": "S", "fraction": 0.999},
            {"disease_state": "E", "fraction": 0.0},
            {"disease_state": "I", "fraction": 0.001},
            {"disease_state": "R", "fraction": 0.0}
        ]
    )
    .build(typology=ModelTypes.DIFFERENCE_EQUATIONS)
)

# Run simulation
simulation = Simulation(model)
results = simulation.run(num_steps=200)

Next Steps¶

Mathematical Expressions - Advanced formulas
Simulations - Running and analyzing models
Examples - Complete model examples