fn_registry/python/functions/datascience/confidence_interval_mean.py

"""Intervalo de confianza de la media (una muestra) o de la diferencia de medias (Welch).

Funcion pura del grupo papers. Calcula el intervalo de confianza (IC) de la media
de una muestra usando la t de Student, o el IC de la diferencia de medias de dos
muestras independientes con el metodo de Welch (sin asumir varianzas iguales).

- Una muestra: ``df = n - 1``, ``se = sd / sqrt(n)`` (sd con ddof=1),
  ``tcrit = t.ppf((1 + confidence) / 2, df)``, ``ci = mean +/- tcrit * se``.
- Dos muestras (Welch): IC de ``mean(data) - mean(other)``, con
  ``se = sqrt(se1^2 + se2^2)`` y grados de libertad de Welch-Satterthwaite.

No lanza excepciones: ante casos degenerados (muestras vacias, ``n < 2``,
varianza cero) devuelve un dict coherente con ``ci_low``/``ci_high``/``se`` en
``nan`` (salvo el sub-caso de varianza cero, donde el IC colapsa al punto) y una
clave ``note`` explicando el caso. Usa ``scipy.stats`` y ``numpy``.
"""

from __future__ import annotations

import math

import numpy as np
from scipy import stats


def confidence_interval_mean(
    data: list, other: list = None, confidence: float = 0.95
) -> dict:
    """Intervalo de confianza de la media o de la diferencia de medias (Welch).

    Si ``other`` es ``None``, calcula el IC de la media de ``data`` con la t de
    Student. Si se proporciona ``other``, calcula el IC de la diferencia
    ``mean(data) - mean(other)`` con el metodo de Welch (no asume varianzas
    iguales) y grados de libertad de Welch-Satterthwaite.

    Es una funcion pura y determinista: no hace I/O ni muta las entradas. No
    lanza excepcion ante datos degenerados; en su lugar devuelve un dict con la
    clave ``note`` y los campos numericos indefinidos a ``nan``.

    Args:
        data: muestra de observaciones numericas (lista de numeros).
        other: segunda muestra independiente. Si se da, el IC es el de la
            diferencia de medias ``mean(data) - mean(other)`` con Welch. Si es
            ``None`` (default), el IC es el de la media de ``data``.
        confidence: nivel de confianza en (0, 1), p.ej. 0.95 para el 95%.

    Returns:
        dict con las claves:
            mean: media de ``data`` (una muestra) o la diferencia
                ``mean(data) - mean(other)`` (dos muestras).
            ci_low: extremo inferior del intervalo de confianza.
            ci_high: extremo superior del intervalo de confianza.
            se: error estandar de la media (o de la diferencia).
            df: grados de libertad de la t (Welch-Satterthwaite si dos muestras).
            confidence: nivel de confianza aplicado (float).
            n: tamano de la muestra (una muestra) o tamano total ``n1 + n2``
                (dos muestras; ademas se incluyen ``n1`` y ``n2``).

        En el caso de dos muestras se incluyen ademas ``n1`` y ``n2``. Casos
        degenerados (muestra vacia, ``n < 2``, etc.) anaden la clave ``note`` y
        dejan ``ci_low``/``ci_high``/``se`` (y a veces ``df``) en ``nan``.
    """
    conf = float(confidence)

    if other is None:
        return _ci_one_sample(data, conf)
    return _ci_welch(data, other, conf)


def _ci_one_sample(data: list, conf: float) -> dict:
    """IC de la media de una sola muestra con la t de Student."""
    arr = np.asarray(list(data), dtype=float)
    n = int(arr.size)

    base = {
        "mean": float("nan"),
        "ci_low": float("nan"),
        "ci_high": float("nan"),
        "se": float("nan"),
        "df": float("nan"),
        "confidence": conf,
        "n": n,
    }

    if n == 0:
        base["note"] = "muestra vacia: media e intervalo indefinidos"
        return base

    mean = float(arr.mean())
    base["mean"] = mean

    if n < 2:
        base["note"] = "n < 2: error estandar y grados de libertad indefinidos"
        return base

    df = n - 1
    base["df"] = float(df)

    sd = float(arr.std(ddof=1))
    se = sd / math.sqrt(n)
    base["se"] = se

    # Varianza cero: el IC colapsa al punto (no es un error).
    if se == 0.0:
        base["ci_low"] = mean
        base["ci_high"] = mean
        base["note"] = "varianza cero: el intervalo colapsa a la media"
        return base

    tcrit = float(stats.t.ppf((1.0 + conf) / 2.0, df))
    margin = tcrit * se
    base["ci_low"] = mean - margin
    base["ci_high"] = mean + margin
    return base


def _ci_welch(data: list, other: list, conf: float) -> dict:
    """IC de la diferencia de medias de dos muestras con el metodo de Welch."""
    a = np.asarray(list(data), dtype=float)
    b = np.asarray(list(other), dtype=float)
    n1 = int(a.size)
    n2 = int(b.size)

    base = {
        "mean": float("nan"),
        "ci_low": float("nan"),
        "ci_high": float("nan"),
        "se": float("nan"),
        "df": float("nan"),
        "confidence": conf,
        "n": n1 + n2,
        "n1": n1,
        "n2": n2,
    }

    if n1 == 0 or n2 == 0:
        base["note"] = "alguna muestra esta vacia: diferencia e intervalo indefinidos"
        return base

    mean1 = float(a.mean())
    mean2 = float(b.mean())
    diff = mean1 - mean2
    base["mean"] = diff

    if n1 < 2 or n2 < 2:
        base["note"] = (
            "n < 2 en alguna muestra: error estandar y grados de libertad indefinidos"
        )
        return base

    sd1 = float(a.std(ddof=1))
    sd2 = float(b.std(ddof=1))
    se1 = sd1 / math.sqrt(n1)
    se2 = sd2 / math.sqrt(n2)
    se = math.sqrt(se1 * se1 + se2 * se2)
    base["se"] = se

    # Ambas varianzas cero: el IC de la diferencia colapsa al punto.
    if se == 0.0:
        base["ci_low"] = diff
        base["ci_high"] = diff
        base["df"] = float("nan")
        base["note"] = "varianza cero en ambas muestras: el intervalo colapsa a la diferencia"
        return base

    # Grados de libertad de Welch-Satterthwaite.
    df = (se1 * se1 + se2 * se2) ** 2 / (
        (se1**4) / (n1 - 1) + (se2**4) / (n2 - 1)
    )
    base["df"] = float(df)

    tcrit = float(stats.t.ppf((1.0 + conf) / 2.0, df))
    margin = tcrit * se
    base["ci_low"] = diff - margin
    base["ci_high"] = diff + margin
    return base