fitness Module

Fitness system for genotype-specific viability, fecundity, and selection.

Overview

The fitness subpackage provides apply_preset_fitness_patch and write_fitness_field for building and applying genotype-specific fitness modifiers.

Complete Module Reference

natal.fitness

Fitness system — fitness patch construction, application, and DSL writing.

This subpackage hosts fitness logic extracted from presets and configurator modules:

• fitness/_patch.py: core fitness patch application (allele scaling, slab scaling, selector-based writes). Uses FitnessPopulationView protocol.
• fitness/_writer.py: Configurator DSL writer — resolves genotype-pattern selectors to ztype indices and writes to config arrays.

apply_preset_fitness_patch

apply_preset_fitness_patch(deps: FitnessPopulationView, patch: PresetFitnessPatch) -> None

Apply a declarative preset fitness patch to population config tensors.

Patch schema (all keys optional): - viability: Dict[genotype_selector, _ViabilityScalingConfig] - fecundity: Dict[genotype_selector, _FecundityScalingConfig] - sexual_selection: Dict[female_selector, Union[float, Dict[male_selector, float]]]

Source code in src/natal/fitness/_patch.py

def apply_preset_fitness_patch(deps: FitnessPopulationView, patch: PresetFitnessPatch) -> None:
    """Apply a declarative preset fitness patch to population config tensors.

    Patch schema (all keys optional):
    - viability: Dict[genotype_selector, _ViabilityScalingConfig]
    - fecundity: Dict[genotype_selector, _FecundityScalingConfig]
    - sexual_selection: Dict[female_selector, Union[float, Dict[male_selector, float]]]
    """
    if not patch:
        return

    all_genotypes = deps.index_registry.index_to_genotype

    # ----------------------------------------------------------------------
    # 1) Selector-based viability patch
    #
    # Input examples:
    # - {"A1|A1": 0.8}
    # - {"A1|A1": {0: 0.9, 1: 0.8}}
    # - {"A1|A1": {"female": 0.9, "male": {0: 0.95}}}
    # ----------------------------------------------------------------------
    viability_patch = patch.get('viability', {})
    for selector, config in viability_patch.items():
        matched = deps.species.resolve_genotype_selectors(
            selector=selector,
            all_genotypes=all_genotypes,
            context='preset.viability',
        )
        for genotype in matched:
            z_indices = deps.index_registry.ztype_indices_for(genotype)

            # scalar: both sexes at default viability age
            if isinstance(config, (int, float)):
                for z_idx in z_indices:
                    deps.config.set_viability_fitness(0, z_idx, float(config))
                    deps.config.set_viability_fitness(1, z_idx, float(config))
                continue

            # age-specific for both sexes: {age: scale}
            config_map = cast(Mapping[object, object], config)
            if _is_simple_age_scale_map(config_map):
                for z_idx in z_indices:
                    for age, scale in config_map.items():
                        deps.config.set_viability_fitness(0, z_idx, float(scale), int(age))
                        deps.config.set_viability_fitness(1, z_idx, float(scale), int(age))
                continue

            # sex-specific: {sex: float | {age: scale}}
            for sex_key, sex_config in config_map.items():
                sex_idx = _normalize_sex_key(_coerce_sex_specifier(sex_key))
                if isinstance(sex_config, (int, float)):
                    for z_idx in z_indices:
                        deps.config.set_viability_fitness(sex_idx, z_idx, float(sex_config))
                elif isinstance(sex_config, Mapping):
                    sex_age_map = cast(Mapping[object, object], sex_config)
                    for z_idx in z_indices:
                        for age, scale in sex_age_map.items():
                            if not isinstance(age, int) or not isinstance(scale, (int, float)):
                                raise TypeError(
                                    f"Invalid viability sex-age config for selector '{selector}', sex '{sex_key}'"
                                )
                            deps.config.set_viability_fitness(sex_idx, z_idx, float(scale), int(age))
                else:
                    raise TypeError(
                        f"Invalid viability sex config for selector '{selector}', sex '{sex_key}': "
                        f"{type(sex_config).__name__}"
                    )

    # ----------------------------------------------------------------------
    # 2) Selector-based fecundity patch
    #
    # Input examples:
    # - {"A1|A1": 0.8}
    # - {"A1|A1": {"female": 0.9, "male": 0.7}}
    # ----------------------------------------------------------------------
    fecundity_patch = patch.get('fecundity', {})
    for selector, config in fecundity_patch.items():
        matched = deps.species.resolve_genotype_selectors(
            selector=selector,
            all_genotypes=all_genotypes,
            context='preset.fecundity',
        )
        for genotype in matched:
            z_indices = deps.index_registry.ztype_indices_for(genotype)

            if isinstance(config, (int, float)):
                for z_idx in z_indices:
                    deps.config.set_fecundity_fitness(0, z_idx, float(config))
                    deps.config.set_fecundity_fitness(1, z_idx, float(config))
                continue

            config_map = cast(Mapping[object, object], config)
            for sex_key, scale in config_map.items():
                sex_idx = _normalize_sex_key(_coerce_sex_specifier(sex_key))
                if not isinstance(scale, (int, float)):
                    raise TypeError(
                        f"Invalid fecundity sex scale for selector '{selector}', sex '{sex_key}'"
                    )
                for z_idx in z_indices:
                    deps.config.set_fecundity_fitness(sex_idx, z_idx, float(scale))

    # ----------------------------------------------------------------------
    # 3) Selector-based sexual-selection patch
    #
    # Input examples:
    # - {"female_selector": 0.9}  # shorthand for all males
    # - {"female_selector": {"male_selector": 1.2}}
    # ----------------------------------------------------------------------
    sexual_selection_patch = patch.get('sexual_selection', {})
    for female_selector, male_config in sexual_selection_patch.items():
        female_matched = deps.species.resolve_genotype_selectors(
            selector=female_selector,
            all_genotypes=all_genotypes,
            context='preset.sexual_selection(female)',
        )

        # Allow shorthand: female_selector -> scalar means all-male targets
        if isinstance(male_config, (int, float)):
            male_map = {'*': float(male_config)}
        else:
            male_map = cast(Mapping[object, object], male_config)

        for male_selector, scale in male_map.items():
            if not isinstance(scale, (int, float)):
                raise TypeError(
                    f"Invalid sexual_selection scale for female selector '{female_selector}'"
                )
            male_matched = deps.species.resolve_genotype_selectors(
                selector=_coerce_selector(male_selector),
                all_genotypes=all_genotypes,
                context='preset.sexual_selection(male)',
            )
            for f_genotype in female_matched:
                f_z_indices = deps.index_registry.ztype_indices_for(f_genotype)
                for m_genotype in male_matched:
                    m_z_indices = deps.index_registry.ztype_indices_for(m_genotype)
                    for f_z in f_z_indices:
                        for m_z in m_z_indices:
                            deps.config.set_sexual_selection_fitness(f_z, m_z, float(scale))

    # ----------------------------------------------------------------------
    # 4) Allele-based patches
    #
    # This layer expands allele-centric config into genotype-level writes:
    # 1) resolve allele name(s) to Gene objects
    # 2) count target copies per genotype (0/1/2)
    # 3) convert copies -> factor according to mode
    # 4) multiply corresponding tensor cells
    # ----------------------------------------------------------------------
    viability_per_allele_patch = patch.get('viability_per_allele', {})
    for allele_name, val in viability_per_allele_patch.items():
        config, mode = _split_config_mode(val)
        if not _is_viability_scaling_config(config):
            raise TypeError(f"Invalid viability_per_allele config for '{allele_name}'")
        _apply_viability_allele_scaling(deps, all_genotypes, allele_name, config, mode)

    fecundity_per_allele_patch = patch.get('fecundity_per_allele', {})
    for allele_name, val in fecundity_per_allele_patch.items():
        config, mode = _split_config_mode(val)
        if not _is_fecundity_scaling_config(config):
            raise TypeError(f"Invalid fecundity_per_allele config for '{allele_name}'")
        _apply_fecundity_allele_scaling(deps, all_genotypes, allele_name, config, mode)

    sexual_selection_per_allele_patch = patch.get('sexual_selection_per_allele', {})
    for allele_name, val in sexual_selection_per_allele_patch.items():
        config, mode = _split_config_mode(val)
        if not _is_sexual_selection_scaling_config(config):
            raise TypeError(f"Invalid sexual_selection_per_allele config for '{allele_name}'")
        _apply_sexual_selection_allele_scaling(deps, all_genotypes, allele_name, config, mode)

    # 5) Zygote fitness patch
    zygote_patch = patch.get('zygote', {})
    for selector, config in zygote_patch.items():
        matched = deps.species.resolve_genotype_selectors(
            selector=selector,
            all_genotypes=all_genotypes,
            context='preset.zygote',
        )
        for genotype in matched:
            z_indices = deps.index_registry.ztype_indices_for(genotype)
            if isinstance(config, (int, float)):
                for z_idx in z_indices:
                    deps.config.set_zygote_viability_fitness(0, z_idx, float(config))
                    deps.config.set_zygote_viability_fitness(1, z_idx, float(config))
            elif isinstance(config, Mapping):
                config_map = cast(Mapping[object, object], config)
                for sex_key, sex_config in config_map.items():
                    sex_idx = _normalize_sex_key(_coerce_sex_specifier(sex_key))
                    if isinstance(sex_config, (int, float)):
                        for z_idx in z_indices:
                            deps.config.set_zygote_viability_fitness(sex_idx, z_idx, float(sex_config))

    # 6) Zygote allele-based fitness patch
    zygote_per_allele_patch = patch.get('zygote_per_allele', {})
    for allele_name, val in zygote_per_allele_patch.items():
        config, mode = _split_config_mode(val)
        if not _is_zygote_viability_scaling_config(config):
            raise TypeError(f"Invalid zygote_per_allele config for '{allele_name}'")
        _apply_zygote_viability_allele_scaling(deps, all_genotypes, allele_name, config, mode)

    # 7) Slab-based fitness patches (per_slab keys — symmetric with per_allele)
    if patch.get('viability_per_slab'):
        _apply_viability_slab_scaling(deps, all_genotypes, patch)
    if patch.get('fecundity_per_slab'):
        _apply_fecundity_slab_scaling(deps, all_genotypes, patch)
    if patch.get('sexual_selection_per_slab'):
        _apply_sexual_selection_slab_scaling(deps, all_genotypes, patch)
    if patch.get('zygote_per_slab'):
        _apply_zygote_slab_scaling(deps, all_genotypes, patch)

