source/src/bgc_retrieval/baselines.py

"""Named retrieval baselines and validation-only ensemble selection."""

from __future__ import annotations

from collections.abc import Mapping, Sequence

import numpy as np
import pandas as pd
import torch
from torch.nn import functional as F


def pfam_jaccard_scores(
    candidate_ids: Sequence[str],
    reference_ids: Sequence[str],
    pfam_sets: Mapping[str, set[str]],
    aggregation: str = "max",
) -> dict[str, float]:
    if aggregation not in {"max", "mean"}:
        raise ValueError("Pfam aggregation must be 'max' or 'mean'")
    result: dict[str, float] = {}
    references = [pfam_sets.get(identifier, set()) for identifier in reference_ids]
    for candidate_id in candidate_ids:
        candidate = pfam_sets.get(candidate_id, set())
        similarities = []
        for reference in references:
            union = candidate | reference
            similarities.append(len(candidate & reference) / len(union) if union else 0.0)
        result[candidate_id] = float(max(similarities) if aggregation == "max" else np.mean(similarities))
    return result


def weighted_pfam_jaccard_scores(
    candidate_ids: Sequence[str],
    reference_ids: Sequence[str],
    pfam_sets: Mapping[str, set[str]],
    weights: Mapping[str, float],
    unknown_weight: float = 1.0,
    aggregation: str = "max",
) -> dict[str, float]:
    """Score BGCs with a learned weighted set Jaccard similarity."""
    if aggregation not in {"max", "mean"}:
        raise ValueError("Weighted Pfam aggregation must be 'max' or 'mean'")

    def value(token: str) -> float:
        return max(0.0, float(weights.get(token, unknown_weight)))

    references = [pfam_sets.get(identifier, set()) for identifier in reference_ids]
    result: dict[str, float] = {}
    for candidate_id in candidate_ids:
        candidate = pfam_sets.get(candidate_id, set())
        similarities = []
        for reference in references:
            union = candidate | reference
            denominator = sum(value(token) for token in union)
            numerator = sum(value(token) for token in candidate & reference)
            similarities.append(numerator / denominator if denominator else 0.0)
        result[candidate_id] = float(
            max(similarities) if aggregation == "max" else np.mean(similarities)
        )
    return result


def aggregate_raw_esm(gene_embeddings: torch.Tensor, aggregation: str = "mean") -> torch.Tensor:
    if gene_embeddings.ndim != 2:
        raise ValueError("Raw ESM aggregation expects [genes, dimension]")
    if aggregation == "mean":
        value = gene_embeddings.mean(dim=0)
    elif aggregation == "max":
        value = gene_embeddings.max(dim=0).values
    else:
        raise ValueError("Raw ESM aggregation must be 'mean' or 'max'")
    return F.normalize(value, p=2, dim=0)


def cosine_scores(
    candidate_ids: Sequence[str],
    reference_ids: Sequence[str],
    embeddings: Mapping[str, torch.Tensor],
    aggregation: str = "mean",
) -> dict[str, float]:
    references = torch.stack([F.normalize(embeddings[item], dim=0) for item in reference_ids])
    candidates = torch.stack([F.normalize(embeddings[item], dim=0) for item in candidate_ids])
    similarities = candidates @ references.T
    if aggregation == "mean":
        values = similarities.mean(dim=1)
    elif aggregation == "max":
        values = similarities.max(dim=1).values
    else:
        raise ValueError("Cosine aggregation must be 'mean' or 'max'")
    return {identifier: float(value) for identifier, value in zip(candidate_ids, values)}


def ensemble_scores(
    cosine: Mapping[str, float],
    jaccard: Mapping[str, float],
    alpha: float,
) -> dict[str, float]:
    if not 0.0 <= alpha <= 1.0:
        raise ValueError("Ensemble alpha must be in [0, 1]")
    if set(cosine) != set(jaccard):
        raise ValueError("Ensemble methods must score identical candidates")
    return {
        identifier: alpha * ((float(cosine[identifier]) + 1.0) / 2.0)
        + (1.0 - alpha) * float(jaccard[identifier])
        for identifier in cosine
    }


def select_ensemble_alpha(
    validation_results: pd.DataFrame,
    primary_metric: str,
) -> float:
    required = {"alpha", "group_id", primary_metric}
    missing = required.difference(validation_results.columns)
    if missing:
        raise ValueError(f"Validation results missing columns: {sorted(missing)}")
    means = validation_results.groupby("alpha")[primary_metric].mean()
    best_value = means.max()
    candidates = sorted(float(alpha) for alpha in means[means == best_value].index)
    return min(candidates, key=lambda alpha: (abs(alpha - 0.5), alpha))