base_population Module

Base population model implementation.

Overview

The base_population module provides the foundational population model that other population types inherit from.

Complete Module Reference

natal.base_population

Base population model helpers and abstractions.

This module provides the abstract base class and utilities for population models (discrete-generation and age-structured). The base class defines common interfaces, evolution methods, history management, and helpers that are implemented by concrete population classes.

This module provides a common abstraction layer for population models while keeping internal state representations compatible with NumPy/Numba kernels.

BasePopulation

BasePopulation(species: Species, name: str = 'Population', hooks: Optional[HookRegistrationMap] = None)

Bases: ABC, Generic[T_State]

Abstract base class for population models.

The base class unifies common behavior for different population model implementations (for example, Wright-Fisher and age-structured non-Wright-Fisher models). It manages the species/genetic architecture, indexing, hook registration, and modifier pipelines.

Attributes:

Name	Type	Description
ALLOWED_EVENTS	List[str]	Event names supported by the hook system.
species	Species	Genetic architecture descriptor for this population.
name	str	Human-readable population name.
tick	int	Current simulation tick.
registry	IndexRegistry	Index registry for genotype/haplotype mappings.
config	PopulationConfig	Active static tensor/config container.
state	T_State	Active population state container.
history	List[Tuple[int, ndarray]]	Recorded state snapshots by tick.

Initialize the base population.

Parameters:

Name	Type	Description	Default
species	Species	Genetic architecture specifying chromosomes, loci, and alleles.	required
name	str	Optional population name (default: "Population").	'Population'
hooks	Optional[HookRegistrationMap]	Optional mapping of event names to hook registrations. Each entry should be a sequence of tuples in the form (func, hook_name, hook_id). Hooks provided here will be registered during initialization.	None

Note

Registry and genotypes are initialized lazily via Template Method. Subclasses must implement _create_registry() and _get_genotypes().

Source code in src/natal/base_population.py

def __init__(
    self,
    species: Species,
    name: str = "Population",
    hooks: Optional[HookRegistrationMap] = None,
):
    """Initialize the base population.

    Args:
        species: Genetic architecture specifying chromosomes, loci, and alleles.
        name: Optional population name (default: "Population").
        hooks: Optional mapping of event names to hook registrations. Each
            entry should be a sequence of tuples in the form ``(func, hook_name, hook_id)``. Hooks
            provided here will be registered during initialization.

    Note:
        Registry and genotypes are initialized lazily via Template Method.
        Subclasses must implement _create_registry() and _get_genotypes().
    """
    self._species = species
    self._name = name
    self._hook_slot = self._derive_hook_slot(name)
    self._tick = 0
    # DELAYED: Registry will be created via _initialize_registry()
    self._index_registry: Optional[IndexRegistry] = None
    self._registry: Optional[IndexRegistry] = None

    # Evolution history as (tick, flattened_array) tuples.
    self._history: List[Tuple[int, np.ndarray]] = []

    # History config
    self.record_every: int = 1
    self.max_history: int = 5000  # Default rolling window size

    # Hooks system: event_name -> [(hook_id, hook_name, hook_func), ...]
    self._hooks: Dict[str, List[HookEntry]] = {
        event: [] for event in self.ALLOWED_EVENTS
    }

    # Unified gamete modifier list.
    self._gamete_modifiers: List[Tuple[int, Optional[str], GameteModifier]] = []

    # Unified zygote modifier list.
    self._zygote_modifiers: List[Tuple[int, Optional[str], ZygoteModifier]] = []

    # Compiled hook descriptors (for Numba-accelerated execution).
    self._compiled_hooks: List[Any] = []  # List[CompiledHookDescriptor]

    # Hook executor (Python-side coordinator for all hook types).
    self._hook_executor: Optional[Any] = None  # HookExecutor

    # Static data container.
    self._config: Optional[PopulationConfig] = None

    # PopulationState container.
    self._state: Optional[T_State] = None

    # Evolution status flag: whether simulation is finished.
    self._finished = False

    # Re-entrancy guard flag.
    self._running = False

    # Observation-based history recording.
    self._observation: Optional[Observation] = None
    self._observation_mask: Optional[np.ndarray] = None

    # Hooks queued for deferred compilation after subclass initialization.
    # Format: [(event_name, func, hook_name, hook_id), ...]
    self._pending_hooks: List[PendingHook] = []

    # Register hooks.
    # If a hook carries @hook metadata, compilation may fail at this stage
    # because IndexRegistry may not be fully initialized yet.
    # Plain functions can be registered immediately; @hook functions are queued.
    hooks_map: HookRegistrationMap = hooks or {}
    if hooks_map:
        for event_name, hooks_list in hooks_map.items():
            for hook_info in hooks_list:
                func, hook_name, hook_id = hook_info

                # Check if function has @hook metadata
                hook_meta = getattr(func, 'meta', None)
                if hook_meta is not None:
                    # Defer compilation until _finalize_hooks() is called
                    self._pending_hooks.append((event_name, func, hook_name, hook_id))
                else:
                    # Plain function, register immediately
                    self.set_hook(event_name, func, hook_id=hook_id, hook_name=hook_name, compile=False)

record_observation property writable

record_observation: Optional[Observation]

The compiled Observation used for observation-mode history.

species property

species: Species

The species/genetic architecture for this population.

name property writable

name: str

The human-readable name of the population.

tick property writable

tick: int

The current simulation tick or generation index.

registry property

registry: IndexRegistry

IndexRegistry instance managing genotype, haplotype, and label indices.

index_registry property

index_registry: IndexRegistry

Public accessor for the internal IndexRegistry.

config property

config: PopulationConfig

Public accessor for compiled population configuration.

state property

state: T_State

Return the current population state container.

Returns:

Name	Type	Description
PopulationState	T_State	The current state object used by the population.

history property

history: List[Tuple[int, ndarray]]

A list of recorded historical states as (tick, flattened_array) tuples.

total_population_size property

total_population_size: int

Total population size (alias of get_total_count).

total_females property

total_females: int

Total number of females (alias of get_female_count).

total_males property

total_males: int

Total number of males (alias of get_male_count).

sex_ratio property

