geographdb-core 0.5.4

//! Takens Delay Embedding and Correlation Dimension estimator.
//!
//! Takens' theorem (1981): if a time series is generated by a deterministic
//! dynamical system with attractor dimension `d`, then `2d+1` delay-embedded
//! observations faithfully reconstruct the attractor — regardless of total
//! sequence length.
//!
//! Applied to KV caches: the minimum number of KV vectors needed for lossless
//! attention reconstruction is bounded by `2d+1`, where `d` is the correlation
//! dimension of the KV sequence treated as a dynamical trajectory.
//!
//! # Workflow
//! 1. Collect a KV sequence as rows of a matrix.
//! 2. Embed with `delay_embedding(series, tau, embed_dim)`.
//! 3. Estimate `d` with `correlation_dimension(points, ...)`.
//! 4. Compute minimum cache size: `min_cache_size_takens(d)`.

// ── Delay embedding ───────────────────────────────────────────────────────────

/// Build a delay-embedded point cloud from a scalar or vector time series.
///
/// Each input row `series[t]` is a feature vector of length `n_features`.
/// The output row `t` is the concatenation of `series[t], series[t-tau],
/// series[t-2*tau], ..., series[t-(embed_dim-1)*tau]`, zero-padded for
/// indices before the start of the series.
///
/// Returns a matrix of shape `[series.len(), n_features * embed_dim]`.
pub fn delay_embedding(series: &[Vec<f32>], tau: usize, embed_dim: usize) -> Vec<Vec<f32>> {
    let n = series.len();
    if n == 0 || embed_dim == 0 {
        return Vec::new();
    }
    let n_features = series[0].len();
    let out_dim = n_features * embed_dim;

    (0..n)
        .map(|t| {
            let mut row = vec![0.0f32; out_dim];
            for lag in 0..embed_dim {
                let src_t = (t as isize) - (lag * tau) as isize;
                if src_t >= 0 {
                    let src = &series[src_t as usize];
                    let offset = lag * n_features;
                    row[offset..offset + n_features].copy_from_slice(src);
                }
                // else: stays zero (zero-pad)
            }
            row
        })
        .collect()
}

// ── Correlation dimension ─────────────────────────────────────────────────────

/// Estimate the correlation dimension of a point cloud using the
/// Grassberger–Procaccia algorithm (1983).
///
/// Sweeps radius `r` from `r_min` to `r_max` (log-uniform, `n_steps` steps),
/// computes `C(r) = fraction of pairs with distance < r`, fits `log C(r) vs
/// log r` by least-squares, and returns the slope (= correlation dimension).
///
/// Uses Euclidean distance. For large point clouds, samples up to
/// `max_pairs` random pairs to keep cost manageable.
pub fn correlation_dimension(points: &[Vec<f32>], r_min: f32, r_max: f32, n_steps: usize) -> f64 {
    let n = points.len();
    if n < 4 || r_min <= 0.0 || r_max <= r_min || n_steps < 2 {
        return 0.0;
    }

    // Build log-uniform radius grid
    let log_rmin = (r_min as f64).ln();
    let log_rmax = (r_max as f64).ln();
    let radii: Vec<f64> = (0..n_steps)
        .map(|i| {
            let t = i as f64 / (n_steps - 1) as f64;
            (log_rmin + t * (log_rmax - log_rmin)).exp()
        })
        .collect();

    // Count pairs within each radius (O(n²) — acceptable for small n)
    let total_pairs = (n * (n - 1) / 2) as f64;
    let mut log_r_fit: Vec<f64> = Vec::with_capacity(n_steps);
    let mut log_c_fit: Vec<f64> = Vec::with_capacity(n_steps);

    for &r in &radii {
        let mut count = 0usize;
        for i in 0..n {
            for j in (i + 1)..n {
                if euclidean_dist_sq(&points[i], &points[j]) < (r * r) as f32 {
                    count += 1;
                }
            }
        }
        // Skip radii where fewer than 2 pairs fall inside — log(ε) at those
        // points is dominated by the floor constant and distorts the slope.
        if count < 2 {
            continue;
        }
        let c = count as f64 / total_pairs;
        log_r_fit.push(r.ln());
        log_c_fit.push(c.ln());
    }

    // Need at least two valid points for a slope
    let m = log_r_fit.len();
    if m < 2 {
        return 0.0;
    }

    // Least-squares slope: d = Σ(log_r * log_c) / Σ(log_r²) (mean-centred)
    let mean_r = log_r_fit.iter().sum::<f64>() / m as f64;
    let mean_c = log_c_fit.iter().sum::<f64>() / m as f64;
    let num: f64 = log_r_fit
        .iter()
        .zip(log_c_fit.iter())
        .map(|(r, c)| (r - mean_r) * (c - mean_c))
        .sum();
    let den: f64 = log_r_fit.iter().map(|r| (r - mean_r).powi(2)).sum();
    if den.abs() < 1e-15 {
        return 0.0;
    }
    (num / den).max(0.0)
}

