discrete_generation_population Module

Discrete generation population simulation model.

Overview

The discrete_generation_population module provides a population implementation without age structure, suitable for simulations that evolve in whole-generation steps.

Complete Module Reference

natal.discrete_generation_population

Discrete-generation population model.

This module provides a lightweight non-overlapping generation model that keeps n_ages=2: - age 0: offspring/zygotes produced in current tick - age 1: reproducing adults

The simulation flow remains split as: first hook -> reproduction -> early hook -> survival -> late hook -> aging

DiscreteGenerationPopulation

DiscreteGenerationPopulation(species: Species, population_config: PopulationConfig, name: Optional[str] = None, initial_individual_count: Optional[Dict[str, Dict[Union[Genotype, str], Union[List[int], Dict[int, int], int, float]]]] = None, hooks: Optional[Dict[str, List[Tuple[Any, Optional[str], Optional[int]]]]] = None)

Bases: BasePopulation[DiscretePopulationState]

Population with strict non-overlapping generations.

Maintains exactly two age classes: - age 0: newly produced offspring - age 1: reproducing adults

Attributes:

Name	Type	Description
state	DiscretePopulationState	Current discrete population state.
config	PopulationConfig	Active normalized configuration with two-age layout.
history	List[Tuple[int, ndarray]]	Flattened snapshots indexed by tick.

Source code in src/natal/discrete_generation_population.py

def __init__(
    self,
    species: Species,
    population_config: PopulationConfig,
    name: Optional[str] = None,
    initial_individual_count: Optional[
        Dict[str, Dict[Union[Genotype, str], Union[List[int], Dict[int, int], int, float]]]
    ] = None,
    hooks: Optional[Dict[str, List[Tuple[Any, Optional[str], Optional[int]]]]] = None,
):
    if name is None:
        name = "DiscreteGenerationPop"

    super().__init__(species, name, hooks=hooks or {})

    self._config = self._normalize_config(population_config)

    self._genotypes_list = species.get_all_genotypes()
    self._haploid_genotypes_list = species.get_all_haploid_genotypes()

    self._initialize_registry()

    n_sexes = self._config_nn.n_sexes
    n_genotypes = self._config_nn.n_genotypes
    n_ages = self._config_nn.n_ages

    self._state = DiscretePopulationState.create(
        n_sexes=n_sexes,
        n_ages=n_ages,
        n_genotypes=n_genotypes,
        n_tick=0,
        individual_count=np.zeros((n_sexes, n_ages, n_genotypes), dtype=np.float64),
    )

    cfg_init_ind = self._config_nn.get_scaled_initial_individual_count()
    if cfg_init_ind.shape == self._state_nn.individual_count.shape:
        self._state_nn.individual_count[:] = cfg_init_ind

    self._history_shape = (
        1 + n_sexes * n_ages * n_genotypes,
    )

    if initial_individual_count is not None:
        self._state_nn.individual_count.fill(0.0)
        self._distribute_initial_population(initial_individual_count)

    self._initial_population_snapshot = (
        self._state_nn.individual_count.copy(),
        None,
        None,
    )

    self._finalize_hooks()

setup classmethod

setup(species: Species, name: str = 'DiscreteGenerationPop', stochastic: bool = True, use_continuous_sampling: bool = False, use_fixed_egg_count: bool = False) -> DiscreteGenerationPopulationBuilder

Create and preconfigure a discrete-generation population builder.

This is a convenience forwarding entry point. Parameter semantics and defaults are the same as DiscreteGenerationPopulationBuilder.setup.

Parameters:

Name	Type	Description	Default
species	Species	Species definition used to initialize the builder.	required
name	str	Population name passed through to builder.setup.	'DiscreteGenerationPop'
stochastic	bool	Whether to use stochastic sampling. Passed through to builder.setup.	True
use_continuous_sampling	bool	If True, use Dirichlet; else Binomial/Multinomial sampling. Passed through to builder.setup.	False
use_fixed_egg_count	bool	If True, egg count is fixed; if False, Poisson distributed. Passed through to builder.setup.	False

Returns:

Type	Description
DiscreteGenerationPopulationBuilder	A configured DiscreteGenerationPopulationBuilder for fluent chaining.

Examples:

DiscreteGenerationPopulation.setup(species).initial_state(...).build()

Source code in src/natal/discrete_generation_population.py

