Skip to main content

lsh_rs2/
hash.rs

1use crate::data::Integer;
2use crate::multi_probe::StepWiseProbe;
3use crate::{data::Numeric, dist::l2_norm, multi_probe::QueryDirectedProbe, utils::create_rng};
4use ndarray::prelude::*;
5use ndarray_rand::rand_distr::{StandardNormal, Uniform};
6use ndarray_rand::RandomExt;
7use num::{traits::NumCast, Float, Zero};
8use std::marker::PhantomData;
9use serde::{Deserialize, Serialize};
10use std::ops::Deref;
11
12/// Implement this trait to create your own custom hashers.
13/// In case of a symmetrical hash function, only `hash_vec_query` needs to be implemented.
14pub trait VecHash<N, K> {
15    /// Create a hash for a query data point.
16    fn hash_vec_query(&self, v: &[N]) -> Vec<K>;
17    /// Create a hash for a data point that is being stored.
18    fn hash_vec_put(&self, v: &[N]) -> Vec<K> {
19        self.hash_vec_query(v)
20    }
21
22    /// If the hasher implements the QueryDirectedProbe trait it should return Some(self)
23    fn as_query_directed_probe(&self) -> Option<&dyn QueryDirectedProbe<N, K>> {
24        None
25    }
26    /// If the hasher implements the StepWiseProbe trait it should return Some(self)
27    fn as_step_wise_probe(&self) -> Option<&dyn StepWiseProbe<N, K>> {
28        None
29    }
30}
31
32/// A family of hashers for the cosine similarity.
33#[derive(Serialize, Deserialize, Clone)]
34pub struct SignRandomProjections<N: Numeric> {
35    ///  Random unit vectors that will lead to the bits of the hash.
36    hyperplanes: Array2<N>,
37}
38
39impl<N: Numeric> SignRandomProjections<N> {
40    ///
41    /// # Arguments
42    ///
43    /// * `k` - Number of hyperplanes used for determining the hash.
44    /// This will also be the hash length.
45    pub fn new(k: usize, dim: usize, seed: u64) -> Self {
46        let mut rng = create_rng(seed);
47        let hp: Array2<f32> = Array::random_using((k, dim), StandardNormal, &mut rng);
48        let hp = hp.mapv(|v| N::from_f32(v).unwrap());
49
50        SignRandomProjections { hyperplanes: hp }
51    }
52
53    fn hash_vec(&self, v: &[N]) -> Vec<i8> {
54        let v = aview1(v);
55        self.hyperplanes
56            .dot(&v)
57            .mapv(|ai| if ai > Zero::zero() { 1 } else { 0 })
58            .to_vec()
59    }
60}
61
62impl<N: Numeric> VecHash<N, i8> for SignRandomProjections<N> {
63    fn hash_vec_query(&self, v: &[N]) -> Vec<i8> {
64        self.hash_vec(v)
65    }
66    fn as_step_wise_probe(&self) -> Option<&dyn StepWiseProbe<N, i8>> {
67        Some(self)
68    }
69}
70
71/// L2 Hasher family. [Read more.](https://arxiv.org/pdf/1411.3787.pdf)
72#[derive(Serialize, Deserialize, Clone)]
73pub struct L2<N = f32, K = i32> {
74    pub a: Array2<N>,
75    pub r: N,
76    pub b: Array1<N>,
77    n_projections: usize,
78    phantom: PhantomData<K>,
79}
80
81impl<N, K> L2<N, K>
82where
83    N: Numeric + Float,
84    K: Integer,
85{
86    pub fn new(dim: usize, r: f32, n_projections: usize, seed: u64) -> Self {
87        let mut rng = create_rng(seed);
88        let a = Array::random_using((n_projections, dim), StandardNormal, &mut rng);
89        let uniform_dist = Uniform::new(0., r);
90        let b = Array::random_using(n_projections, uniform_dist, &mut rng);
91
92        // cast to generic
93        let a = a.mapv(|v| N::from_f32(v).unwrap());
94        let b = b.mapv(|v| N::from_f32(v).unwrap());
95        let r = N::from_f32(r).unwrap();
96
97        L2 {
98            a,
99            r,
100            b,
101            n_projections,
102            phantom: PhantomData,
103        }
104    }
105
106    pub(crate) fn hash_vec(&self, v: &[N]) -> Array1<N> {
107        ((self.a.dot(&aview1(v)) + &self.b) / self.r).mapv(|x| x.floor())
108    }
109
110    fn hash_and_cast_vec(&self, v: &[N]) -> Vec<K> {
111        let div_r = N::from_i8(1).unwrap() / self.r;
112        // not DRY. we don't call hash_vec to save function call.
113        ((self.a.dot(&aview1(v)) + &self.b) * div_r)
114            .mapv(|x| {
115                let hp = NumCast::from(x.floor())
116                    .expect("Hash value doesnt fit in the Hash primitive type");
117                hp
118            })
119            .to_vec()
120    }
121}
122
123impl<N, K> VecHash<N, K> for L2<N, K>
124where
125    N: Numeric + Float,
126    K: Integer,
127{
128    fn hash_vec_query(&self, v: &[N]) -> Vec<K> {
129        self.hash_and_cast_vec(v)
130    }
131
132    fn as_query_directed_probe(&self) -> Option<&dyn QueryDirectedProbe<N, K>> {
133        Some(self)
134    }
135}
136
137/// Maximum Inner Product Search. [Read more.](https://papers.nips.cc/paper/5329-asymmetric-lsh-alsh-for-sublinear-time-maximum-inner-product-search-mips.pdf)
138#[derive(Serialize, Deserialize, Clone)]
139pub struct MIPS<N, K = i32> {
140    U: N,
141    M: N,
142    m: usize,
143    dim: usize,
144    hasher: L2<N, K>,
145}
146
147impl<N, K> MIPS<N, K>
148where
149    N: Numeric + Float,
150    K: Integer,
151{
152    pub fn new(dim: usize, r: f32, U: N, m: usize, n_projections: usize, seed: u64) -> Self {
153        let l2 = L2::new(dim + m, r, n_projections, seed);
154        MIPS {
155            U,
156            M: Zero::zero(),
157            m,
158            dim,
159            hasher: l2,
160        }
161    }
162
163    pub fn fit(&mut self, v: &[Vec<N>]) {
164        // TODO: add fit to vechash trait?
165        let mut max_l2 = Zero::zero();
166        for x in v.iter() {
167            let l2 = l2_norm(x);
168            if l2 > max_l2 {
169                max_l2 = l2
170            }
171        }
172        self.M = max_l2
173    }
174
175    pub fn tranform_put(&self, x: &[N]) -> Vec<N> {
176        let mut x_new = Vec::with_capacity(x.len() + self.m);
177
178        if self.M == Zero::zero() {
179            panic!("MIPS is not fitted")
180        }
181
182        // shrink norm such that l2 norm < U < 1.
183        for x_i in x.iter().cloned() {
184            x_new.push(x_i / self.M * self.U)
185        }
186
187        let norm_sq = l2_norm(&x_new).powf(N::from_f32(2.).unwrap());
188        for i in 1..(self.m + 1) {
189            x_new.push(norm_sq.powf(N::from_usize(i).unwrap()))
190        }
191        x_new
192    }
193
194    pub fn transform_query(&self, x: &[N]) -> Vec<N> {
195        let mut x_new = Vec::with_capacity(x.len() + self.m);
196
197        // normalize query to have l2 == 1.
198        let l2 = l2_norm(x);
199        for x_i in x.iter().cloned() {
200            x_new.push(x_i / l2)
201        }
202
203        let half = N::from_f32(0.5).unwrap();
204        for _ in 0..self.m {
205            x_new.push(half)
206        }
207        x_new
208    }
209}
210
211impl<N, K> VecHash<N, K> for MIPS<N, K>
212where
213    N: Numeric + Float,
214    K: Integer,
215{
216    fn hash_vec_query(&self, v: &[N]) -> Vec<K> {
217        let q = self.transform_query(v);
218        self.hasher.hash_vec_query(&q)
219    }
220
221    fn hash_vec_put(&self, v: &[N]) -> Vec<K> {
222        let p = self.tranform_put(v);
223        self.hasher.hash_vec_query(&p)
224    }
225}
226
227impl<N, K> Deref for MIPS<N, K>
228where
229    N: Numeric,
230    K: Integer,
231{
232    type Target = L2<N, K>;
233
234    fn deref(&self) -> &Self::Target {
235        &self.hasher
236    }
237}
238
239/// A hash family for the [Jaccard Index](https://en.wikipedia.org/wiki/Jaccard_index)
240/// /// The generic integer N, needs to be able to hold the number of dimensions.
241/// so a `u8` with a vector of > 255 dimensions will cause a `panic`.
242#[derive(Serialize, Deserialize, Clone)]
243pub struct MinHash<N = u8, K = i32> {
244    pub pi: Array2<N>,
245    n_projections: usize,
246    phantom: PhantomData<K>,
247}
248
249impl<N, K> MinHash<N, K>
250where
251    N: Integer,
252    K: Integer,
253{
254    pub fn new(n_projections: usize, dim: usize, seed: u64) -> Self {
255        let mut pi = Array::zeros((n_projections, dim));
256        let mut rng = create_rng(seed);
257
258        for row in 0..n_projections {
259            // randomly permute the indexes of vector that should be hashed.
260            // So a vector of length 4 could have the following random pi permutation:
261            // [3, 2, 4, 1]
262            // We start counting from 1, as we want to multiply with these pi vectors and take the
263            // lowest non zero output
264            let permutation_idx = rand::seq::index::sample(&mut rng, dim, dim)
265                .into_iter()
266                .map(|idx| N::from_usize(idx + 1).expect("could not cast idx to generic"))
267                .collect::<Vec<_>>();
268            let mut slice = pi.slice_mut(s![row, ..]);
269            slice += &aview1(&permutation_idx);
270        }
271        MinHash {
272            pi,
273            n_projections,
274            phantom: PhantomData,
275        }
276    }
277}
278
279impl<N, K> VecHash<N, K> for MinHash<N, K>
280where
281    N: Integer,
282    K: Integer,
283{
284    fn hash_vec_query(&self, v: &[N]) -> Vec<K> {
285        let a = &self.pi * &aview1(v);
286        let init = K::from_usize(self.n_projections).expect("could not cast to K");
287        let hash = a.map_axis(Axis(1), |view| {
288            view.into_iter().fold(init, |acc, v| {
289                if *v > Zero::zero() {
290                    let v = K::from(*v).expect("could not cast N to K");
291                    if v < acc {
292                        v
293                    } else {
294                        acc
295                    }
296                } else {
297                    acc
298                }
299            })
300        });
301        hash.to_vec()
302    }
303}
304
305#[cfg(test)]
306mod test {
307    use super::*;
308
309    #[test]
310    fn test_l2() {
311        // Only test if it runs
312        let l2 = <L2>::new(5, 2.2, 7, 1);
313        // two close vector
314        let h1 = l2.hash_vec_query(&[1., 2., 3., 1., 3.]);
315        let h2 = l2.hash_vec_query(&[1.1, 2., 3., 1., 3.1]);
316
317        // a distant vec
318        let h3 = l2.hash_vec_query(&[10., 10., 10., 10., 10.1]);
319
320        println!("close: {:?} distant: {:?}", (&h1, &h2), &h3);
321        assert_eq!(h1, h2);
322        assert_ne!(h1, h3);
323    }
324
325    #[test]
326    fn test_minhash() {
327        let n_projections = 3;
328        let h = <MinHash>::new(n_projections, 5, 0);
329        let hash = h.hash_vec_query(&[1, 0, 1, 0, 1]);
330        assert_eq!(hash.len(), n_projections)
331    }
332}