Simulation Models

GPU Agent-Based Model

W.I.N.G.S. GPU-Accelerated Simulation Engine

Replaces the per-beetle Python loop with fully vectorized PyTorch tensor operations. Designed for NVIDIA L40S (48 GB VRAM) but will run on any CUDA device (or CPU as fallback).

Scaling strategy

The original ABM is O(F·M) per time-step for the mating search because every female is checked against every male. At N = 20 000 that means ~100 M distance evaluations per step – trivial on an L40S but wasteful in memory for N > 50 000.

This module provides two mating backends:

1. brute – compute the full female × male distance matrix on GPU. Memory: O(F·M) floats ≈ 400 MB at N = 20 000. Fast and simple.

2. cell_list – partition the toroidal grid into cells of side length ≥ mating_distance, then only check the 3×3 neighbourhood of each female’s cell. Memory: O(N · k) where k is the mean number of neighbours. Scales to N > 100 000.

Usage

>>> from gpu_simulation import GPUSimulation, SimConfig
>>> cfg = SimConfig(initial_population=20_000, max_population=25_000,
...                 grid_size=500, wolbachia_effects={'cytoplasmic_incompatibility': True,
...                 'male_killing': False, 'increased_exploration_rate': True,
...                 'increased_eggs': True, 'reduced_eggs': False})
>>> sim = GPUSimulation(cfg)
>>> for day in range(365):
...     sim.step_one_day()          # 24 hourly sub-steps
...     print(sim.get_infection_rate())

Author: Adapted from WINGS ABM (Geurten et al.)

class wings.models.gpu_abm.SimConfig(initial_population=50, max_population=20000, max_eggs=800000, infected_fraction=0.1, male_to_female_ratio=0.5, grid_size=500, mating_distance=5.0, wolbachia_effects=<factory>, egg_laying_max=15, ci_strength=1.0, fecundity_increase_factor=1.2, fecundity_decrease_factor=0.85, male_offspring_rate=0.1, exploration_rate_boost=1.4, mating_cooldown_female=48, mating_cooldown_male=5, multiple_mating=True, egg_hatching_age=552, life_expectancy_min=6720, life_expectancy_max=10800, initial_age_min=889, initial_age_max=2500, levy_alpha=1.5, mortality_mode='cannibalism', mortality_beta=2.0, cannibalism_rate=6e-07, mating_backend='cell_list', device='cuda', seed=None)

Bases: object

Configuration dataclass for the GPU ABM simulation.

All simulation parameters are set here and passed to GPUSimulation. Immutable after construction.

Parameters:

• initial_population (int)

• max_population (int)

• max_eggs (int)

• infected_fraction (float)

• male_to_female_ratio (float)

• grid_size (int)

• mating_distance (float)

• wolbachia_effects (Dict[str, bool])

• egg_laying_max (int)

• ci_strength (float)

• fecundity_increase_factor (float)

• fecundity_decrease_factor (float)

• male_offspring_rate (float)

• exploration_rate_boost (float)

• mating_cooldown_female (int)

• mating_cooldown_male (int)

• multiple_mating (bool)

• egg_hatching_age (int)

• life_expectancy_min (int)

• life_expectancy_max (int)

• initial_age_min (int)

• initial_age_max (int)

• levy_alpha (float)

• mortality_mode (str)

• mortality_beta (float)

• cannibalism_rate (float)

• mating_backend (str)

• device (str)

• seed (int | None)

initial_population

Starting number of adult beetles.

Type:

int

max_population

Carrying capacity (K).

Type:

int

max_eggs

Egg buffer cap.

Type:

int

grid_size

Side length of the toroidal grid.

Type:

int

ci_strength

CI intensity (0–1).

Type:

float

mortality_mode

Density-dependent mortality type. One of "none", "logistic", "cannibalism", "contest".

Type:

str

mortality_beta

Exponent for density dependence.

Type:

float

cannibalism_rate

Egg cannibalism rate at N=K.

Type:

float

mating_backend

"brute" (O(F·M) distance matrix) or "cell_list" (O(N·k) spatial hashing).

Type:

str

device

"cuda" or "cpu".

Type:

str

seed

Random seed for reproducibility.

Type:

int, optional

wolbachia_effects

Boolean toggles for CI, MK, ER, IE, RE.

Type:

dict

