geographdb-core 0.5.4

//! Sparse attention via octree nearest-neighbor search and causal time masking.
//!
//! Each token occupies a position in 3D concept space. For each query token the
//! octree finds the `k` spatially closest key tokens in O(n log n) rather than
//! O(n²). A causal mask then removes keys whose `time_step` lies strictly after
//! the query, producing an auto-regressive attention pattern. Attention weights
//! are computed as a softmax over negative squared distances scaled by a
//! temperature parameter.

use crate::spatial::octree::Octree;
use crate::storage::data_structures::NodePoint;
use glam::Vec3;
use std::collections::HashMap;

// ── Public Types ──────────────────────────────────────────────────────────────

/// A token embedded in 3D concept space with a causal time index.
pub struct Token {
    pub id: u64,
    /// Position in concept space; distance determines attention affinity.
    pub position: Vec3,
    /// Auto-regressive time index; queries only attend to keys with `time_step
    /// <= query.time_step`.
    pub time_step: u64,
}

/// Attention output for a single query token.
pub struct AttentionResult {
    pub query_id: u64,
    /// Causally-valid neighbors sorted by descending weight: `(key_id, weight)`.
    /// Weights are softmax-normalised and sum to 1.0 (empty when no valid key
    /// exists).
    pub attended: Vec<(u64, f32)>,
}

/// Aggregate sparsity statistics over all query results.
pub struct AttentionStats {
    pub n_queries: usize,
    pub mean_attended: f32,
    pub min_attended: usize,
    pub max_attended: usize,
    /// n_queries * n_keys — the dense baseline cost.
    pub total_pairs: usize,
    /// Total (query, key) pairs with nonzero weight.
    pub attended_pairs: usize,
    /// `1 - attended_pairs / total_pairs`; 1.0 means fully sparse.
    pub sparsity: f32,
}

// ── Helpers ───────────────────────────────────────────────────────────────────

fn to_node_point(t: &Token) -> NodePoint {
    NodePoint {
        id: t.id,
        x: t.position.x,
        y: t.position.y,
        z: t.position.z,
    }
}

/// Build an octree from `keys` whose bounding box also encompasses all `query`
/// positions. This ensures every query falls inside the octree bounds so
/// `query_knn` uses a radius large enough to find neighbors.
fn build_key_octree(queries: &[Token], keys: &[Token]) -> Octree {
    use crate::spatial::octree::BoundingBox;

    let all: Vec<Vec3> = queries
        .iter()
        .chain(keys.iter())
        .map(|t| t.position)
        .collect();

    let mut min = all[0];
    let mut max = all[0];
    for &p in &all {
        min = min.min(p);
        max = max.max(p);
    }

    // Generous padding so the KNN radius covers the full space.
    let span = (max - min).length().max(1.0);
    let pad = Vec3::splat(span * 0.25);
    let bounds = BoundingBox::new(min - pad, max + pad);

    let mut octree = Octree::new(bounds);
    for key in keys {
        octree.insert(to_node_point(key));
    }
    octree
}

fn softmax_weights(neg_dist_sq: &[f32]) -> Vec<f32> {
    let max = neg_dist_sq
        .iter()
        .cloned()
        .fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = neg_dist_sq.iter().map(|&v| (v - max).exp()).collect();
    let sum: f32 = exps.iter().sum();
    exps.iter().map(|&e| e / sum).collect()
}

// ── Public API ────────────────────────────────────────────────────────────────

