Rotasim

https://github.com/starsimhub/rotasim
Admin
Rotasim is an agent-based model for simulating rotavirus genetic diversity and vaccination...
Tokens:26,167
Snippets:189
Trust Score:8.9
License:MIT
Update:5 months ago
Show doc for...
Tokens:
Raw
### Basic Simulation Setup and Run

Source: https://github.com/starsimhub/rotasim/blob/main/CLAUDE.md

Set up and run a RotaABM simulation with specified agent count, time limit, and verbosity. Custom Rota-specific keyword arguments can be passed.

```python
import rotasim as rs

sim = rs.Sim(
    n_agents=50000,
    timelimit=10,
    verbose=True,
    rota_kwargs={"vaccination_time": 5, "time_to_equilibrium": 2}
)
sim.run()
```

--------------------------------

### Basic Simulation Setup

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Initializes and runs a basic simulation using a predefined scenario. Access simulation results like total diseases, scenario name, and initial strains.

```python
import rotasim as rs

# Simple simulation using predefined scenario
sim = rs.Sim(scenario='simple')  # Two-strain scenario (G1P8, G2P4)
sim.run()

# Access results
print(f"Simulation completed with {len(sim.diseases)} total diseases")
print(f"Scenario: {sim.scenario}")
print(f"Initial strains: {sim.initial_strains}")
```

--------------------------------

### Import Libraries and Set Up Plotting

Source: https://github.com/starsimhub/rotasim/blob/main/examples/rotasim_demo.ipynb

Imports necessary libraries for simulation and plotting, and configures plotting styles. Ensure these libraries are installed before running.

```python
# Import required libraries
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import starsim as ss
import rotasim as rs

# Set up plotting
plt.style.use('default')
sns.set_palette('husl')
plt.rcParams['figure.figsize'] = (10, 6)

print(f"Success: Rotasim v2 loaded successfully")
print(f"Success: Ready for multi-strain rotavirus modeling")
```

--------------------------------

### Full Pipeline Example

Source: https://context7.com/starsimhub/rotasim/llms.txt

Demonstrates a comprehensive simulation pipeline using Rotasim, integrating custom connectors, interventions, and analyzers for multi-strain epidemiological research.

```APIDOC
## Combined Full-Pipeline Example

### Description
Rotasim's primary use case is multi-strain epidemiological research: running scenario-based simulations to study strain competition, genetic reassortment, and the impact of vaccination programs on rotavirus genetic diversity. Researchers can rapidly iterate across built-in scenarios or define custom strain configurations, attach multiple analyzers for comprehensive output, and export results to CSV for downstream analysis.

Integration with the broader Starsim ecosystem means Rotasim supports any standard `ss.Analyzer`, `ss.Intervention`, `ss.Demographics`, or `ss.Network` alongside its specialized components. The scenario override system enables parameter sweeps without modifying base scenario definitions, making it straightforward to run calibration studies or sensitivity analyses across fitness values, transmission rates, immunity parameters, and vaccination coverage levels.

### Method
```python
# Setup custom connectors and interventions
immunity = RotaImmunityConnector(...)
vax = RotaVaccination(...)

# Initialize and run the simulation
sim = rs.Sim(
    scenario='realistic_competition',
    base_beta=0.1,
    connectors=[immunity],
    interventions=[vax],
    analyzers=[rs.StrainStats(), rs.EventStats(), rs.AgeStats()],
    n_agents=50000,
    start='2018-01-01',
    stop='2033-01-01',
    dt=ss.days(1),
    verbose=1
)
sim.run()

# Export results and print summaries
sim.analyzers['strainstats'].to_df().to_csv('strains.csv', index=False)
sim.analyzers['eventstats'].to_df().to_csv('events.csv', index=False)
sim.analyzers['agestats'].to_df().to_csv('ages.csv', index=False)
sim.interventions['rotavaccination'].print_vaccination_summary()
sim.print_strain_summary()
```

### Request Example
```python
import starsim as ss
import rotasim as rs
from rotasim import RotaVaccination, RotaImmunityConnector

# Full pipeline: custom immunity + vaccination + comprehensive analyzers
immunity = RotaImmunityConnector(
    homotypic_immunity_efficacy=0.9,
    partial_heterotypic_immunity_efficacy=0.5,
    complete_heterotypic_immunity_efficacy=0.3,
)

vax = RotaVaccination(
    start_date='2022-01-01',
    n_doses=2,
    G_antigens=[1, 2],
    P_antigens=[8, 4],
    dose_effectiveness=[0.6, 0.8],
    uptake_dist=ss.bernoulli(p=0.85),
    verbose=False
)

sim = rs.Sim(
    scenario='realistic_competition',
    base_beta=0.1,
    connectors=[immunity],         # Custom immunity (no default reassortment)
    interventions=[vax],
    analyzers=[rs.StrainStats(), rs.EventStats(), rs.AgeStats()],
    n_agents=50000,
    start='2018-01-01',
    stop='2033-01-01',
    dt=ss.days(1),
    verbose=1
)
sim.run()

# Export all results
sim.analyzers['strainstats'].to_df().to_csv('strains.csv', index=False)
sim.analyzers['eventstats'].to_df().to_csv('events.csv', index=False)
sim.analyzers['agestats'].to_df().to_csv('ages.csv', index=False)

# Vaccination summary
sim.interventions['rotavaccination'].print_vaccination_summary()

