population_state Module

Population state management and data structures.

Overview

The population_state module defines the runtime state of population simulations, including individual tracking, age structure, and genetic composition.

Complete Module Reference

natal.population_state

Population state containers based on NamedTuple.

These containers keep scalar metadata immutable while allowing in-place mutation of NumPy array contents, which remains compatible with Numba kernels.

PopulationState

Bases: NamedTuple

Age-structured state container.

Scalars are immutable (use _replace to rebuild); array values remain mutable in-place.

Attributes:

Name	Type	Description
n_tick	int	Current simulation time step.
individual_count	NDArray[float64]	Array of shape (n_sexes, n_ages, n_genotypes) – counts of individuals per sex, age, and genotype.
sperm_storage	NDArray[float64]	Array of shape (n_ages, n_genotypes, n_genotypes) – stored sperm counts per female age, female genotype, and male genotype.

create classmethod

create(n_genotypes: int, n_sexes: Optional[int] = None, n_ages: int = 2, n_tick: int = 0, individual_count: Optional[NDArray[float64]] = None, sperm_storage: Optional[NDArray[float64]] = None) -> PopulationState

Create a PopulationState with optionally provided arrays.

If arrays are not provided, they are initialised to zeros.

Parameters:

Name	Type	Description	Default
n_genotypes	int	Number of diploid genotype types.	required
n_sexes	Optional[int]	Number of sexes (defaults to 2 if not given).	None
n_ages	int	Number of age classes (default 2).	2
n_tick	int	Initial tick value (default 0).	0
individual_count	Optional[NDArray[float64]]	Optional array (n_sexes, n_ages, n_genotypes).	None
sperm_storage	Optional[NDArray[float64]]	Optional array (n_ages, n_genotypes, n_genotypes).	None

Returns:

Type	Description
PopulationState	A new PopulationState instance.

Raises:

Type	Description
AssertionError	If dimensions are invalid or provided arrays have wrong shape.

Source code in src/natal/population_state.py

@classmethod
def create(
    cls,
    n_genotypes: int,
    n_sexes: Optional[int] = None,
    n_ages: int = 2,
    n_tick: int = 0,
    individual_count: Optional[NDArray[np.float64]] = None,
    sperm_storage: Optional[NDArray[np.float64]] = None,
) -> PopulationState:
    """Create a PopulationState with optionally provided arrays.

    If arrays are not provided, they are initialised to zeros.

    Args:
        n_genotypes: Number of diploid genotype types.
        n_sexes: Number of sexes (defaults to 2 if not given).
        n_ages: Number of age classes (default 2).
        n_tick: Initial tick value (default 0).
        individual_count: Optional array (n_sexes, n_ages, n_genotypes).
        sperm_storage: Optional array (n_ages, n_genotypes, n_genotypes).

    Returns:
        A new PopulationState instance.

    Raises:
        AssertionError: If dimensions are invalid or provided arrays have wrong shape.
    """
    if n_sexes is None:
        n_sexes = 2
    assert n_genotypes > 0, "n_genotypes must be positive"
    assert n_ages > 0, "n_ages must be positive"
    assert n_tick >= 0, "n_tick must be non-negative"

    if individual_count is None:
        ind = np.zeros((n_sexes, n_ages, n_genotypes), dtype=np.float64)
    else:
        expected_shape = (n_sexes, n_ages, n_genotypes)
        assert individual_count.shape == expected_shape, (
            f"Invalid shape for individual_count: expected {expected_shape}, got {individual_count.shape}"
        )
        ind = individual_count.astype(np.float64)

    if sperm_storage is None:
        sperm = np.zeros((n_ages, n_genotypes, n_genotypes), dtype=np.float64)
    else:
        expected_shape = (n_ages, n_genotypes, n_genotypes)
        assert sperm_storage.shape == expected_shape, (
            f"Invalid shape for sperm_storage: expected {expected_shape}, got {sperm_storage.shape}"
        )
        sperm = sperm_storage.astype(np.float64)

    return cls(n_tick=int(n_tick), individual_count=ind, sperm_storage=sperm)

get_count

get_count(sex: int, age: int, genotype_index: int) -> float

Retrieve the count of individuals for a specific category.

Parameters:

Name	Type	Description	Default
sex	int	Sex index.	required
age	int	Age class index.	required
genotype_index	int	Diploid genotype index.	required

Returns:

Type	Description
float	The count (float).

Source code in src/natal/population_state.py

def get_count(self, sex: int, age: int, genotype_index: int) -> float:
    """Retrieve the count of individuals for a specific category.

    Args:
        sex: Sex index.
        age: Age class index.
        genotype_index: Diploid genotype index.

    Returns:
        The count (float).
    """
    return self.individual_count[sex, age, genotype_index]