infected_fraction

Initial infection prevalence.

Type:

float

initial_population: int = 50

max_population: int = 20000

max_eggs: int = 800000

infected_fraction: float = 0.1

male_to_female_ratio: float = 0.5

grid_size: int = 500

mating_distance: float = 5.0

wolbachia_effects: Dict[str, bool]

egg_laying_max: int = 15

ci_strength: float = 1.0

fecundity_increase_factor: float = 1.2

fecundity_decrease_factor: float = 0.85

male_offspring_rate: float = 0.1

exploration_rate_boost: float = 1.4

mating_cooldown_female: int = 48

mating_cooldown_male: int = 5

multiple_mating: bool = True

egg_hatching_age: int = 552

life_expectancy_min: int = 6720

life_expectancy_max: int = 10800

initial_age_min: int = 889

initial_age_max: int = 2500

levy_alpha: float = 1.5

mortality_mode: str = 'cannibalism'

mortality_beta: float = 2.0

cannibalism_rate: float = 6e-07

mating_backend: str = 'cell_list'

device: str = 'cuda'

seed: int | None = None

class wings.models.gpu_abm.PopulationState(device)

Bases: object

Holds the entire population (adults + eggs) as flat GPU tensors.

We use a Structure-of-Arrays layout so that every operation (move, age, filter, mate) is a single vectorized kernel call instead of a Python loop over beetles.

Attributes (all tensors on device):

x, y : float32 [N] – positions infected : bool [N] – Wolbachia status is_male : bool [N] – sex (True = male) age : int32 [N] – current age in hours max_life : int32 [N] – sampled life expectancy (hours) last_mate : int32 [N] – sim-time of last mating event is_egg : bool [N] – True while in egg stage (not yet hatched)

Parameters:

device (torch.device)

property n: int

property n_adults: int

property n_eggs: int

append(**kwargs)

Concatenate new individuals (given as keyword tensors) onto the state.

keep(mask)

Keep only individuals where mask is True (boolean tensor [N]).

Parameters:

mask (torch.Tensor)

subsample(max_n)

If n > max_n, randomly keep max_n individuals (GPU grim_reaper).

Parameters:

max_n (int)

class wings.models.gpu_abm.GPUSimulation(cfg)

Bases: object

GPU-accelerated agent-based model of Wolbachia spread.

Replaces per-beetle Python loops with fully vectorised PyTorch tensor operations. All beetle state (position, sex, infection, age, mating cooldown) is stored as contiguous GPU tensors.

Two mating backends are supported:

• brute: Full F×M distance matrix. Fast for N < 20,000.

• cell_list: Spatial hashing into grid cells. Scales to N > 100,000.

The egg pipeline models Tribolium’s 23-day (552-hour) development period as a ring buffer of daily cohorts.

Parameters:

• config (SimConfig) – Simulation configuration.

• cfg (SimConfig)

sim_time: int

infection_history: List[float]

population_history: List[int]

step()

Advance the simulation by one hour.

Sequence:

1. Age all adults; remove those exceeding max lifespan.

2. Move adults (Lévy flight; ER beetles get 1.4× step size).

3. Find mating pairs within mating distance.

4. Reproduce (apply CI, MK, IE/RE).

5. Add new eggs to the pipeline.

6. Hatch eggs from 23 days ago.

7. Apply density-dependent egg mortality (cannibalism).

8. Record population size and infection rate.

step_one_day()

Run 24 hourly steps, recording stats only at the end of the day.

get_infection_rate()

Return type:

float

get_population_size()

Return type:

int

get_sex_ratio()

Returns counts of adult males and females (infected and uninfected).

Return type:

dict

export_history_csv(path)

Write the simulation time series to a CSV file.

Columns: Population Size, Infection Rate (one row per recorded time point — typically daily).

Parameters:

path (str) – Output CSV file path.

wings.models.gpu_abm.run_experiment(cfg, n_days=365, verbose=True)

Run a complete simulation experiment.

Creates a GPUSimulation from the config, runs for n_days × 24 hours, and returns the simulation object with its recorded history.

Parameters:

• config (SimConfig) – Simulation parameters.

• n_days (int) – Duration in days. Defaults to 365.

• cfg (SimConfig)

• verbose (bool)

Returns:

Completed simulation with history.

Return type:

GPUSimulation