sex_ratio: float

Return the female-to-male ratio, or np.inf when male count is zero.

is_finished property

is_finished: bool

Whether the population is marked as finished (finish=True).

set_observations

set_observations(groups: GroupsInput, *, collapse_age: bool = False) -> None

Register observation groups and immediately compile the binary mask.

Once set, the mask is passed to simulation kernels to record observation-aggregated snapshots (compressed format) instead of raw flattened state.

Parameters:

Name	Type	Description	Default
groups	GroupsInput	Observation groups passed to ObservationFilter. When None, one group per genotype index is used.	required
collapse_age	bool	Whether to collapse the age axis during projection. The stored kernel mask is always 4-D; collapse_age is recorded as metadata and respected by export functions.	False

Source code in src/natal/base_population.py

def set_observations(self, groups: GroupsInput, *, collapse_age: bool = False) -> None:
    """Register observation groups and immediately compile the binary mask.

    Once set, the mask is passed to simulation kernels to record
    observation-aggregated snapshots (compressed format) instead of raw
    flattened state.

    Args:
        groups: Observation groups passed to ``ObservationFilter``.
            When ``None``, one group per genotype index is used.
        collapse_age: Whether to collapse the age axis during projection.
            The stored kernel mask is always 4-D; collapse_age is recorded
            as metadata and respected by export functions.
    """
    from natal.observation import ObservationFilter

    obs_filter = ObservationFilter(self.index_registry)
    self._observation = obs_filter.build_filter(
        diploid_genotypes=self.species,
        groups=groups,
        collapse_age=bool(collapse_age),
    )
    self._observation_mask = self._build_observation_mask(self._observation)

refresh_modifier_maps

refresh_modifier_maps() -> None

Public wrapper that rebuilds modifier maps from registered modifiers.

Source code in src/natal/base_population.py

def refresh_modifier_maps(self) -> None:
    """Public wrapper that rebuilds modifier maps from registered modifiers."""
    self._refresh_modifier_maps()

add_gamete_modifier

add_gamete_modifier(modifier: GameteModifier, name: Optional[str] = None, modifier_id: Optional[int] = None, refresh: bool = True) -> None

Register a gamete-level modifier.

Parameters:

Name	Type	Description	Default
modifier	GameteModifier	A GameteModifier callable or object.	required
name	Optional[str]	Optional human-readable name for debugging.	None
modifier_id	Optional[int]	Optional numeric priority used for ordering.	None

Source code in src/natal/base_population.py

def add_gamete_modifier(
    self,
    modifier: GameteModifier,
    name: Optional[str] = None,
    modifier_id: Optional[int] = None,
    refresh: bool = True,
) -> None:
    """Register a gamete-level modifier.

    Args:
        modifier: A ``GameteModifier`` callable or object.
        name: Optional human-readable name for debugging.
        modifier_id: Optional numeric priority used for ordering.
    """
    resolved_id = self._resolve_modifier_id(modifier_id, self._gamete_modifiers)
    self._gamete_modifiers.append((resolved_id, name, modifier))
    self._gamete_modifiers.sort(key=lambda x: x[0])
    if refresh:
        self._refresh_modifier_maps()

add_zygote_modifier

add_zygote_modifier(modifier: ZygoteModifier, name: Optional[str] = None, modifier_id: Optional[int] = None, refresh: bool = True) -> None

Register a zygote-level modifier.

Parameters:

Name	Type	Description	Default
modifier	ZygoteModifier	A ZygoteModifier callable or object.	required
name	Optional[str]	Optional human-readable name for debugging.	None
modifier_id	Optional[int]	Optional numeric priority used for ordering.	None

Source code in src/natal/base_population.py

def add_zygote_modifier(
    self,
    modifier: ZygoteModifier,
    name: Optional[str] = None,
    modifier_id: Optional[int] = None,
    refresh: bool = True,
) -> None:
    """Register a zygote-level modifier.

    Args:
        modifier: A ``ZygoteModifier`` callable or object.
        name: Optional human-readable name for debugging.
        modifier_id: Optional numeric priority used for ordering.
    """
    resolved_id = self._resolve_modifier_id(modifier_id, self._zygote_modifiers)
    self._zygote_modifiers.append((resolved_id, name, modifier))
    self._zygote_modifiers.sort(key=lambda x: x[0])
    if refresh:
        self._refresh_modifier_maps()

set_zygote_modifier

set_zygote_modifier(modifier: ZygoteModifier, modifier_id: Optional[int] = None, modifier_name: Optional[str] = None) -> None

Register a zygote modifier with an optional priority.

Parameters:

Name	Type	Description	Default
modifier	ZygoteModifier	A ZygoteModifier instance or callable.	required
modifier_id	Optional[int]	Numeric priority (lower values execute earlier). If omitted an id will be auto-assigned.	None
modifier_name	Optional[str]	Optional name for debugging.	None

Source code in src/natal/base_population.py

def set_zygote_modifier(
    self,
    modifier: ZygoteModifier,
    modifier_id: Optional[int] = None,
    modifier_name: Optional[str] = None
) -> None:
    """Register a zygote modifier with an optional priority.

    Args:
        modifier: A ``ZygoteModifier`` instance or callable.
        modifier_id: Numeric priority (lower values execute earlier). If omitted
            an id will be auto-assigned.
        modifier_name: Optional name for debugging.
    """
    if not callable(modifier):
        raise TypeError("Zygote modifier must be callable")

    resolved_id = self._resolve_modifier_id(modifier_id, self._zygote_modifiers)

    # Add and sort by priority id.
    self._zygote_modifiers.append((resolved_id, modifier_name, modifier))
    self._zygote_modifiers.sort(key=lambda x: x[0])

set_gamete_modifier

set_gamete_modifier(modifier: GameteModifier, modifier_id: Optional[int] = None, modifier_name: Optional[str] = None) -> None

Register a gamete modifier with optional priority and name.

Source code in src/natal/base_population.py