add_count

add_count(sex: int, age: int, genotype_index: int, count: float) -> None

Add to the count of individuals for a specific category.

Parameters:

Name	Type	Description	Default
sex	int	Sex index.	required
age	int	Age class index.	required
genotype_index	int	Diploid genotype index.	required
count	float	Amount to add (can be negative).	required

Source code in src/natal/population_state.py

def add_count(self, sex: int, age: int, genotype_index: int, count: float) -> None:
    """Add to the count of individuals for a specific category.

    Args:
        sex: Sex index.
        age: Age class index.
        genotype_index: Diploid genotype index.
        count: Amount to add (can be negative).
    """
    self.individual_count[sex, age, genotype_index] += count

set_count

set_count(sex: int, age: int, genotype_index: int, count: float) -> None

Set the count of individuals for a specific category.

Parameters:

Name	Type	Description	Default
sex	int	Sex index.	required
age	int	Age class index.	required
genotype_index	int	Diploid genotype index.	required
count	float	New count.	required

Source code in src/natal/population_state.py

def set_count(self, sex: int, age: int, genotype_index: int, count: float) -> None:
    """Set the count of individuals for a specific category.

    Args:
        sex: Sex index.
        age: Age class index.
        genotype_index: Diploid genotype index.
        count: New count.
    """
    self.individual_count[sex, age, genotype_index] = count

get_stored_sperm

get_stored_sperm(age: int, female_genotype_index: int, male_genotype_index: int) -> float

Retrieve stored sperm count for a given combination.

Parameters:

Name	Type	Description	Default
age	int	Age class of the female.	required
female_genotype_index	int	Female genotype index.	required
male_genotype_index	int	Male genotype index.	required

Returns:

Type	Description
float	Stored sperm count.

Source code in src/natal/population_state.py

def get_stored_sperm(self, age: int, female_genotype_index: int, male_genotype_index: int) -> float:
    """Retrieve stored sperm count for a given combination.

    Args:
        age: Age class of the female.
        female_genotype_index: Female genotype index.
        male_genotype_index: Male genotype index.

    Returns:
        Stored sperm count.
    """
    return self.sperm_storage[age, female_genotype_index, male_genotype_index]

set_stored_sperm

set_stored_sperm(age: int, female_genotype_index: int, male_genotype_index: int, count: float) -> None

Add to stored sperm count (in‑place addition).

Parameters:

Name	Type	Description	Default
age	int	Age class of the female.	required
female_genotype_index	int	Female genotype index.	required
male_genotype_index	int	Male genotype index.	required
count	float	Amount to add (can be negative).	required

Source code in src/natal/population_state.py

def set_stored_sperm(self, age: int, female_genotype_index: int, male_genotype_index: int, count: float) -> None:
    """Add to stored sperm count (in‑place addition).

    Args:
        age: Age class of the female.
        female_genotype_index: Female genotype index.
        male_genotype_index: Male genotype index.
        count: Amount to add (can be negative).
    """
    self.sperm_storage[age, female_genotype_index, male_genotype_index] += count

flatten_all

flatten_all() -> NDArray[np.float64]

Flatten the entire state into a single 1D array.

The order is: tick, then individual_count flattened (row‑major), then sperm_storage flattened.

Returns:

Type	Description
NDArray[float64]	1D array of floats.

Source code in src/natal/population_state.py

def flatten_all(self) -> NDArray[np.float64]:
    """Flatten the entire state into a single 1D array.

    The order is: tick, then individual_count flattened (row‑major),
    then sperm_storage flattened.

    Returns:
        1D array of floats.
    """
    tick_arr = np.array([float(self.n_tick)], dtype=np.float64)
    return np.concatenate((tick_arr, self.individual_count.flatten(), self.sperm_storage.flatten()))

DiscretePopulationState

Bases: NamedTuple

Discrete‑generation state container (no sperm storage).

Attributes:

Name	Type	Description
n_tick	int	Current simulation time step.
individual_count	NDArray[float64]	Array of shape (n_sexes, n_ages, n_genotypes) – counts of individuals per sex, age, and genotype.

create classmethod

create(n_sexes: int, n_ages: int, n_genotypes: int, n_tick: int = 0, individual_count: Optional[NDArray[float64]] = None) -> DiscretePopulationState

Create a DiscretePopulationState with optionally provided array.

Parameters:

Name	Type	Description	Default
n_sexes	int	Number of sexes.	required
n_ages	int	Number of age classes.	required
n_genotypes	int	Number of diploid genotype types.	required
n_tick	int	Initial tick value (default 0).	0
individual_count	Optional[NDArray[float64]]	Optional array (n_sexes, n_ages, n_genotypes); if None, filled with zeros.	None

Returns:

