laser.measles.utils

laser.measles.utils

This module provides utility functions for the laser-measles project.

Functions:

Name	Description
calc_distances	np.ndarray, longitudes: np.ndarray, verbose: bool = False) -> np.ndarray: Calculate the pairwise distances between points given their latitudes and longitudes.
calc_capacity	np.uint32, nticks: np.uint32, cbr: np.float32, verbose: bool = False) -> np.uint32: Calculate the population capacity after a given number of ticks based on a constant birth rate.
seed_infections_randomly	int = 100) -> None: Seed initial infections in random locations at the start of the simulation.
seed_infections_in_patch	int, ninfections: int = 100) -> None: Seed initial infections in a specific location at the start of the simulation.
cast_type	Cast a value to a specified data type.
select_implementation	bool = True) -> callable: Select between numpy and numba implementations based on availability and preference.
dual_implementation	callable, numba_func: callable) -> callable: Decorator to create function selector that chooses between numpy and numba implementations.

laser.measles.utils.StateArray

Bases: ndarray

A numpy array wrapper that provides attribute access to state compartments.

This class allows accessing state compartments by name (e.g., states.S, states.I, states.R) while maintaining full numpy array functionality and backward compatibility with numeric indexing (e.g., states[0], states[1]).

Example

states = StateArray(source_array=np.zeros((3, 100)), state_names=["S", "I", "R"], state_axis=0) states.S[0] = 1000 # Set susceptible population in patch 0 prevalence = states.I / states.sum(axis=0) # Calculate prevalence states[0] += births # Numeric indexing still works N = states.sum(axis=states.state_axis) # Sum over state axis to get total population per patch

laser.measles.utils.StateArray.state_axis property

Get the axis index for the state compartments.

laser.measles.utils.StateArray.state_names property

Return the list of state compartment names.

laser.measles.utils.StateArray.get_state_index(name)

Get the numeric index for a state compartment name.

laser.measles.utils.assert_row_vector(vec)

Assert that vec is a scalar, 1-D array, or 2-D row vector.

laser.measles.utils.calc_capacity(population, nticks, cbr, verbose=False)

Calculate the population capacity after a given number of ticks based on a constant birth rate (CBR).

Args:

population (np.uint32): The initial population.
nticks (np.uint32): The number of ticks (time steps) to simulate.
cbr (np.float32): The constant birth rate per 1000 people per year.
verbose (bool, optional): If True, prints detailed population growth information. Defaults to False.

Returns:

np.uint32: The estimated population capacity after the given number of ticks.

laser.measles.utils.calc_distances(latitudes, longitudes, verbose=False)

Calculate the pairwise distances between points given their latitudes and longitudes.

Parameters:

latitudes (np.ndarray): A 1-dimensional array of latitudes.
longitudes (np.ndarray): A 1-dimensional array of longitudes with the same shape as latitudes.
verbose (bool, optional): If True, prints the upper left corner of the distance matrix. Default is False.

Returns:

np.ndarray: A 2-dimensional array where the element at [i, j] represents the distance in km between the i-th and j-th points.

Raises:

AssertionError: If latitudes is not 1-dimensional or if latitudes and longitudes do not have the same shape.

laser.measles.utils.cast_type(a, dtype, round=False)

Cast a value to a specified data type. Note that this casting truncates by default.

Parameters:

Name	Type	Description	Default
a	Any	The value to cast	required
dtype	DTypeLike	The target data type	required

Returns:

Type	Description
Any	The value cast to the specified data type

laser.measles.utils.coerce_utf8_date_column(df, date_column)

Cast a Utf8 (string) date column on a polars DataFrame to Datetime.

Used by SIACalendarParams validators across all three model variants to silently fix the laser-measles #215 footgun: users build the SIA schedule with date strings (e.g. ["2024-06-15"]), polars defaults to Utf8, and the per-tick filter in SIACalendarProcess.__call__ later raises an opaque InvalidOperationError when polars compares the string column to a datetime.datetime.

Parameters:

Name	Type	Description	Default
df	DataFrame	The polars DataFrame holding the schedule.	required
date_column	str	Name of the column to coerce.	required

Returns:

Type	Description
DataFrame	The same DataFrame if the column is missing or already typed; a new
DataFrame	DataFrame with the named column cast to Datetime if it was Utf8.
DataFrame	Mis-formatted date strings raise a clear ValueError from polars
DataFrame	at this call rather than at a much later filter site.

The format pinning ("%Y-%m-%d") is intentional — schedules with non-ISO date strings should fail loud at coercion time rather than being silently misinterpreted.

laser.measles.utils.dual_implementation(numpy_func, numba_func)

Decorator to create function selector that chooses between numpy and numba implementations.

Parameters:

Name	Type	Description	Default
numpy_func	Callable	The numpy implementation function.	required
numba_func	Callable	The numba implementation function.	required

Returns:

Type	Description
Callable	A wrapper function that selects the appropriate implementation.

laser.measles.utils.get_laserframe_properties(laserframe)

Get the scalar and vector properties of a laserframe that are numpy arrays.

laser.measles.utils.seed_infections_in_patch(model, ipatch, ninfections=100)

Seed initial infections in a specific patch of the population at the start of the simulation. This function randomly selects individuals from the specified patch and sets their infection timer to the mean infection duration, effectively marking them as infected. The process continues until the desired number of initial infections is reached.

Args:

model: The simulation model containing the population and parameters.
ipatch (int): The identifier of the patch where infections should be seeded.
ninfections (int, optional): The number of initial infections to seed. Defaults to 100.

Returns:

None

laser.measles.utils.seed_infections_randomly(model, ninfections=100)

Seed initial infections in random locations at the start of the simulation. This function randomly selects individuals from the population and seeds them with an infection, based on the specified number of initial infections.

Args:

model: The simulation model containing the population and parameters.
ninfections (int, optional): The number of initial infections to seed.
                             Defaults to 100.

Returns:

None

laser.measles.utils.select_implementation(numpy_func, numba_func, use_numba=True)

Select between numpy and numba implementations based on availability and preference.

Parameters:

Name	Type	Description	Default
numpy_func	Callable	The numpy implementation function.	required
numba_func	Callable	The numba implementation function.	required
use_numba	bool	Whether to prefer numba implementation if available.	True

Returns:

Type	Description
Callable	The selected function implementation.