# Strain summary
sim.print_strain_summary()
```

### Response
#### Success Response
- The simulation runs and generates CSV files for strain, event, and age statistics.
- Vaccination and strain summaries are printed to the console.

#### Response Example
```
# Example output from print_vaccination_summary()
Vaccination Summary:
- Total doses administered: 85000
- Total individuals vaccinated: 42500

# Example output from print_strain_summary()
=== Rotasim Strain Summary ===
Total diseases: 16
Initial strains: 4
Dormant reassortants: 12
```
```

--------------------------------

### Rotasim Full Pipeline Example with Customization

Source: https://context7.com/starsimhub/rotasim/llms.txt

Sets up and runs a comprehensive Rotasim simulation pipeline, including custom immunity connectors, vaccination interventions, and multiple analyzers. Use for detailed epidemiological studies and parameter sweeps.

```python
import starsim as ss
import rotasim as rs
from rotasim import RotaVaccination, RotaImmunityConnector

# Full pipeline: custom immunity + vaccination + comprehensive analyzers
immunity = RotaImmunityConnector(
    homotypic_immunity_efficacy=0.9,
    partial_heterotypic_immunity_efficacy=0.5,
    complete_heterotypic_immunity_efficacy=0.3,
)

vax = RotaVaccination(
    start_date='2022-01-01',
    n_doses=2,
    G_antigens=[1, 2],
    P_antigens=[8, 4],
    dose_effectiveness=[0.6, 0.8],
    uptake_dist=ss.bernoulli(p=0.85),
    verbose=False
)

sim = rs.Sim(
    scenario='realistic_competition',
    base_beta=0.1,
    connectors=[immunity],         # Custom immunity (no default reassortment)
    interventions=[vax],
    analyzers=[rs.StrainStats(), rs.EventStats(), rs.AgeStats()],
    n_agents=50000,
    start='2018-01-01',
    stop='2033-01-01',
    dt=ss.days(1),
    verbose=1
)
sim.run()

# Export all results
sim.analyzers['strainstats'].to_df().to_csv('strains.csv', index=False)
sim.analyzers['eventstats'].to_df().to_csv('events.csv', index=False)
sim.analyzers['agestats'].to_df().to_csv('ages.csv', index=False)

# Vaccination summary
sim.interventions['rotavaccination'].print_vaccination_summary()

# Strain summary
sim.print_strain_summary()
# === Rotasim Strain Summary ===
# Total diseases: 16
# Initial strains: 4
# Dormant reassortants: 12
```

--------------------------------

### V2 Core Components Setup

Source: https://github.com/starsimhub/rotasim/blob/main/CLAUDE.md

Instantiate individual Rotavirus strains with G,P parameters and connect them using RotaImmunityConnector and RotaReassortment for cross-strain interactions. This pattern is suitable for setting up complex simulations with genetic diversity.

```python
rota_g1p8 = Rotavirus(G=1, P=8, backbone=(1,1))  # Auto-named "G1P8_11"
rota_g2p4 = Rotavirus(G=2, P=4, backbone=(1,1))  # Auto-named "G2P4_11"

# Connectors auto-detect Rotavirus instances
immunity = RotaImmunityConnector()
reassortment = RotaReassortment()

sim = ss.Sim(
    diseases=[rota_g1p8, rota_g2p4],
    connectors=[immunity, reassortment],
    networks='random'
)
```

--------------------------------

### Install RotaABM

Source: https://github.com/starsimhub/rotasim/blob/main/CLAUDE.md

Install the RotaABM package in development mode. This makes the package importable in your Python environment.

```bash
pip install -e .
```

--------------------------------

### Define and Simulate Rotavirus Strains with Connectors

Source: https://github.com/starsimhub/rotasim/blob/main/planning/rotasim_v2_architecture_plan.md

This example demonstrates setting up multiple Rotavirus strains, configuring their transmission parameters, and integrating immunity and reassortment connectors within a standard Starsim simulation. Ensure all possible G,P reassortments are generated at simulation start as Starsim cannot add new diseases mid-simulation.

```python
from starsim import Sim
from starsim.connectors import Connector

# Placeholder for Rotavirus class and helper function
class Rotavirus(ss.Infection):
    def __init__(self, G, P, init_prev):
        super().__init__(name=f"Rota_G{G}P{P}")
        self.G = G
        self.P = P
        self.pars.init_prev = init_prev
        self.pars.beta = 0.1 # Default beta, will be overridden

class RotaImmunityConnector(ss.Connector):
    def __init__(self):
        super().__init__(name='rota_immunity')

    def call(self, sim):
        # Placeholder for immunity logic
        pass

class RotaReassortment(ss.Connector):
    def __init__(self):
        super().__init__(name='rota_reassortment')

    def set_prognoses(self, sim):
        # Placeholder for reassortment logic
        pass

def generate_gp_reassortments(initial_strains):
    # Placeholder for generating all G,P pairs
    # Example: return [(1,8), (2,4), (3,6), (1,4), (2,8)]
    return initial_strains # Simplified for example

# Define initial G,P strain combinations (backbone removed for performance)
initial_strains = [(1,8), (2,4), (3,6)]

# Auto-generate ALL possible G,P reassortments at simulation start
# (Required: Starsim cannot add new diseases mid-simulation)
all_gp_pairs = generate_gp_reassortment(initial_strains)  # ~20-70 combinations

# Create Rotavirus instances for all possible G,P combinations
base_beta = 0.1  # Global transmission scaling
diseases = []
for G, P in all_gp_pairs:
    # Most reassortants start with zero prevalence
    init_prev = 0.01 if (G,P) in initial_strains else 0.0
    diseases.append(Rotavirus(G=G, P=P, init_prev=init_prev))

# Strain-specific fitness multipliers for easy parameter tuning
x_beta_lookup = {
    (1,8): 1.0,   # G1P8 baseline fitness
    (2,4): 0.8,   # G2P4 reduced fitness  
    (3,6): 1.2,   # G3P6 enhanced fitness
}
for disease in diseases:
    x_beta = x_beta_lookup.get((disease.G, disease.P), 1.0)
    disease.pars.beta = base_beta * x_beta

# Connectors auto-detect Rotavirus instances - no manual listing needed
immunity = RotaImmunityConnector()  # Finds all Rotavirus diseases automatically
reassortment = RotaReassortment()   # Uses set_prognoses() for co-infected hosts

# Standard Starsim simulation - no custom Sim class needed
sim = ss.Sim(
    diseases=diseases,  # All possible G,P combinations
    connectors=[immunity, reassortment],
    networks='random',
    n_agents=10000
)

sim.run()

```

