.claude

Imported from clorton/laser-malaria (.claude/AGENTS.md). Install upstream with npx skills add clorton/laser-malaria --skill .claude. Copyright stays with the author.

LASER Framework Development Guide for AI Agents

This document provides comprehensive information about the LASER framework for AI agents assisting with development.

Framework Overview

LASER (Light Agent Spatial modeling for ERadication) is a high-performance, stochastic agent-based simulation framework for modeling the spread of infectious diseases. When developing LASER models, follow the architecture and style guidelines below.

LASER Framework API Summary

Overview

LASER (Light Agent Spatial modeling for ERadication) is a high-performance, stochastic agent-based simulation framework for modeling the spread of infectious diseases. It emphasizes spatial structure, age demographics, and modular disease logic through Python components.

Architecture

The framework represents systems as mutable dataframes where:

• Each column corresponds to a numerical property (e.g., node_id, age, infection_status, infectious_timer)
• Each row represents an individual agent
• Components process properties and run at each timestep to update the dataframe
• Can be implemented in NumPy (default), Numba, or C with optimizations

laser-core Module

LaserFrame Class

Core dataframe class similar to Pandas DataFrame that manages agent and node data.

Key Methods:

• add_scalar_property(name: str, dtype=np.uint32, default=0) - Add scalar property to agents
• add_vector_property(name: str, length: int, dtype=np.uint32, default=0) - Add vector property with specified length
• add(count: int) -> tuple[int, int] - Add new agents, returns (start_index, end_index)
• count - Property returning current count (equivalent to len())
• capacity - Property returning total capacity for dynamic properties
• save_snapshot() - Save LaserFrame to HDF5 file
• load_snapshot() - Load LaserFrame from HDF5 file
• sort() - Sort numpy arrays based on provided indices
• squash() - Compress dataframe by removing inactive agents

Demographics Module (laser_core.demographics)

pyramid module:

• Supports initialization of agents with plausible initial ages
• AliasedDistribution class - Generate samples using Vose alias method

Population distribution functions:

• distribute_population_tapered() - Distribute population across nodes with tapering
• distribute_population_skewed() - Calculate population distribution with rural/urban split

Migration Module (laser_core.migration)

Provides various migration network models:

• gravity() - Gravity model for migration based on population and distance
• competing_destinations() - Gravity model adjusted for destination competition
• stouffer() - Stouffer's intervening opportunities model
• radiation() - Radiation model for human mobility
• distance() - Calculate great-circle distance using Haversine formula

Utility Classes

• SortedQueue - High-performance priority queue for scheduled events
• PropertySet - Smart dictionary with dot-notation access to keys

laser-generic Module

Built on top of laser-core, provides epidemiological modeling tools.

Pre-built Disease Models

• SI (Susceptible-Infected)
• SIS (Susceptible-Infected-Susceptible)
• SIR (Susceptible-Infected-Recovered)
• SIRS (Susceptible-Infected-Recovered-Susceptible)
• SEIR (Susceptible-Exposed-Infected-Recovered)
• SEIRS (Susceptible-Exposed-Infected-Recovered-Susceptible)

Key Feature Modules

1. Disease Transmission - Models infection dynamics between agents
2. Vital Dynamics - Components for births and mortality modeling
3. Immunization - Routine immunization tracking and campaigns
4. Importation - Disease case importation mechanisms
5. Distributions - Numba-compatible probability distributions

Model Structure

The Model object is the central data structure encapsulating:

• Agent Population - Represented as a LaserFrame
• Node Information - Spatial/geographic data
• Components - Modular behavior units with:
  • Initialization function
  • Step function (runs each timestep)

Performance Optimizations

Key design principles for efficiency:

• Preallocated Memory - All arrays allocated at initialization
• Sequential Array Access - Process data sequentially to optimize CPU cache
• No Runtime Allocation - Eliminates dynamic memory allocation during simulation
• SIMD/OpenMP Support - Can leverage parallel processing optimizations

Installation

uv pip install laser-generic

Usage Pattern

1. Create LaserFrame for agents and nodes
2. Add properties (scalar/vector) for model state
3. Define components with init and step functions
4. Compose components into complete model
5. Run simulation by calling step functions each timestep

