qualia_core_db/solvers/learning/kg_embedding/
predict.rs

1//! Link prediction: rank candidate entities for an incomplete triple, and the
2//! standard ranking metrics (mean rank, MRR, Hits@k). This is the **cheap, always-on**
3//! path — given a trained [`EmbeddingTable`], answer "which tail best completes
4//! `(h, r, ?)`" by scoring candidates and ranking by plausibility.
5
6use super::{EmbeddingTable, KgEmbeddingError};
7
8/// Whether to use the *filtered* ranking protocol (Bordes et al. 2013): known-true
9/// triples other than the target are removed from the candidate set before ranking,
10/// so a model is not penalised for ranking another genuine answer above the target.
11pub enum RankFilter<'a> {
12    /// Raw ranking — score against all candidates.
13    Raw,
14    /// Filtered — exclude these (already-known-true) tail indices from competing.
15    Known(&'a [usize]),
16}
17
18/// Rank of the true tail `t` among `candidates` for `(h, r, ?)`. Rank 1 is best.
19/// The rank is `1 + (#candidates scoring strictly higher than the true tail)`; ties
20/// are broken pessimistically by also counting equal-scoring *different* candidates
21/// at half weight is avoided — we use the strict-greater convention (optimistic ties),
22/// which is the common reporting choice. Fails closed on bad indices.
23pub fn rank_tail(
24    table: &EmbeddingTable,
25    h: usize,
26    r: usize,
27    true_t: usize,
28    candidates: &[usize],
29    filter: RankFilter,
30) -> Result<usize, KgEmbeddingError> {
31    let target = table.score(h, r, true_t)?;
32    let known: &[usize] = match filter {
33        RankFilter::Raw => &[],
34        RankFilter::Known(k) => k,
35    };
36    let mut rank = 1usize;
37    for &c in candidates {
38        if c == true_t {
39            continue;
40        }
41        if known.contains(&c) {
42            continue; // filtered out — a genuine answer, not a distractor
43        }
44        let s = table.score(h, r, c)?;
45        if s > target {
46            rank += 1;
47        }
48    }
49    Ok(rank)
50}
51
52/// Mean rank over a set of `(h, r, t)` test triples, each ranked against `candidates`
53/// (typically all entities). Lower is better.
54pub fn mean_rank(
55    table: &EmbeddingTable,
56    triples: &[(usize, usize, usize)],
57    candidates: &[usize],
58) -> Result<f64, KgEmbeddingError> {
59    if triples.is_empty() {
60        return Err(KgEmbeddingError::InsufficientData);
61    }
62    let mut sum = 0.0;
63    for &(h, r, t) in triples {
64        sum += rank_tail(table, h, r, t, candidates, RankFilter::Raw)? as f64;
65    }
66    Ok(sum / triples.len() as f64)
67}
68
69/// Mean reciprocal rank (MRR) — `mean(1/rank)`. Higher is better, in `(0, 1]`.
70pub fn mean_reciprocal_rank(
71    table: &EmbeddingTable,
72    triples: &[(usize, usize, usize)],
73    candidates: &[usize],
74) -> Result<f64, KgEmbeddingError> {
75    if triples.is_empty() {
76        return Err(KgEmbeddingError::InsufficientData);
77    }
78    let mut sum = 0.0;
79    for &(h, r, t) in triples {
80        let rk = rank_tail(table, h, r, t, candidates, RankFilter::Raw)?;
81        sum += 1.0 / rk as f64;
82    }
83    Ok(sum / triples.len() as f64)
84}
85
86/// Hits@k — fraction of test triples whose true tail ranks within the top `k`.
87pub fn hits_at_k(
88    table: &EmbeddingTable,
89    triples: &[(usize, usize, usize)],
90    candidates: &[usize],
91    k: usize,
92) -> Result<f64, KgEmbeddingError> {
93    if triples.is_empty() || k == 0 {
94        return Err(KgEmbeddingError::InsufficientData);
95    }
96    let mut hits = 0usize;
97    for &(h, r, t) in triples {
98        if rank_tail(table, h, r, t, candidates, RankFilter::Raw)? <= k {
99            hits += 1;
100        }
101    }
102    Ok(hits as f64 / triples.len() as f64)
103}
104
105#[cfg(test)]
106mod tests {
107    use super::super::score::ScoreModel;
108    use super::*;
109
110    /// Build a tiny TransE table by hand where the geometry is exactly right:
111    /// entity 0 + relation 0 = entity 1. Then (0, 0, 1) must rank 1.
112    fn hand_table() -> EmbeddingTable {
113        let mut t = EmbeddingTable::zeros(ScoreModel::TransE { p: 2 }, 2, 3, 1).unwrap();
114        t.entity_mut(0).copy_from_slice(&[0.0, 0.0]);
115        t.entity_mut(1).copy_from_slice(&[1.0, 0.0]);
116        t.entity_mut(2).copy_from_slice(&[5.0, 5.0]); // distractor, far away
117        t.relation_mut(0).copy_from_slice(&[1.0, 0.0]);
118        t
119    }
120
121    #[test]
122    fn true_tail_ranks_first() {
123        let t = hand_table();
124        let rank = rank_tail(&t, 0, 0, 1, &[0, 1, 2], RankFilter::Raw).unwrap();
125        assert_eq!(rank, 1);
126    }
127
128    #[test]
129    fn metrics_on_a_perfect_table() {
130        let t = hand_table();
131        let test = [(0usize, 0usize, 1usize)];
132        let cands = [0, 1, 2];
133        assert!((mean_rank(&t, &test, &cands).unwrap() - 1.0).abs() < 1e-12);
134        assert!((mean_reciprocal_rank(&t, &test, &cands).unwrap() - 1.0).abs() < 1e-12);
135        assert!((hits_at_k(&t, &test, &cands, 1).unwrap() - 1.0).abs() < 1e-12);
136    }
137
138    #[test]
139    fn filtered_protocol_excludes_known_true() {
140        // Make entity 2 also a good answer; filtering it should keep rank 1.
141        let mut t = hand_table();
142        t.entity_mut(2).copy_from_slice(&[1.0, 0.0]); // now also exactly h+r
143        let raw = rank_tail(&t, 0, 0, 1, &[0, 1, 2], RankFilter::Raw).unwrap();
144        let filt = rank_tail(&t, 0, 0, 1, &[0, 1, 2], RankFilter::Known(&[2])).unwrap();
145        assert!(raw >= 1 && filt == 1, "raw {raw} filt {filt}");
146    }
147
148    #[test]
149    fn empty_test_set_fails_closed() {
150        let t = hand_table();
151        assert_eq!(
152            mean_rank(&t, &[], &[0, 1, 2]).unwrap_err(),
153            KgEmbeddingError::InsufficientData
154        );
155    }
156}
qualia_core_db/solvers/learning/kg_embedding/predict.rs

qualia_core_db/solvers/learning/kg_embedding/
predict.rs
