use std::borrow::Cow;
use bytemuck::{Pod, Zeroable};
use rand::Rng;
use super::{two_means_binary_quantized as two_means, Manhattan};
use crate::distance::Distance;
use crate::node::Leaf;
use crate::parallel::ImmutableSubsetLeafs;
use crate::spaces::simple::dot_product_binary_quantized;
use crate::unaligned_vector::{self, BinaryQuantized, UnalignedVector};
#[derive(Debug, Clone)]
pub enum BinaryQuantizedManhattan {}
#[repr(C)]
#[derive(Pod, Zeroable, Debug, Clone, Copy)]
pub struct NodeHeaderBinaryQuantizedManhattan {
bias: f32,
}
impl Distance for BinaryQuantizedManhattan {
const DEFAULT_OVERSAMPLING: usize = 3;
type Header = NodeHeaderBinaryQuantizedManhattan;
type VectorCodec = unaligned_vector::BinaryQuantized;
fn name() -> &'static str {
"binary quantized manhattan"
}
fn new_header(_vector: &UnalignedVector<Self::VectorCodec>) -> Self::Header {
NodeHeaderBinaryQuantizedManhattan { bias: 0.0 }
}
fn built_distance(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
manhattan_distance_binary_quantized(&p.vector, &q.vector)
}
fn normalized_distance(d: f32, dimensions: usize) -> f32 {
d.max(0.0) / dimensions as f32
}
fn norm_no_header(v: &UnalignedVector<Self::VectorCodec>) -> f32 {
let ones = v
.as_bytes()
.iter()
.map(|b| b.count_ones() as i32 - b.count_zeros() as i32)
.sum::<i32>() as f32;
ones.sqrt()
}
fn init(_node: &mut Leaf<Self>) {}
fn create_split<'a, R: Rng>(
children: &'a ImmutableSubsetLeafs<Self>,
rng: &mut R,
) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>> {
let [node_p, node_q] = two_means::<Self, Manhattan, R>(rng, children, false)?;
let vector: Vec<f32> =
node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
let mut normal = Leaf {
header: NodeHeaderBinaryQuantizedManhattan { bias: 0.0 },
vector: UnalignedVector::from_slice(&vector),
};
Self::normalize(&mut normal);
Ok(Cow::Owned(normal.vector.into_owned()))
}
fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
p.header.bias + dot_product_binary_quantized(&p.vector, &q.vector)
}
fn margin_no_header(
p: &UnalignedVector<Self::VectorCodec>,
q: &UnalignedVector<Self::VectorCodec>,
) -> f32 {
dot_product_binary_quantized(p, q)
}
}
fn manhattan_distance_binary_quantized(
u: &UnalignedVector<BinaryQuantized>,
v: &UnalignedVector<BinaryQuantized>,
) -> f32 {
let ret =
u.as_bytes().iter().zip(v.as_bytes()).map(|(u, v)| (u ^ v).count_ones()).sum::<u32>() * 2;
ret as f32
}