steiner-tree 0.0.5

// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: GPL-3.0-or-later

//! Compute a hash of multi-dimensional points such that points which are close
//! likely map to the same or similar hash.
//!
//! When the distance metric is the Hamming distance (number of different bits)
//! a simple locality sensitive hash is implemented by selecting some of the bits.

use bitvec::prelude::*;
use itertools::Itertools;
use std::cmp::Ordering;
use std::hash::Hash;
use std::marker::PhantomData;

/// Define a locality sensitive hash function together with
/// a distance metric in the hash domain and a distance metric in the hash image.
/// The two distance metrics and the hash function **must satisfy** `domain_distance(p1, p2) >= image_distance(hash(p1), hash(p2))`.
pub trait LocalitySensitiveHasher {
    type Domain;
    type Image: Eq + Ord;
    /// Distance value between inputs to the hasher.
    type DomainDistance: Ord;
    /// Distance value between hashes.
    type ImageDistance: Ord;

    fn hash(&self, i: &Self::Domain) -> Self::Image;

    /// Compute the distance between two input points.
    fn domain_distance(&self, a: &Self::Domain, b: &Self::Domain) -> Self::DomainDistance;

    /// Compute the distance between two hashes.
    /// The distance between two hashes `h1` and `h2` gives a lower bound of the between
    /// the two set of points which hash to `h1` and `h2`.
    fn image_distance(&self, a: &Self::Image, b: &Self::Image) -> Self::ImageDistance;
}

/// A point in hamming space.
/// This is typically implemented by a sequence of bits.
/// Used to generalize `BitSlice` and machine integers.
pub trait HammingPoint {
    /// Type of coordinate.
    type Coord: Copy + Eq;

    /// Length of the bit sequence.
    fn len(&self) -> usize;

    /// Get the value of a bit.
    fn get(&self, index: usize) -> Self::Coord;

    /// Extract bits at `indices` from the value and pack them into an `usize`.
    fn sub_sample<B>(&self, indices: &[usize]) -> B
    where
        B: BuildHammingPoint<Coord = Self::Coord>,
    {
        debug_assert!(
            indices.len() <= 64,
            "cannot store more than 64 bits in a usize"
        );

        indices
            .into_iter()
            .map(|&i| self.get(i))
            .fold(B::empty(), |acc, x| acc.concatenated(x))
    }

    /// Compute distance between the two strings.
    fn distance(&self, other: &Self) -> u32 {
        let len = self.len().max(other.len());
        (0..len).filter(|&i| self.get(i) != other.get(i)).count() as u32
    }
}

impl HammingPoint for usize {
    type Coord = bool;

    fn len(&self) -> usize {
        64
    }

    fn get(&self, index: usize) -> bool {
        ((*self >> index) & 0b1) == 1
    }
}

impl HammingPoint for BitSlice {
    type Coord = bool;

    fn len(&self) -> usize {
        self.len()
    }

    fn get(&self, index: usize) -> bool {
        if index < self.len() {
            self[index]
        } else {
            false
        }
    }
}

impl HammingPoint for BitVec {
    type Coord = bool;

    fn len(&self) -> usize {
        self.len()
    }

    fn get(&self, index: usize) -> bool {
        if index < self.len() {
            self[index]
        } else {
            false
        }
    }
}

/// Construct a point in Hamming space by appending elements.
pub trait BuildHammingPoint {
    /// Type of coordinate.
    type Coord: Copy + Eq;

    /// Create empty string.
    fn empty() -> Self;

    /// Append an element.
    fn push(&mut self, element: Self::Coord);

    /// Append an element.
    fn concatenated(mut self, element: Self::Coord) -> Self
    where
        Self: Sized,
    {
        self.push(element);
        self
    }
}

impl BuildHammingPoint for usize {
    type Coord = bool;

    fn empty() -> Self {
        0
    }

    fn push(&mut self, element: Self::Coord) {
        *self = (*self << 1) | (element as usize)
    }
}

/// Locality sensitive hasher for bit strings using the hamming distance.
/// Hashing works by sub-sampling the bits, i.e. selecting some bits from the bit string.
/// This is equivalent to a projection onto a lower dimension.
pub struct HammingLSH<HashInput, HashOutput> {
    sub_sample_indices: Vec<usize>,
    _input_type: PhantomData<HashInput>,
    _output_type: PhantomData<HashOutput>,
}