@classmethod
def setup(
    cls,
    species: Species,
    name: str = "DiscreteGenerationPop",
    stochastic: bool = True,
    use_continuous_sampling: bool = False,
    use_fixed_egg_count: bool = False,
) -> DiscreteGenerationPopulationBuilder:
    """Create and preconfigure a discrete-generation population builder.

    This is a convenience forwarding entry point. Parameter semantics and
    defaults are the same as ``DiscreteGenerationPopulationBuilder.setup``.

    Args:
        species: Species definition used to initialize the builder.
        name: Population name passed through to ``builder.setup``.
        stochastic: Whether to use stochastic sampling. Passed through to ``builder.setup``.
        use_continuous_sampling: If True, use Dirichlet; else Binomial/Multinomial sampling.
            Passed through to ``builder.setup``.
        use_fixed_egg_count: If True, egg count is fixed; if False, Poisson distributed.
            Passed through to ``builder.setup``.

    Returns:
        A configured ``DiscreteGenerationPopulationBuilder`` for fluent chaining.

    Examples:
        ``DiscreteGenerationPopulation.setup(species).initial_state(...).build()``
    """
    from natal.population_builder import DiscreteGenerationPopulationBuilder

    builder = DiscreteGenerationPopulationBuilder(species)
    builder.setup(
        name=name,
        stochastic=stochastic,
        use_continuous_sampling=use_continuous_sampling,
        use_fixed_egg_count=use_fixed_egg_count,
    )
    return builder

run

run(n_steps: int = 1, record_every: Optional[int] = None, finish: bool = False, clear_history_on_start: bool = False) -> DiscreteGenerationPopulation

Run multi-step evolution using optimized simulation kernels.

Parameters:

Name	Type	Description	Default
n_steps	int	Number of steps to evolve.	1
record_every	Optional[int]	Interval for recording snapshots. If None, uses self.record_every. If 0, no snapshots are recorded.	None
finish	bool	Whether to mark the population as finished after the run.	False
clear_history_on_start	bool	Whether to clear existing history before starting.	False

Returns:

Name	Type	Description
DiscreteGenerationPopulation	DiscreteGenerationPopulation	Self for chaining.

Raises:

Type	Description
RuntimeError	If the population is already finished and cannot continue.

Source code in src/natal/discrete_generation_population.py

def run(
    self,
    n_steps: int = 1,
    record_every: Optional[int] = None,
    finish: bool = False,
    clear_history_on_start: bool = False,
) -> DiscreteGenerationPopulation:
    """Run multi-step evolution using optimized simulation kernels.

    Args:
        n_steps: Number of steps to evolve.
        record_every: Interval for recording snapshots.
            If None, uses self.record_every. If 0, no snapshots are recorded.
        finish: Whether to mark the population as finished after the run.
        clear_history_on_start: Whether to clear existing history before starting.

    Returns:
        DiscreteGenerationPopulation: Self for chaining.

    Raises:
        RuntimeError: If the population is already finished and cannot continue.
    """
    if self._finished:
        raise RuntimeError(
            f"Population '{self.name}' has finished. "
            "Cannot run() again after finish=True."
        )

    if record_every is None:
        record_every = self.record_every

    if self.should_use_python_dispatch():
        return self._run_python_dispatch(
            n_steps=n_steps,
            record_every=record_every,
            finish=finish,
            clear_history_on_start=clear_history_on_start,
        )

    hooks = self.get_compiled_event_hooks()

    assert hooks.registry is not None, "hooks.registry should always be initialized"

    obs_mask = self._observation_mask
    n_obs = len(self._observation.labels) if self._observation is not None else 0

    if hooks.run_discrete_fn is not None:
        final_state_tuple, history_new, was_stopped = hooks.run_discrete_fn(
            state=self._state_nn,
            config=self._config_nn,
            registry=hooks.registry,
            n_ticks=n_steps,
            record_interval=record_every,
            observation_mask=obs_mask,
            n_obs_groups=n_obs,
        )
    else:
        from natal.kernels.simulation_kernels import run_discrete_with_hooks

        final_state_tuple, history_new, was_stopped = run_discrete_with_hooks(
            state=self._state_nn,
            config=self._config_nn,
            registry=hooks.registry,
            first_hook=hooks.first,
            early_hook=hooks.early,
            late_hook=hooks.late,
            n_ticks=n_steps,
            record_interval=record_every,
            observation_mask=obs_mask,
            n_obs_groups=n_obs,
        )

    self._state = DiscretePopulationState(
        n_tick=int(final_state_tuple[1]),
        individual_count=final_state_tuple[0],
    )
    self._tick = int(final_state_tuple[1])

    self._process_kernel_history(history_new, clear_history_on_start)

    if was_stopped:
        self._finished = True
        self.trigger_event("finish")
    elif finish:
        self.finish_simulation()

    return self

run_tick

run_tick() -> DiscreteGenerationPopulation

Execute a single simulation tick.

Overrides BasePopulation.run_tick to use the accelerated run() pipeline which correctly handles compiled hooks.

Source code in src/natal/discrete_generation_population.py