/// Run sparse causal attention.
///
/// For each query the octree returns at most `k` nearest key candidates, then
/// causal masking removes keys with `time_step > query.time_step`. Weights are
/// softmax-normalised over the remaining candidates.
///
/// `temperature` scales the negative distances before softmax: higher values
/// produce a flatter (more uniform) distribution; lower values sharpen toward
/// the single nearest token.
pub fn sparse_attention(
    queries: &[Token],
    keys: &[Token],
    k: usize,
    temperature: f32,
) -> Vec<AttentionResult> {
    if queries.is_empty() {
        return Vec::new();
    }
    if keys.is_empty() {
        return queries
            .iter()
            .map(|q| AttentionResult {
                query_id: q.id,
                attended: Vec::new(),
            })
            .collect();
    }

    let octree = build_key_octree(queries, keys);

    let key_time: HashMap<u64, u64> = keys.iter().map(|t| (t.id, t.time_step)).collect();
    let temp = temperature.max(1e-6);

    queries
        .iter()
        .map(|query| {
            let neighbors = octree.query_knn(query.position, k);

            // Apply causal mask: only keys at or before query's time step.
            let causal: Vec<(u64, f32)> = neighbors
                .into_iter()
                .filter(|(np, _)| {
                    key_time
                        .get(&np.id)
                        .is_some_and(|&ts| ts <= query.time_step)
                })
                .map(|(np, dist_sq)| (np.id, dist_sq))
                .collect();

            if causal.is_empty() {
                return AttentionResult {
                    query_id: query.id,
                    attended: Vec::new(),
                };
            }

            let neg_dists: Vec<f32> = causal.iter().map(|(_, d)| -d / temp).collect();
            let weights = softmax_weights(&neg_dists);

            let attended = causal
                .iter()
                .zip(weights)
                .map(|((id, _), w)| (*id, w))
                .collect();

            AttentionResult {
                query_id: query.id,
                attended,
            }
        })
        .collect()
}

