heyting 0.12.0

//! The tranz temporal-model adapter: trained timestamp-set hops.
//!
//! Wraps a `tranz::temporal::TemporalScorer` (a trained TComplEx) as an
//! [`AtomicScorer`], with the same virtual-relation registration pattern as
//! [`TemporalKg`](crate::TemporalKg): register `(base relation, TimeSet)`
//! and hop through the returned id anywhere a relation id goes. A hop's
//! degree is existential over the set — `max over τ ∈ S` of the calibrated
//! atom degree — which is the trained counterpart of `TemporalKg` taking
//! the best admitted fact.

use crate::query::AtomicScorer;
use crate::temporal::TimeSet;

/// Wraps a trained temporal scorer as an [`AtomicScorer`] with
/// timestamp-set virtual relations.
///
/// Base relation ids `0..num_relations` hop with no temporal constraint
/// (the whole axis); [`windowed`](Self::windowed) registers time-scoped
/// views at fresh ids. Energies map to degrees via
/// `sigmoid(-score / temperature)`, exactly as
/// [`PointModel`](crate::adapters::PointModel).
pub struct TemporalPointModel<S> {
    /// The wrapped tranz temporal scorer.
    pub model: S,
    /// Sigmoid temperature; `1.0` is the plain logistic map.
    pub temperature: f32,
    virtuals: Vec<(usize, TimeSet)>,
}

impl<S: tranz::temporal::TemporalScorer> TemporalPointModel<S> {
    /// Wrap a scorer with the plain sigmoid map (temperature `1.0`).
    pub fn new(model: S) -> Self {
        Self::with_temperature(model, 1.0)
    }

    /// Wrap a scorer with a temperature; non-finite or non-positive values
    /// fall back to `1.0` (the map must stay strictly monotone).
    pub fn with_temperature(model: S, temperature: f32) -> Self {
        let temperature = if temperature.is_finite() && temperature > 0.0 {
            temperature
        } else {
            1.0
        };
        Self {
            model,
            temperature,
            virtuals: Vec::new(),
        }
    }

    /// Register a time-scoped view of `relation` over the timestamp set
    /// and return its virtual relation id. `None` if `relation` is not a
    /// base relation or the set's axis does not match the model's.
    pub fn windowed(&mut self, relation: usize, times: TimeSet) -> Option<usize> {
        if relation >= self.model.num_relations()
            || times.num_timestamps() != self.model.num_timestamps()
        {
            return None;
        }
        self.virtuals.push((relation, times));
        Some(self.model.num_relations() + self.virtuals.len() - 1)
    }

    /// Resolve a (possibly virtual) relation id to `(base, TimeSet)`.
    fn resolve(&self, relation: usize) -> Option<(usize, TimeSet)> {
        let n = self.model.num_relations();
        if relation < n {
            Some((relation, TimeSet::all(self.model.num_timestamps())))
        } else {
            self.virtuals.get(relation - n).cloned()
        }
    }
}

impl<S: tranz::temporal::TemporalScorer> AtomicScorer for TemporalPointModel<S> {
    fn num_entities(&self) -> usize {
        self.model.num_entities()
    }

    fn project(&self, anchor: usize, relation: usize) -> Vec<f32> {
        let n = self.model.num_entities();
        let Some((base, times)) = self.resolve(relation) else {
            return vec![0.0; n]; // unknown virtual id, as in TemporalKg.
        };
        if anchor >= n {
            return vec![0.0; n];
        }
        // Existential over the set: keep the best (lowest) energy per
        // entity across τ ∈ S, then map once (sigmoid is monotone).
        let mut best: Option<Vec<f32>> = None;
        for tau in times.iter() {
            let scores = self.model.score_all_tails(anchor, base, tau);
            best = Some(match best {
                None => scores,
                Some(mut acc) => {
                    for (a, s) in acc.iter_mut().zip(&scores) {
                        if *s < *a {
                            *a = *s;
                        }
                    }
                    acc
                }
            });
        }
        match best {
            None => vec![0.0; n], // empty set: no admissible timestamp.
            Some(energies) => energies
                .iter()
                .map(|&e| sigmoid(-e / self.temperature))
                .collect(),
        }
    }
}