def set_gamete_modifier(
    self,
    modifier: GameteModifier,
    modifier_id: Optional[int] = None,
    modifier_name: Optional[str] = None
) -> None:
    """Register a gamete modifier with optional priority and name."""
    if not callable(modifier):
        raise TypeError("Gamete modifier must be callable")

    resolved_id = self._resolve_modifier_id(modifier_id, self._gamete_modifiers)

    # Add and sort by priority id.
    self._gamete_modifiers.append((resolved_id, modifier_name, modifier))
    self._gamete_modifiers.sort(key=lambda x: x[0])

apply_preset

apply_preset(preset: GeneticPreset) -> None

Apply a genetic preset to this population.

This is the preferred API for registering presets. The preset's gamete modifiers, zygote modifiers, and fitness effects are registered in the correct order.

Parameters:

Name	Type	Description	Default
preset	GeneticPreset	A GeneticPreset instance (e.g., HomingDrive or custom preset).	required

Examples:

>>> from natal.genetic_presets import HomingDrive
>>> drive = HomingDrive(
...     name="MyDrive",
...     drive_allele="Drive",
...     target_allele="WT",
...     drive_conversion_rate=0.95
... )
>>> population.apply_preset(drive)

See Also

:class:natal.genetic_presets.GeneticPreset - Base class for creating custom presets :class:natal.genetic_presets.HomingDrive - Built-in gene drive preset

Source code in src/natal/base_population.py

def apply_preset(self, preset: GeneticPreset) -> None:
    """Apply a genetic preset to this population.

    This is the preferred API for registering presets. The preset's
    gamete modifiers, zygote modifiers, and fitness effects are
    registered in the correct order.

    Args:
        preset: A GeneticPreset instance (e.g., HomingDrive or custom preset).

    Examples:
        >>> from natal.genetic_presets import HomingDrive
        >>> drive = HomingDrive(
        ...     name="MyDrive",
        ...     drive_allele="Drive",
        ...     target_allele="WT",
        ...     drive_conversion_rate=0.95
        ... )
        >>> population.apply_preset(drive)

    See Also:
        :class:`natal.genetic_presets.GeneticPreset` - Base class for creating custom presets
        :class:`natal.genetic_presets.HomingDrive` - Built-in gene drive preset
    """
    from natal.genetic_presets import apply_preset_to_population
    apply_preset_to_population(self, preset)

builder classmethod

builder(species: Species) -> Any

Create a builder for this population type.

This is the recommended way to construct populations with presets.

Parameters:

Name	Type	Description	Default
species	Species	Genetic architecture for the population.	required

Returns:

Type	Description
Any	A builder instance for this population type.

Examples:

>>> pop = (AgeStructuredPopulation.builder(species)
...     .set_age_structure(n_ages=10)
...     .add_preset(HomingModificationDrive(...))
...     .build())

Source code in src/natal/base_population.py

@classmethod
def builder(cls, species: Species) -> Any:
    """Create a builder for this population type.

    This is the recommended way to construct populations with presets.

    Args:
        species: Genetic architecture for the population.

    Returns:
        A builder instance for this population type.

    Examples:
        >>> pop = (AgeStructuredPopulation.builder(species)
        ...     .set_age_structure(n_ages=10)
        ...     .add_preset(HomingModificationDrive(...))
        ...     .build())
    """
    raise NotImplementedError(f"{cls.__name__} must implement builder()")

initialize_config

initialize_config() -> None

Initialize static lookup tensors used by the population model.

This prepares precomputed maps such as gametes_to_zygote_map and genotype_to_gametes_map and wraps high-level modifiers so they can be applied at tensor-level during simulation steps.

Note

Ensures registry is initialized before proceeding.

Source code in src/natal/base_population.py

def initialize_config(self) -> None:
    """Initialize static lookup tensors used by the population model.

    This prepares precomputed maps such as ``gametes_to_zygote_map`` and
    ``genotype_to_gametes_map`` and wraps high-level modifiers so they can
    be applied at tensor-level during simulation steps.

    Note:
        Ensures registry is initialized before proceeding.
    """
    # ✅ Ensure registry is initialized
    if self._index_registry is None:
        self._initialize_registry()

    # Retrieve all possible haploid and diploid genotypes.
    haploid_genotypes: List[HaploidGenotype] = self.species.get_all_haploid_genotypes()
    diploid_genotypes: List[Genotype] = self.species.get_all_genotypes()

    n_hg = len(haploid_genotypes)
    n_genotypes = len(diploid_genotypes)
    n_glabs = 1  # TODO: configure based on use case.

    # Create static configuration container.
    self._config = build_population_config(
        n_genotypes=n_genotypes,
        n_haploid_genotypes=n_hg,
        n_sexes=2, # TODO
        n_glabs=n_glabs
    )

    # Convert high-level modifiers into tensor-level wrappers.
    gamete_modifier_funcs, zygote_modifier_funcs = self._build_modifier_wrappers(
        haploid_genotypes=haploid_genotypes,
        diploid_genotypes=diploid_genotypes,
        n_glabs=n_glabs
    )

    # Initialize gametes_to_zygote_map and genotype_to_gametes_map.
    gametes_to_zygote_map = initialize_zygote_map(
        haploid_genotypes=haploid_genotypes,
        diploid_genotypes=diploid_genotypes,
        n_glabs=n_glabs,
        zygote_modifiers=zygote_modifier_funcs,
    )

    genotype_to_gametes_map = initialize_gamete_map(
        diploid_genotypes=diploid_genotypes,
        haploid_genotypes=haploid_genotypes,
        n_glabs=n_glabs,
        gamete_modifiers=gamete_modifier_funcs,
    )

    self._config = self._config._replace(
        gametes_to_zygote_map=gametes_to_zygote_map,
        genotype_to_gametes_map=genotype_to_gametes_map,
    )

register_gamete_labels

register_gamete_labels(labels: Optional[Sequence[str]]) -> None

Register gamete labels in the IndexRegistry.

Parameters:

Name	Type	Description	Default
labels	Optional[Sequence[str]]	Sequence of string labels to register. Labels must be unique in the provided sequence. Existing labels are ignored.	required

Source code in src/natal/base_population.py

