E-Value Methods¶

class EBH(StatefulMethodMixin):
    """Batch FDR control with e-values via e-Benjamini-Hochberg.

    References:
        Wang, R. and Ramdas, A. (2022). False discovery rate control with
        e-values. Journal of the Royal Statistical Society: Series B.
        Author code: https://github.com/ruoduwang/e-BH
    """

    error_rate = "FDR"

    def __init__(self, alpha: float):
        check_alpha(alpha)
        self.target_level = float(alpha)
        self._num_hypotheses = 0
        self._num_batches = 0
        self._last_test_level: float | None = None
        self._last_rejection_threshold: float | None = None
        self._last_num_rejections = 0

    @property
    def num_hypotheses(self) -> int:
        return self._num_hypotheses

    @property
    def num_tests(self) -> int:
        return self.num_hypotheses

    @property
    def num_batches(self) -> int:
        return self._num_batches

    @property
    def last_test_level(self) -> float | None:
        return self._last_test_level

    @property
    def last_rejection_threshold(self) -> float | None:
        return self._last_rejection_threshold

    @property
    def current_threshold(self) -> float | None:
        return self.last_rejection_threshold

    @property
    def current_k(self) -> int:
        return self._last_num_rejections

    def test_batch(self, e_values: Sequence[float]) -> list[bool]:
        """Test one batch of e-values and return rejection decisions."""
        return list(self.test_batch_detail(e_values).rejected)

    def test_batch_detail(self, e_values: Sequence[float]) -> BatchDecision:
        """Test one batch of e-values and return immutable decision details."""
        raw_values = list(e_values)
        if not raw_values:
            return BatchDecision(
                rejected=(),
                values=(),
                rejection_threshold=self.last_rejection_threshold,
                batch_index=self.num_batches,
                test_level=self.last_test_level,
                error_rate=self.error_rate,
                metadata={"num_rejections": 0},
            )
        check_e_values(raw_values)
        values = [float(value) for value in raw_values]
        decisions = e_bh(values, self.target_level)
        self._last_num_rejections = sum(decisions)
        self._last_test_level = self.target_level
        self._last_rejection_threshold = (
            len(values) / (self.target_level * self._last_num_rejections)
            if self._last_num_rejections
            else math.inf
        )
        self._num_batches += 1
        self._num_hypotheses += len(values)
        return BatchDecision(
            rejected=tuple(decisions),
            values=tuple(values),
            rejection_threshold=self.last_rejection_threshold,
            batch_index=self.num_batches,
            test_level=self.last_test_level,
            error_rate=self.error_rate,
            metadata={"num_rejections": self._last_num_rejections},
        )

class ELond(StatefulMethodMixin):
    """Online FDR control for e-values with e-LOND.

    e-LOND uses the same test levels as p-value LOND but rejects when the
    incoming e-value exceeds the reciprocal test level. Valid e-values give FDR
    control under arbitrary dependence.

    References:
        Xu, Z. and Ramdas, A. (2024). Online multiple testing with e-values.
        Proceedings of AISTATS 2024.
        Author code: https://github.com/neilzxu/evalue-omt
    """

    error_rate = "FDR"

    def __init__(
        self,
        alpha: float,
        gamma_seq: AbstractGammaSequence | None = None,
    ):
        check_alpha(alpha)
        self.target_level = float(alpha)
        self._num_hypotheses = 0
        self.num_reject = 0
        self._current_level: float | None = None
        self._last_rejection_threshold: float | None = None
        self.seq = gamma_seq or DefaultLondGammaSequence(c=0.07720838)

    @property
    def num_hypotheses(self) -> int:
        return self._num_hypotheses

    @property
    def num_tests(self) -> int:
        return self.num_hypotheses

    @property
    def last_test_level(self) -> float | None:
        return self._current_level

    @property
    def current_level(self) -> float | None:
        return self.last_test_level

    @property
    def last_rejection_threshold(self) -> float | None:
        return self._last_rejection_threshold

    @property
    def current_threshold(self) -> float | None:
        return self.last_rejection_threshold

    def test_one(self, e_value: float) -> bool:
        """Test a single e-value and return whether it is rejected."""
        return self.test_one_detail(e_value).rejected

    def test_one_detail(self, e_value: float) -> TestDecision:
        """Test a single e-value and return immutable decision details."""
        check_e_value(e_value)
        self._num_hypotheses += 1

        gamma_t = self._calc_gamma(self.num_hypotheses)
        self._current_level = self.target_level * gamma_t * (self.num_reject + 1)
        self._last_rejection_threshold = (
            math.inf if self._current_level <= 0 else 1.0 / self._current_level
        )

        rejected = float(e_value) >= self._last_rejection_threshold
        if rejected:
            self.num_reject += 1
        return TestDecision(
            rejected=bool(rejected),
            value=float(e_value),
            rejection_threshold=self.last_rejection_threshold,
            index=self.num_hypotheses,
            test_level=self.last_test_level,
            error_rate=self.error_rate,
        )

    def _calc_gamma(self, index: int) -> float:
        try:
            return float(self.seq.calc_gamma(index, alpha=1.0))
        except TypeError:
            return float(self.seq.calc_gamma(index))