Type	Description
DiscretePopulationState	A new DiscretePopulationState instance.

Raises:

Type	Description
AssertionError	If dimensions are invalid or array shape mismatch.

Source code in src/natal/population_state.py

@classmethod
def create(
    cls,
    n_sexes: int,
    n_ages: int,
    n_genotypes: int,
    n_tick: int = 0,
    individual_count: Optional[NDArray[np.float64]] = None,
) -> DiscretePopulationState:
    """Create a DiscretePopulationState with optionally provided array.

    Args:
        n_sexes: Number of sexes.
        n_ages: Number of age classes.
        n_genotypes: Number of diploid genotype types.
        n_tick: Initial tick value (default 0).
        individual_count: Optional array (n_sexes, n_ages, n_genotypes);
            if None, filled with zeros.

    Returns:
        A new DiscretePopulationState instance.

    Raises:
        AssertionError: If dimensions are invalid or array shape mismatch.
    """
    assert n_sexes > 0, "n_sexes must be positive"
    assert n_ages > 0, "n_ages must be positive"
    assert n_genotypes > 0, "n_genotypes must be positive"
    assert n_tick >= 0, "n_tick must be non-negative"

    if individual_count is None:
        ind = np.zeros((n_sexes, n_ages, n_genotypes), dtype=np.float64)
    else:
        expected_shape = (n_sexes, n_ages, n_genotypes)
        assert individual_count.shape == expected_shape, (
            f"Invalid shape for individual_count: expected {expected_shape}, got {individual_count.shape}"
        )
        ind = individual_count.astype(np.float64)

    return cls(n_tick=int(n_tick), individual_count=ind)

flatten_all

flatten_all() -> NDArray[np.float64]

Flatten the entire state into a single 1D array.

The order is: tick, then individual_count flattened (row‑major).

Returns:

Type	Description
NDArray[float64]	1D array of floats.

Source code in src/natal/population_state.py

def flatten_all(self) -> NDArray[np.float64]:
    """Flatten the entire state into a single 1D array.

    The order is: tick, then individual_count flattened (row‑major).

    Returns:
        1D array of floats.
    """
    tick_arr = np.array([float(self.n_tick)], dtype=np.float64)
    return np.concatenate((tick_arr, self.individual_count.flatten()))

to_plain_population_state

to_plain_population_state(state: PopulationState, copy: bool = True) -> PlainPopulationState

Convert a PopulationState to a plain (copied) instance.

Parameters:

Name	Type	Description	Default
state	PopulationState	Input PopulationState.	required
copy	bool	If True, arrays are deep‑copied; otherwise they are referenced.	True

Returns:

Type	Description
PlainPopulationState	A new PopulationState (or the same arrays if copy=False).

Source code in src/natal/population_state.py

def to_plain_population_state(state: PopulationState, copy: bool = True) -> PlainPopulationState:
    """Convert a PopulationState to a plain (copied) instance.

    Args:
        state: Input PopulationState.
        copy: If True, arrays are deep‑copied; otherwise they are referenced.

    Returns:
        A new PopulationState (or the same arrays if copy=False).
    """
    ind = state.individual_count.copy() if copy else state.individual_count
    sperm = state.sperm_storage.copy() if copy else state.sperm_storage
    return PopulationState(n_tick=int(state.n_tick), individual_count=ind, sperm_storage=sperm)

to_plain_discrete_population_state

to_plain_discrete_population_state(state: DiscretePopulationState, copy: bool = True) -> PlainDiscretePopulationState

Convert a DiscretePopulationState to a plain (copied) instance.

Parameters:

Name	Type	Description	Default
state	DiscretePopulationState	Input DiscretePopulationState.	required
copy	bool	If True, the array is deep‑copied; otherwise it is referenced.	True

Returns:

Type	Description
PlainDiscretePopulationState	A new DiscretePopulationState (or the same array if copy=False).

Source code in src/natal/population_state.py

def to_plain_discrete_population_state(
    state: DiscretePopulationState,
    copy: bool = True,
) -> PlainDiscretePopulationState:
    """Convert a DiscretePopulationState to a plain (copied) instance.

    Args:
        state: Input DiscretePopulationState.
        copy: If True, the array is deep‑copied; otherwise it is referenced.

    Returns:
        A new DiscretePopulationState (or the same array if copy=False).
    """
    ind = state.individual_count.copy() if copy else state.individual_count
    return DiscretePopulationState(n_tick=int(state.n_tick), individual_count=ind)

from_plain_population_state

from_plain_population_state(plain: PlainPopulationState) -> PopulationState

Convert a plain PopulationState back (arrays are referenced, not copied).

Source code in src/natal/population_state.py

