jaxabm.analysis.SensitivityAnalysis

class jaxabm.analysis.SensitivityAnalysis(model_factory, param_ranges, metrics_of_interest, num_samples=100, seed=0)

Bases: object

Perform sensitivity analysis on model parameters.

This class provides tools for analyzing how changes in model parameters affect model outputs, using efficient sampling techniques and sensitivity indices calculation.

model_factory

Function to create model instances

param_ranges

Dictionary mapping parameter names to (min, max) ranges

metrics_of_interest

List of metric names to analyze

num_samples

Number of parameter samples to generate

key

JAX random key

samples

Generated parameter samples

results

Analysis results (populated after run())

__init__(model_factory, param_ranges, metrics_of_interest, num_samples=100, seed=0)

Initialize sensitivity analysis.

Parameters:

• model_factory (TypeVar(ModelFactory, bound= Callable[..., Model])) – Function to create model instances

• param_ranges (Dict[str, Tuple[float, float]]) – Dictionary mapping parameter names to (min, max) ranges

• metrics_of_interest (List[str]) – List of metric names to analyze

• num_samples (int) – Number of parameter samples to generate

• seed (int) – Random seed

Methods

__init__(model_factory, param_ranges, ...[, ...])	Initialize sensitivity analysis.
plot([metric, ax])	Plot sensitivity analysis results.
plot_indices([figsize])	Plot the sensitivity indices.
run([verbose])	Run sensitivity analysis.
sobol_indices()	Calculate sensitivity indices for each parameter and metric.

run(verbose=True)

Run sensitivity analysis.

Parameters:

verbose (bool) – Whether to print progress information

Return type:

Dict[str, Array]

Returns:

Dictionary mapping metric names to arrays of results

sobol_indices()

Calculate sensitivity indices for each parameter and metric.

This is a simplified implementation that calculates correlation-based indices as a proxy for Sobol indices. For a full Sobol analysis, specialized sampling would be required.

Return type:

Dict[str, Dict[str, float]]

Returns:

Dictionary mapping metric names to dictionaries of parameter name -> sensitivity index

plot(metric=None, ax=None, **kwargs)

Plot sensitivity analysis results.

Parameters:

• metric – Metric to plot. If None, plot sobol indices for all metrics.

• ax – Matplotlib axis to use for plotting.

• **kwargs – Additional keyword arguments to pass to plotting function.

Returns:

Matplotlib axis.

plot_indices(figsize=(10, 6))

Plot the sensitivity indices.

Parameters:

figsize (Tuple[int, int]) – Figure size as (width, height)

Return type:

Any

Returns:

Matplotlib figure and axes