File I/O

• No required input file format - flexible data loading
• HDF5 preferred for large output files
• User responsible for collecting and writing output data

LASER Model and Component Implementation Style Guide

Overview

LASER (Location-Aware Stochastic Epidemic Realization) is a framework for building compartmental epidemiological models with geographic structure and agent-based simulation capabilities. This guide provides instructions for implementing LASER-style models and components based on the architecture from laser-generic.

Core Architecture Principles

1. Model Structure

The LASER model follows a specific initialization and execution pattern:

class Model:
    def __init__(self, scenario, params, birthrates=None, name="model_name",
                 skip_capacity=False, states=None, additional_states=None):
        """
        Initialize model with:
        - scenario: GeoDataFrame containing population, initial states, and geometry
        - params: PropertySet with simulation parameters (nticks, beta, etc.)
        - birthrates: Optional birth rates per patch per tick
        - states: Compartment names (default: {"S", "E", "I", "R"})
        """

2. Key Model Components

LaserFrame Data Structures

• model.people: Agent-level properties (LaserFrame)
• model.nodes: Node-level aggregate counts (LaserFrame)

Network Construction

Models must build a gravity-based migration network:

# Calculate distances between node centroids
dist_matrix = distance(lats, longs, lats, longs)

# Build gravity network with configurable parameters
network = gravity(population, dist_matrix, k=k, a=a, b=b, c=c)
network = row_normalizer(network, normalization_factor)

Component Registration

model.components = [component1, component2, ...]

3. Component Implementation Pattern

Each component follows a consistent structure:

class ComponentName:
    def __init__(self, model, **kwargs):
        """
        Initialize component with:
        - Reference to model
        - Add properties to model.people (agent-level)
        - Add properties to model.nodes (node-level aggregates)
        """
        self.model = model
        # Add agent properties
        model.people.add_property("property_name", dtype=np.int32)
        # Add node properties for each time step
        model.nodes.add_property("StateCount", shape=(params.nticks, num_nodes))

    def step(self, tick):
        """
        Execute component logic for current tick.
        Updates agent states and node aggregates.
        """
        # Process agents
        # Update node counts
        pass

    def plot(self, figure=None, **kwargs):
        """
        Generate visualization for this component.
        Yields matplotlib figures.
        """
        pass

4. State Management Rules

Flow Initialization

Before each tick, preserve current state values:

def _initialize_flows(self, tick):
    for state in self.states:
        if (prop := getattr(self.nodes, state, None)) is not None:
            # state(t+1) = state(t) + state(t)
            prop[tick + 1, :] = prop[tick, :]

Agent Property Conventions

• state: Integer encoding of compartment (0=S, 1=E, 2=I, 3=R, etc.)
• nodeid: Geographic location of agent
• timer: Countdown for state transitions (e.g., infection duration)

5. Numba Optimization

Use Numba JIT compilation for performance-critical code:

import numba as nb

@nb.njit(nogil=True, parallel=True, cache=True)
def process_agents(state, timer, nodeid, ...):
    """
    Process agent states in parallel.
    - nogil=True: Release Python GIL
    - parallel=True: Enable parallelization
    - cache=True: Cache compiled functions
    """
    for i in nb.prange(len(state)):
        # Agent processing logic
        pass

6. Validation Framework

Implement validation decorators for consistency checking:

def validate(func):
    """Decorator to wrap step methods with validation."""
    def wrapper(self, tick):
        if self.model.validating:
            self.prevalidate_step(tick)
        result = func(self, tick)
        if self.model.validating:
            self.postvalidate_step(tick)
        return result
    return wrapper

class Component:
    @validate
    def step(self, tick):
        # Component logic
        pass

7. Component Types and Responsibilities

State Initialization Components

• Susceptible: Initialize S compartment from scenario data
• Exposed: Initialize E compartment (for SEIR models)
• Infectious: Initialize I compartment variants

Transmission Components

• TransmissionSI: SI/SIS dynamics (no recovery or with reinfection)
• TransmissionSE: SEIR exposure dynamics
• TransmissionSIR: SIR with immunity

Transition Components

• Recovery: I→R transitions with configurable duration
• Waning: R→S immunity loss

Demographic Components