/// Compute sparsity statistics over a batch of attention results.
pub fn attention_stats(results: &[AttentionResult], n_keys: usize) -> AttentionStats {
    let n_queries = results.len();
    if n_queries == 0 {
        return AttentionStats {
            n_queries: 0,
            mean_attended: 0.0,
            min_attended: 0,
            max_attended: 0,
            total_pairs: 0,
            attended_pairs: 0,
            sparsity: 1.0,
        };
    }

    let counts: Vec<usize> = results.iter().map(|r| r.attended.len()).collect();
    let attended_pairs: usize = counts.iter().sum();
    let min_attended = *counts.iter().min().unwrap();
    let max_attended = *counts.iter().max().unwrap();
    let mean_attended = attended_pairs as f32 / n_queries as f32;
    let total_pairs = n_queries * n_keys;
    let sparsity = if total_pairs == 0 {
        1.0
    } else {
        1.0 - attended_pairs as f32 / total_pairs as f32
    };

    AttentionStats {
        n_queries,
        mean_attended,
        min_attended,
        max_attended,
        total_pairs,
        attended_pairs,
        sparsity,
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use glam::Vec3;

    fn tok(id: u64, x: f32, y: f32, z: f32, t: u64) -> Token {
        Token {
            id,
            position: Vec3::new(x, y, z),
            time_step: t,
        }
    }

    #[test]
    fn test_empty_queries_returns_empty() {
        let result = sparse_attention(&[], &[tok(0, 0.0, 0.0, 0.0, 0)], 3, 1.0);
        assert!(result.is_empty());
    }

    #[test]
    fn test_empty_keys_returns_empty_attended() {
        let queries = vec![tok(0, 0.0, 0.0, 0.0, 0)];
        let result = sparse_attention(&queries, &[], 3, 1.0);
        assert_eq!(result.len(), 1);
        assert!(result[0].attended.is_empty());
    }

    #[test]
    fn test_single_pair_weight_is_one() {
        let queries = vec![tok(0, 0.0, 0.0, 0.0, 5)];
        let keys = vec![tok(1, 0.1, 0.0, 0.0, 3)];
        let result = sparse_attention(&queries, &keys, 1, 1.0);
        assert_eq!(result[0].attended.len(), 1);
        let (id, w) = result[0].attended[0];
        assert_eq!(id, 1);
        assert!(
            (w - 1.0).abs() < 1e-5,
            "single key must get weight 1.0, got {w}"
        );
    }

    #[test]
    fn test_causal_mask_excludes_future_keys() {
        // Query at time 2; only key 1 (time 1) should be attended; key 2 (time 3) masked.
        let queries = vec![tok(0, 0.0, 0.0, 0.0, 2)];
        let keys = vec![tok(1, 0.1, 0.0, 0.0, 1), tok(2, 0.2, 0.0, 0.0, 3)];
        let result = sparse_attention(&queries, &keys, 5, 1.0);
        let ids: Vec<u64> = result[0].attended.iter().map(|(id, _)| *id).collect();
        assert!(ids.contains(&1), "past key must be attended");
        assert!(!ids.contains(&2), "future key must be masked");
    }

    #[test]
    fn test_same_time_step_is_allowed() {
        let queries = vec![tok(0, 0.0, 0.0, 0.0, 3)];
        let keys = vec![tok(1, 0.1, 0.0, 0.0, 3)];
        let result = sparse_attention(&queries, &keys, 1, 1.0);
        assert_eq!(
            result[0].attended.len(),
            1,
            "equal time_step should not be masked"
        );
    }

    #[test]
    fn test_k_limits_attended_count() {
        let queries = vec![tok(0, 0.5, 0.5, 0.5, 10)];
        let keys: Vec<Token> = (1u64..=8)
            .map(|i| tok(i, i as f32 * 0.1, 0.0, 0.0, i))
            .collect();
        let result = sparse_attention(&queries, &keys, 3, 1.0);
        assert!(
            result[0].attended.len() <= 3,
            "attended count must not exceed k"
        );
    }

    #[test]
    fn test_weights_sum_to_one() {
        let queries = vec![tok(0, 0.0, 0.0, 0.0, 10)];
        let keys: Vec<Token> = (1u64..=5)
            .map(|i| tok(i, i as f32 * 0.2, 0.0, 0.0, i))
            .collect();
        let result = sparse_attention(&queries, &keys, 5, 1.0);
        let sum: f32 = result[0].attended.iter().map(|(_, w)| w).sum();
        assert!(
            (sum - 1.0).abs() < 1e-5,
            "weights must sum to 1.0, got {sum}"
        );
    }

    #[test]
    fn test_closer_token_gets_higher_weight() {
        let queries = vec![tok(0, 0.0, 0.0, 0.0, 10)];
        let keys = vec![
            tok(1, 0.1, 0.0, 0.0, 1), // close
            tok(2, 5.0, 0.0, 0.0, 2), // far
        ];
        let result = sparse_attention(&queries, &keys, 5, 1.0);
        let w1 = result[0]
            .attended
            .iter()
            .find(|(id, _)| *id == 1)
            .map(|(_, w)| *w)
            .expect("close key must be attended");
        let w2 = result[0]
            .attended
            .iter()
            .find(|(id, _)| *id == 2)
            .map(|(_, w)| *w)
            .expect("far key must be attended");
        assert!(
            w1 > w2,
            "closer token must have higher weight ({w1:.4} vs {w2:.4})"
        );
    }

    #[test]
    fn test_attention_stats_fields() {
        let queries: Vec<Token> = (0..4).map(|i| tok(i, i as f32, 0.0, 0.0, 10)).collect();
        let keys: Vec<Token> = (10..14)
            .map(|i| tok(i, (i - 10) as f32 * 0.5, 0.0, 0.0, 5))
            .collect();
        let results = sparse_attention(&queries, &keys, 2, 1.0);
        let stats = attention_stats(&results, keys.len());
        assert_eq!(stats.n_queries, 4);
        assert_eq!(stats.total_pairs, 16);
        assert!(stats.sparsity >= 0.0 && stats.sparsity <= 1.0);
        assert!(stats.attended_pairs <= stats.total_pairs);
    }
}