/// Minimum number of KV vectors needed to reconstruct the attractor, per
/// Takens' theorem: `2 * ceil(d) + 1`.
pub fn min_cache_size_takens(dim_estimate: f64) -> usize {
    2 * dim_estimate.ceil() as usize + 1
}

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

fn euclidean_dist_sq(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| (x - y).powi(2)).sum()
}

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

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

    // ── delay_embedding ───────────────────────────────────────────────────────

    #[test]
    fn test_delay_embedding_shape() {
        let series: Vec<Vec<f32>> = (0..10).map(|i| vec![i as f32, i as f32 + 1.0]).collect();
        let embedded = delay_embedding(&series, 1, 3);
        assert_eq!(embedded.len(), 10);
        assert_eq!(embedded[0].len(), 6); // 2 features × 3 lags
    }

    #[test]
    fn test_delay_embedding_tau1_lag0_matches_original() {
        let series: Vec<Vec<f32>> = (0..5).map(|i| vec![i as f32]).collect();
        let embedded = delay_embedding(&series, 1, 2);
        // First lag should always equal the current time step
        for (t, row) in embedded.iter().enumerate() {
            assert_eq!(row[0], t as f32, "lag-0 mismatch at t={t}");
        }
    }

    #[test]
    fn test_delay_embedding_zero_pads_before_start() {
        let series: Vec<Vec<f32>> = (0..4).map(|i| vec![i as f32]).collect();
        let embedded = delay_embedding(&series, 1, 3);
        // t=0: lags [t=0, t=-1, t=-2] → [0.0, 0.0, 0.0]
        assert_eq!(embedded[0], vec![0.0, 0.0, 0.0]);
        // t=1: lags [t=1, t=0, t=-1] → [1.0, 0.0, 0.0]
        assert_eq!(embedded[1], vec![1.0, 0.0, 0.0]);
        // t=2: lags [t=2, t=1, t=0] → [2.0, 1.0, 0.0]
        assert_eq!(embedded[2], vec![2.0, 1.0, 0.0]);
    }

    #[test]
    fn test_delay_embedding_empty_series() {
        let embedded = delay_embedding(&[], 1, 3);
        assert!(embedded.is_empty());
    }

    // ── correlation_dimension ─────────────────────────────────────────────────

    /// A 1D line → correlation dimension ≈ 1.0
    #[test]
    fn test_correlation_dim_line() {
        // 50 evenly-spaced points on [0, 1]
        let points: Vec<Vec<f32>> = (0..50).map(|i| vec![i as f32 / 49.0]).collect();
        let d = correlation_dimension(&points, 0.01, 1.0, 20);
        assert!(
            (d - 1.0).abs() < 0.3,
            "line should have corr dim ≈ 1.0, got {d:.3}"
        );
    }

    /// Points on a unit circle → correlation dimension ≈ 1.0 (1D manifold)
    #[test]
    fn test_correlation_dim_circle() {
        let points: Vec<Vec<f32>> = (0..60)
            .map(|i| {
                let theta = 2.0 * std::f32::consts::PI * i as f32 / 60.0;
                vec![theta.cos(), theta.sin()]
            })
            .collect();
        let d = correlation_dimension(&points, 0.05, 1.5, 20);
        assert!(
            (d - 1.0).abs() < 0.4,
            "circle should have corr dim ≈ 1.0, got {d:.3}"
        );
    }

    /// Points on a 2D surface → correlation dimension ≈ 2.0
    #[test]
    fn test_correlation_dim_plane() {
        // 100 points on [0,1]² grid
        let points: Vec<Vec<f32>> = (0..10)
            .flat_map(|i| (0..10).map(move |j| vec![i as f32 / 9.0, j as f32 / 9.0]))
            .collect();
        let d = correlation_dimension(&points, 0.05, 1.0, 20);
        assert!(
            (d - 2.0).abs() < 0.5,
            "plane should have corr dim ≈ 2.0, got {d:.3}"
        );
    }

    /// Too few points → returns 0.0 without panic
    #[test]
    fn test_correlation_dim_degenerate() {
        let d = correlation_dimension(&[vec![1.0f32], vec![2.0f32]], 0.1, 1.0, 10);
        assert_eq!(d, 0.0);
    }

    // ── min_cache_size_takens ─────────────────────────────────────────────────

    #[test]
    fn test_min_cache_size_formula() {
        assert_eq!(min_cache_size_takens(1.0), 3); // 2*1+1
        assert_eq!(min_cache_size_takens(2.0), 5); // 2*2+1
        assert_eq!(min_cache_size_takens(64.0), 129); // 2*64+1
        assert_eq!(min_cache_size_takens(1.5), 5); // 2*ceil(1.5)+1 = 2*2+1
    }

    #[test]
    fn test_min_cache_size_fractional_rounds_up() {
        // d=0.7 → ceil=1 → 2*1+1 = 3
        assert_eq!(min_cache_size_takens(0.7), 3);
        // d=2.1 → ceil=3 → 2*3+1 = 7
        assert_eq!(min_cache_size_takens(2.1), 7);
    }
}