Wright-Fisher Model

W.I.N.G.S. — Fixed-size discrete-generation Wolbachia spread model.

Wright-Fisher style: N adults are sampled each generation from the offspring pool. This decouples infection-frequency dynamics from population-size dynamics, providing a clean baseline for measuring the effect of CI, male-killing, increased exploration, and increased eggs on Wolbachia invasion.

Biology

Tribolium castaneum generation time ≈ 30 days at 30 °C (Pointer et al. 2021, Heredity). One year ≈ 12 generations.

Mapping spatial effects to well-mixed

The ABM’s “increased exploration rate” gives infected ♀ a 1.4× mating distance. In a well-mixed cage of 50 beetles, the spatial advantage is translated as a mating-probability advantage: infected ♀ mate with probability 1.0 while uninfected ♀ mate with probability 1 / exploration_rate_boost ≈ 0.714. This represents the empirical observation that more-active females encounter males more frequently even in small arenas (Pai & Bhatt 1995, J. Stored Prod. Res.).

Expected CI dynamics (sanity check)

With full CI (strength 1.0) and initial freq p = 0.5:

Infected ♀ (freq p) mate with any ♂ → offspring all infected. Uninfected ♀ (freq 1-p) × uninfected ♂ (freq 1-p) → offspring uninfected. Uninfected ♀ × infected ♂ → NO offspring (CI).

New p ≈ p / (p + (1-p)²)

Gen 0: p=0.50 → Gen 1: p≈0.67 → Gen 2: p≈0.86 → Gen 3: p≈0.97 → fixation. Expect fixation within 4–6 generations with stochastic variation.

Usage

# Single run python fixed_generation_sim.py –ci –seed 42 –output result.csv

# All 16 combos × 200 reps (takes ~30 seconds on one CPU core) python fixed_generation_sim.py –run-all –nreps 200

# Quick test python fixed_generation_sim.py –run-all –nreps 2 –outdir ./test_fixed

wings.models.wfm.simulate(N=50, max_generations=12, seed=42, ci=False, mk=False, er=False, ie=False, ci_strength=1.0, egg_laying_max=15, male_offspring_rate=0.1, fecundity_increase_factor=1.35, er_mating_advantage=1.4, initial_infection_freq=0.5)

Run one fixed-size discrete-generation simulation.

Parameters

Nint

Fixed adult population size per generation.

max_generationsint

Maximum generations to simulate (≈ 1 year at 12).

seedint

RNG seed for reproducibility.

ci, mk, er, iebool

Wolbachia effect toggles.

ci_strengthfloat

Probability that each egg from ♂I × ♀U cross is killed (0–1).

egg_laying_maxint

Maximum clutch size per mating (uniform draw from 1..max).

male_offspring_ratefloat

Probability offspring is male when MK active and mother infected.

fecundity_increase_factorfloat

Clutch multiplier for infected mothers when IE active.

er_mating_advantagefloat

Mating probability ratio (infected / uninfected) when ER active.

initial_infection_freqfloat

Starting Wolbachia frequency (0–1).

Returns

list of (population_size: int, infection_rate: float)

One entry per generation, padded to max_generations + 1.

Parameters:

• N (int)

• max_generations (int)

• seed (int)

• ci (bool)

• mk (bool)

• er (bool)

• ie (bool)

• ci_strength (float)

• egg_laying_max (int)

• male_offspring_rate (float)

• fecundity_increase_factor (float)

• er_mating_advantage (float)

• initial_infection_freq (float)

wings.models.wfm.save_csv(history, path)

Save simulation history as CSV matching gpu_simulation.py format.

wings.models.wfm.combo_id(ci, mk, er, ie)

Compute the 4-bit combo index consistent with the GPU submit script.

wings.models.wfm.combo_label(ci, mk, er, ie)

wings.models.wfm.make_filename(ci, mk, er, ie, rep)

Build filename matching the original ABM format for ingest compatibility.

wings.models.wfm.run_all(args)

Run all 16 combos × nreps replicates.

wings.models.wfm.main()

CLI entry point for the Wright-Fisher model.

Supports single-run mode and --run-all for the full 16-combination sweep with multiple replicates.

CPU Agent-Based Model (Legacy)

Original CPU-based agent-based model (per-beetle Python loops).