--------------------------------

### Assemble a simulation manually with custom Rotavirus instances

Source: https://context7.com/starsimhub/rotasim/llms.txt

A simulation can be assembled manually by providing a list of custom `Rotavirus` disease instances and connectors (like `RotaImmunityConnector`) to the `ss.Sim` constructor. This approach offers more control over the simulation setup compared to using predefined scenarios.

```python
# Assemble simulation manually using ss.Sim
immunity = RotaImmunityConnector()
sim = ss.Sim(
    diseases=[g1p8, g2p4],
    connectors=[immunity],
    networks='random',
    n_agents=10000,
    dt=ss.days(1),
    start='2020-01-01',
    stop='2022-01-01'
)
sim.run()

# After run: check infection states
print(f"G1P8 currently infected: {g1p8.infected.sum()}")
print(f"G2P4 currently infected: {g2p4.infected.sum()}")
print(f"G1P8 total recoveries (last step): {g1p8.results.new_recovered[-1]}")
```

--------------------------------

### Monitor Cross-Strain Interactions

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Track immunity and reassortment events during a simulation to analyze cross-strain dynamics. This setup is useful for emergence studies.

```python
import rotasim as rs
from rotasim.analyzers import EventStats

# Track immunity and reassortment events
analyzer = EventStats()
sim = rs.Sim(
    scenario='emergence_scenario',  # Weak background for emergence studies
    analyzers=[analyzer],
    stop='2035-01-01',  # 15-year simulation
    dt=rs.days(1),     # Daily timesteps
    verbose=1  # Show summary information
)

sim.run()

# Analyze cross-strain dynamics
events = analyzer.to_df()
print(f"Immunity waning events: {events['wanings'].sum()}")
print(f"Reassortment events: {events['reassortments'].sum()}")
print(f"Peak coinfected agents: {events['coinfected_agents'].max()}")
print(f"Final scenario used: {sim.final_scenario['strains']}")
```

--------------------------------

### Listing Available Scenarios

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Retrieves and prints a list of all available built-in simulation scenarios and their descriptions. Demonstrates how to initialize simulations with different scenarios.

```python
import rotasim as rs

# List all available scenarios
scenarios = rs.list_scenarios()
for name, description in scenarios.items():
    print(f"{name}: {description}")

# Use different scenarios
sim_baseline = rs.Sim(scenario='baseline')           # 3 common global strains
sim_diverse = rs.Sim(scenario='high_diversity')      # 12 strains with varied fitness
sim_competition = rs.Sim(scenario='realistic_competition')  # G1P8 dominant with competition
```

--------------------------------

### Run Quick Test

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Execute a quick test by navigating to the tests directory and running the simple.py script.

```bash
# Run from tests directory
cd tests
python simple.py
```

--------------------------------

### Get Rotasim strain summary after simulation

Source: https://context7.com/starsimhub/rotasim/llms.txt

After running a simulation, the `get_strain_summary()` method on the `Sim` object provides statistics on total diseases, active strains, and dormant reassortants.

```python
# --- Strain summary ---
summary = sim4.get_strain_summary()
print(f"Total diseases: {summary['total_diseases']}")
print(f"Active strains: {len(summary['active_strains'])}")
print(f"Dormant reassortants: {len(summary['dormant_strains'])}")
```

--------------------------------

### Run Simple Simulation

Source: https://github.com/starsimhub/rotasim/blob/main/CLAUDE.md

Execute a basic simulation script located in the 'tests' directory. Results will be saved to a 'results' folder.

```bash
cd tests
python simple.py
```

--------------------------------

### Manage Simulation Scenarios with Scenario Utilities

Source: https://context7.com/starsimhub/rotasim/llms.txt

Provides functions to list, retrieve, validate, and apply overrides to built-in simulation scenarios. Demonstrates listing available scenarios and their descriptions.

```python
import rotasim as rs
from rotasim.utils import (
    list_scenarios, get_scenario,
    validate_scenario, apply_scenario_overrides
)

# List all scenarios
for name, desc in list_scenarios().items():
    print(f"  {name}: {desc}")
```

--------------------------------

### Implement Multi-Dose Vaccination with RotaVaccination

Source: https://context7.com/starsimhub/rotasim/llms.txt

