use crate::prob::{clamp_prob, sigmoid};
use crate::{error::invalid_input, ScoringResult};
#[derive(Debug, Clone, Copy)]
pub struct LegacyCompositePriorTransform {
pub alpha: f64,
pub beta: f64,
pub base_rate: Option<f64>,
}
impl Default for LegacyCompositePriorTransform {
fn default() -> Self {
Self {
alpha: 1.0,
beta: 0.0,
base_rate: None,
}
}
}
impl LegacyCompositePriorTransform {
pub fn new(alpha: f64, beta: f64, base_rate: Option<f64>) -> ScoringResult<Self> {
if !alpha.is_finite() || !beta.is_finite() {
return Err(invalid_input(format!(
"legacy transform alpha and beta must be finite, got alpha={alpha}, beta={beta}"
)));
}
if let Some(br) = base_rate {
if !br.is_finite() || !(0.0..1.0).contains(&br) || br == 0.0 {
return Err(invalid_input(format!(
"base_rate must be finite and in (0, 1), got {br}"
)));
}
}
Ok(Self {
alpha,
beta,
base_rate,
})
}
#[inline]
pub fn score_signal(&self, score: f64) -> f64 {
sigmoid(self.alpha * (score - self.beta))
}
#[inline]
pub fn tf_prior(tf: f64) -> f64 {
0.2 + 0.7 * (tf / 10.0).min(1.0)
}
#[inline]
pub fn norm_prior(doc_len_ratio: f64) -> f64 {
0.3 + 0.6 * (1.0 - ((doc_len_ratio - 0.5).abs() * 2.0).min(1.0))
}
#[inline]
pub fn composite_prior(tf: f64, doc_len_ratio: f64) -> f64 {
let p_tf = Self::tf_prior(tf);
let p_norm = Self::norm_prior(doc_len_ratio);
(0.7 * p_tf + 0.3 * p_norm).clamp(0.1, 0.9)
}
pub fn no_match_floor() -> f64 {
Self::composite_prior(0.0, 1.0)
}
pub fn combined_score(score_signal: f64, prior: f64, base_rate: Option<f64>) -> f64 {
let l = score_signal;
let p = prior;
let num = l * p;
let denom = num + (1.0 - l) * (1.0 - p);
let mut result = clamp_prob(num / denom);
if let Some(br) = base_rate {
let n2 = result * br;
let d2 = n2 + (1.0 - result) * (1.0 - br);
result = clamp_prob(n2 / d2);
}
result
}
pub fn transform_score(&self, score: f64, tf: f64, doc_len_ratio: f64) -> f64 {
let l = self.score_signal(score);
let prior = Self::composite_prior(tf, doc_len_ratio);
Self::combined_score(l, prior, self.base_rate)
}
pub fn heuristic_upper_bound(&self, bm25_upper_bound: f64, p_max: f64) -> f64 {
let l_max = self.score_signal(bm25_upper_bound);
Self::combined_score(l_max, p_max, self.base_rate)
}
}
#[cfg(test)]
mod no_match_floor_tests {
use super::*;
#[test]
fn no_match_floor_matches_zero_evidence_composite_bias() {
let floor = LegacyCompositePriorTransform::no_match_floor();
assert!((floor - 0.23).abs() < 1e-12, "{floor}");
}
#[test]
fn matched_scores_stay_above_the_no_match_floor_under_defaults() {
let transform = LegacyCompositePriorTransform::default();
let floor = LegacyCompositePriorTransform::no_match_floor();
for score_tenths in 0..=100 {
let score = f64::from(score_tenths) / 10.0;
for tf in 1..=10 {
for ratio_tenths in 0..=30 {
let ratio = f64::from(ratio_tenths) / 10.0;
let combined_score = transform.transform_score(score, f64::from(tf), ratio);
assert!(
combined_score >= floor,
"combined_score {combined_score} fell below floor {floor} \
at score={score} tf={tf} ratio={ratio}",
);
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::prob::{logit, sigmoid};
fn approx_eq(a: f64, b: f64, eps: f64) {
assert!((a - b).abs() < eps, "expected {a} ~ {b} within {eps}");
}
#[test]
fn score_signal_at_beta_is_half() {
let t = LegacyCompositePriorTransform::new(1.0, 5.0, None).unwrap();
approx_eq(t.score_signal(5.0), 0.5, 1e-12);
}
#[test]
fn constructor_rejects_non_finite_parameters_and_invalid_base_rates() {
assert!(LegacyCompositePriorTransform::new(f64::NAN, 0.0, None).is_err());
assert!(LegacyCompositePriorTransform::new(1.0, f64::INFINITY, None).is_err());
assert!(LegacyCompositePriorTransform::new(1.0, 0.0, Some(0.0)).is_err());
assert!(LegacyCompositePriorTransform::new(1.0, 0.0, Some(1.0)).is_err());
assert!(LegacyCompositePriorTransform::new(1.0, 0.0, Some(f64::NAN)).is_err());
}
#[test]
fn tf_prior_floor_and_ceiling() {
approx_eq(LegacyCompositePriorTransform::tf_prior(0.0), 0.2, 1e-12);
approx_eq(LegacyCompositePriorTransform::tf_prior(10.0), 0.9, 1e-12);
approx_eq(LegacyCompositePriorTransform::tf_prior(100.0), 0.9, 1e-12);
}
#[test]
fn norm_prior_peaks_at_half() {
approx_eq(LegacyCompositePriorTransform::norm_prior(0.5), 0.9, 1e-12);
approx_eq(LegacyCompositePriorTransform::norm_prior(0.0), 0.3, 1e-12);
approx_eq(LegacyCompositePriorTransform::norm_prior(1.0), 0.3, 1e-12);
approx_eq(LegacyCompositePriorTransform::norm_prior(2.0), 0.3, 1e-12);
}
#[test]
fn composite_prior_in_clamp_window() {
for tf in [0.0, 1.0, 5.0, 50.0] {
for r in [0.0, 0.5, 1.0, 2.0] {
let p = LegacyCompositePriorTransform::composite_prior(tf, r);
assert!((0.1..=0.9).contains(&p), "p={p} for tf={tf} r={r}");
}
}
}
#[test]
fn combined_score_matches_three_term_logit_form() {
let l = 0.7;
let p = 0.4;
let br = 0.3;
let two_step = LegacyCompositePriorTransform::combined_score(l, p, Some(br));
let logit_form = sigmoid(logit(l) + logit(p) + logit(br));
approx_eq(two_step, logit_form, 1e-9);
}
#[test]
fn combined_score_without_base_rate_matches_logit_form() {
let l = 0.6;
let p = 0.3;
let two_step = LegacyCompositePriorTransform::combined_score(l, p, None);
let logit_form = sigmoid(logit(l) + logit(p));
approx_eq(two_step, logit_form, 1e-9);
}
#[test]
fn combined_score_is_monotone_in_score_signal() {
let p = 0.5;
let prev = LegacyCompositePriorTransform::combined_score(0.1, p, None);
let mut last = prev;
for l in [0.2, 0.3, 0.5, 0.7, 0.9] {
let cur = LegacyCompositePriorTransform::combined_score(l, p, None);
assert!(
cur > last,
"combined_score should rise with L: {last} -> {cur}"
);
last = cur;
}
}
#[test]
fn transform_score_pipeline() {
let t = LegacyCompositePriorTransform::new(1.0, 0.0, None).unwrap();
let p_zero = t.transform_score(0.0, 0.0, 1.0);
approx_eq(
p_zero,
LegacyCompositePriorTransform::composite_prior(0.0, 1.0),
1e-12,
);
}
#[test]
fn heuristic_upper_bound_dominates_actual_score() {
let t = LegacyCompositePriorTransform::new(1.0, 0.0, None).unwrap();
let upper = t.heuristic_upper_bound(5.0, 0.9);
for tf in [0.0, 1.0, 10.0] {
for r in [0.1, 0.5, 1.0, 2.0] {
let actual = t.transform_score(5.0, tf, r);
assert!(actual <= upper + 1e-12, "actual {actual} > upper {upper}");
}
}
}
}