Source code for orgmatt.metabolism.intake_fraction

# -*- coding: utf-8 -*-
# SPDX-FileCopyrightText: 2022-2023 Tanguy Fardet
# SPDX-License-Identifier: GPL-3.0-or-later
# orgmatt/metabolism/intake_fraction.py

import numpy as np

from .._utils.age import _array_numeric_age
from .._utils.tools import is_arraylike
from ..typing import NumericOrArray
from ..units import Quantity, ureg




def intake_fraction_age(
    age: NumericOrArray | str,
    adult_consumption: Quantity,
    sex: str | None = None,
) -> NumericOrArray:
    '''
    Fraction of the adult intake consumed at a given age.

    Parameters
    ----------
    age : float, int, str or array
        Age of the individual(s) in years. If a string is provided, it must be
        a group name, an integer or a range, e.g. ("4" or "16-20").
    adult_consumption : Quantity
        Average adult intake (mass of nitrogen, phosphorus, or potassium).
    sex : str, optional (default: average value)
        Sex of the indivual(s).
    '''
    numeric_age = _array_numeric_age(age)

    if isinstance(sex, str) and sex != "mixed":
        sex = [sex]
    else:
        sex = ["male", "female"]

    res = None

    if adult_consumption.is_compatible_with("gN"):
        res = _nitrogen_intake_fraction_age(numeric_age, adult_consumption, sex)
    elif adult_consumption.is_compatible_with("gP"):
        res = _phosphorus_intake_fraction_age(
            numeric_age, adult_consumption, sex
        )
    elif adult_consumption.is_compatible_with("gK"):
        res = _potassium_intake_fraction_age(
            numeric_age, adult_consumption, sex
        )
    else:
        raise ValueError(
            f"Unsupported `adult_consumption`: {adult_consumption}."
        )

    if is_arraylike(age):
        return res

    return res[0]



def _nitrogen_intake_fraction_age(
    age: np.ndarray, adult_consumption: Quantity, sex: list[str]
) -> NumericOrArray:
    min_child_frac = (2.4 * ureg("gN") / adult_consumption.to("gN")).m

    age_half = 2
    age_peak = 55
    age_diff_sex = 5

    frac = np.zeros(len(age))

    for s in sex:
        age_plateau = 14 if s == "female" else 18
        frac_plateau = 0.83 if s == "female" else 1.1
        frac_peak = 0.9 if s == "female" else 1.15
        frac_100 = 0.725 if s == "female" else 0.95

        keep = age <= age_half
        frac[keep] += (
            min_child_frac + age[keep] * (0.5 - min_child_frac) / age_half
        )

        keep = (age_half < age) & (age <= age_diff_sex)
        frac[keep] += 0.5 + (age[keep] - age_half) * 0.15 / (
            age_diff_sex - age_half
        )

        keep = (age_diff_sex < age) & (age <= age_plateau)
        frac[keep] += 0.65 + (age[keep] - age_diff_sex) * (
            frac_plateau - 0.65
        ) / (age_plateau - age_diff_sex)

        keep = (age_plateau < age) & (age < age_peak)
        frac[keep] += frac_plateau + (age[keep] - age_plateau) * (
            frac_peak - frac_plateau
        ) / (age_peak - age_plateau)

        keep = age >= age_peak
        frac[keep] += frac_peak + (age[keep] - age_peak) * (
            frac_100 - frac_peak
        ) / (100 - age_peak)

    return frac / len(sex)


def _phosphorus_intake_fraction_age(
    age: np.ndarray, adult_consumption: Quantity, sex: list[str]
) -> NumericOrArray:
    min_child_frac = (375 * ureg("mgP") / adult_consumption.to("mgP")).m

    age_65pc = 2
    age_peak = 55
    age_diff_sex = 5

    frac = np.zeros(len(age))

    for s in sex:
        age_plateau = 13 if s == "female" else 16
        frac_plateau = 0.88 if s == "female" else 1.1
        frac_peak = 0.89 if s == "female" else 1.12
        frac_100 = 0.71 if s == "female" else 0.925

        keep = age <= age_65pc
        frac[keep] += (
            min_child_frac + age[keep] * (0.65 - min_child_frac) / age_65pc
        )

        keep = (age_65pc < age) & (age <= age_diff_sex)
        frac[keep] += 0.65 + (age[keep] - age_65pc) * 0.12 / (
            age_diff_sex - age_65pc
        )

        keep = (age_diff_sex < age) & (age <= age_plateau)
        frac[keep] += 0.77 + (age[keep] - age_diff_sex) * (
            frac_plateau - 0.77
        ) / (age_plateau - age_diff_sex)

        keep = (age_plateau < age) & (age < age_peak)
        frac[keep] += frac_plateau + (age[keep] - age_plateau) * (
            frac_peak - frac_plateau
        ) / (age_peak - age_plateau)

        keep = age >= age_peak
        frac[keep] += frac_peak + (age[keep] - age_peak) * (
            frac_100 - frac_peak
        ) / (100 - age_peak)

    return frac / len(sex)


def _potassium_intake_fraction_age(
    age: np.ndarray, adult_consumption: Quantity, sex: list[str]
) -> NumericOrArray:
    min_child_frac = (800 * ureg("mgK") / adult_consumption.to("mgK")).m

    age_57pc = 2
    age_peak = 55
    age_diff_sex = 5

    frac = np.zeros(len(age))

    for s in sex:
        age_plateau = 13 if s == "female" else 10
        frac_plateau = 0.8 if s == "female" else 0.875
        frac_peak = 0.94 if s == "female" else 1.15
        frac_100 = 0.825 if s == "female" else 1

        keep = age <= age_57pc
        frac[keep] += (
            min_child_frac + age[keep] * (0.57 - min_child_frac) / age_57pc
        )

        keep = (age_57pc < age) & (age <= age_diff_sex)
        frac[keep] += 0.57 + (age[keep] - age_57pc) * 0.11 / (
            age_diff_sex - age_57pc
        )

        keep = (age_diff_sex < age) & (age <= age_plateau)
        frac[keep] += 0.68 + (age[keep] - age_diff_sex) * (
            frac_plateau - 0.68
        ) / (age_plateau - age_diff_sex)

        keep = (age_plateau < age) & (age < age_peak)
        frac[keep] += frac_plateau + (age[keep] - age_plateau) * (
            frac_peak - frac_plateau
        ) / (age_peak - age_plateau)

        keep = age >= age_peak
        frac[keep] += frac_peak + (age[keep] - age_peak) * (
            frac_100 - frac_peak
        ) / (100 - age_peak)

    return frac / len(sex)