Configure RotaVaccination for a multi-dose vaccination program, specifying start date, dose intervals, antigen targets, effectiveness, age eligibility, uptake, and waning rates.

```python
import starsim as ss
import rotasim as rs
from rotasim import RotaVaccination

# --- Simple 2-dose G1P8 vaccine ---
vax = RotaVaccination(
    start_date='2023-01-01',
    n_doses=2,
    dose_interval=ss.days(28),
    G_antigens=[1],
    P_antigens=[8],
    dose_effectiveness=[0.6, 0.8],
    min_age=ss.days(42),    # 6 weeks
    max_age=ss.days(365),   # 1 year
    uptake_dist=ss.bernoulli(p=0.8),
    waning_rate_dist=ss.lognorm_ex(mean=365),  # 1-year mean waning
    homotypic_efficacy=1.0,
    partial_heterotypic_efficacy=0.6,
    complete_heterotypic_efficacy=0.3,
    verbose=True
)

sim = rs.Sim(
    scenario='realistic_competition',
    interventions=[vax],
    n_agents=50000,
    start='2020-01-01',
    stop='2033-01-01',
    dt=ss.days(1)
)
sim.run()

# Vaccination summary
vax_obj = sim.interventions['rotavaccination']
vax_obj.print_vaccination_summary()
# === RotaVaccination Summary ===
# Total agents: 50000
# Received any dose: ...
# Completed schedule: ...

summary = vax_obj.get_vaccination_summary()
print(f"Completion rate: {100 * summary['completed_schedule'] / summary['total_agents']:.1f}%")
```

--------------------------------

### Run Performance Benchmarks

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Test the performance of the simulation with a large number of strains.

```bash
# Test performance with large strain numbers
python tests/test_performance.py
```

--------------------------------

### Initialize and run Rotasim simulation with predefined scenario

Source: https://context7.com/starsimhub/rotasim/llms.txt

Use the `rs.Sim` class to initialize a simulation with a predefined scenario like 'realistic_competition'. You can specify parameters such as base transmission rate, number of agents, and simulation duration. The simulation automatically generates all possible reassortant strain combinations.

```python
import rotasim as rs
import starsim as ss

# --- Predefined scenario ---
sim = rs.Sim(
    scenario='realistic_competition',  # G1P8 dominant with competition
    base_beta=0.1,
    n_agents=50000,
    start='2020-01-01',
    stop='2030-01-01',
    dt=ss.days(1)
)
sim.run()

print(sim)
# Sim(scenario=realistic_competition, strains=4, n_agents=50000)
print(f"Total diseases (active + dormant): {len(sim.diseases)}")
print(f"Initial strains: {sim.initial_strains}")
# [(1, 8), (2, 4), (3, 8), (4, 8)]
```

--------------------------------

### Set Up Basic Multi-Strain Simulation

Source: https://github.com/starsimhub/rotasim/blob/main/examples/rotasim_demo.ipynb

Configures a multi-strain Rotasim simulation with specified population size, network type, strain scenario, and simulation duration. Includes analyzers for tracking strain dynamics, events, and age distribution.

```python
# Create a basic multi-strain simulation
print("=== Setting Up Multi-Strain Simulation ===\n")

# Create simulation with analyzers for data collection
sim = rs.Sim(
    # Population and network
    n_agents=5000,
    networks='random', # Randomly connected network
    
    # Strain configuration using unified scenario API
    scenario='simple',  # Two-strain scenario (G1P8, G2P4) with equal fitness and prevalence
    
    # Simulation parameters
    start='2020-01-01',
    stop='2020-06-01',
    dt=ss.days(1),   # Daily timesteps
    verbose=0,       # Quiet operation

    # Data collection
    analyzers=[
        rs.StrainStats(),  # Track strain dynamics
        rs.EventStats(),   # Track simulation events
        rs.AgeStats()      # Track age distribution
    ]
)

# Display simulation setup
print(f"\nSimulation Configuration:")
print(f"   Population: {sim.pars.n_agents:,} agents")
print(f"   Duration: {sim.pars.start} to {sim.pars.stop}")
print(f"   Time step: {sim.pars.dt}")
print(f"   Scenario: {sim.scenario}")

# Show strain details
strain_summary = sim.get_strain_summary()
print(f"\nStrain Details:")
print(f"   Total possible strains: {strain_summary['total_diseases']}")
print(f"   Active at start: {len(strain_summary['active_strains'])}")
print(f"   Dormant reassortants: {len(strain_summary['dormant_strains'])}")

print(f"\n   Active strains:")
for strain in strain_summary['active_strains']:
    print(f"     • {strain['name']}: G{strain['G']}P{strain['P']}")

print(f"\n   Dormant reassortants (can emerge during simulation):")
for strain in strain_summary['dormant_strains']:
    print(f"     • {strain['name']}: G{strain['G']}P{strain['P']}")
```

--------------------------------

### Parameter Sensitivity Analysis

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Perform parameter sensitivity analysis by running simulations with different scenarios and cross-protection levels. Collect and print results to understand parameter impact.

