File size: 4,637 Bytes
c87881a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """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))
|