E-Value Methods¶

class EBH(AbstractBatchingTest):
    """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
    """

    def __init__(self, alpha: float):
        super().__init__(alpha)
        self._last_num_rejections = 0

    @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."""
        raw_values = list(e_values)
        if not raw_values:
            return []
        decisions = e_bh(raw_values, self.target_level)
        self._last_num_rejections = sum(decisions)
        threshold = (
            len(raw_values) / (self.target_level * self._last_num_rejections)
            if self._last_num_rejections
            else math.inf
        )
        self._set_test_level(self.target_level, rejection_threshold=threshold)
        self._advance_batch(len(raw_values))
        return decisions

    def test_batch_detail(self, e_values: Sequence[float]) -> BatchDecision:
        """Return details using the same validation and state as ``test_batch``."""
        return super().test_batch_detail(e_values)

class ELond(AbstractSequentialTest):
    """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
    """

    def __init__(
        self,
        alpha: float,
        gamma_seq: AbstractGammaSequence | None = None,
    ):
        super().__init__(alpha)
        self.num_reject = 0
        self.seq = gamma_seq or DefaultLondGammaSequence(c=0.07720838)

    @classmethod
    def _restore_snapshot_state(cls, state: dict[str, Any]) -> dict[str, Any]:
        # Schema-1 snapshots before the shared state used this private name.
        if "_current_level" in state:
            state["_last_test_level"] = state.pop("_current_level")
        return state

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

    def test_one(self, e_value: float) -> bool:
        """Test a single e-value and return whether it is rejected."""
        check_e_value(e_value)
        gamma_t = float(self.seq.calc_gamma(self.num_hypotheses + 1, alpha=1.0))
        level = self.target_level * gamma_t * (self.num_reject + 1)
        threshold = math.inf if level <= 0 else 1.0 / level
        self._set_test_level(level, rejection_threshold=threshold)
        self._advance_hypotheses()

        rejected = float(e_value) >= threshold
        if rejected:
            self.num_reject += 1
        return rejected

    def test_one_detail(self, e_value: float) -> TestDecision:
        """Return details using the same validation and state as ``test_one``."""
        return super().test_one_detail(e_value)

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))
        self.log_current = _advance_log_capital(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)
        self._weights, self._log_weights = _prepare_weights(
            len(self.log_likelihood_ratios),
            weights,
            mismatch_message="weights must match likelihood-ratio components.",
        )
        self.component_log_currents = [0.0] * len(self.log_likelihood_ratios)

    @property
    def weights(self) -> list[float]:
        """Return the normalized fixed prior without exposing mutable state."""
        return list(self._weights)

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

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

    def update(self, observation: Any) -> float:
        updated = self.component_log_currents.copy()
        for idx, (log_weight, log_lr) in enumerate(
            zip(self._log_weights, self.log_likelihood_ratios)
        ):
            if log_weight == -math.inf:
                continue
            value = float(log_lr(observation))
            updated[idx] = _advance_log_capital(updated[idx], value)
        self.component_log_currents = updated
        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,
    ):
        check_positive_integer(n, "n")
        if batch_size is not None:
            check_positive_integer(batch_size, "batch_size")
        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 = None if batch_size is None else int(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 if self.batch_size is not None else self.n
        check_positive_integer(size, "size")
        size = int(size)
        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