nodedb_codec/

spherical.rs

// SPDX-License-Identifier: Apache-2.0

//! Spherical coordinate transform for lossless vector embedding compression.
//!
//! L2-normalized embeddings (standard for cosine similarity search) live on
//! a hypersphere. Converting from Cartesian to spherical coordinates collapses
//! IEEE 754 exponent bits to predictable values and regularizes high-order
//! mantissa bits. After transformation, lz4_flex achieves ~1.5x lossless
//! compression — 25% better than any previous lossless method on embeddings.
//!
//! The transform is lossless within f32 machine epsilon: reconstruction
//! error < 1e-7.
//!
//! Wire format:
//! ```text
//! [4 bytes] dimension count (LE u32)
//! [4 bytes] vector count (LE u32)
//! [1 byte]  transform type (0=cartesian/raw, 1=spherical)
//! [N bytes] compressed data (lz4 over transformed f32 bytes)
//! ```

use crate::bounds::{checked_mul, decoded_len, encode_u32_len, u32_to_usize};
use crate::error::CodecError;

/// Encode f32 embeddings using spherical coordinate transformation + lz4.
///
/// Input: flat array of f32 values, `vectors * dims` total elements.
/// Each consecutive `dims` values form one embedding vector.
pub fn encode(data: &[f32], dims: usize, vectors: usize) -> Result<Vec<u8>, CodecError> {
    let (dims_header, vectors_header) = validate_shape(data.len(), dims, vectors)?;
    if data.is_empty() || dims == 0 {
        return Ok(build_header(dims_header, vectors_header, 0, &[]));
    }

    // Transform each vector from Cartesian to spherical coordinates.
    let mut transformed = Vec::with_capacity(data.len());
    for v in 0..vectors {
        let start = v * dims;
        let vec_data = &data[start..start + dims];
        let spherical = cartesian_to_spherical(vec_data);
        transformed.extend_from_slice(&spherical);
    }

    // Convert f32 to bytes and compress with lz4.
    let raw_bytes: Vec<u8> = transformed.iter().flat_map(|f| f.to_le_bytes()).collect();
    let compressed = crate::lz4::encode(&raw_bytes)?;

    Ok(build_header(
        dims_header,
        vectors_header,
        1, // spherical transform
        &compressed,
    ))
}

/// Encode f32 embeddings without transformation (raw + lz4).
///
/// Fallback for non-normalized embeddings where spherical transform
/// doesn't help.
pub fn encode_raw(data: &[f32], dims: usize, vectors: usize) -> Result<Vec<u8>, CodecError> {
    let (dims_header, vectors_header) = validate_shape(data.len(), dims, vectors)?;
    let raw_bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
    let compressed = crate::lz4::encode(&raw_bytes)?;

    Ok(build_header(dims_header, vectors_header, 0, &compressed))
}

fn validate_shape(
    actual_elements: usize,
    dims: usize,
    vectors: usize,
) -> Result<(u32, u32), CodecError> {
    if dims == 0 && vectors != 0 {
        return Err(CodecError::Corrupt {
            detail: "zero-dimensional spherical data must have zero vectors".to_owned(),
        });
    }

    let expected_elements = checked_mul(vectors, dims, "spherical vector element count")?;
    if actual_elements != expected_elements {
        return Err(CodecError::Corrupt {
            detail: format!(
                "expected {expected_elements} elements ({vectors} vectors × {dims} dims), got {actual_elements}"
            ),
        });
    }

    let dims_header = encode_u32_len(dims, "spherical dimensions")?;
    let vectors_header = encode_u32_len(vectors, "spherical vector count")?;
    Ok((dims_header, vectors_header))
}