def register_gamete_labels(self, labels: Optional[Sequence[str]]) -> None:
    """
    Register gamete labels in the IndexRegistry.

    Args:
        labels: Sequence of string labels to register. Labels must be
            unique in the provided sequence. Existing labels are ignored.
    """
    if not hasattr(self, "_index_registry") or self._index_registry is None:
        raise RuntimeError("IndexRegistry not initialized; cannot register gamete labels")

    if labels is None:
        return

    # Normalize and validate input
    try:
        seq = list(labels)
    except Exception as e:
        raise TypeError("labels must be a sequence of strings") from e

    # Ensure provided labels are unique
    if len(set(seq)) != len(seq):
        raise ValueError("labels must be unique")

    # Register each string label if not already present
    for lab in seq:
        if lab not in self._index_registry.glab_to_index:
            self._index_registry.register_gamete_label(lab)

run_tick abstractmethod

run_tick() -> BasePopulation[T_State]

Execute one simulation tick.

Typical sequence: 1. Check termination and re-entrancy guards. 2. Trigger first hooks. 3. Run reproduction step. 4. Trigger early hooks. 5. Run survival step. 6. Trigger late hooks. 7. Run aging step. 8. Increment tick and clear running flag.

If any hook returns RESULT_STOP, remaining steps are skipped and the population is marked as finished.

Returns:

Type	Description
BasePopulation[T_State]	BasePopulation[T_State]: self for chaining.

Raises:

Type	Description
RuntimeError	If the population is finished or already running.

Source code in src/natal/base_population.py

@abstractmethod
def run_tick(self) -> BasePopulation[T_State]:
    """Execute one simulation tick.

    Typical sequence:
    1. Check termination and re-entrancy guards.
    2. Trigger ``first`` hooks.
    3. Run reproduction step.
    4. Trigger ``early`` hooks.
    5. Run survival step.
    6. Trigger ``late`` hooks.
    7. Run aging step.
    8. Increment tick and clear running flag.

    If any hook returns ``RESULT_STOP``, remaining steps are skipped and
    the population is marked as finished.

    Returns:
        BasePopulation[T_State]: ``self`` for chaining.

    Raises:
        RuntimeError: If the population is finished or already running.
    """
    pass

step

step() -> BasePopulation[T_State]

Alias for BasePopulation.run_tick()

Source code in src/natal/base_population.py

def step(self) -> BasePopulation[T_State]:
    """Alias for `BasePopulation.run_tick()`"""
    return self.run_tick()

get_total_count abstractmethod

get_total_count() -> int

Return the total number of individuals in the population.

Source code in src/natal/base_population.py

@abstractmethod
def get_total_count(self) -> int:
    """Return the total number of individuals in the population."""
    pass

get_female_count abstractmethod

get_female_count() -> int

Return the total number of female individuals.

Source code in src/natal/base_population.py

@abstractmethod
def get_female_count(self) -> int:
    """Return the total number of female individuals."""
    pass

get_male_count abstractmethod

get_male_count() -> int

Return the total number of male individuals.

Source code in src/natal/base_population.py

@abstractmethod
def get_male_count(self) -> int:
    """Return the total number of male individuals."""
    pass

create_observation

create_observation(*, groups: Optional[GroupsInput] = None, collapse_age: bool = False) -> Observation

Create a compiled observation from the current population schema.

Parameters:

Name	Type	Description	Default
groups	Optional[GroupsInput]	Observation groups passed to ObservationFilter. When None, one group per genotype index is used.	None
collapse_age	bool	Whether observation collapses the age axis.	False

Returns:

Type	Description
Observation	Compiled Observation object that can be reused across states.

Source code in src/natal/base_population.py

def create_observation(
    self,
    *,
    groups: Optional[GroupsInput] = None,
    collapse_age: bool = False,
) -> Observation:
    """Create a compiled observation from the current population schema.

    Args:
        groups: Observation groups passed to ``ObservationFilter``.
            When ``None``, one group per genotype index is used.
        collapse_age: Whether observation collapses the age axis.

    Returns:
        Compiled ``Observation`` object that can be reused across states.
    """
    from natal.observation import ObservationFilter

    obs_filter = ObservationFilter(self.index_registry)
    return obs_filter.build_filter(
        diploid_genotypes=self.species,
        groups=groups,
        collapse_age=collapse_age,
    )

output_current_state

output_current_state(*, observation: Optional[Observation] = None, groups: Optional[GroupsInput] = None, collapse_age: bool = False, include_zero_counts: bool = False, output_path: Optional[Union[str, Path]] = None, indent: int = 2) -> Dict[str, Any]

Export the current population state with observation rules applied.

This method integrates observation with state translation and can optionally write the JSON payload to a file.

Parameters:

Name	Type	Description	Default
observation	Optional[Observation]	Optional prebuilt observation object. When provided, groups and collapse_age are ignored.	None
groups	Optional[GroupsInput]	Observation groups passed to ObservationFilter. When None, one group per genotype index is used.	None
collapse_age	bool	Whether observation rule generation collapses age axis.	False
include_zero_counts	bool	Whether to keep zero-valued entries.	False
output_path	Optional[Union[str, Path]]	Optional JSON file path. When provided, the payload is written to this file as UTF-8 JSON.	None
indent	int	Indentation used when writing JSON.	2

Returns:

Type	Description
Dict[str, Any]	A dictionary with observation metadata and observed counts.

Source code in src/natal/base_population.py

def output_current_state(
    self,
    *,
    observation: Optional[Observation] = None,
    groups: Optional[GroupsInput] = None,
    collapse_age: bool = False,
    include_zero_counts: bool = False,
    output_path: Optional[Union[str, Path]] = None,
    indent: int = 2,
) -> Dict[str, Any]:
    """Export the current population state with observation rules applied.

    This method integrates observation with state translation and can
    optionally write the JSON payload to a file.

    Args:
        observation: Optional prebuilt observation object. When provided,
            ``groups`` and ``collapse_age`` are ignored.
        groups: Observation groups passed to ``ObservationFilter``.
            When ``None``, one group per genotype index is used.
        collapse_age: Whether observation rule generation collapses age axis.
        include_zero_counts: Whether to keep zero-valued entries.
        output_path: Optional JSON file path. When provided, the payload is
            written to this file as UTF-8 JSON.
        indent: Indentation used when writing JSON.

    Returns:
        A dictionary with observation metadata and observed counts.
    """
    return _output_current_state(
        self,
        observation=observation,
        groups=groups,
        collapse_age=collapse_age,
        include_zero_counts=include_zero_counts,
        output_path=output_path,
        indent=indent,
    )