write_fitness_field

write_fitness_field(config: PopulationConfig | DiscretePopulationConfig, field_name: str, patch: Mapping[str, float | Mapping[str, float]], mode: str, *, species: Species, registry: IndexRegistry, all_genotypes: list[Genotype]) -> None

Resolve genotype-pattern strings and write into a fitness tensor.

field_name is one of "viability", "fecundity", "sexual_selection", or "zygote_viability". patch is a dict mapping genotype selectors to fitness values, with optional sex-keyed or genotype-keyed nesting.

The function detects the format of patch and dispatches to one of four branches. See inline comments for the detection rules.

Supported formats::

{genotype: val}                                        # scalar → both sexes, all ages
{genotype: {"female": val, "male": val}}               # per-selector sex-keyed
{genotype: {0: val, 1: val}}                           # per-selector age-keyed
{genotype: {"female": {0: val}}}                       # per-selector sex+age keyed
{"female": {genotype: val}, "male": {...}}             # top-level sex-keyed
{female_g: {male_g: val}}                              # sexual_selection pair format

Source code in src/natal/fitness/_writer.py

def write_fitness_field(
    config: PopulationConfig | DiscretePopulationConfig,
    field_name: str,
    patch: Mapping[str, float | Mapping[str, float]],
    mode: str,
    *,
    species: Species,
    registry: IndexRegistry,
    all_genotypes: list[Genotype],
) -> None:
    """Resolve genotype-pattern strings and write into a fitness tensor.

    *field_name* is one of ``"viability"``, ``"fecundity"``,
    ``"sexual_selection"``, or ``"zygote_viability"``.  *patch* is a
    dict mapping genotype selectors to fitness values, with optional
    sex-keyed or genotype-keyed nesting.

    The function detects the format of *patch* and dispatches to one of
    four branches.  See inline comments for the detection rules.

    Supported formats::

        {genotype: val}                                        # scalar → both sexes, all ages
        {genotype: {"female": val, "male": val}}               # per-selector sex-keyed
        {genotype: {0: val, 1: val}}                           # per-selector age-keyed
        {genotype: {"female": {0: val}}}                       # per-selector sex+age keyed
        {"female": {genotype: val}, "male": {...}}             # top-level sex-keyed
        {female_g: {male_g: val}}                              # sexual_selection pair format
    """
    # ══════════════════════════════════════════════════════════════════════
    # BRANCH 1: top-level sex-keyed
    #   {"female": {genotype: val}, "male": {genotype: val}}
    #
    # Detection: ALL top-level keys are "female" or "male".
    # Action: iterate sex→genotype_dict, delegate each to _write_fitness_field_flat.
    # ══════════════════════════════════════════════════════════════════════
    if patch and all(k in ("female", "male") for k in patch):
        # ---- guard: every value must itself be a dict {genotype: val} ----
        if not all(isinstance(v, Mapping) for v in patch.values()):
            raise TypeError(
                "sex-keyed fitness dict values must be genotype→value mappings"
            )
        sex_patch: Mapping[str, Mapping[str | tuple[Genotype | str, str], float]]
        sex_patch = patch  # type: ignore[assignment]  # Mapping key invariance; narrowed by branch guard
        # ---- write female slice, then male slice ----
        for sex_key, geno_dict in sex_patch.items():
            sex_idx = 0 if sex_key == "female" else 1
            _write_fitness_field_flat(
                config, field_name, geno_dict, mode,
                sex_idx=sex_idx,
                species=species, registry=registry,
                all_genotypes=all_genotypes,
            )
        return

    # ══════════════════════════════════════════════════════════════════════
    # BRANCH 2: sexual_selection — nested female→male pair format
    #   {female_g: {male_g: val}}
    #
    # Detection: field is "sexual_selection" AND any value is a Mapping.
    # Action: resolve female & male selectors independently,
    #         then write into each [f_idx, m_idx] cell.
    # ══════════════════════════════════════════════════════════════════════
    if field_name == "sexual_selection":
        arr = config.sexual_selection_fitness          # shape: (g, g) — [female_idx, male_idx]
        has_nested = any(isinstance(v, Mapping) for v in patch.values())
        if has_nested:
            for female_selector, male_map in patch.items():           # outer: female genotype key
                if not isinstance(male_map, Mapping):                 # guard: must be {male_g: val}
                    raise TypeError(
                        "Mixed scalar/nested format in sexual_selection. "
                        "When using nested female→male pairs, all values "
                        "must be dicts mapping male selectors to values."
                    )
                for male_selector, value in male_map.items():         # inner: male genotype key → float
                    # ---- resolve both selectors to genotype indices ----
                    matched_f = species.resolve_genotype_selectors(
                        selector=female_selector,
                        all_genotypes=all_genotypes,
                        context="sexual_selection (female)",
                    )
                    matched_m = species.resolve_genotype_selectors(
                        selector=male_selector,
                        all_genotypes=all_genotypes,
                        context="sexual_selection (male)",
                    )
                    # ---- write every female×male combination ----
                    for f_geno in matched_f:
                        for f_z in registry.ztype_indices_for(f_geno):
                            for m_geno in matched_m:
                                for m_z in registry.ztype_indices_for(m_geno):
                                    val = float(value)
                                    if mode == "replace":
                                        arr[f_z, m_z] = val
                                    else:
                                        arr[f_z, m_z] *= val
            return

        # ═══════════════════════════════════════════════════════════════
        # BRANCH 3: sexual_selection — flat male-keyed format
        #   {male_g: val}
        #
        # Detection: field is "sexual_selection" AND no nested values.
        # Action: resolve male selector, write value to ALL female rows
        #         of the matched male column (arr[:, m_idx]).
        # ═══════════════════════════════════════════════════════════════
        for male_selector, value in patch.items():
            if isinstance(value, Mapping):                            # guard: mixed scalar/nested
                raise TypeError(
                    "Mixed scalar/nested format in sexual_selection. "
                    "When using nested female→male pairs, all values "
                    "must be dicts."
                )
            # ---- resolve the male genotype to an index ----
            matched_m = species.resolve_genotype_selectors(
                selector=male_selector,
                all_genotypes=all_genotypes,
                context="sexual_selection (male)",
            )
            for m_geno in matched_m:
                for m_z in registry.ztype_indices_for(m_geno):
                    val = float(value)
                    if mode == "replace":
                        arr[:, m_z] = val        # broadcast: all females × this male
                    else:
                        arr[:, m_z] *= val
        return

    # ══════════════════════════════════════════════════════════════════════
    # BRANCH 4: per-selector resolution (viability / fecundity / zygote)
    #   {genotype: val}
    #   {genotype: {"female": val, "male": val}}           — sex-keyed
    #   {genotype: {0: val, 1: val}}                       — age-keyed
    #   {genotype: {"female": {0: val}}}                   — sex+age keyed
    #
    # Detection: everything not caught by branches 1-3.
    # Each selector value may be:
    #   - scalar → apply to both sexes, all ages
    #   - Mapping → inspect the first key to decide the format
    # ══════════════════════════════════════════════════════════════════════
    for selector, value in patch.items():
        if isinstance(value, Mapping):
            # Inspect the first key to determine the nesting structure.
            first_key = next(iter(value.keys()))
            if isinstance(first_key, int) and not isinstance(first_key, bool):  # type: ignore[unnecessary-isinstance] — bool ⊂ int in Python
                # ---- age-keyed: {genotype: {0: val, 1: val}} ----
                for age_key, age_val in value.items():          # type: ignore[var-unknown]  # Mapping values are Any without explicit type params
                    if age_val is None:                         # type: ignore[unnecessary-comparison]  # user may pass {age: None} to skip
                        continue
                    age = int(age_key)
                    for sex_idx in (0, 1):
                        _write_fitness_field_flat(
                            config, field_name,
                            {selector: float(age_val)}, mode,
                            sex_idx=sex_idx, age_idx=age,
                            species=species, registry=registry,
                            all_genotypes=all_genotypes,
                        )
            elif first_key in ("female", "male"):
                # ---- sex-keyed: {genotype: {"female": val, "male": val}} ----
                for sex_key, sex_val in value.items():
                    sex_idx = 0 if sex_key == "female" else 1
                    if isinstance(sex_val, Mapping):
                        # ---- sex+age keyed: {genotype: {"female": {0: val}}} ----
                        for age_key, age_val in sex_val.items():          # type: ignore[var-unknown]  # Mapping values are Any without explicit type params
                            if age_val is None:
                                continue
                            _write_fitness_field_flat(
                                config, field_name,
                                {selector: float(age_val)}, mode,  # type: ignore[arg-type]  # age_val is Unknown from unparameterized Mapping
                                sex_idx=sex_idx, age_idx=int(age_key),  # type: ignore[arg-type]  # age_key is Unknown from unparameterized Mapping
                                species=species, registry=registry,
                                all_genotypes=all_genotypes,
                            )
                    else:
                        # ---- simple sex-keyed (existing behavior) ----
                        _write_fitness_field_flat(
                            config, field_name,
                            {selector: float(sex_val)}, mode,
                            sex_idx=sex_idx,
                            species=species, registry=registry,
                            all_genotypes=all_genotypes,
                        )
            else:
                raise TypeError(
                    f"Unrecognised key in fitness value dict: {first_key!r}. "
                    f"Expected 'female'/'male' (sex-keyed) or int (age-keyed)."
                )
        else:
            # ---- scalar format: {genotype: val} → apply to both sexes, all ages ----
            for sex_idx in (0, 1):
                _write_fitness_field_flat(
                    config, field_name,
                    {selector: float(value)}, mode,
                    sex_idx=sex_idx,
                    species=species, registry=registry,
                    all_genotypes=all_genotypes,
                )