• Births: Add new susceptible agents
• Deaths: Remove agents based on mortality rates
• Migration: Move agents between nodes

8. Visualization Standards

Components should provide plotting methods that:

• Accept optional Figure object for multi-panel layouts
• Support basemap overlays for geographic visualization
• Yield figures for PDF generation
• Display both temporal dynamics and spatial distributions

def plot(self, figure=None, basemap_provider=None, **kwargs):
    fig = plt.figure(figsize=(12, 9), dpi=200) if figure is None else figure

    # Plot logic

    yield fig  # Allow caller to save/display

9. Parameter Conventions

Standard parameter names in params:

• nticks: Total simulation time steps
• beta: Transmission rate
• gravity_k, gravity_a, gravity_b, gravity_c: Migration parameters
• prng_seed: Random number generator seed
• capacity_safety_factor: Agent allocation multiplier

10. Error Handling and Assertions

Include assertions to validate data consistency:

assert dist_matrix.shape == (num_nodes, num_nodes), "Distance matrix shape mismatch"
assert np.all(timer[state == RECOVERED] == 0), "Recovered agents should have timer=0"

Implementation Checklist

When implementing a new LASER model or component:

• Inherit from appropriate base class or follow Model/Component pattern
• Initialize LaserFrame properties in __init__
• Implement step() method with tick parameter
• Use Numba JIT for performance-critical loops
• Add validation methods if maintaining invariants
• Include plotting method for visualization
• Document all parameters and state transitions
• Use consistent state encoding (0=S, 1=E, 2=I, 3=R)
• Handle edge cases (empty nodes, zero population)
• Test with both single and multi-node scenarios

Example: Minimal SIR Component

import numba as nb
import numpy as np

class TransmissionSIR:
    """SIR transmission dynamics."""

    def __init__(self, model, beta=0.3, recovery_rate=0.1):
        self.model = model
        self.beta = beta
        self.recovery_rate = recovery_rate

        # Ensure required states exist
        if not hasattr(model.people, "state"):
            model.people.add_property("state", dtype=np.int32)

    @nb.njit(nogil=True, parallel=True, cache=True)
    def _transmit(self, state, nodeid, node_I, node_N, beta, dt):
        """Calculate new infections."""
        for i in nb.prange(len(state)):
            if state[i] == 0:  # Susceptible
                force_of_infection = beta * node_I[nodeid[i]] / node_N[nodeid[i]]
                if np.random.random() < force_of_infection * dt:
                    state[i] = 2  # Become Infectious
        return state

    def step(self, tick):
        # Get current counts
        node_S = self.model.nodes.S[tick, :]
        node_I = self.model.nodes.I[tick, :]
        node_N = node_S + node_I + self.model.nodes.R[tick, :]

        # Transmit infections
        self.model.people.state[:] = self._transmit(
            self.model.people.state,
            self.model.people.nodeid,
            node_I,
            node_N,
            self.beta,
            1.0  # dt
        )

        # Update node counts for next tick
        # (Actual implementation would aggregate from agent states)

Best Practices

1. Modularity: Keep components focused on single responsibilities
2. Performance: Use Numba for loops over agents
3. Validation: Add checks to catch inconsistencies early
4. Documentation: Include docstrings with parameter descriptions
5. Testing: Verify conservation laws (population totals)
6. Visualization: Provide meaningful plots for debugging
7. Configurability: Expose key parameters through constructor
8. Compatibility: Follow LaserFrame property conventions

This style guide ensures consistency and interoperability when building LASER epidemiological models.

Additional Development Guidelines

Code Quality Standards

• Use double quotes for strings
• Run ruff on code for formatting and linting
• Write docstrings for all classes and functions using triple double quotes
• Use pathlib.Path for all file and disk operations rather than the os package
• Write unittests in the GWT pattern - given a particular setup or scenario, when the code is executed or functions are called, then we expect and assert specific test results

When Working with LASER Components

1. Always follow the component implementation pattern described above
2. Use Numba optimizations for performance-critical code
3. Maintain consistent state encoding across components
4. Implement validation methods for debugging
5. Provide visualization methods for all components
6. Document all parameters and their expected ranges
7. Test with both single-node and multi-node scenarios