```python
import rotasim as rs
import numpy as np

# Test scenario and immunity parameter sensitivity
scenarios = ['simple', 'baseline', 'realistic_competition']
cross_protection_levels = [0.2, 0.4, 0.6, 0.8]
results = []

for scenario in scenarios:
    for cross_prot in cross_protection_levels:
        immunity = rs.RotaImmunityConnector(
            partial_heterotypic_immunity_efficacy=cross_prot
        )
        
        sim = rs.Sim(
            scenario=scenario,
            connectors=[immunity],
            dt=rs.days(1),
            stop='2030-01-01',
            verbose=0
        )
        sim.run()
        
        results.append({
            'scenario': scenario,
            'cross_protection': cross_prot,
            'total_diseases': len(sim.diseases),
            'initial_strains': len(sim.initial_strains)
        })

print("Scenario and cross-protection sensitivity:")
for result in results:
    print(f"Scenario: {result['scenario']}, Protection: {result['cross_protection']:.1f}, "
          f"Total diseases: {result['total_diseases']}, Initial: {result['initial_strains']}")
```

--------------------------------

### Monitor Population Age Distribution with AgeStats

Source: https://context7.com/starsimhub/rotasim/llms.txt

Initializes and runs a simulation using the AgeStats analyzer to track the age distribution of the population over time across nine standard age bins. The resulting DataFrame is printed and saved to a CSV.

```python
import rotasim as rs

age_analyzer = rs.AgeStats()

sim = rs.Sim(
    scenario='baseline',
    analyzers=[age_analyzer],
    n_agents=15000,
    stop='2027-01-01',
    dt=rs.days(1)
)
sim.run()

aa = sim.analyzers['agestats']
df = aa.to_df()

print(df.columns.tolist())
# ['timevec', '0-2', '2-4', '4-6', '6-12', '12-24', '24-36', '36-48', '48-60', '60+']

df.to_csv('age_distribution.csv', index=False)

# Check age structure at final timestep
print(df.iloc[-1])
```

--------------------------------

### Run Performance Tests

Source: https://github.com/starsimhub/rotasim/blob/main/CLAUDE.md

Execute performance testing scripts for the RotaABM simulation.

```bash
python tests/test_performance.py
```

--------------------------------

### Configure Custom Protection Factors

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Set up a Rotasim simulation with custom asymmetric protection factors using RotaImmunityConnector. This allows for fine-tuning cross-protection and waning rates.

```python
import rotasim as rs
from rotasim import RotaImmunityConnector

# Asymmetric protection scenarios
strong_homotypic = RotaImmunityConnector(
    homotypic_immunity_efficacy=0.95,       # Very strong same-strain
    partial_heterotypic_immunity_efficacy=0.4,   # Moderate cross-protection
    complete_heterotypic_immunity_efficacy=0.1,  # Minimal heterotypic
    full_waning_rate=rs.freqperyear(365/180)     # 6-month waning
)

sim = rs.Sim(
    scenario='balanced_competition',  # Use built-in 4-strain scenario
    connectors=[strong_homotypic],
    base_beta=0.12,
    verbose=2  # Detailed output
)
```

--------------------------------

### Run the Multi-Strain Simulation

Source: https://github.com/starsimhub/rotasim/blob/main/examples/rotasim_demo.ipynb

Executes the configured Rotasim simulation and confirms completion. This step may take some time to complete.

```python
# Run the simulation
print("Running simulation...")
print("   (This may take 30-60 seconds)\n")

sim.run()

print("Success: Simulation completed successfully!")
print(f"\nSimulation Results:")
print(f"   Total timesteps: {len(sim.timevec)}")
print(f"   Analyzers available: {list(sim.analyzers.keys())}")
```

--------------------------------

### Run Integration Tests

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Execute integration tests for multi-strain architecture and unified scenarios using pytest. Also, test the scenario system.

```bash
# Test multi-strain architecture and unified scenarios
cd tests
python -m pytest test_integration.py -v
python test_utils.py  # Test scenario system
```

--------------------------------

### Manage Rotasim Scenarios Programmatically

Source: https://context7.com/starsimhub/rotasim/llms.txt

Demonstrates how to retrieve, validate, and modify Rotasim scenarios. Use this for custom scenario configurations or parameter sweeps.

```python
sc = get_scenario('high_diversity')
print(sc['strains'].keys())  # dict_keys with 12 (G,P) tuples
```

```python
custom = {
    'strains': {
        (1, 8): {'fitness': 1.0, 'prevalence': 0.015},
        (2, 4): {'fitness': 0.8, 'prevalence': 0.010},
    },
    'default_fitness': 0.3
}
validated = validate_scenario(custom)   # Returns scenario or raises ValueError
```

```python
modified = apply_scenario_overrides(
    validated,
    override_fitness=0.9,                       # Set all fitness to 0.9
    override_prevalence={(1, 8): 0.02},         # Override only G1P8 prevalence
    override_strains={(4, 8): {'fitness': 0.7, 'prevalence': 0.005}}  # Add G4P8
)
print(modified['strains'])
```

```python
sim = rs.Sim(scenario=modified, n_agents=10000)
sim.run()
```

--------------------------------

### Analyze Strain Proportions and Counts with StrainStats

Source: https://context7.com/starsimhub/rotasim/llms.txt

Initializes and runs a simulation using the StrainStats analyzer to track the proportion and count of infected agents per strain over time. Results are saved to a CSV and a programmatic summary is generated.