class EProcess(Protocol):
    """Protocol for anytime-valid e-processes."""

    @property
    def current(self) -> float:
        """Current e-value."""
        ...

    def update(self, observation: Any) -> float:
        """Update with one observation and return the current e-value."""
        ...

    def reset(self) -> None:
        """Reset the process to its initial e-value."""
        ...

class LikelihoodRatioEProcess:
    """Log-space likelihood-ratio e-process.

    ``log_likelihood_ratio`` must return log ``f_alt(x) / f_null(x)`` for each
    observation. The caller is responsible for the model assumptions that make
    this an e-process under the null.
    """

    def __init__(self, log_likelihood_ratio: Callable[[Any], float]):
        self.log_likelihood_ratio = log_likelihood_ratio
        self.log_current = 0.0

    @property
    def current(self) -> float:
        return _exp_or_inf(self.log_current)

    def update(self, observation: Any) -> float:
        increment = float(self.log_likelihood_ratio(observation))
        if math.isnan(increment):
            raise ValueError("log likelihood ratio must not be NaN.")
        self.log_current += increment
        return self.current

    def reset(self) -> None:
        self.log_current = 0.0

class MixtureLikelihoodRatioEProcess:
    """Fixed-prior mixture of likelihood-ratio e-processes.

    Each component accumulates its own likelihood-ratio process over time. The
    reported value is the weighted arithmetic mixture of those cumulative
    component processes, not a product of per-observation mixture likelihoods.
    The caller is responsible for the null and alternative model assumptions.
    """

    def __init__(
        self,
        log_likelihood_ratios: Sequence[Callable[[Any], float]],
        weights: Sequence[float] | None = None,
    ):
        if not log_likelihood_ratios:
            raise ValueError("at least one likelihood-ratio component is required.")
        self.log_likelihood_ratios = list(log_likelihood_ratios)
        if weights is None:
            self.weights = [1.0 / len(self.log_likelihood_ratios)] * len(
                self.log_likelihood_ratios
            )
        else:
            if len(weights) != len(self.log_likelihood_ratios):
                raise ValueError("weights must match likelihood-ratio components.")
            check_nonnegative_weights(weights)
            total = float(sum(weights))
            self.weights = [float(weight) / total for weight in weights]
        self.component_log_currents = [0.0] * len(self.log_likelihood_ratios)

    @property
    def current(self) -> float:
        return _exp_or_inf(self.log_current)

    @property
    def log_current(self) -> float:
        terms = []
        for weight, component_log_current in zip(
            self.weights, self.component_log_currents
        ):
            if weight == 0:
                terms.append(-math.inf)
            else:
                terms.append(math.log(weight) + component_log_current)
        return _logsumexp(terms)

    def update(self, observation: Any) -> float:
        for idx, (weight, log_lr) in enumerate(
            zip(self.weights, self.log_likelihood_ratios)
        ):
            if weight == 0:
                continue
            value = float(log_lr(observation))
            if math.isnan(value):
                raise ValueError("log likelihood ratio must not be NaN.")
            self.component_log_currents[idx] += value
            if math.isnan(self.component_log_currents[idx]):
                raise ValueError("cumulative log likelihood ratio must not be NaN.")
        return _exp_or_inf(self.log_current)

    def reset(self) -> None:
        self.component_log_currents = [0.0] * len(self.log_likelihood_ratios)

