Source code for knowledgespaces.estimation.discrepancy

"""State inclusion rules for discrepancy-based estimation, as in pks.

Independent implementations of the distance rules, with exact feasibility
under fixed-zero error parameters and log-scaled hyperbolic weights.
"""

from dataclasses import dataclass
from operator import index
from typing import Literal

import numpy as np




@dataclass(frozen=True)
class MDOptions:
    """Discrepancy assignment options shared by fitting and refitting.

    ``minimum`` includes feasible states with d(R,K) <= d_min + radius.
    ``hypblc1`` weights all feasible states by (1+d-d_min)**(-exponent);
    ``hypblc2`` uses (1+d)**(-exponent). Hyperbolic rules are available
    for non-iterative MD only, with radius=0. The exponent must be positive.
    Radius is an excess distance from the response, not a distance between
    a state and the nearest state. Defaults recover ordinary MD/MDML.
    These are the pks ``incradius`` and ``blimMD`` inclusion conventions.
    """

    rule: Literal["minimum", "hypblc1", "hypblc2"] = "minimum"
    radius: int = 0
    exponent: float = 1.0

    def __post_init__(self) -> None:
        if self.rule not in ("minimum", "hypblc1", "hypblc2"):
            raise ValueError("rule must be 'minimum', 'hypblc1', or 'hypblc2'.")
        try:
            radius = index(self.radius)
        except TypeError as error:
            raise ValueError("radius must be a nonnegative integer.") from error
        if isinstance(self.radius, (bool, np.bool_)) or radius < 0:
            raise ValueError("radius must be a nonnegative integer.")
        object.__setattr__(self, "radius", radius)
        if not np.isfinite(self.exponent) or self.exponent <= 0:
            raise ValueError("exponent must be finite and positive.")
        if self.rule != "minimum" and radius:
            raise ValueError("Hyperbolic rules require radius=0.")
        if self.rule == "minimum" and self.exponent != 1:
            raise ValueError("exponent applies only to hyperbolic rules.")



def _validate_options(options: MDOptions | None, method: str) -> MDOptions:
    if options is not None and not isinstance(options, MDOptions):
        raise TypeError("discrepancy must be an MDOptions instance.")
    if options is not None and method == "ML":
        raise ValueError("Discrepancy options apply only to MD/MDML.")
    spec = options or MDOptions()
    if method != "MD" and spec.rule != "minimum":
        raise ValueError("Hyperbolic inclusion is supported only by MD.")
    return spec


def _discrepancy_weights(
    responses: np.ndarray,
    membership: np.ndarray,
    counts: np.ndarray,
    options: MDOptions,
    feasible: np.ndarray | None = None,
) -> np.ndarray:
    distance = responses @ (1 - membership.T) + (1 - responses) @ membership.T
    if feasible is not None:
        distance[~feasible] = np.inf
    minimum = distance.min(axis=1, keepdims=True)
    possible = np.isfinite(minimum[:, 0])
    if np.any((counts > 0) & ~possible):
        raise ValueError("Observed responses are impossible under the fixed parameters.")
    # Arbitrary harmless assignment for impossible zero-frequency rows.
    distance[~possible] = 0
    minimum[~possible] = 0
    if options.rule == "minimum":
        # Hamming distances never exceed the item count; avoid conversion
        # overflow for an arbitrarily large integer radius.
        radius = min(options.radius, responses.shape[1])
        return (np.isfinite(distance) & (distance <= minimum + radius)).astype(float)
    offset = distance - minimum if options.rule == "hypblc1" else distance
    log_base = np.log1p(offset)
    # A rowwise constant cancels on normalization; subtract before multiplying
    # to avoid all weights underflowing for very large exponents.
    log_base -= log_base.min(axis=1, keepdims=True)
    with np.errstate(over="ignore"):
        return np.exp(-options.exponent * log_base)