class wings.models.cpu.beetle.Beetle(position, infected, sex, environment, age=0, mating_cooldown=48)

Bases: object

A single beetle agent in the CPU-based ABM.

Represents an individual Tribolium beetle with spatial position, infection status, sex, age, and mating state. Movement follows a Lévy flight distribution, with infected beetles optionally moving further (increased exploration rate).

position

(x, y) coordinates on the toroidal grid.

Type:

tuple

infected

Whether the beetle carries Wolbachia.

Type:

bool

sex

'male' or 'female'.

Type:

str

age

Current age in hours.

Type:

int

environment

Reference to the simulation environment.

Type:

Environment

max_life_expectancy

Maximum lifespan in hours.

Type:

float

mating_cooldown

Hours between successive matings.

Type:

int

last_mating_time

Hour of the most recent mating event.

Type:

int

grid_size

Side length of the simulation grid.

Type:

int

__init__(position, infected, sex, environment, age=0, mating_cooldown=48)

Initialise a beetle agent.

Parameters:

• position (tuple) – Initial (x, y) grid coordinates.

• infected (bool) – Wolbachia infection status.

• sex (str) – 'male' or 'female'.

• environment (Environment) – The simulation environment.

• age (int) – Starting age in hours. Defaults to 0.

• mating_cooldown (int) – Minimum hours between matings. Defaults to 48 (2 days).

generate_life_expectancy()

Generates the beetle’s life expectancy (in hours) based on a uniform distribution. Roughly between 9 and 15 months.

levy_flight_step()

Performs a movement step based on a Lévy flight pattern. The step size follows a Pareto distribution, and direction is random (0 to 2π).

move()

Move the beetle using a Lévy flight step.

Step length is drawn from a power-law distribution. Infected beetles with the increased_exploration_rate phenotype receive a 1.4× movement multiplier, increasing their effective mating radius. The grid wraps toroidally.

update_last_mating_time(current_time)

Updates the last mating time of the beetle to the current time (hours).

can_mate(current_time)

Determines whether the beetle can mate based on current time and its mating cooldown. Males have a shorter cooldown (10% of the female cooldown period).

class wings.models.cpu.environment.Environment(size, initial_population, wolbachia_effects, infected_fraction=0.1, max_population=50, max_eggs=40, male_to_female_ratio=0.5, param_set=None, ci_strength=1.0, multiple_mating=True, use_gpu=False)

Bases: object

The simulation arena for the CPU-based ABM.

Manages a population of Beetle agents on a toroidal grid, handles hourly time-stepping (movement, mating, ageing, mortality), and records population size and infection rate over time.

grid_size

Side length of the toroidal grid.

Type:

int

population

Current adult beetles.

Type:

list[Beetle]

eggs

Developing eggs (hatch after ~552 hours).

Type:

list

wolbachia_effects

Boolean toggles for CI, MK, ER, IE, RE.

Type:

dict

max_population

Carrying capacity.

Type:

int

max_eggs

Maximum egg buffer size.

Type:

int

population_size

Time series of population counts.

Type:

list[int]

infection_history

Time series of infection rates.

Type:

list[float]

infected_fraction

Current infection prevalence.

Type:

float

__init__(size, initial_population, wolbachia_effects, infected_fraction=0.1, max_population=50, max_eggs=40, male_to_female_ratio=0.5, param_set=None, ci_strength=1.0, multiple_mating=True, use_gpu=False)

Initialise the simulation environment.

Creates the initial population with the given infection fraction, balanced sex ratio, and optional stochastic parameters.

Parameters:

• size (int) – Grid side length (grid is size × size).

• initial_population (int) – Number of starting beetles.

• wolbachia_effects (dict) – Boolean toggles, e.g. {'cytoplasmic_incompatibility': True, 'male_killing': False, ...}.

• infected_fraction (float) – Proportion of initial population carrying Wolbachia. Defaults to 0.1.

• max_population (int) – Carrying capacity. Defaults to 50.

• max_eggs (int) – Egg buffer cap. Defaults to 40.

• male_to_female_ratio (float) – Sex ratio. Defaults to 0.5.

• param_set (ParameterSet, optional) – Stochastic parameter source.

• ci_strength (float) – CI intensity (0–1). Defaults to 1.0.

• multiple_mating (bool) – Allow females to mate multiple times per cycle. Defaults to True.