/// Numerically stable logistic sigmoid.
fn sigmoid(x: f32) -> f32 {
    if x >= 0.0 {
        1.0 / (1.0 + (-x).exp())
    } else {
        let e = x.exp();
        e / (1.0 + e)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{answer_query, Godel, Query, QueryConfig};

    /// 3 entities, 1 relation, 4 timestamps. The fact (0, r, 1) holds only
    /// at τ = 1 (low energy); (0, r, 2) holds only at τ = 3.
    struct Planted;
    impl tranz::temporal::TemporalScorer for Planted {
        fn score(&self, h: usize, _r: usize, t: usize, tau: usize) -> f32 {
            match (h, t, tau) {
                (0, 1, 1) => -8.0,
                (0, 2, 3) => -8.0,
                _ => 8.0,
            }
        }
        fn num_entities(&self) -> usize {
            3
        }
        fn num_relations(&self) -> usize {
            1
        }
        fn num_timestamps(&self) -> usize {
            4
        }
    }

    #[test]
    fn windowed_hops_scope_answers_existentially() {
        let mut m = TemporalPointModel::new(Planted);
        let early = m.windowed(0, TimeSet::before(2, 4)).unwrap();
        let late = m.windowed(0, TimeSet::after(2, 4)).unwrap();
        let cfg = QueryConfig::default();

        let s = answer_query::<Godel>(&m, &Query::anchor(0, early), &cfg);
        assert!(s[1] > 0.99, "entity 1 holds before τ=2: {s:?}");
        assert!(s[2] < 0.01, "entity 2 does not: {s:?}");

        let s = answer_query::<Godel>(&m, &Query::anchor(0, late), &cfg);
        assert!(s[2] > 0.99 && s[1] < 0.01, "{s:?}");

        // The base relation is unconstrained: both facts admit.
        let s = answer_query::<Godel>(&m, &Query::anchor(0, 0), &cfg);
        assert!(s[1] > 0.99 && s[2] > 0.99, "{s:?}");
    }

    /// The non-contiguous set intervals cannot represent: NOT-during
    /// (complement of a between) admits both rays.
    #[test]
    fn complement_set_hops_work() {
        let mut m = TemporalPointModel::new(Planted);
        let not_mid = m
            .windowed(0, TimeSet::between(2, 2, 4).complement())
            .unwrap();
        let s = m.project(0, not_mid);
        assert!(
            s[1] > 0.99 && s[2] > 0.99,
            "τ=1 and τ=3 both admitted: {s:?}"
        );
    }

    #[test]
    fn invalid_registrations_and_ids_are_safe() {
        let mut m = TemporalPointModel::new(Planted);
        assert_eq!(m.windowed(5, TimeSet::all(4)), None, "unknown relation");
        assert_eq!(m.windowed(0, TimeSet::all(9)), None, "axis mismatch");
        let empty = m.windowed(0, TimeSet::empty(4)).unwrap();
        assert!(m.project(0, empty).iter().all(|&d| d == 0.0));
        assert!(m.project(0, 99).iter().all(|&d| d == 0.0), "unknown id");
        assert!(m.project(99, 0).iter().all(|&d| d == 0.0), "bad anchor");
    }

    #[test]
    fn temperature_spreads_but_preserves_ranking() {
        let sharp = TemporalPointModel::new(Planted);
        let mut soft = TemporalPointModel::with_temperature(Planted, 5.0);
        let _ = soft.windowed(0, TimeSet::all(4));
        let a = sharp.project(0, 0);
        let b = soft.project(0, 0);
        assert!(b[1] < a[1] && b[1] > 0.5, "softened but still positive");
        assert!(b[0] > a[0] && b[0] < 0.5);
    }
}