Source code for knowledgespaces.estimation.observation_masks

"""Empirical independent observation masks and their explicitly scoped Pearson statistic."""

from __future__ import annotations

import numpy as np
from scipy.special import logsumexp

from knowledgespaces._patterns import _positive_integer
from knowledgespaces.estimation.incomplete import IncompleteBLIMEstimate, IncompleteResponseMatrix


def _empirical_mask_distribution(data: IncompleteResponseMatrix) -> tuple[np.ndarray, np.ndarray]:
    masks, inverse = np.unique(data.observed, axis=0, return_inverse=True)
    weights = np.bincount(inverse, weights=data.effective_counts)
    keep = weights > 0
    return masks[keep], weights[keep] / weights.sum()




def sample_observation_masks(
    data: IncompleteResponseMatrix,
    n_respondents: int,
    *,
    seed: int | np.random.Generator | None = None,
    max_memory_bytes: int = 512_000_000,
) -> np.ndarray:
    """Sample entire observed masks independently from their weighted empirical law.

    True means observed. Columns follow data.items; rows are newly sampled
    individuals. Counts weight masks even for aggregated response patterns;
    repeated masks pool their weights, zero-weight masks are never sampled.
    Fractional nonnegative source weights are allowed: these specify a
    discrete sampling distribution, not a fractional simulated sample size.
    Joint item missingness is retained; masks are not sampled cell by cell.

    Applying these draws independently of complete responses models MCAR,
    not generic MAR. This function does not infer a missingness mechanism.
    """
    n_respondents = _positive_integer(n_respondents, "n_respondents")
    max_memory_bytes = _positive_integer(max_memory_bytes, "max_memory_bytes")
    if (
        8 * (data.n_patterns * (data.n_items + 4) + n_respondents * (data.n_items + 1))
        > max_memory_bytes
    ):
        raise MemoryError("Observation mask sampling exceeds max_memory_bytes.")
    masks, probabilities = _empirical_mask_distribution(data)
    rng = np.random.default_rng(seed)
    result = masks[rng.choice(len(masks), n_respondents, p=probabilities)]
    result.flags.writeable = False
    return result





def independent_mask_pearson(
    estimate: IncompleteBLIMEstimate,
    data: IncompleteResponseMatrix,
    *,
    max_memory_bytes: int = 512_000_000,
) -> float:
    """Pearson X2 for an explicitly independent empirical-mask response model.

    For each observed mask O, fit its mass as N_O/N. Joint cell probabilities
    are P(O) P_theta(R_O), giving X2 = sum_r n_r^2 / [N_O(r) P_theta(R_O)] - N.
    Sum over unique value-and-mask patterns; unseen cells are included by
    the normalization identity. This also equals the sum of conditional
    Pearson statistics across fixed exogenous mask strata, with sizes N_O.

    This factorization requires independent masks (MCAR) or an exogenous
    fixed design. Ignorable MAR alone does NOT justify it. No automatic
    degrees of freedom or chi-squared reference distribution is supplied.
    All-missing strata contribute zero. Fractional weights produce only a
    descriptive statistic; respondent resampling requires integer counts.
    Impossible observed cells yield infinity. Columns align by item label.
    """
    max_memory_bytes = _positive_integer(max_memory_bytes, "max_memory_bytes")
    if 8 * data.n_patterns * (data.n_items + 6) > max_memory_bytes:
        raise MemoryError("Independent-mask Pearson exceeds max_memory_bytes.")
    aggregated = data.aggregate()
    _, inverse = np.unique(aggregated.observed, axis=0, return_inverse=True)
    counts = aggregated.effective_counts
    mask_counts = np.bincount(inverse, weights=counts)
    log_probability = estimate.predict(
        aggregated.patterns,
        items=aggregated.items,
        max_memory_bytes=max_memory_bytes,
    ).log_probabilities
    log_sum = logsumexp(2 * np.log(counts) - np.log(mask_counts[inverse]) - log_probability)
    # expm1 avoids subtracting nearly equal O(N) numbers for a fitted table.
    with np.errstate(over="ignore"):
        result = aggregated.n_respondents * np.expm1(log_sum - np.log(aggregated.n_respondents))
    return max(0.0, float(result))