• use_gpu (bool) – Whether to attempt GPU acceleration. Defaults to False.

initialize_population(initial_population)

Create the initial population of beetles.

Distributes beetles uniformly on the grid with the specified infection fraction and sex ratio.

generate_position_in_central_third()

Generates a random (x, y) position within the central third region of the grid.

run_simulation_step()

Advance the simulation by one hour.

Sequence per step:

1. Move all beetles (Lévy flight).

2. Identify mating pairs within mating distance.

3. Reproduce (with CI/MK/IE/RE effects).

4. Age all beetles; remove dead.

5. Hatch eggs that have completed development.

6. Enforce carrying capacity via random culling.

7. Record population size and infection rate.

grim_reaper(target_list, max_size)

Ensures the list does not exceed max_size by randomly removing surplus elements. Returns a pruned list (or the original list if within limit).

retire_old_beetles()

Removes beetles that have exceeded their life expectancy from the population.

check_infection_status()

Calculates the current infected fraction of the population and logs it.

check_for_mating()

Checks each female for mating opportunities with nearby males. Allows multiple matings if enabled. Produces offspring eggs for each successful mating.

is_within_mating_distance(female, male)

Determines if two beetles are within mating distance. If ‘increased_exploration_rate’ is in effect and the female is infected, expands mating range by 40%.

update_population_arrays()

Updates cached tensors of population positions and infection statuses for GPU-based reproduction. Called after any change in the population when GPU is in use.

class wings.models.cpu.reproduction.Reproduction(environment)

Bases: object

Handles beetle mating, offspring production, and Wolbachia effects.

Implements cytoplasmic incompatibility (CI), male killing (MK), increased egg laying (IE), and reduced egg laying (RE) as modular toggles on the reproduction pipeline.

grid_size

Side length of the simulation grid.

Type:

int

wolbachia_effects

Boolean toggles for each effect.

Type:

dict

environment

Reference to the simulation environment.

Type:

Environment

egg_laying_max

Maximum eggs per clutch (default: 15).

Type:

int

__init__(environment)

Initialise the reproduction handler.

Parameters:

environment (Environment) – The simulation environment, from which grid size and effect toggles are read.

mate(female, male)

Attempt mating between a female and a male beetle.

Applies CI logic: if the male is infected and the female is not, offspring viability is reduced by ci_strength.

Parameters:

• female (Beetle) – The female beetle.

• male (Beetle) – The male beetle.

Returns:

Offspring beetles (may be empty if CI kills the brood or population cap is reached).

Return type:

list[Beetle]

generate_offspring(female)

Generates offspring (as Beetle objects) for a female beetle after a mating event (CPU mode). Applies male-killing effect and maternal transmission of infection.

determine_offspring_count(female)

Determines how many eggs a female will lay from a mating event. If Wolbachia infection affects fecundity (increase or reduction), adjust the base egg count.

get_nearby_position(position)

Generates a new position (within 1 unit in x and y) near the given position. Uses toroidal wrapping if the position goes out of bounds.

batch_mating_events(sim, female_indices, male_indices)

Vectorized offspring generation for multiple mating pairs (GPU mode). sim (int): Index of the simulation batch. female_indices (List[int]): Population indices of mothers. male_indices (List[int]): Population indices of fathers. Returns a dict with offspring attributes (positions, infected, sex, age, life) for all offspring.

class wings.models.cpu.parameters.ParameterSet

Bases: object

Stochastic physiological and behavioural parameters.

Each attribute is sampled from a biologically plausible range on instantiation, introducing inter-individual variability into the simulation.

fecundity_increase_factor

Multiplier (>1) for Wolbachia-induced fecundity increase. Range: 1.1–1.3.

Type:

float

fecundity_decrease_factor

Multiplier (<1) for Wolbachia-induced fecundity cost. Range: 0.8–0.9.

Type:

float

ci_strength

Cytoplasmic incompatibility strength. Sampled from {0.5, 0.75, 1.0}.

Type:

float

infected_male_advantage

Sperm competition advantage factor for infected males. Range: 0.7–0.9.

Type:

float

male_offspring_rate

Fraction of offspring that are male under male killing. Range: 0.1–0.2.

Type:

float

female_mating_interval

Base mating cooldown in hours. Fixed at 48 (2 days).

Type:

int