output_history

output_history(*, observation: Optional[Observation] = None, groups: Optional[GroupsInput] = None, collapse_age: bool = False, include_zero_counts: bool = False, history: Optional[ndarray] = None, output_path: Optional[Union[str, Path]] = None, indent: int = 2) -> Dict[str, Any]

Export the observation history for this population.

Parameters:

Name	Type	Description	Default
observation	Optional[Observation]	Optional prebuilt observation object. When provided, groups and collapse_age are ignored.	None
groups	Optional[GroupsInput]	Observation groups passed to ObservationFilter. When None, one group per genotype index is used.	None
collapse_age	bool	Whether observation rule generation collapses age axis.	False
include_zero_counts	bool	Whether to keep zero-valued entries.	False
history	Optional[ndarray]	Optional flattened history array. When omitted, the population history is fetched from get_history().	None
output_path	Optional[Union[str, Path]]	Optional JSON file path. When provided, the payload is written to this file as UTF-8 JSON.	None
indent	int	Indentation used when writing JSON.	2

Returns:

Type	Description
Dict[str, Any]	A dictionary containing observation metadata and per-snapshot outputs.

Source code in src/natal/base_population.py

def output_history(
    self,
    *,
    observation: Optional[Observation] = None,
    groups: Optional[GroupsInput] = None,
    collapse_age: bool = False,
    include_zero_counts: bool = False,
    history: Optional[np.ndarray] = None,
    output_path: Optional[Union[str, Path]] = None,
    indent: int = 2,
) -> Dict[str, Any]:
    """Export the observation history for this population.

    Args:
        observation: Optional prebuilt observation object. When provided,
            ``groups`` and ``collapse_age`` are ignored.
        groups: Observation groups passed to ``ObservationFilter``.
            When ``None``, one group per genotype index is used.
        collapse_age: Whether observation rule generation collapses age axis.
        include_zero_counts: Whether to keep zero-valued entries.
        history: Optional flattened history array. When omitted, the
            population history is fetched from ``get_history()``.
        output_path: Optional JSON file path. When provided, the payload is
            written to this file as UTF-8 JSON.
        indent: Indentation used when writing JSON.

    Returns:
        A dictionary containing observation metadata and per-snapshot outputs.
    """
    return _output_history(
        self,
        observation=observation,
        groups=groups,
        collapse_age=collapse_age,
        include_zero_counts=include_zero_counts,
        history=history,
        output_path=output_path,
        indent=indent,
    )

finish_simulation

finish_simulation() -> None

End simulation, trigger the finish event, and lock the population.

This method may be called by hooks for early termination. After calling it, step(), run_tick(), and run() cannot run again.

Raises:

Type	Description
RuntimeError	If the population is already finished.

Examples:

>>> def check_extinction(pop):
...     if pop.get_total_count() == 0:
...         print("Population extinct, finishing simulation.")
...         pop.finish_simulation()
>>> pop.set_hook('late', check_extinction)

Source code in src/natal/base_population.py

def finish_simulation(self) -> None:
    """
    End simulation, trigger the ``finish`` event, and lock the population.

    This method may be called by hooks for early termination.
    After calling it, ``step()``, ``run_tick()``, and ``run()`` cannot run again.

    Raises:
        RuntimeError: If the population is already finished.

    Examples:
        >>> def check_extinction(pop):
        ...     if pop.get_total_count() == 0:
        ...         print("Population extinct, finishing simulation.")
        ...         pop.finish_simulation()
        >>> pop.set_hook('late', check_extinction)
    """
    if self._finished:
        raise RuntimeError(
            f"Population '{self.name}' has already finished."
        )

    self._finished = True
    self.trigger_event("finish")

run abstractmethod

run(n_steps: int, record_every: Optional[int] = None, finish: bool = False) -> BasePopulation[Any]

Run multi-step evolution.

Source code in src/natal/base_population.py

@abstractmethod
def run(
    self,
    n_steps: int,
    record_every: Optional[int] = None,
    finish: bool = False
) -> BasePopulation[Any]:
    """
    Run multi-step evolution.
    """
    pass

reset abstractmethod

reset() -> None

Reset the population to its initial state.

Source code in src/natal/base_population.py

@abstractmethod
def reset(self) -> None:
    """Reset the population to its initial state."""
    pass

compute_allele_frequencies

compute_allele_frequencies() -> Dict[str, float]

Compute frequencies of all alleles in the population, normalized per locus.

Returns:

Type	Description
Dict[str, float]	Dict[str, float]: Mapping {allele_name: frequency}.
Dict[str, float]	Frequencies are per-locus proportions in the range [0.0, 1.0].

Source code in src/natal/base_population.py

def compute_allele_frequencies(self) -> Dict[str, float]:
    """
    Compute frequencies of all alleles in the population, normalized per locus.

    Returns:
        Dict[str, float]: Mapping ``{allele_name: frequency}``.
        Frequencies are per-locus proportions in the range ``[0.0, 1.0]``.
    """
    if self._state is None or self._registry is None:
        return {}

    # 1. Initialize counters.
    allele_counts: Dict[str, float] = {}
    locus_totals: Dict[str, float] = {}  # locus_name -> total_count

    for chromosome in self.species.chromosomes:
        for locus in chromosome.loci:
            locus_totals[locus.name] = 0.0
            for gene in locus.alleles:
                allele_counts[gene.name] = 0.0

    # 2. Aggregate genotype counts.
    # individual_count shape: (n_sexes, n_ages, n_genotypes)
    # Sum over sex and age to get total count per genotype.
    genotype_counts = self._state.individual_count.sum(axis=(0, 1))

    registry = self._registry
    for g_idx, count in enumerate(genotype_counts):
        if count <= 0:
            continue

        genotype = registry.index_to_genotype[g_idx]
        for chrom in self.species.chromosomes:
            for locus in chrom.loci:
                mat, pat = genotype.get_alleles_at_locus(locus)
                for allele in (mat, pat):
                    if allele is not None:
                        allele_counts[allele.name] += count
                        locus_totals[locus.name] += count

    # 3. Compute frequencies.
    frequencies: Dict[str, float] = {}
    for allele_name, count in allele_counts.items():
        # Lookup the locus total for this allele.
        # We do not keep a direct fast gene->locus reverse index here,
        # so we safely resolve via species.gene_index.
        gene = self.species.gene_index.get(allele_name)
        if gene and locus_totals[gene.locus.name] > 0:
            frequencies[allele_name] = count / locus_totals[gene.locus.name]
        else:
            frequencies[allele_name] = 0.0

    return frequencies