```python
import rotasim as rs

strain_analyzer = rs.StrainStats()

sim = rs.Sim(
    scenario='realistic_competition',
    analyzers=[strain_analyzer],
    n_agents=20000,
    start='2020-01-01',
    stop='2025-01-01',
    dt=rs.days(1)
)
sim.run()

# Access from sim.analyzers (recommended — Starsim may copy analyzer instances)
s a = sim.analyzers['strainstats']
df = sa.to_df()
print(df.head())
# timevec  (1, 8) proportion  (1, 8) count  (2, 4) proportion  ... 

df.to_csv('strain_proportions.csv', index=False)

# Programmatic summary
summary = sa.get_strain_summary()
for strain, stats in summary['strain_stats'].items():
    print(f"{strain}: max_proportion={stats['max_proportion']:.3f}, "
          f"mean_count={stats['mean_count']:.0f}, "
          f"active_timesteps={stats['total_timesteps_active']}")
```

--------------------------------

### List Available Rotasim Scenarios

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Retrieve and print a list of all available built-in scenarios in Rotasim, along with their descriptions.

```python
import rotasim as rs

# List all built-in scenarios
print(rs.list_scenarios())
```

--------------------------------

### Run Rotasim Simulation and Export Results

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Sets up analyzers to track strain statistics and events, runs a simulation, and exports the results to CSV files. Examine immunity waning events and the final scenario used.

```python
import rotasim as rs

analyzers = [
    rs.StrainStats(),    # Strain proportions and counts
    rs.EventStats()      # Births, deaths, infections, waning events
]

sim = rs.Sim(
    scenario='realistic_competition',  # Use built-in scenario
    analyzers=analyzers,
    dt=rs.days(1),  # Daily timesteps
    stop='2030-01-01',  # 10-year simulation
    verbose=1
)
sim.run()

# Export results
strain_stats = analyzers[0].to_df()
event_stats = analyzers[1].to_df()

strain_stats.to_csv("strain_proportions.csv", index=False)
event_stats.to_csv("simulation_events.csv", index=False)

# Examine immunity waning events
waning_events = event_stats['wanings'].sum() 
print(f"Total immunity waning events: {waning_events}")
print(f"Final scenario: {sim.final_scenario}")
```

--------------------------------

### Track Simulation Events with EventStats

Source: https://context7.com/starsimhub/rotasim/llms.txt

Sets up and runs a simulation using the EventStats analyzer to collect data on births, deaths, recoveries, infections, wanings, reassortments, and infected/coinfected agents. Key event data is printed and saved to a CSV.

```python
import rotasim as rs

event_analyzer = rs.EventStats()

sim = rs.Sim(
    scenario='emergence_scenario',
    analyzers=[event_analyzer],
    n_agents=25000,
    stop='2035-01-01',
    dt=rs.days(1),
    verbose=1
)
sim.run()

ea = sim.analyzers['eventstats']
df = ea.to_df()

print(df[['timevec', 'new_infections', 'recoveries', 'coinfected_agents',
          'reassortments', 'wanings']].tail(10))

print(f"Total new infections: {df['new_infections'].sum()}")
print(f"Peak coinfected agents: {df['coinfected_agents'].max()}")
print(f"Total reassortment events: {df['reassortments'].sum()}")

df.to_csv('simulation_events.csv', index=False)
```

--------------------------------

### Defining Custom Simulation Scenarios

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Creates a simulation using a completely custom scenario defined by a dictionary. Specifies strains with their fitness and prevalence, and sets a default fitness for dormant reassortants.

```python
import rotasim as rs

# Define custom scenario
custom_scenario = {
    'strains': {
        (1, 8): {'fitness': 1.0, 'prevalence': 0.015},
        (2, 4): {'fitness': 0.8, 'prevalence': 0.010},
        (3, 6): {'fitness': 0.9, 'prevalence': 0.005}
    },
    'default_fitness': 0.3  # For dormant reassortants
}

sim = rs.Sim(
    scenario=custom_scenario,
    base_beta=0.1,
    n_agents=25000
)
sim.run()
```

--------------------------------

### Customizing Simulation Scenarios

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Overrides default scenario parameters like prevalence, fitness, and transmission rate. Allows adding new strains to existing scenarios and viewing strain summaries.

```python
import rotasim as rs

# Override scenario parameters
sim = rs.Sim(
    scenario='baseline',
    override_prevalence=0.02,  # Set all strains to 2% prevalence
    override_fitness={(1,8): 0.95, (2,4): 0.8},  # Override specific strain fitness
    base_beta=0.15,  # Adjust base transmission rate
    n_agents=50000
)
sim.run()

# Add new strain to existing scenario
sim = rs.Sim(
    scenario='baseline', 
    override_strains={(9,6): {'fitness': 0.7, 'prevalence': 0.003}},  # Add G9P6
    verbose=True
)

# View strain summary
summary = sim.get_strain_summary()
print(f"Total diseases: {summary['total_diseases']}")
print(f"Active strains: {len(summary['active_strains'])}")
print(f"Dormant reassortants: {len(summary['dormant_strains'])}")
```

--------------------------------

### Run a Simple Rotavirus Simulation

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Initializes and runs a basic Rotasim simulation with predefined G1P8 and G2P4 strains. Displays the number of diseases simulated.

```python
import rotasim as rs

# Run a simple simulation
sim = rs.Sim(scenario='simple')  # G1P8 and G2P4 strains
sim.run()
print(f"Simulation completed with {len(sim.diseases)} diseases")
```

--------------------------------

### Importing Analysis Tools

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Imports specific analysis tools from the rotasim.analyzers module, such as StrainStats and EventStats, for post-simulation data processing.

```python
import rotasim as rs
from rotasim.analyzers import StrainStats, EventStats
```