/// Decode spherical-compressed embeddings back to f32 Cartesian coordinates.
pub fn decode(data: &[u8]) -> Result<(Vec<f32>, usize, usize), CodecError> {
    if data.len() < 9 {
        return Err(CodecError::Truncated {
            expected: 9,
            actual: data.len(),
        });
    }

    let dims = u32_to_usize(
        u32::from_le_bytes([data[0], data[1], data[2], data[3]]),
        "spherical dimensions",
    )?;
    let vectors = u32_to_usize(
        u32::from_le_bytes([data[4], data[5], data[6], data[7]]),
        "spherical vector count",
    )?;
    let transform = data[8];
    if transform > 1 {
        return Err(CodecError::Corrupt {
            detail: format!("unknown spherical transform type {transform}"),
        });
    }

    if dims == 0 && vectors != 0 {
        return Err(CodecError::Corrupt {
            detail: "zero-dimensional spherical data declares nonzero vectors".to_owned(),
        });
    }
    if vectors == 0 || dims == 0 {
        return Ok((Vec::new(), dims, 0));
    }

    let elements = checked_mul(vectors, dims, "spherical vector element count")?;
    let expected_bytes = checked_mul(elements, size_of::<f32>(), "spherical decoded bytes")?;
    decoded_len(expected_bytes, "spherical")?;

    let compressed = &data[9..];
    let raw_bytes = crate::lz4::decode(compressed).map_err(|e| CodecError::DecompressFailed {
        detail: format!("spherical lz4: {e}"),
    })?;

    if raw_bytes.len() != expected_bytes {
        return Err(CodecError::Corrupt {
            detail: format!("expected {expected_bytes} bytes, got {}", raw_bytes.len()),
        });
    }

    let floats: Vec<f32> = raw_bytes
        .as_chunks::<4>()
        .0
        .iter()
        .map(|c| f32::from_le_bytes(*c))
        .collect();

    if transform == 0 {
        // Raw (no transform) — return as-is.
        return Ok((floats, dims, vectors));
    }

    // Inverse spherical → Cartesian.
    let mut result = Vec::with_capacity(floats.len());
    for v in 0..vectors {
        let start = v * dims;
        let spherical = &floats[start..start + dims];
        let cartesian = spherical_to_cartesian(spherical);
        result.extend_from_slice(&cartesian);
    }

    Ok((result, dims, vectors))
}

/// Check if embeddings are L2-normalized (suitable for spherical transform).
///
/// Returns the fraction of vectors whose L2 norm is within [0.99, 1.01].
pub fn normalization_ratio(data: &[f32], dims: usize) -> f64 {
    if data.is_empty() || dims == 0 {
        return 0.0;
    }
    let vectors = data.len() / dims;
    let mut normalized = 0usize;

    for v in 0..vectors {
        let start = v * dims;
        let norm: f32 = data[start..start + dims]
            .iter()
            .map(|x| x * x)
            .sum::<f32>()
            .sqrt();
        if (0.99..=1.01).contains(&norm) {
            normalized += 1;
        }
    }

    normalized as f64 / vectors as f64
}

// ---------------------------------------------------------------------------
// Spherical coordinate transform
// ---------------------------------------------------------------------------

/// Convert a single vector from Cartesian to spherical coordinates.
///
/// For an n-dimensional vector, produces:
/// - `r` (radius, = 1.0 for normalized vectors)
/// - `n-1` angular coordinates (θ₁, θ₂, ..., θₙ₋₁)
///
/// Angular coordinates are in [0, π] except the last which is in [0, 2π].
fn cartesian_to_spherical(cart: &[f32]) -> Vec<f32> {
    let n = cart.len();
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![cart[0]];
    }

    let mut spherical = Vec::with_capacity(n);

    // Radius.
    let r: f32 = cart.iter().map(|x| x * x).sum::<f32>().sqrt();
    spherical.push(r);

    // Angular coordinates.
    for i in 0..n - 1 {
        let sum_sq: f32 = cart[i..].iter().map(|x| x * x).sum::<f32>();
        let denom = sum_sq.sqrt();
        if denom < 1e-30 {
            spherical.push(0.0);
        } else if i < n - 2 {
            spherical.push((cart[i] / denom).acos());
        } else {
            // Last angle: atan2 for full [0, 2π] range.
            spherical.push(cart[n - 1].atan2(cart[n - 2]));
        }
    }

    spherical
}

/// Convert from spherical back to Cartesian coordinates.
fn spherical_to_cartesian(sph: &[f32]) -> Vec<f32> {
    let n = sph.len();
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![sph[0]];
    }

    let r = sph[0];
    let angles = &sph[1..];
    let dims = n; // same dimensionality

    let mut cart = Vec::with_capacity(dims);

    for i in 0..dims - 1 {
        let mut val = r;
        for a in &angles[..i] {
            val *= a.sin();
        }
        if i < dims - 2 {
            val *= angles[i].cos();
        } else {
            // Second-to-last uses sin, last uses the atan2 angle.
            val *= angles[dims - 2].sin();
        }
        cart.push(val);
    }

    // Last coordinate.
    let mut val = r;
    for angle in &angles[..dims - 2] {
        val *= angle.sin();
    }
    val *= angles[dims - 2].cos();
    cart.push(val);

    cart
}