set_hook

set_hook(event_name: str, func: HookCallback, hook_id: Optional[int] = None, hook_name: Optional[str] = None, compile: bool = True, deme_selector: Optional[DemeSelector] = None) -> None

Register an event hook with optional automatic compilation.

When compile=True and the function carries @hook metadata, it enters the DSL compilation pipeline: - declarative hook -> CSR plan in HookProgram (kernel executable) - selector hook -> py_wrapper or njit_fn (mode dependent) - numba hook -> njit_fn

Plain Python functions are still registered in traditional _hooks for backward-compatible execution.

Parameters:

Name	Type	Description	Default
event_name	str	Event name (must exist in ALLOWED_EVENTS).	required
func	HookCallback	Callback function, supported forms include: - plain function: func(population) - declarative @hook function: returns [Op.scale(...), ...] - selector @hook(selectors={...}) function	required
hook_id	Optional[int]	Numeric execution priority (optional, auto-assigned if omitted). Lower IDs execute first.	None
hook_name	Optional[str]	Optional human-readable name for debugging.	None
compile	bool	Whether to try compiling @hook-decorated functions (default True).	True
deme_selector	Optional[DemeSelector]	Optional deme selector. - None: keep panmictic default behavior (no explicit selector override) - non-None: passed into hook registration for spatial filtering	None

Raises:

Type	Description
ValueError	If event does not exist or hook_id is already in use.

Examples:

>>> # Plain function (backward compatible)
>>> pop.set_hook('first', lambda p: print(f'Step {p.tick}'))
>>>
>>> # Declarative @hook function (auto-compiled)
>>> @hook()
>>> def reduce_juveniles():
...     return [Op.scale(genotypes='AA', ages=[0, 1], factor=0.9)]
>>> pop.set_hook('early', reduce_juveniles)
>>>
>>> # Selector @hook function (auto-compiled)
>>> @hook(selectors={'target': 'AA'})
>>> def release(pop, target):
...     pop.state.individual_count[1, 2, target] += 100
>>> pop.set_hook('first', release)

Source code in src/natal/base_population.py

def set_hook(
    self,
    event_name: str,
    func: HookCallback,
    hook_id: Optional[int] = None,
    hook_name: Optional[str] = None,
    compile: bool = True,
    deme_selector: Optional[DemeSelector] = None,
) -> None:
    """
    Register an event hook with optional automatic compilation.

    When ``compile=True`` and the function carries ``@hook`` metadata,
    it enters the DSL compilation pipeline:
    - declarative hook -> CSR plan in HookProgram (kernel executable)
    - selector hook -> ``py_wrapper`` or ``njit_fn`` (mode dependent)
    - numba hook -> ``njit_fn``

    Plain Python functions are still registered in traditional ``_hooks``
    for backward-compatible execution.

    Args:
        event_name: Event name (must exist in ``ALLOWED_EVENTS``).
        func: Callback function, supported forms include:
              - plain function: ``func(population)``
              - declarative ``@hook`` function: returns ``[Op.scale(...), ...]``
              - selector ``@hook(selectors={...})`` function
        hook_id: Numeric execution priority (optional, auto-assigned if omitted).
                 Lower IDs execute first.
        hook_name: Optional human-readable name for debugging.
        compile: Whether to try compiling ``@hook``-decorated functions (default ``True``).
        deme_selector: Optional deme selector.
            - ``None``: keep panmictic default behavior (no explicit selector override)
            - non-``None``: passed into hook registration for spatial filtering

    Raises:
        ValueError: If event does not exist or hook_id is already in use.

    Examples:
        >>> # Plain function (backward compatible)
        >>> pop.set_hook('first', lambda p: print(f'Step {p.tick}'))
        >>>
        >>> # Declarative @hook function (auto-compiled)
        >>> @hook()
        >>> def reduce_juveniles():
        ...     return [Op.scale(genotypes='AA', ages=[0, 1], factor=0.9)]
        >>> pop.set_hook('early', reduce_juveniles)
        >>>
        >>> # Selector @hook function (auto-compiled)
        >>> @hook(selectors={'target': 'AA'})
        >>> def release(pop, target):
        ...     pop.state.individual_count[1, 2, target] += 100
        >>> pop.set_hook('first', release)
    """
    if event_name not in self.ALLOWED_EVENTS:
        raise ValueError(f"Event '{event_name}' not in {self.ALLOWED_EVENTS}")

    # BasePopulation itself is panmictic. Non-wildcard deme selectors are
    # interpreted by SpatialPopulation orchestration and should not be
    # consumed here.
    if deme_selector is not None and deme_selector != "*":
        warnings.warn(
            "BasePopulation ignores non-'*' deme_selector. "
            "Apply deme selection through SpatialPopulation-level logic instead.",
            UserWarning,
            stacklevel=2,
        )
        deme_selector = None

    # Check if function has @hook metadata and should be compiled
    hook_meta = getattr(func, 'meta', None)
    if is_numba_enabled() and hook_meta is None:
        raise TypeError(
            "Python-layer hooks are not allowed when Numba is enabled. "
            "Use @hook(...) with a compilable body or disable Numba."
        )

    if compile and hook_meta is not None:
        # Use the hook's register method with event override
        register_fn = getattr(func, 'register', None)
        if register_fn is not None:
            # Panmictic path: do not force any selector override.
            if deme_selector is None:
                register_fn(self, event_override=event_name)
            else:
                register_fn(self, event_override=event_name, deme_selector_override=deme_selector)
            # Compiled hooks are stored in _compiled_hooks.
            # Only selector-mode hooks with py_wrapper are mirrored to _hooks.
            self._hook_executor = None
            return

    # Traditional registration (no compilation)
    actual_name = hook_name or getattr(func, '__name__', None)

    current_ids = [hid for hid, _, _ in self._hooks[event_name]]

    if hook_id is None:
        hook_id = (max(current_ids) + 1) if current_ids else 0

    if hook_id in current_ids:
        raise ValueError(f"hook_id {hook_id} already exists in event '{event_name}'")

    self._hooks[event_name].append((hook_id, actual_name, func))
    # Sort by hook ID to preserve execution order.
    self._hooks[event_name].sort(key=lambda x: x[0])
    self._hook_executor = None

