geographdb-core 0.5.4

//! Parallel transport on the geometric graph.
//!
//! Standard transformer attention is a scalar-weighted identity transport over a
//! fully-connected graph. Here the transport matrix `A_uv` is a Rodrigues rotation
//! derived directly from the 3D positions of tokens on the measured activation
//! manifold. The rotation angle and incoming Ollivier-Ricci curvature combine into
//! a geometric transition score with zero learned parameters.

use crate::algorithms::four_d::GraphNode4D;
use crate::algorithms::ricci::{ollivier_ricci, RicciEdge};
use glam::{Mat3, Vec3};
use std::collections::HashMap;

/// Rodrigues rotation matrix that maps unit vector `from` onto unit vector `to`.
///
/// Returns the identity when the vectors are parallel, anti-parallel, or either is
/// degenerate. The rotation axis is `from × to`; the angle is `acos(from · to)`.
pub fn rodrigues_rotation(from: Vec3, to: Vec3) -> Mat3 {
    let a = from.normalize();
    let b = to.normalize();
    if !a.is_finite() || !b.is_finite() {
        return Mat3::IDENTITY;
    }

    let v = a.cross(b);
    let s = v.length();
    if s < 1e-6 {
        // Parallel or anti-parallel: identity is the correct limiting rotation.
        return Mat3::IDENTITY;
    }

    let c = a.dot(b).clamp(-1.0, 1.0);
    let k = v / s; // unit rotation axis

    // Skew-symmetric cross-product matrix for axis k.
    let skew = Mat3::from_cols(
        Vec3::new(0.0, k.z, -k.y),
        Vec3::new(-k.z, 0.0, k.x),
        Vec3::new(k.y, -k.x, 0.0),
    );

    // R = I + sin(θ) K + (1 - cos(θ)) K^2, with sin(θ) = s and cos(θ) = c.
    Mat3::IDENTITY + s * skew + (1.0 - c) * (skew * skew)
}

/// Tangent vector along a directed edge (normalised displacement).
pub fn edge_tangent(graph: &[GraphNode4D], from_idx: usize, to_idx: usize) -> Vec3 {
    let from_pos = graph[from_idx].position();
    let to_pos = graph[to_idx].position();
    let dir = to_pos - from_pos;
    let len = dir.length();
    if len < 1e-6 {
        // Degenerate edge: use a finite fallback direction rather than NaN.
        Vec3::X
    } else {
        dir / len
    }
}

/// Angle between the incoming tangent (prev → curr) and outgoing tangent
/// (curr → next). Returns a value in [0, π].
pub fn transport_angle(
    graph: &[GraphNode4D],
    prev_idx: usize,
    curr_idx: usize,
    next_idx: usize,
) -> f32 {
    let t_in = edge_tangent(graph, prev_idx, curr_idx);
    let t_out = edge_tangent(graph, curr_idx, next_idx);
    let cos_angle = t_in.dot(t_out).clamp(-1.0, 1.0);
    cos_angle.acos()
}

/// Build a lookup from undirected edge to Ollivier-Ricci curvature.
pub fn curvature_map(graph: &[GraphNode4D], alpha: f32) -> HashMap<(u64, u64), f32> {
    let mut map = HashMap::new();
    for RicciEdge {
        src,
        dst,
        curvature,
        ..
    } in ollivier_ricci(graph, alpha)
    {
        let key = if src <= dst { (src, dst) } else { (dst, src) };
        map.insert(key, curvature);
    }
    map
}

/// Curvature of the undirected edge between two node IDs.
fn lookup_curvature(curvature: Option<&HashMap<(u64, u64), f32>>, a: u64, b: u64) -> f32 {
    curvature
        .and_then(|m| {
            let key = if a <= b { (a, b) } else { (b, a) };
            m.get(&key).copied()
        })
        .unwrap_or(0.0)
}

/// Geometric transition score for moving curr → next given that we arrived at
/// curr from prev.
///
/// `lambda` controls how sharply the score decays with rotation angle.
/// `curvature` is optional; when present, the incoming Ricci curvature κ(prev,curr)
/// modulates the score via `exp(-kappa * κ)` (negative curvature = hyperbolic =
/// encouraged).
pub fn transport_score(
    graph: &[GraphNode4D],
    prev_idx: usize,
    curr_idx: usize,
    next_idx: usize,
    lambda: f32,
    curvature: Option<&HashMap<(u64, u64), f32>>,
) -> f32 {
    let angle = transport_angle(graph, prev_idx, curr_idx, next_idx);
    let kappa = lookup_curvature(curvature, graph[prev_idx].id, graph[curr_idx].id);

    // Rotation smoothness: small angle → high score.
    let angle_factor = (-lambda * angle).exp();

    // Curvature factor: negative Ricci curvature (hyperbolic) increases score,
    // positive curvature (spherical) decreases it.
    let curvature_factor = (-kappa).exp();

    angle_factor * curvature_factor
}

