"""
This module provides utility functions for the laser-core project.
Functions:
calc_capacity(birthrates: np.ndarray, initial_pop: np.ndarray, safety_factor: float = 1.0) -> np.ndarray:
Calculate the population capacity after a given number of ticks based on a constant birth rate.
grid(shape: Tuple[int, int], fill_value: float = 0.0) -> np.ndarray:
Create a 2D grid (numpy array) of the specified shape, filled with the given value.
"""
import geopandas as gpd
import numpy as np
from shapely.geometry import Polygon
[docs]
def calc_capacity(
birthrates: np.ndarray,
initial_pop: np.ndarray,
safety_factor: float = 1.0,
*,
deathrates: np.ndarray | None = None,
mortality_safety_factor: float = 1.0,
) -> np.ndarray:
"""
Estimate the required capacity (number of agents) to model a population given birthrates over time.
By default (``deathrates=None``) this estimates a **cumulative-births bound**: enough slots to
hold every agent that will ever be born. When ``deathrates`` are provided, the estimate becomes
a **peak-living bound** appropriate for simulations that reclaim dead-agent slots via
`LaserFrame.squash`. The peak-living bound credits a configurable fraction of the projected
deaths against births so the returned capacity holds the maximum simultaneous living population
rather than every agent ever alive.
Args:
birthrates (np.ndarray): 2D array of shape (nsteps, nnodes) representing birthrates (CBR)
per 1,000 individuals per year. Each value in ``[0, 100]``.
initial_pop (np.ndarray): 1D array of length nnodes representing the initial population at
each node. Non-negative.
safety_factor (float): Headroom multiplier on the births-driven growth term, ``1 +
safety_factor * (sqrt(exp_mu_t) - 1)``. Must be in ``[0, 6]``. Default 1.0.
deathrates (np.ndarray, optional): 2D array with the **same shape as birthrates**,
representing crude death rates (CDR) per 1,000 individuals per year. Each value in
``[0, 100]``. Default ``None`` (no mortality — same behavior as previous versions).
Keyword-only.
mortality_safety_factor (float): Controls how much of the projected mortality is credited
against births when ``deathrates`` is given. Only ``1 / (1 + mortality_safety_factor)``
of the death sum is credited; the remainder is held back as headroom against a
lower-mortality (faster-growing) realization. Must be in ``[0, 6]``. Default 1.0
(credit half of deaths, hold half back). Set to ``0.0`` to credit deaths fully (tightest
bound). Ignored when ``deathrates`` is ``None``. Keyword-only.
Returns:
np.ndarray: 1D ``uint32`` array of length ``nnodes`` representing the estimated required
capacity (number of agents) at each node. Per-node estimates that exceed the maximum
representable ``uint32`` value (``2**32 - 1``) are clamped to that maximum before being
cast, so the returned array is always safe to cast or sum without overflow. When
``deathrates`` is provided, each estimate is also floored at the corresponding
``initial_pop`` value — a net-shrinking projection cannot report below the starting
population (whose peak living count occurs at t=0).
"""
# Validate birthrates shape against initial_pop shape
_, nnodes = birthrates.shape
assert len(initial_pop) == nnodes, f"Number of nodes in birthrates ({nnodes}) and initial_pop length ({len(initial_pop)}) must match"
assert np.all(
initial_pop >= 0
), f"Initial populations must be >= 0. Found {','.join(map(str, initial_pop[np.where(initial_pop < 0)]))} in initial_pop"
# Validate birthrates values, must be >= 0 and <= 100
assert np.all(birthrates >= 0.0), "All birthrate values must be non-negative"
assert np.all(birthrates <= 100.0), "All birthrate values must be less than or equal to 100"
# Validate safety_factor
assert 0 <= safety_factor <= 6, f"safety_factor must be between 0 and 6, got {safety_factor}"
# Validate mortality_safety_factor unconditionally for consistency with safety_factor —
# the actual mortality math only fires when deathrates is provided.
assert 0 <= mortality_safety_factor <= 6, f"mortality_safety_factor must be between 0 and 6, got {mortality_safety_factor}"
# Validate deathrates if provided
if deathrates is not None:
assert deathrates.shape == birthrates.shape, f"deathrates shape {deathrates.shape} must match birthrates shape {birthrates.shape}"
assert np.all(deathrates >= 0.0), "All deathrate values must be non-negative"
assert np.all(deathrates <= 100.0), "All deathrate values must be less than or equal to 100"
# Convert CBR to daily growth rate
# CBR = births per 1,000 individuals per year
# CBR / 1000 = births per individual per year
# Growth = (1 + CBR / 1000) per individual per year
# Daily growth = (1 + CBR / 1000) ^ (1/365)
# Daily growth rate = (1 + CBR / 1000) ^ (1/365) - 1
# Note, sticking with "lamda" here a) for consistency with modern Greek and Unicode, b) to avoid confusion with Python keyword "lambda"
lamda = (1.0 + birthrates / 1000) ** (1.0 / 365) - 1.0
# Geometric Brownian motion approximation for population growth (https://en.wikipedia.org/wiki/Geometric_Brownian_motion#Properties)
# E(P_t) = P_0 * exp(mu * t)
# where mu is the daily growth rate and t is the number of time steps (days)
# Since we may have time-varying rates, add up daily growth rates over all time steps
# Consider alternative: np.prod(1 + lamda, axis=0)
# For 0 <= CBR <= 40, difference is negligible (< 1:1e6)
exp_mu_t = np.exp(lamda.sum(axis=0))
safety_multiplier = 1 + safety_factor * (np.sqrt(exp_mu_t) - 1)
estimates = initial_pop * safety_multiplier * exp_mu_t
if deathrates is not None:
# Daily per-individual death rates, same Geometric Brownian approximation as for births.
lamda_d = (1.0 + deathrates / 1000) ** (1.0 / 365) - 1.0
# Underestimate mortality: credit only `1 / (1 + mortality_safety_factor)` of the death
# rate against births and hold the rest back as headroom. mortality_safety_factor=0
# credits deaths fully (tightest bound); larger values credit fewer deaths (looser bound).
death_credit = 1.0 / (1.0 + mortality_safety_factor)
# Peak-across-time net growth: the *end-of-sim* cumulative exponent
# (sum_b - death_credit*sum_d) can be smaller than the cumulative exponent reached at an
# intermediate tick. An early high-CBR-low-CDR window followed by a later high-CDR-low-CBR
# window peaks well above either endpoint. The capacity bound must hold the peak living
# population, so we take the max across time of the cumulative net exponent — floored at 0
# because growth begins at exp(0) = 1 (the initial population is always representable).
net_daily = lamda - death_credit * lamda_d
# Inner floor: clamp the cumulative-exponent peak at 0 (growth begins at exp(0) = 1 —
# the initial population is always representable). For net-shrinking scenarios this
# PRESERVES the births-variance headroom from safety_multiplier (estimates ≈
# initial_pop * safety_multiplier), which is the semantically correct bound.
peak_exponent = np.maximum(np.max(np.cumsum(net_daily, axis=0), axis=0), 0.0)
# Note: safety_multiplier stays anchored to the births-only end-of-sim growth (exp_mu_t),
# which equals the births-only peak because CBR >= 0 makes that cumulative monotonic.
estimates = initial_pop * safety_multiplier * np.exp(peak_exponent)
# Outer floor: defensive backstop at initial_pop. Intentionally redundant with the
# inner floor — they only differ when safety_factor > 0 (the inner floor keeps the
# multiplier's headroom, the outer floor would strip it). Both are kept so a future
# refactor that drops either one in isolation still won't under-allocate.
estimates = np.maximum(estimates, initial_pop)
# Clamp to uint32.max before casting so per-node values that overflow uint32 saturate
# at the maximum representable value rather than wrapping around.
estimates = np.round(estimates)
estimates = np.minimum(estimates, np.iinfo(np.uint32).max).astype(np.uint32)
return estimates
[docs]
def grid(M=5, N=5, node_size_degs=0.08983, population_fn=None, origin_x=0, origin_y=0, states=None):
"""
Create an MxN grid of cells anchored at (lat, long) with populations and geometries.
By default all nodes are initialized with the full population in the first state (e.g., "S" for susceptible).
Args:
M (int): Number of rows (north-south).
N (int): Number of columns (east-west).
node_size_degs (float): Size of each cell in decimal degrees (default 0.08983 ≈ 10km at the equator).
population_fn (callable): Function(row, col) returning population for a cell. Default is uniform random between 1,000 and 100,000.
origin_x (float): longitude of the origin in decimal degrees (bottom-left corner) -180 <= origin_x < 180.
origin_y (float): latitude of the origin in decimal degrees (bottom-left corner) -90 <= origin_y < 90.
states (list): List of state names to initialize in the GeoDataFrame. Default is ["S", "E", "I", "R"].
Returns:
scenario (GeoDataFrame): Columns are nodeid, population, geometry.
"""
if M < 1:
raise ValueError("M must be >= 1")
if N < 1:
raise ValueError("N must be >= 1")
if node_size_degs <= 0:
raise ValueError("node_size_degs must be > 0")
if node_size_degs > 1.0:
raise ValueError("node_size_degs must be <= 1.0")
if not (-180 <= origin_x < 180):
raise ValueError("origin_x must be -180 <= origin_x < 180")
if not (-90 <= origin_y < 90):
raise ValueError("origin_y must be -90 <= origin_y < 90")
if population_fn is None:
def population_fn(row: int, col: int) -> int:
return int(np.random.uniform(1_000, 100_000))
states = states or ["S", "E", "I", "R"]
cells = []
nodeid = 0
for row in range(M):
for col in range(N):
# TODO - use latitude sensitive conversion of km to degrees
x0 = origin_x + col * node_size_degs
y0 = origin_y + row * node_size_degs
x1 = x0 + node_size_degs
y1 = y0 + node_size_degs
poly = Polygon(
[
(x0, y0), # SW
(x1, y0), # SE
(x1, y1), # NE
(x0, y1), # NW
(x0, y0), # Close polygon in SW
]
)
population = int(population_fn(row, col))
if population < 0:
raise ValueError(f"population_fn returned negative population {population} for row {row}, col {col}")
cells.append({"nodeid": nodeid, "population": population, "geometry": poly})
nodeid += 1
gdf = gpd.GeoDataFrame(cells, columns=["nodeid", "population", "geometry"], crs="EPSG:4326")
for state in states:
gdf[state] = 0
gdf[states[0]] = gdf.population # All state[0] (susceptible?) by default
return gdf
[docs]
def initialize_population(grid, initial: list | np.ndarray, states=None):
"""
Initialize the population states in the grid based on the initial state counts provided.
Provide integer values to set the exact counts for each state at each node.
Alternatively, provide fractional values between 0.0 and 1.0 to set proportions of the population for each state at each node.
In the latter case, the first state in the states list will be computed as the remainder of the population after assigning the other states.
Args:
grid (GeoDataFrame): The grid GeoDataFrame with population and state columns.
initial (list or np.ndarray): A list or array of shape (1|nnodes, nstates) representing the initial counts for each state at each node. If the shape is (1, nstates), the same initial state distribution will be applied to all nodes.
states (list): List of state names corresponding to the columns in the grid. Default is ["S", "E", "I", "R"].
Returns:
GeoDataFrame: The updated grid with initialized population states.
"""
states = states or ["S", "E", "I", "R"]
nnodes = len(grid)
nstates = len(states)
if isinstance(initial, list):
initial = np.array(initial)
# Convert 1D array to 2D by reshaping to (1, nstates)
if len(initial.shape) == 1:
if initial.shape[0] != nstates:
raise ValueError(f"Initial state array with shape {initial.shape} does not match expected number of states ({nstates})")
initial = initial.reshape(1, nstates)
elif len(initial.shape) != 2:
raise ValueError(f"Initial state array must be 1D or 2D, got shape {initial.shape}")
if initial.shape[0] == 1:
# Broadcast single row to all nodes
initial = np.broadcast_to(initial, (nnodes, nstates))
if initial.shape != (nnodes, nstates):
raise ValueError(f"Initial state array shape {initial.shape} does not match expected shape ({nnodes}, {nstates})")
# If all values are integral values, use them as counts
total = np.zeros((nnodes,), dtype=np.int32)
if np.all(np.mod(initial, 1) == 0):
initial = initial.astype(np.int32)
for index, state in enumerate(states):
grid[state] = initial[:, index]
total += initial[:, index]
assert np.all(total == grid.population), "Sum of initial states does not equal population at some nodes"
elif np.all((initial >= 0.0) & (initial <= 1.0)):
# If any rows sum to > 1.0, raise error
row_sums = initial.sum(axis=1)
if np.any(row_sums > 1.0):
raise ValueError("Initial state proportions sum to more than 1.0 at some nodes")
# Handle fractional values as proportions
for index, state in enumerate(states):
if index == 0:
continue # Susceptible will be computed as remainder
grid[state] = np.round(initial[:, index] * grid.population).astype(np.int32)
total += grid[state]
grid[states[0]] = grid.population - total
# Double check no negatives in the remainder state.
if np.any(grid[states[0]] < 0):
raise ValueError(f"Computed {states[0]} counts are negative at some nodes")
else:
raise ValueError("Initial state array must contain either all integer counts or all proportions between 0.0 and 1.0")
return grid