trigger_event

trigger_event(event_name: str, deme_id: int = -1) -> int

    Trigger an event and execute all registered hooks.

    Execution order:
    1. CSR operations (Numba fast path)
    2. ``njit_fn`` hooks (user-defined Numba functions)
    3. ``py_wrapper`` hooks (Python wrapper functions)

Parameters:

Name	Type	Description	Default
event_name	str	Event name to trigger.	required
deme_id	int	Deme index. Default -1 for non-spatial populations.	-1

Returns:

Name	Type	Description
int	int	RESULT_CONTINUE (0) to continue, RESULT_STOP (1) to stop.

Note

• Prefer HookExecutor (unified three-layer coordination).
• If executor is not built, fall back to traditional _hooks (Python callbacks only).
• In accelerated run(), core events are mostly executed by kernels; trigger_event is used mainly for explicit events (for example finish) and compatibility paths.

Examples:

>>> result = pop.trigger_event('first')  # Executes all 'first' hooks

>>> if result == RESULT_STOP:
...     print("Simulation stopped by hook")

Source code in src/natal/base_population.py

def trigger_event(self, event_name: str, deme_id: int = -1) -> int:
    """
            Trigger an event and execute all registered hooks.

            Execution order:
            1. CSR operations (Numba fast path)
            2. ``njit_fn`` hooks (user-defined Numba functions)
            3. ``py_wrapper`` hooks (Python wrapper functions)

    Args:
                    event_name: Event name to trigger.
                    deme_id: Deme index. Default -1 for non-spatial populations.

    Returns:
                    int: ``RESULT_CONTINUE`` (0) to continue, ``RESULT_STOP`` (1) to stop.

    Note:
                    - Prefer HookExecutor (unified three-layer coordination).
                    - If executor is not built, fall back to traditional ``_hooks``
                        (Python callbacks only).
                    - In accelerated ``run()``, core events are mostly executed by kernels;
                        ``trigger_event`` is used mainly for explicit events (for example ``finish``)
                        and compatibility paths.

    Examples:
                    >>> result = pop.trigger_event('first')  # Executes all 'first' hooks
        >>> if result == RESULT_STOP:
        ...     print("Simulation stopped by hook")
    """
    from natal.hook_dsl import RESULT_CONTINUE

    # Prefer HookExecutor when available.
    if self._hook_executor is not None:
        from natal.hook_dsl import EVENT_ID_MAP
        event_id = EVENT_ID_MAP.get(event_name)
        if event_id is not None:
            result = self._hook_executor.execute_event(event_id, self, self.tick, deme_id=deme_id)
            return result

    # Fallback to traditional _hooks for compatibility.
    for _, _, hook in self._hooks.get(event_name, []):
        hook(self)

    return RESULT_CONTINUE

get_hooks

get_hooks(event_name: str) -> List[HookEntry]

Get all registered hooks for a specific event.

Parameters:

Name	Type	Description	Default
event_name	str	Event name.	required

Returns:

Type	Description
List[HookEntry]	List of tuples [(hook_id, hook_name, hook_func), ...].

Source code in src/natal/base_population.py

def get_hooks(self, event_name: str) -> List[HookEntry]:
    """
    Get all registered hooks for a specific event.

    Args:
        event_name: Event name.

    Returns:
        List of tuples ``[(hook_id, hook_name, hook_func), ...]``.
    """
    return list(self._hooks.get(event_name, []))

remove_hook

remove_hook(event_name: str, hook_id: int) -> bool

Remove a specific hook from an event.

Parameters:

Name	Type	Description	Default
event_name	str	Event name.	required
hook_id	int	Hook ID.	required

Returns:

Type	Description
bool	True if removed successfully, otherwise False.

Source code in src/natal/base_population.py

def remove_hook(self, event_name: str, hook_id: int) -> bool:
    """
    Remove a specific hook from an event.

    Args:
        event_name: Event name.
        hook_id: Hook ID.

    Returns:
        True if removed successfully, otherwise False.
    """
    if event_name not in self._hooks:
        return False

    original_len = len(self._hooks[event_name])
    self._hooks[event_name] = [(hid, name, func) for hid, name, func in self._hooks[event_name]
                                if hid != hook_id]
    self._hook_executor = None
    return len(self._hooks[event_name]) < original_len

has_python_hooks

has_python_hooks() -> bool

Return whether any Python-layer hooks are currently registered.

Source code in src/natal/base_population.py

def has_python_hooks(self) -> bool:
    """Return whether any Python-layer hooks are currently registered."""
    hooks_map = cast(Dict[str, List[HookEntry]], getattr(self, "_hooks", {}))
    return any(len(entries) > 0 for entries in hooks_map.values())

has_mixed_hook_types

has_mixed_hook_types() -> bool

Return whether any event mixes declarative/njit/python hook types.

Source code in src/natal/base_population.py

def has_mixed_hook_types(self) -> bool:
    """Return whether any event mixes declarative/njit/python hook types."""
    for event_name in self.ALLOWED_EVENTS:
        kinds: set[str] = set()
        for desc in self.get_compiled_hooks(event_name):
            if getattr(desc, "plan", None) is not None:
                kinds.add("declarative")
            if getattr(desc, "njit_fn", None) is not None:
                kinds.add("njit")
            if getattr(desc, "py_wrapper", None) is not None and getattr(desc, "njit_fn", None) is None:
                kinds.add("python")
        if len(kinds) > 1:
            return True
    return False