--------------------------------

### Print Rotasim v2 Simulation Summary

Source: https://github.com/starsimhub/rotasim/blob/main/examples/rotasim_demo.ipynb

This Python code prints a summary of the completed Rotasim v2 simulation, detailing the generated data and suggesting next steps for scaling up research.

```python
# Final simulation summary
print("🎉 Rotasim v2 Demonstration Complete!\n")

print(f"📊 This simulation generated:")
print(f"   • {len(strain_df)} timesteps of strain dynamics data")
print(f"   • {len(proportion_cols)} strain proportion timeseries")
print(f"   • {len([col for col in event_df.columns if col != 'timevec'])} event type timeseries")
print(f"   • {len([col for col in age_df.columns if col != 'timevec'])} age group timeseries")
print(f"   • 3 CSV files ready for downstream analysis")

print(f"\n🚀 Ready to scale up for production research!")
print(f"   • Increase n_agents for better statistical power")
print(f"   • Extend simulation duration for long-term dynamics")
print(f"   • Add more initial strains for higher genetic diversity")
print(f"   • Include vaccination interventions")
print(f"   • Calibrate parameters using real-world data")
```

--------------------------------

### Step Method for RotaImmunityConnector

Source: https://github.com/starsimhub/rotasim/blob/main/planning/rotasim_v2_architecture_plan.md

The main step method that orchestrates the vectorized immunity waning and cross-immunity protection calculations.

```python
    def step(self):
        # 1. Vectorized immunity waning
        self._apply_waning()
        
        # 2. Vectorized cross-immunity protection calculation
        self._update_cross_immunity()
```

--------------------------------

### Define a fully custom Rotasim scenario with a dictionary

Source: https://context7.com/starsimhub/rotasim/llms.txt

A completely custom simulation scenario can be defined by passing a dictionary to the `scenario` argument of `rs.Sim`. This dictionary allows explicit definition of strains with their fitness and prevalence, as well as a default fitness for dormant reassortants.

```python
# --- Fully custom scenario dict ---
sim4 = rs.Sim(
    scenario={
        'strains': {
            (1, 8): {'fitness': 1.0, 'prevalence': 0.015},
            (2, 4): {'fitness': 0.8, 'prevalence': 0.010},
            (3, 6): {'fitness': 0.9, 'prevalence': 0.005},
        },
        'default_fitness': 0.3  # Fitness for dormant reassortants
    },
    base_beta=0.1,
    n_agents=25000
)
sim4.run()
```

--------------------------------

### Scenario Utilities

Source: https://context7.com/starsimhub/rotasim/llms.txt

Functions to inspect, retrieve, validate, and modify built-in scenarios. Built-in scenarios include: `simple`, `baseline`, `realistic_competition`, `balanced_competition`, `high_diversity`, `low_diversity`, `emergence_scenario`.

```APIDOC
## Scenario Utilities

### Description
Functions to inspect, retrieve, validate, and modify built-in scenarios. Built-in scenarios: `simple`, `baseline`, `realistic_competition`, `balanced_competition`, `high_diversity`, `low_diversity`, `emergence_scenario`.

### Functions
- `list_scenarios()`: Lists all available scenarios and their descriptions.
- `get_scenario(name)`: Retrieves a specific scenario by name.
- `validate_scenario(scenario)`: Validates a given scenario.
- `apply_scenario_overrides(scenario, overrides)`: Applies overrides to a scenario.

### Usage
```python
import rotasim as rs
from rotasim.utils import (
    list_scenarios, get_scenario,
    validate_scenario, apply_scenario_overrides
)

# List all scenarios
for name, desc in list_scenarios().items():
    print(f"  {name}: {desc}")
```
```

--------------------------------

### Visualize Strain Dynamics Over Time

Source: https://github.com/starsimhub/rotasim/blob/main/examples/rotasim_demo.ipynb

Generates two plots: one showing strain proportions over time and another showing total infection counts over time. It also prints summary statistics for strain proportions.

```python
# Create strain dynamics visualization
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))

# Plot 1: Strain proportions over time
time_points = np.arange(len(strain_df))
colors = sns.color_palette('husl', len(proportion_cols))

for i, col in enumerate(proportion_cols):
    # Extract strain name from column (e.g., "(1, 8) proportion" -> "G1P8")
    strain_tuple = col.split(' proportion')[0]
    parts = strain_tuple.strip('()').split(', ')
    g, p = parts[0], parts[1] 
    strain_name = f"G{g}P{p}"
    
    ax1.plot(time_points, strain_df[col], 
            label=strain_name, color=colors[i], linewidth=2)

ax1.set_xlabel('Time (days)')
ax1.set_ylabel('Strain Proportion')
ax1.set_title('Rotavirus Strain Proportions Over Time')
ax1.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
ax1.grid(True, alpha=0.3)
ax1.set_ylim(0, 1)

# Plot 2: Total infection counts
for i, col in enumerate(count_cols):
    # Extract strain name
    strain_tuple = col.split(' count')[0]
    parts = strain_tuple.strip('()').split(', ')
    g, p = parts[0], parts[1]
    strain_name = f"G{g}P{p}"
    
    ax2.plot(time_points, strain_df[col], 
            label=strain_name, color=colors[i], linewidth=2)

ax2.set_xlabel('Time (days)')
ax2.set_ylabel('Infected Count')
ax2.set_title('Rotavirus Strain Infection Counts Over Time')
ax2.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# Summary statistics
print(f"\nStrain Dynamics Summary:")
for col in proportion_cols:
    strain_tuple = col.split(' proportion')[0]
    parts = strain_tuple.strip('()').split(', ')
    g, p = parts[0], parts[1]
    strain_name = f"G{g}P{p}"
    
    max_prop = strain_df[col].max()
    mean_prop = strain_df[col].mean()
    
    print(f"   {strain_name}: max={max_prop:.3f}, mean={mean_prop:.3f}")
```

