jaxabm.Model

class jaxabm.Model(parameters=None, seed=None)

Bases: object

Base class for agent-based models in JaxABM.

This class provides an AgentPy-like interface for creating and running agent-based models.

Example

```python class MyModel(Model):

def setup(self):

self.agents = self.add_agents(10, MyAgent) self.env.temperature = 25.0

def step(self):

# Agents are stepped automatically by default # Add additional model logic here self.env.temperature += 0.1

# Record data self.record(‘temperature’, self.env.temperature)

model = MyModel(parameters) results = model.run() ```

__init__(parameters=None, seed=None)

Initialize model.

Parameters:

• parameters (Optional[Dict[str, Any]]) – Model parameters.

• seed (Optional[int]) – Random seed.

Methods

__init__([parameters, seed])	Initialize model.
add_agents(n, agent_class[, name])	Add agents to the model.
batch_run(parameter_ranges[, repetitions])	Run model with multiple parameter combinations.
compute_metrics(env_state, agent_states, ...)	Compute model metrics.
end()	Execute code at the end of a simulation.
get_agent(collection_name, agent_id)	Get an agent instance by collection name and ID.
record(name, value)	Record data for later analysis.
run([steps])	Run the model for the specified number of steps.
setup()	Set up model.
step()	Execute a single time step.
update_state(env_state, agent_states, ...)	Update model environment state.

setup()

Set up model.

Override this method to set up agents and environment.

Return type:

None

step()

Execute a single time step.

Override this method to define model behavior. By default, it steps all agent lists.

Return type:

None

end()

Execute code at the end of a simulation.

Override this method to define behavior at the end of a simulation.

Return type:

None

update_state(env_state, agent_states, model_params, key)

Update model environment state.

This method is called by the JAX model to update the environment state based on agent states. By default, it returns the environment state unchanged. Override this method to define custom state update logic.

Parameters:

• env_state (Dict[str, Any]) – Current environment state.

• agent_states (Dict[str, Dict[str, Any]]) – Current agent states by collection.

• model_params (Dict[str, Any]) – Model parameters.

• key (Array) – JAX random key.

Return type:

Dict[str, Any]

Returns:

Updated environment state.

compute_metrics(env_state, agent_states, model_params)

Compute model metrics.

This method is called by the JAX model to compute metrics from model state. By default, it returns an empty dictionary. Override this method to define custom metrics.

Parameters:

• env_state (Dict[str, Any]) – Current environment state.

• agent_states (Dict[str, Dict[str, Any]]) – Current agent states by collection.

• model_params (Dict[str, Any]) – Model parameters.

Return type:

Dict[str, Any]

Returns:

Dictionary of metrics.

add_agents(n, agent_class, name=None, **kwargs)

Add agents to the model.

Parameters:

• n (int) – Number of agents to add.

• agent_class (Type[Agent]) – Agent class to use.

• name (Optional[str]) – Name for this agent collection.

• **kwargs – Parameters to pass to the agents.

Return type:

AgentList

Returns:

AgentList of created agents.

get_agent(collection_name, agent_id)

Get an agent instance by collection name and ID.

This allows accessing agent instances for calling custom methods.

Parameters:

• collection_name (str) – Name of the agent collection.

• agent_id (int) – ID of the agent.

Return type:

Optional[Agent]

Returns:

Agent instance if found, None otherwise.

record(name, value)

Record data for later analysis.

Parameters:

• name (str) – Name for the recorded data.

• value (Any) – Value to record.

Return type:

None

run(steps=None)

Run the model for the specified number of steps.

Parameters:

steps (Optional[int]) – Number of steps to run.

Return type:

Results

Returns:

Results object containing simulation results.

batch_run(parameter_ranges, repetitions=1)

Run model with multiple parameter combinations.

Parameters:

• parameter_ranges (Dict[str, List[Any]]) – Dictionary mapping parameter names to lists of values.

• repetitions (int) – Number of repetitions for each parameter combination.

Return type:

Dict[str, Results]

Returns:

Dictionary mapping parameter combinations to Results objects.