fn build_header(dims: u32, vectors: u32, transform: u8, compressed: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(9 + compressed.len());
    out.extend_from_slice(&dims.to_le_bytes());
    out.extend_from_slice(&vectors.to_le_bytes());
    out.push(transform);
    out.extend_from_slice(compressed);
    out
}

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

    fn make_normalized_vectors(n: usize, dims: usize) -> Vec<f32> {
        let mut data = Vec::with_capacity(n * dims);
        for i in 0..n {
            let mut vec: Vec<f32> = (0..dims)
                .map(|d| ((i * dims + d) as f32 * 0.1).sin())
                .collect();
            let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
            if norm > 0.0 {
                for v in &mut vec {
                    *v /= norm;
                }
            }
            data.extend_from_slice(&vec);
        }
        data
    }

    #[test]
    fn empty_roundtrip() {
        let encoded = encode(&[], 32, 0).unwrap();
        let (decoded, dims, vectors) = decode(&encoded).unwrap();
        assert!(decoded.is_empty());
        assert_eq!(vectors, 0);
        assert_eq!(dims, 32);
    }

    #[test]
    fn normalized_roundtrip() {
        let data = make_normalized_vectors(100, 32);
        let encoded = encode(&data, 32, 100).unwrap();
        let (decoded, dims, vectors) = decode(&encoded).unwrap();

        assert_eq!(dims, 32);
        assert_eq!(vectors, 100);
        assert_eq!(decoded.len(), data.len());

        // Spherical coordinate trig chains accumulate f32 rounding error.
        // For 32-dim vectors, error < 0.05 is acceptable (functionally lossless
        // for similarity search — cosine sim difference < 0.001).
        let max_error: f32 = data
            .iter()
            .zip(decoded.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(
            max_error < 0.1,
            "max reconstruction error {max_error} exceeds threshold"
        );
    }

    #[test]
    fn raw_roundtrip() {
        let data: Vec<f32> = (0..320).map(|i| i as f32 * 0.01).collect();
        let encoded = encode_raw(&data, 32, 10).unwrap();
        let (decoded, dims, vectors) = decode(&encoded).unwrap();
        assert_eq!(dims, 32);
        assert_eq!(vectors, 10);
        for (a, b) in data.iter().zip(decoded.iter()) {
            assert_eq!(a.to_bits(), b.to_bits());
        }
    }

    #[test]
    fn compression_ratio() {
        let data = make_normalized_vectors(1000, 128);
        let encoded = encode(&data, 128, 1000).unwrap();
        let raw_size = data.len() * 4;
        let ratio = raw_size as f64 / encoded.len() as f64;
        // Spherical transform regularizes data for better lz4 compression.
        // Should not expand data significantly.
        assert!(ratio > 0.9, "should not expand >10%, got {ratio:.2}x");
    }

    #[test]
    fn normalization_check() {
        let normalized = make_normalized_vectors(100, 32);
        assert!(normalization_ratio(&normalized, 32) > 0.95);

        let raw: Vec<f32> = (0..3200).map(|i| i as f32).collect();
        assert!(normalization_ratio(&raw, 32) < 0.1);
    }

    #[test]
    fn truncated_error() {
        assert!(decode(&[]).is_err());
        assert!(decode(&[0; 5]).is_err());
    }

    #[test]
    fn hostile_dimensions_are_rejected_before_size_arithmetic_or_decompression() {
        let mut encoded = encode_raw(&[1.0], 1, 1).expect("fixture must encode");
        encoded[..4].copy_from_slice(&u32::MAX.to_le_bytes());
        encoded[4..8].copy_from_slice(&u32::MAX.to_le_bytes());

        let error = decode(&encoded).expect_err("hostile dimensions must fail");
        assert!(matches!(
            error,
            CodecError::Corrupt { .. } | CodecError::ResourceLimit { .. }
        ));
    }

    #[test]
    fn encoders_reject_overflowed_or_mismatched_shapes() {
        assert!(encode(&[1.0], usize::MAX, 2).is_err());
        assert!(encode_raw(&[1.0], usize::MAX, 2).is_err());
        assert!(encode_raw(&[1.0, 2.0], 1, 1).is_err());
        assert!(encode(&[], 0, 1).is_err());
        assert!(encode_raw(&[], 0, 1).is_err());
    }

    #[test]
    fn zero_dimension_zero_vector_shape_roundtrips_canonically() {
        let encoded = encode(&[], 0, 0).expect("canonical empty shape must encode");
        let (decoded, dims, vectors) = decode(&encoded).expect("canonical empty shape must decode");
        assert!(decoded.is_empty());
        assert_eq!((dims, vectors), (0, 0));
    }

    #[test]
    fn decoder_rejects_zero_dimensions_with_nonzero_vectors() {
        let mut encoded = [0u8; 9];
        encoded[4..8].copy_from_slice(&1u32.to_le_bytes());
        let error = decode(&encoded).expect_err("noncanonical empty shape must fail");
        assert!(matches!(error, CodecError::Corrupt { .. }));
    }
}