Model Calibration

JaxABM provides state-of-the-art parameter calibration methods to find optimal model parameters that match target outcomes. The calibration system supports traditional optimization methods as well as cutting-edge reinforcement learning approaches.

Overview

Parameter calibration is the process of finding model parameters that produce outputs matching desired target metrics. JaxABM’s calibration system is designed for:

Flexibility: Multiple optimization algorithms
Performance: JAX-accelerated computations
Robustness: Handles complex, non-linear parameter spaces
Advanced Methods: Reinforcement learning for challenging problems

Quick Example

import jaxabm as jx

# Create calibrator
calibrator = jx.analysis.ModelCalibrator(
    model_factory=your_model_factory,
    initial_params={'param1': 0.1, 'param2': 0.5},
    target_metrics={'output1': 2.5, 'output2': 1.0},
    method='actor_critic'  # RL method
)

# Find optimal parameters
best_params = calibrator.calibrate()

Available Methods

Traditional Optimization

Method	Description	Best For	Performance
`adam`	Gradient-based optimization	Smooth landscapes	Fast
`pso`	Particle Swarm Optimization	Multi-modal problems	Good
`es`	Evolution Strategies	Noisy functions	Robust
`cem`	Cross-Entropy Method	High-dimensional spaces	Excellent

Reinforcement Learning Methods ⭐

Method	Description	Best For	Performance
`actor_critic`	Actor-Critic algorithms	Complex dynamics	Excellent
`policy_gradient`	Policy gradient methods	Continuous spaces	Very Good
`q_learning`	Deep Q-Learning	Discrete actions	Good
`dqn`	Deep Q-Networks	Large state spaces	Very Good

Note

New in v0.1.0: All RL methods have been significantly improved with enhanced numerical stability and convergence guarantees.

Method Selection Guide

Choose your calibration method based on your problem characteristics:

For Economic Models: actor_critic or cem

Handle complex parameter interactions
Robust to noisy evaluations
Good convergence properties

For Epidemiological Models: pso or policy_gradient

Navigate multi-modal fitness landscapes
Handle discrete parameter changes well

For Social Network Models: es or dqn

Robust to evaluation noise
Handle large parameter spaces

For Quick Prototyping: adam or pso

Fast convergence for testing
Good general-purpose methods

Advanced Features

Multi-Objective Calibration

calibrator = jx.analysis.ModelCalibrator(
    model_factory=model_factory,
    initial_params=params,
    target_metrics={
        'accuracy': 0.95,
        'efficiency': 0.8,
        'robustness': 0.9
    },
    metrics_weights={
        'accuracy': 2.0,     # Higher priority
        'efficiency': 1.0,
        'robustness': 1.5
    },
    method='actor_critic'
)

Parameter Bounds and Constraints

calibrator = jx.analysis.ModelCalibrator(
    model_factory=model_factory,
    initial_params={'growth': 0.02, 'volatility': 0.1},
    target_metrics={'gdp': 2.5, 'unemployment': 5.0},
    param_bounds={
        'growth': (0.0, 0.1),        # Economic constraints
        'volatility': (0.01, 0.5)    # Realistic ranges
    },
    method='policy_gradient'
)

Robust Evaluation

calibrator = jx.analysis.ModelCalibrator(
    model_factory=model_factory,
    initial_params=params,
    target_metrics=targets,
    num_evaluation_runs=5,    # Multiple runs for robustness
    evaluation_steps=100,     # Longer simulations
    method='actor_critic'
)

Detailed Guides

Reinforcement Learning Calibration

Examples

Best Practices

Start Simple: Begin with traditional methods before using RL
Parameter Scaling: Normalize parameters to similar ranges
Target Scaling: Use relative errors for different metric scales
Evaluation Budget: Balance accuracy vs computational cost
Convergence Monitoring: Use early stopping and convergence criteria
Reproducibility: Set random seeds for consistent results

Performance Tips

Use jit compilation for model functions
Batch multiple evaluations when possible
Consider GPU acceleration for large models
Monitor memory usage with large agent populations
Use appropriate tolerance levels for convergence