def from_plain_population_state(plain: PlainPopulationState) -> PopulationState:
    """Convert a plain PopulationState back (arrays are referenced, not copied)."""
    return PopulationState(
        n_tick=int(plain.n_tick),
        individual_count=plain.individual_count,
        sperm_storage=plain.sperm_storage,
    )

from_plain_discrete_population_state

from_plain_discrete_population_state(plain: PlainDiscretePopulationState) -> DiscretePopulationState

Convert a plain DiscretePopulationState back (array is referenced).

Source code in src/natal/population_state.py

def from_plain_discrete_population_state(plain: PlainDiscretePopulationState) -> DiscretePopulationState:
    """Convert a plain DiscretePopulationState back (array is referenced)."""
    return DiscretePopulationState(
        n_tick=int(plain.n_tick),
        individual_count=plain.individual_count,
    )

parse_flattened_state

parse_flattened_state(flat_array: NDArray[float64], n_sexes: Union[int, integer], n_ages: Union[int, integer], n_genotypes: Union[int, integer], copy: bool = True) -> PopulationState

Reconstruct a PopulationState from a flattened array.

The flattened array must be in the format produced by flatten_all().

Parameters:

Name	Type	Description	Default
flat_array	NDArray[float64]	1D array containing tick, individual_count, sperm_storage.	required
n_sexes	Union[int, integer]	Number of sexes.	required
n_ages	Union[int, integer]	Number of age classes.	required
n_genotypes	Union[int, integer]	Number of diploid genotype types.	required
copy	bool	If True, arrays are deep‑copied; otherwise they are viewed.	True

Returns:

Type	Description
PopulationState	A PopulationState instance.

Source code in src/natal/population_state.py

def parse_flattened_state(
    flat_array: NDArray[np.float64],
    n_sexes: Union[int, np.integer],
    n_ages: Union[int, np.integer],
    n_genotypes: Union[int, np.integer],
    copy: bool = True,
) -> PopulationState:
    """Reconstruct a PopulationState from a flattened array.

    The flattened array must be in the format produced by ``flatten_all()``.

    Args:
        flat_array: 1D array containing tick, individual_count, sperm_storage.
        n_sexes: Number of sexes.
        n_ages: Number of age classes.
        n_genotypes: Number of diploid genotype types.
        copy: If True, arrays are deep‑copied; otherwise they are viewed.

    Returns:
        A PopulationState instance.
    """
    n_tick = int(flat_array[0])
    end = 1 + n_sexes * n_ages * n_genotypes
    individual_count = flat_array[1:end].reshape((n_sexes, n_ages, n_genotypes))
    sperm_storage = flat_array[end:].reshape((n_ages, n_genotypes, n_genotypes))

    if copy:
        individual_count = individual_count.copy()
        sperm_storage = sperm_storage.copy()

    return PopulationState(
        n_tick=n_tick,
        individual_count=individual_count,
        sperm_storage=sperm_storage,
    )

parse_flattened_discrete_state

parse_flattened_discrete_state(flat_array: NDArray[float64], n_sexes: Union[int, integer], n_ages: Union[int, integer], n_genotypes: Union[int, integer], copy: bool = True) -> DiscretePopulationState

Reconstruct a DiscretePopulationState from a flattened array.

The flattened array must be in the format produced by flatten_all().

Parameters:

Name	Type	Description	Default
flat_array	NDArray[float64]	1D array containing tick and individual_count.	required
n_sexes	Union[int, integer]	Number of sexes.	required
n_ages	Union[int, integer]	Number of age classes.	required
n_genotypes	Union[int, integer]	Number of diploid genotype types.	required
copy	bool	If True, the array is deep‑copied; otherwise it is viewed.	True

Returns:

Type	Description
DiscretePopulationState	A DiscretePopulationState instance.

Source code in src/natal/population_state.py

def parse_flattened_discrete_state(
    flat_array: NDArray[np.float64],
    n_sexes: Union[int, np.integer],
    n_ages: Union[int, np.integer],
    n_genotypes: Union[int, np.integer],
    copy: bool = True,
) -> DiscretePopulationState:
    """Reconstruct a DiscretePopulationState from a flattened array.

    The flattened array must be in the format produced by ``flatten_all()``.

    Args:
        flat_array: 1D array containing tick and individual_count.
        n_sexes: Number of sexes.
        n_ages: Number of age classes.
        n_genotypes: Number of diploid genotype types.
        copy: If True, the array is deep‑copied; otherwise it is viewed.

    Returns:
        A DiscretePopulationState instance.
    """
    n_tick = int(flat_array[0])
    individual_count = flat_array[1:].reshape((n_sexes, n_ages, n_genotypes))

    if copy:
        individual_count = individual_count.copy()

    return DiscretePopulationState(
        n_tick=n_tick,
        individual_count=individual_count,
    )