def run_tick(self) -> DiscreteGenerationPopulation:
    """Execute a single simulation tick.

    Overrides BasePopulation.run_tick to use the accelerated run() pipeline
    which correctly handles compiled hooks.
    """
    return self.run(n_steps=1, record_every=self.record_every)

reset

reset() -> None

Reset the population to its initial state.

Source code in src/natal/discrete_generation_population.py

def reset(self) -> None:
    """Reset the population to its initial state."""
    self._tick = 0
    self._history = []
    self._finished = False
    if hasattr(self, '_initial_population_snapshot'):
        ind_copy, _, _ = self._initial_population_snapshot

        # Recreate state with initial data
        self._state = DiscretePopulationState.create(
            n_sexes=self._config_nn.n_sexes,
            n_ages=self._config_nn.n_ages,
            n_genotypes=self._config_nn.n_genotypes,
            n_tick=0,
            individual_count=ind_copy.copy(),
        )

export_state

export_state() -> Tuple[NDArray[np.float64], Optional[NDArray[np.float64]]]

Export the current state and optional history.

Returns:

Type	Description
NDArray[float64]	A tuple of (state_flat, history) where state_flat contains
Optional[NDArray[float64]]	tick and individual counts, and history is either None or a
Tuple[NDArray[float64], Optional[NDArray[float64]]]	stacked history array.

Source code in src/natal/discrete_generation_population.py

def export_state(self) -> Tuple[NDArray[np.float64], Optional[NDArray[np.float64]]]:
    """Export the current state and optional history.

    Returns:
        A tuple of ``(state_flat, history)`` where ``state_flat`` contains
        tick and individual counts, and ``history`` is either ``None`` or a
        stacked history array.
    """
    state_flat = self._state_nn.flatten_all()
    history = self.get_history() if self._history else None
    return state_flat, history

export_config

export_config() -> PopulationConfig

Export the current population configuration.

Returns:

Type	Description
PopulationConfig	The active PopulationConfig used by the population.

Source code in src/natal/discrete_generation_population.py

def export_config(self) -> PopulationConfig:
    """Export the current population configuration.

    Returns:
        The active ``PopulationConfig`` used by the population.
    """
    return self._config_nn

import_config

import_config(config: PopulationConfig) -> None

Import a population configuration into the discrete model.

Parameters:

Name	Type	Description	Default
config	PopulationConfig	Configuration object to install.	required

Source code in src/natal/discrete_generation_population.py

def import_config(self, config: PopulationConfig) -> None:
    """Import a population configuration into the discrete model.

    Args:
        config: Configuration object to install.
    """
    self._config = self._normalize_config(config)

import_state

import_state(state: Union[DiscretePopulationState, NDArray[float64], Dict[str, ndarray]], history: Optional[NDArray[float64]] = None) -> None

Import state and optional history records.

Parameters:

Name	Type	Description	Default
state	Union[DiscretePopulationState, NDArray[float64], Dict[str, ndarray]]	DiscretePopulationState, flattened state array, or a mapping containing individual_count.	required
history	Optional[NDArray[float64]]	Optional 2D history array previously returned by export_state.	None

Source code in src/natal/discrete_generation_population.py

def import_state(
    self,
    state: Union[DiscretePopulationState, NDArray[np.float64], Dict[str, np.ndarray]],
    history: Optional[NDArray[np.float64]] = None,
) -> None:
    """Import state and optional history records.

    Args:
        state: ``DiscretePopulationState``, flattened state array, or a mapping
            containing ``individual_count``.
        history: Optional 2D history array previously returned by ``export_state``.
    """
    assert isinstance(state, (np.ndarray, DiscretePopulationState, dict)), \
        "state must be a DiscretePopulationState, flattened ndarray, or dict"
    if isinstance(state, np.ndarray):
        state_obj = parse_flattened_discrete_state(
            state,
            n_sexes=self._config_nn.n_sexes,
            n_ages=self._config_nn.n_ages,
            n_genotypes=self._config_nn.n_genotypes,
        )
    elif isinstance(state, DiscretePopulationState):
        state_obj = state
    else:
        state_obj = DiscretePopulationState(
            n_tick=int(state.get("n_tick", self._tick)),
            individual_count=np.asarray(state["individual_count"], dtype=np.float64),
        )

    self._state = DiscretePopulationState(
        n_tick=int(state_obj.n_tick),
        individual_count=state_obj.individual_count.copy(),
    )
    self._tick = int(state_obj.n_tick)

    if history is not None and history.shape[0] > 0:
        self.clear_history()
        for row_idx in range(history.shape[0]):
            flat = history[row_idx, :]
            tick = int(flat[0])
            self._history.append((tick, flat.copy()))