impl<HashInput, HashOutput> HammingLSH<HashInput, HashOutput>
where
    HashInput: HammingPoint,
    HashOutput: BuildHammingPoint<Coord = HashInput::Coord> + Copy + Eq + Hash,
{
    /// Find the a given number of bit indices which are most distinguishing for the given points.
    pub fn new(num_sub_samples: usize, points: &[HashInput]) -> Self {
        assert!(num_sub_samples <= 64);

        // Greedily find bit indices which yield highest entropy when used for subsampling.
        let mut indices = Vec::new();
        let mut tmp_indices = Vec::new();

        // Find maximum length of the bit strings.
        let max_len = points.iter().map(|p| p.len()).max().unwrap_or(0);

        // Find high-entropy bit positions.
        for _ in 0..num_sub_samples {
            // Find the hightest-entropy bit which is not selected yet.
            let best_index = (0..max_len)
                // Skip indices which are already used.
                .filter(|i| !indices.contains(i))
                // Compute entropy of hashed points.
                .map(|i| {
                    // Add new index.
                    tmp_indices.push(i);
                    // Compute hashes of all points.
                    let hashed_points = points
                        .iter()
                        .map(|p| p.sub_sample::<HashOutput>(&tmp_indices));

                    // Compute entropy of hashed points.
                    let entropy = compute_entropy(hashed_points);

                    // Restore indices.
                    tmp_indices.pop();

                    (i, entropy)
                })
                // Find index which yields the highest entropy.
                .max_by(|(_, entropy1), (_, entropy2)| {
                    if entropy1 > entropy2 {
                        Ordering::Greater
                    } else {
                        Ordering::Less
                    }
                })
                .map(|(i, _)| i)
                .unwrap();

            // Use found index.
            indices.push(best_index);
            tmp_indices.push(best_index);
        }

        assert_eq!(indices.len(), num_sub_samples);

        indices.sort_unstable();

        Self {
            sub_sample_indices: indices,
            _input_type: Default::default(),
            _output_type: Default::default(),
        }
    }
}

impl<B> LocalitySensitiveHasher for HammingLSH<B, usize>
where
    B: HammingPoint<Coord = bool>,
{
    type Domain = B;
    type Image = usize;
    type DomainDistance = u32;
    type ImageDistance = u32;

    fn hash(&self, i: &Self::Domain) -> Self::Image {
        i.sub_sample(&self.sub_sample_indices)
    }

    fn domain_distance(&self, a: &Self::Domain, b: &Self::Domain) -> Self::DomainDistance {
        a.distance(b)
    }

    fn image_distance(&self, a: &Self::Image, b: &Self::Image) -> Self::ImageDistance {
        (a ^ b).count_ones()
    }
}

#[test]
fn test_high_entropy_bit_selection() {
    let points = [0b000, 0b010];
    let lsh: HammingLSH<_, usize> = HammingLSH::new(1, &points);
    assert_eq!(lsh.sub_sample_indices, [1]);

    let points = [0b000, 0b010, 0b110];
    let lsh: HammingLSH<_, usize> = HammingLSH::new(2, &points);
    assert_eq!(lsh.sub_sample_indices, [1, 2]);
}

#[test]
fn test_bitvec() {
    let n = 0b101usize;
    let bits = n.view_bits::<Lsb0>();
    assert!(bits[0]);
    assert!(!bits[1]);
    assert!(bits[2]);
}

/// Compute empirical entropy of a sequence.
///
/// Entropy `H(X) = -Sum( p(x_i) * log( p( x_i ) )`
/// where `p(x)` is the probability of `x`.
fn compute_entropy<I>(elements: I) -> f64
where
    I: Iterator,
    I::Item: Hash + Eq + Copy,
{
    let counts = elements.counts();
    let total_count = counts.values().sum::<usize>() as f64;

    // Empirical probabilities.
    let probabilities = counts.into_values().map(|c| (c as f64) / total_count);

    let entropy = -probabilities.map(|p| p * p.log2()).sum::<f64>();

    entropy
}

#[test]
fn test_compute_entropy() {
    let elements = [1, 1, 1, 1];
    assert_eq!(compute_entropy(elements.iter()), 0.0);
    let elements = [1, 2];
    assert_eq!(compute_entropy(elements.iter()), 1.0);
    let elements = [1, 2, 2, 1];
    assert_eq!(compute_entropy(elements.iter()), 1.0);
    let elements = [1, 2, 2, 1, 1];
    assert!(compute_entropy(elements.iter()) < 1.0);
}