class BettingEProcess:
    """Simple betting e-process from multiplicative betting factors.

    For each score, the factor is ``1 + stake * score``. The caller is
    responsible for using scores and predictable stakes that make the factor
    conditionally expectation-bounded by one under the null.
    """

    def __init__(self, stake: float | Callable[[float], float]):
        self.stake = stake
        self.log_current = 0.0

    @property
    def current(self) -> float:
        return _exp_or_inf(self.log_current)

    def update(self, score: float) -> float:
        stake = self.stake(score) if callable(self.stake) else self.stake
        stake = float(stake)
        score = float(score)
        factor = 1.0 + stake * score
        if not math.isfinite(factor) or factor < 0:
            raise ValueError("betting factor must be finite and nonnegative.")
        if factor == 0:
            self.log_current = -math.inf
        elif self.log_current != -math.inf:
            self.log_current += math.log(factor)
        return self.current

    def reset(self) -> None:
        self.log_current = 0.0

class GaussianEValueGenerator:
    """Generate labeled Gaussian likelihood-ratio e-values."""

    def __init__(
        self,
        n: int,
        pi0: float,
        alt_mean: float = 3.0,
        null_mean: float = 0.0,
        std: float = 1.0,
        batch_size: Optional[int] = None,
        seed: int = 1,
    ):
        if n <= 0:
            raise ValueError("n must be positive.")
        if not 0 <= pi0 <= 1:
            raise ValueError("pi0 must be in [0, 1].")
        if std <= 0:
            raise ValueError("std must be positive.")
        self.n = int(n)
        self.pi0 = float(pi0)
        self.alt_mean = float(alt_mean)
        self.null_mean = float(null_mean)
        self.std = float(std)
        self.batch_size = batch_size
        self.rng = np.random.RandomState(seed)
        self.current_idx = 0

        n_null = int(self.n * self.pi0)
        n_alt = self.n - n_null
        null_samples = self.rng.normal(self.null_mean, self.std, n_null)
        alt_samples = self.rng.normal(self.alt_mean, self.std, n_alt)
        samples = np.concatenate([null_samples, alt_samples])
        labels = np.concatenate([np.zeros(n_null), np.ones(n_alt)]).astype(bool)
        order = self.rng.permutation(self.n)

        self.samples = samples[order]
        self.labels = labels[order]
        self.e_values = gaussian_likelihood_ratio_e_values(
            self.samples,
            self.alt_mean,
            null_mean=self.null_mean,
            std=self.std,
        )

    def sample_one(self) -> Tuple[float, bool]:
        """Sample one ``(e_value, is_alternative)`` pair."""
        if self.current_idx >= self.n:
            raise StopIteration("All samples have been generated.")
        e_value = self.e_values[self.current_idx]
        label = self.labels[self.current_idx]
        self.current_idx += 1
        return float(e_value), bool(label)

    def sample_batch(self, size: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
        """Sample a batch of e-values and labels."""
        if size is None:
            size = self.batch_size or self.n
        if self.current_idx >= self.n:
            raise StopIteration("All samples have been generated.")
        end_idx = min(self.current_idx + size, self.n)
        e_values = self.e_values[self.current_idx : end_idx]
        labels = self.labels[self.current_idx : end_idx]
        self.current_idx = end_idx
        return e_values, labels

    def reset(self) -> None:
        """Reset the generator to the beginning."""
        self.current_idx = 0

    @property
    def remaining(self) -> int:
        """Number of samples remaining."""
        return self.n - self.current_idx