should_use_python_dispatch

should_use_python_dispatch() -> bool

Return whether this population should run with Python event dispatch.

Policy

• When Numba is disabled, any registered hook type uses Python dispatch so py/declarative/njit hooks share one sequential path.
• When Numba is enabled, only mixed hook-type timelines fall back to Python dispatch; homogeneous compiled timelines stay on kernel wrappers.

Source code in src/natal/base_population.py

def should_use_python_dispatch(self) -> bool:
            """Return whether this population should run with Python event dispatch.

            Policy:
                    - When Numba is disabled, any registered hook type uses Python
                        dispatch so py/declarative/njit hooks share one sequential path.
                    - When Numba is enabled, only mixed hook-type timelines fall back
                        to Python dispatch; homogeneous compiled timelines stay on kernel
                        wrappers.
            """
            if not is_numba_enabled():
                    return self.has_python_hooks() or len(self.get_compiled_hooks()) > 0
            return self.has_mixed_hook_types()

ensure_hook_executor

ensure_hook_executor() -> None

Build HookExecutor lazily for Python event-dispatch paths.

Source code in src/natal/base_population.py

def ensure_hook_executor(self) -> None:
    """Build HookExecutor lazily for Python event-dispatch paths."""
    if self._hook_executor is None:
        self._hook_executor = self._build_hook_executor()

register_compiled_hook

register_compiled_hook(desc: Any) -> None

Public wrapper for registering compiled hooks.

Source code in src/natal/base_population.py

def register_compiled_hook(self, desc: Any) -> None:
    """Public wrapper for registering compiled hooks."""
    self._register_compiled_hook(desc)

get_compiled_hooks

get_compiled_hooks(event: Optional[str] = None) -> List[Any]

Get compiled hook descriptors, optionally filtered by event.

Parameters:

Name	Type	Description	Default
event	Optional[str]	Optional event name to filter by.	None

Returns:

Type	Description
List[Any]	List of CompiledHookDescriptor sorted by priority.

Source code in src/natal/base_population.py

def get_compiled_hooks(self, event: Optional[str] = None) -> List[Any]:
    """Get compiled hook descriptors, optionally filtered by event.

    Args:
        event: Optional event name to filter by.

    Returns:
        List of CompiledHookDescriptor sorted by priority.
    """
    hooks = cast(List[Any], getattr(self, "_compiled_hooks", []))
    if event is not None:
        hooks = [h for h in hooks if h.event == event]
    return sorted(hooks, key=lambda h: h.priority)

register_declarative_hook

register_declarative_hook(event: str, ops: List[Any], priority: int = 0, name: str = 'declarative_hook') -> Any

Register a declarative hook from a list of operations.

This is an alternative to using the @hook decorator.

Parameters:

Name	Type	Description	Default
event	str	Event name ('first', 'early', 'late', 'finish')	required
ops	List[Any]	List of HookOp operations (from Op.scale, Op.add, etc.)	required
priority	int	Execution priority (lower = earlier)	0
name	str	Hook name for debugging	'declarative_hook'

Returns:

Name	Type	Description
CompiledHookDescriptor	Any	The compiled descriptor

Examples:

>>> from natal.hook_dsl import Op
>>> pop.register_declarative_hook(
...     event='early',
...     ops=[
...         Op.scale(genotypes='AA', ages=[0, 1], factor=0.9),
...         Op.add(genotypes='*', ages=0, delta=50, when='tick % 10 == 0'),
...     ],
...     name='juvenile_control'
... )

Source code in src/natal/base_population.py

def register_declarative_hook(
    self,
    event: str,
    ops: List[Any],
    priority: int = 0,
    name: str = "declarative_hook"
) -> Any:
    """Register a declarative hook from a list of operations.

    This is an alternative to using the @hook decorator.

    Args:
        event: Event name ('first', 'early', 'late', 'finish')
        ops: List of HookOp operations (from Op.scale, Op.add, etc.)
        priority: Execution priority (lower = earlier)
        name: Hook name for debugging

    Returns:
        CompiledHookDescriptor: The compiled descriptor

    Examples:
        >>> from natal.hook_dsl import Op
        >>> pop.register_declarative_hook(
        ...     event='early',
        ...     ops=[
        ...         Op.scale(genotypes='AA', ages=[0, 1], factor=0.9),
        ...         Op.add(genotypes='*', ages=0, delta=50, when='tick % 10 == 0'),
        ...     ],
        ...     name='juvenile_control'
        ... )
    """
    from natal.hook_dsl import compile_declarative_hook
    desc = compile_declarative_hook(
        ops,
        self,
        event,
        priority=priority,
        name=name,
    )
    self._register_compiled_hook(desc)
    return desc

get_compiled_event_hooks

get_compiled_event_hooks() -> CompiledEventHooks

Get compiled hooks for use with generated kernel wrappers.

This method collects all registered hooks and compiles them into Numba-friendly combined functions, one per event.

Returns:

Name	Type	Description
CompiledEventHooks	CompiledEventHooks	Container with combined @njit hooks per event. Access via .first, .early, .late, .finish

Examples:

>>> hooks = pop.get_compiled_event_hooks()
>>> hooks.run_fn is not None
True

Source code in src/natal/base_population.py

def get_compiled_event_hooks(self) -> CompiledEventHooks:
    """Get compiled hooks for use with generated kernel wrappers.

    This method collects all registered hooks and compiles them into
    Numba-friendly combined functions, one per event.

    Returns:
        CompiledEventHooks: Container with combined @njit hooks per event.
                            Access via .first, .early, .late, .finish

    Examples:
        >>> hooks = pop.get_compiled_event_hooks()
        >>> hooks.run_fn is not None
        True
    """
    from natal.hook_dsl import CompiledEventHooks

    registry = self._build_hook_program()
    return CompiledEventHooks.from_compiled_hooks(
        self._compiled_hooks,
        registry=registry,
        include_spatial_wrappers=False,
    )