--------------------------------

### Configure and Run RotaVaccination Campaign

Source: https://context7.com/starsimhub/rotasim/llms.txt

Sets up a pentavalent-like 3-dose multi-strain rotavirus vaccination campaign with specified antigen targets, dose effectiveness, and uptake distribution. The campaign is then applied to a simulation.

```python
pentavalent = RotaVaccination(
    start_date='2025-06-01',
    end_date='2030-12-31',           # Time-limited campaign
    n_doses=3,
    G_antigens=[1, 2, 3, 4],
    P_antigens=[8, 4, 6],
    dose_effectiveness=[0.5, 0.7, 0.85],
    uptake_dist=ss.bernoulli(p=0.9),
)
sim2 = rs.Sim(scenario='high_diversity', interventions=[pentavalent], n_agents=30000)
sim2.run()
```

--------------------------------

### Testing Cross-Strain Protection Scenarios

Source: https://github.com/starsimhub/rotasim/blob/main/README.md

Compares simulation outcomes with different levels of cross-strain protection. Initializes `RotaImmunityConnector` with varying efficacy values and runs simulations using a 'high_diversity' scenario.

```python
import rotasim as rs
from rotasim import RotaImmunityConnector

# Test different protection scenarios
high_cross_protection = RotaImmunityConnector(
    homotypic_immunity_efficacy=0.9,
    partial_heterotypic_immunity_efficacy=0.8,  # High cross-protection
    complete_heterotypic_immunity_efficacy=0.6
)

low_cross_protection = RotaImmunityConnector(
    homotypic_immunity_efficacy=0.9,
    partial_heterotypic_immunity_efficacy=0.3,  # Low cross-protection
    complete_heterotypic_immunity_efficacy=0.1
)

# Compare strain diversity outcomes
for immunity, label in [(high_cross_protection, "High"), (low_cross_protection, "Low")]:
    sim = rs.Sim(
        scenario='high_diversity',  # Use built-in scenario
        connectors=[immunity],
        n_agents=25000,
        dt=rs.days(1),  # Daily timesteps
        verbose=0
    )
    sim.run()
    print(f"{label} cross-protection: simulation completed")
```

--------------------------------

### Performance Critical Points: Bitmask Operations

Source: https://github.com/starsimhub/rotasim/blob/main/planning/PHASE_1_IMPLEMENTATION_TODO.md

Notes on using `astype(int)` for bitwise operations and `astype(float)` for storing in FloatArr for performance.

```python
1. **Bitmask Operations**: Use `astype(int)` for bitwise ops, `astype(float)` to store in FloatArr
```

--------------------------------

### Run Pytest Tests

Source: https://github.com/starsimhub/rotasim/blob/main/CLAUDE.md

Execute all unit tests for the RotaABM project using pytest. Individual test files can also be run by specifying their path.

```bash
python -m pytest tests/
```

```bash
python -m pytest tests/test_rotasim.py
```

```bash
python -m pytest tests/test_analyzers.py
```

```bash
python -m pytest tests/test_calibration.py
```

--------------------------------

### Initialize RotaImmunityConnector

Source: https://github.com/starsimhub/rotasim/blob/main/planning/rotasim_v2_architecture_plan.md

Initializes the RotaImmunityConnector with immunity rates and waning parameters. It also auto-detects Rotavirus diseases and creates bitmask mappings for unique G and P types.

```python
class RotaImmunityConnector(ss.Connector):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
        self.define_pars(
            homotypic_immunity_rate=0.5,
            partial_heterotypic_immunity_rate=0.5,
            complete_heterotypic_immunity_rate=0.5,
            waning_rate=365/273,  # omega parameter
        )
        
    def init_post(self, sim):
        """Auto-detect Rotavirus diseases and create bitmask mappings"""
        super().init_post(sim)
        self.rota_diseases = [d for d in sim.diseases.values() 
                            if isinstance(d, Rotavirus)]
        
        # Create bit mappings for unique G,P types in simulation
        unique_G = sorted(set(d.G for d in self.rota_diseases))
        unique_P = sorted(set(d.P for d in self.rota_diseases))
        
        self.G_to_bit = {g: i for i, g in enumerate(unique_G)}
        self.P_to_bit = {p: i for i, p in enumerate(unique_P)}
        
        # Pre-compute disease-specific bitmasks for fast lookup
        self.disease_G_masks = {}
        self.disease_P_masks = {}
        for disease in self.rota_diseases:
            self.disease_G_masks[disease.name] = 1 << self.G_to_bit[disease.G]
            self.disease_P_masks[disease.name] = 1 << self.P_to_bit[disease.P]
```

--------------------------------

### Plot Simulation Results

Source: https://github.com/starsimhub/rotasim/blob/main/CLAUDE.md

Generate plots from simulation results. Ensure simulations have been run prior to executing this script.

```bash
python data_analysis/plot_results.py
```
Tokens:
Raw