/// Graph attention scores over transported neighbor features.
///
/// Returns a map from candidate graph index to `(raw_score, softmax_alpha)`.
/// With identity W_Q and W_K:
///   Q   = h_current (the current node's feature vector)
///   K_u = R(t_in, t_out) · h_u
///   V_u = R(t_in, t_out) · h_u
///
/// The raw score is Q · K_u / sqrt(dim). `dim` is taken as 3 because the
/// current graph uses 3D positions as node features.
pub fn graph_attention_scores(
    graph: &[GraphNode4D],
    current_idx: usize,
    prev_idx: Option<usize>,
    candidate_indices: &[usize],
) -> HashMap<usize, (f32, f32)> {
    // Use unit-normalised positions as the first test features (W_Q = W_K = I).
    // This makes the attention score a cosine of the angle between the current
    // node's direction in space and the transported neighbour direction.
    let h_current = graph[current_idx].position().normalize_or(Vec3::X);
    let sqrt_d = 3.0f32.sqrt();
    let t_in = prev_idx.map(|p| edge_tangent(graph, p, current_idx));

    let mut raw_scores = Vec::with_capacity(candidate_indices.len());
    for &next_idx in candidate_indices {
        let h_u = graph[next_idx].position().normalize_or(Vec3::X);
        let transported = if let Some(t_in) = t_in {
            let t_out = edge_tangent(graph, current_idx, next_idx);
            let rotation = rodrigues_rotation(t_in, t_out);
            rotation * h_u
        } else {
            h_u
        };
        let score = h_current.dot(transported.normalize_or(Vec3::X)) / sqrt_d;
        raw_scores.push((next_idx, score));
    }

    // Softmax over raw scores for the alpha distribution.
    let max_score = raw_scores
        .iter()
        .map(|(_, s)| *s)
        .fold(f32::NEG_INFINITY, f32::max);
    let exp_scores: Vec<f32> = raw_scores
        .iter()
        .map(|(_, s)| (s - max_score).exp())
        .collect();
    let sum_exp: f32 = exp_scores.iter().sum();

    let mut result = HashMap::with_capacity(raw_scores.len());
    for ((idx, raw), e) in raw_scores.iter().zip(exp_scores.iter()) {
        let alpha = if sum_exp > 0.0 { *e / sum_exp } else { 0.0 };
        result.insert(*idx, (*raw, alpha));
    }
    result
}

/// Score all successors of `curr_idx` given arrival from `prev_idx`.
pub fn successor_transport_scores(
    graph: &[GraphNode4D],
    prev_idx: usize,
    curr_idx: usize,
    lambda: f32,
    curvature: Option<&HashMap<(u64, u64), f32>>,
) -> HashMap<u64, f32> {
    let mut scores = HashMap::new();
    for edge in &graph[curr_idx].successors {
        let next_idx = graph
            .iter()
            .position(|n| n.id == edge.dst)
            .expect("successor node not found in graph");
        scores.insert(
            edge.dst,
            transport_score(graph, prev_idx, curr_idx, next_idx, lambda, curvature),
        );
    }
    scores
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::algorithms::four_d::{GraphProperties, TemporalEdge};
    use crate::GraphNode4D;

    fn make_node(id: u64, x: f32, y: f32, z: f32) -> GraphNode4D {
        GraphNode4D {
            id,
            x,
            y,
            z,
            begin_ts: 0,
            end_ts: u64::MAX,
            properties: GraphProperties::default(),
            successors: Vec::new(),
        }
    }

    #[test]
    fn rodrigues_maps_x_to_y() {
        let r = rodrigues_rotation(Vec3::X, Vec3::Y);
        let mapped = r * Vec3::X;
        assert!((mapped - Vec3::Y).length() < 1e-4);
    }

    #[test]
    fn rodrigues_identity_for_parallel() {
        let r = rodrigues_rotation(Vec3::X, Vec3::X);
        assert_eq!(r, Mat3::IDENTITY);
    }

    #[test]
    fn straight_line_zero_transport_angle() {
        let mut a = make_node(1000, 0.0, 0.0, 0.0);
        let mut b = make_node(2000, 1.0, 0.0, 0.0);
        let c = make_node(3000, 2.0, 0.0, 0.0);
        a.successors.push(TemporalEdge {
            dst: 2000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        b.successors.push(TemporalEdge {
            dst: 3000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        let graph = vec![a, b, c];

        let angle = transport_angle(&graph, 0, 1, 2);
        assert!(angle.abs() < 1e-4);
    }

    #[test]
    fn right_angle_transport_score_lower() {
        let mut a = make_node(1000, 0.0, 0.0, 0.0);
        let mut b = make_node(2000, 1.0, 0.0, 0.0);
        let c = make_node(3000, 1.0, 1.0, 0.0);
        let d = make_node(4000, 2.0, 0.0, 0.0);
        a.successors.push(TemporalEdge {
            dst: 2000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        b.successors.push(TemporalEdge {
            dst: 3000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        b.successors.push(TemporalEdge {
            dst: 4000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        let graph = vec![a, b, c, d];

        // b -> d continues straight; b -> c turns 90°.
        let score_straight = transport_score(&graph, 0, 1, 3, 1.0, None);
        let score_turn = transport_score(&graph, 0, 1, 2, 1.0, None);
        assert!(score_straight > score_turn);
    }
}