Skip to content

Commit

Permalink
feat: add f16 support, example and docs to show usage with half crate
Browse files Browse the repository at this point in the history
  • Loading branch information
sdd committed Feb 18, 2024
1 parent 2e2b989 commit cc5f942
Show file tree
Hide file tree
Showing 9 changed files with 72 additions and 15 deletions.
14 changes: 12 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -91,21 +91,27 @@ optional = true
version = "3"
optional = true

[dependencies.half]
version = "2.3"
optional = true
features = ["num-traits"]

[dependencies.rkyv]
version = "0.7"
optional = true
default-features = false
features = ["alloc", "copy_unsafe", "size_64"]

[features]
default = ["tracing"]
f16 = ["half"]
global_allocate = []
immutable = []
serialize = ["serde", "serde/derive", "serde_derive", "serde_with", "fixed/serde"]
serialize_rkyv = ["rkyv"]
simd = []
tracing = ["dep:tracing", "tracing-subscriber"]
default = ["tracing"]
test_utils = ["rand", "rand_chacha", "rayon"]
tracing = ["dep:tracing", "tracing-subscriber"]

[package.metadata.docs.rs]
all-features = true
Expand Down Expand Up @@ -206,3 +212,7 @@ required-features = ["serialize_rkyv"]
name = "build-float-doctest-tree"
path = "examples/build-float-doctest-tree.rs"
required-features = ["serialize_rkyv"]

[[example]]
name = "half"
path = "examples/half.rs"
6 changes: 4 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ Kiddo is ideal for super-fast spatial / geospatial lookups and nearest-neighbour
Kiddo provides:
- Its standard floating point k-d tree, exposed as [`kiddo::KdTree`](`crate::KdTree`)
- **integer / fixed point support** via the [`Fixed`](https://docs.rs/fixed/latest/fixed/) library;
- **`f16` support** via the [`half`](https://docs.rs/half/latest/half/) library;
- **instant zero-copy deserialization** and serialization via [`Rkyv`](https://docs.rs/rkyv/latest/rkyv/) ([`Serde`](https://docs.rs/serde/latest/serde/) still available).
- An [`ImmutableKdTree`](`immutable::float::kdtree::ImmutableKdTree`) with space and performance advantages over the standard
k-d tree, for situations where the tree does not need to be modified after creation
Expand All @@ -26,7 +27,7 @@ Kiddo provides:
Add `kiddo` to `Cargo.toml`
```toml
[dependencies]
kiddo = "3.0.0"
kiddo = "4.2.0"
```

Add points to kdtree and query nearest n points with distance function
Expand Down Expand Up @@ -70,6 +71,7 @@ The Kiddo crate exposes the following features. Any labelled as **(NIGHTLY)** ar
* `serialize_rkyv` - zero-copy serialization / deserialization via [`Rkyv`](https://docs.rs/rkyv/latest/rkyv/)
* `global_allocate` **(NIGHTLY)** - When enabled Kiddo will use the unstable allocator_api feature within [`ImmutableKdTree`](`immutable::float::kdtree::ImmutableKdTree`) to get a slight performance improvement when allocating space for leaves.
* `simd` **(NIGHTLY)** - enables some hand-written SIMD intrinsic code within [`ImmutableKdTree`](`immutable::float::kdtree::ImmutableKdTree`) that may improve performance (currently only on the nearest_one method when using `f64`)
* `f16` - enables usage of `f16` from the `half` crate for float trees.

## v3.x

Expand All @@ -82,7 +84,7 @@ use kiddo::distance::squared_euclidean;
let result = kdtree.nearest_one(&[0f64, 0f64], &squared_euclidean);
```

Now in v3, you'll need to switch to this syntax:
Now for v3 onwards, you'll need to switch to this syntax:

```
use kiddo::SquaredEuclidean;
Expand Down
37 changes: 37 additions & 0 deletions examples/half.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
#[cfg(feature = "half")]
fn main() -> Result<(), Box<dyn std::error::Error>> {
use half::f16;
use kiddo::{KdTree, SquaredEuclidean};
use num_traits::real::Real;

// build and serialize small tree for ArchivedKdTree doctests
let mut tree: KdTree<f16, 3> = KdTree::new();
tree.add(
&[f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(5.0)],
100,
);
tree.add(
&[f16::from_f32(2.0), f16::from_f32(3.0), f16::from_f32(6.0)],
101,
);

let nearest = tree.nearest_one::<SquaredEuclidean>(&[
f16::from_f32(1.0),
f16::from_f32(2.0),
f16::from_f32(5.1),
]);

println!("Nearest: {:?}", &nearest);

assert!((nearest.distance - f16::from_f32(0.01)).abs() < f16::EPSILON);
assert_eq!(nearest.item, 100);
Ok(())
}

#[cfg(not(feature = "half"))]
fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("Activate the 'half' feature to run this example properly");
println!("Try this: cargo run --example half --features=half");

Ok(())
}
9 changes: 6 additions & 3 deletions src/float/kdtree.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
//! Floating point k-d tree, for use when the co-ordinates of the points being stored in the tree
//! are floats. f64 or f32 are supported currently.
//! are floats. f64 or f32 are supported currently, or [`f16`](https://docs.rs/half/latest/half/struct.f16.html)
//! if the `f16` feature is enabled

use az::{Az, Cast};
use divrem::DivCeil;
Expand All @@ -16,7 +17,8 @@ use serde::{Deserialize, Serialize};

/// Axis trait represents the traits that must be implemented
/// by the type that is used as the first generic parameter, `A`,
/// on the float [`KdTree`]. This will be [`f64`] or [`f32`].
/// on the float [`KdTree`]. This will be [`f64`] or [`f32`],
/// or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled
pub trait Axis: FloatCore + Default + Debug + Copy + Sync + Send + std::ops::AddAssign {
/// returns absolute diff between two values of a type implementing this trait
fn saturating_dist(self, other: Self) -> Self;
Expand Down Expand Up @@ -47,7 +49,8 @@ impl<T: FloatCore + Default + Debug + Copy + Sync + Send + std::ops::AddAssign>
/// Floating point k-d tree
///
/// For use when the co-ordinates of the points being stored in the tree
/// on the float [`KdTree`]. This will be [`f64`] or [`f32`].
/// on the float [`KdTree`]. This will be [`f64`] or [`f32`],
/// or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled
///
/// A convenient type alias exists for KdTree with some sensible defaults set: [`kiddo::KdTree`](`crate::KdTree`).

Expand Down
3 changes: 2 additions & 1 deletion src/float/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
//! Floating point k-d tree, for use when the co-ordinates of the points being stored in the tree
//! are floats. [`f64`] or [`f32`] are the types that are supported for use as co-ordinates.
//! are floats. [`f64`] or [`f32`] are the types that are supported for use as co-ordinates,
//! or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled

#[doc(hidden)]
pub mod construction;
Expand Down
4 changes: 2 additions & 2 deletions src/immutable/float/kdtree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
//! expense of not being able to modify the contents of the tree after its initial
//! construction, and longer construction times - perhaps prohibitively so.
//! As with the vanilla tree, [`f64`] or [`f32`] are supported currently for co-ordinate
//! values.
//! values, or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled

use az::{Az, Cast};
use ordered_float::OrderedFloat;
Expand All @@ -28,7 +28,7 @@ use serde::{Deserialize, Serialize};
/// expense of not being able to modify the contents of the tree after its initial
/// construction, and longer construction times.
/// As with the vanilla tree, [`f64`] or [`f32`] are supported currently for co-ordinate
/// values.
/// values, or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled
///
/// A convenient type alias exists for ImmutableKdTree with some sensible defaults set: [`kiddo::ImmutableKdTree`](`crate::ImmutableKdTree`).
#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
Expand Down
3 changes: 2 additions & 1 deletion src/immutable/float/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@
//! Likewise for [`f32`] based trees, up to a few million nodes.
//!
//! As per the other Kiddo float-type trees, points being stored
//! in the tree must be floats ([`f64`] or [`f32`] are supported currently).
//! in the tree must be floats ([`f64`] or [`f32`] are supported currently,
//! or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled).

#[doc(hidden)]
pub mod construction;
Expand Down
3 changes: 2 additions & 1 deletion src/immutable/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@
//! Likewise for [`f32`] based trees, up to a few million nodes.
//!
//! As per the other Kiddo float-type trees, points being stored
//! in the tree must be floats ([`f64`] or [`f32`] are supported currently).
//! in the tree must be floats ([`f64`] or [`f32`] are supported currently,
//! or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled).
#[doc(hidden)]
pub(crate) mod common;
pub mod float;
8 changes: 5 additions & 3 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
//! Kiddo provides:
//! - Its standard floating-point k-d tree, exposed as [`kiddo::KdTree`](`crate::KdTree`)
//! - **integer / fixed point support** via the [`Fixed`](https://docs.rs/fixed/latest/fixed/) library;
//! - **`f16` support** via the [`half`](https://docs.rs/half/latest/half/) library;
//! - **instant zero-copy deserialization** and serialization via [`Rkyv`](https://docs.rs/rkyv/latest/rkyv/) ([`Serde`](https://docs.rs/serde/latest/serde/) still available).
//! - An [`ImmutableKdTree`](`immutable::float::kdtree::ImmutableKdTree`) with space and performance advantages over the standard
//! k-d tree, for situations where the tree does not need to be modified after creation
Expand All @@ -32,7 +33,7 @@
//! Add `kiddo` to `Cargo.toml`
//! ```toml
//! [dependencies]
//! kiddo = "3.0.0"
//! kiddo = "4.2.0"
//! ```
//!
//! ## Usage
Expand Down Expand Up @@ -76,6 +77,7 @@
//! * **serialize_rkyv** - zero-copy serialization / deserialization via [`Rkyv`](https://docs.rs/rkyv/latest/rkyv/)
//! * `global_allocate` **(NIGHTLY)** - When enabled Kiddo will use the unstable allocator_api feature within [`ImmutableKdTree`](`immutable::float::kdtree::ImmutableKdTree`) to get a slight performance improvement when allocating space for leaves.
//! * `simd` **(NIGHTLY)** - enables some hand written SIMD and pre-fetch intrinsics code within [`ImmutableKdTree`](`immutable::float::kdtree::ImmutableKdTree`) that may improve performance (currently only on nearest_one with `f64`)
//! * `f16` - enables usage of `f16` from the `half` crate for float trees.

#[macro_use]
extern crate doc_comment;
Expand Down Expand Up @@ -107,7 +109,7 @@ pub mod float_leaf_simd;

/// A floating-point k-d tree with default parameters.
///
/// `A` is the floating point type (`f32` or `f64`).
/// `A` is the floating point type (`f32` or `f64`, or `f16` if the `f16` feature is enabled).
/// `K` is the number of dimensions. See [`KdTree`](`float::kdtree::KdTree`) for details of how to use.
///
/// To manually specify more advanced parameters, use [`KdTree`](`float::kdtree::KdTree`) directly.
Expand All @@ -116,7 +118,7 @@ pub type KdTree<A, const K: usize> = float::kdtree::KdTree<A, u64, K, 32, u32>;

/// An immutable floating-point k-d tree with default parameters.
///
/// `A` is the floating point type (`f32` or `f64`).
/// `A` is the floating point type (`f32` or `f64`, or `f16` if the `f16` feature is enabled).
/// `K` is the number of dimensions. See [`ImmutableKdTree`](`immutable::float::kdtree::ImmutableKdTree`) for details of how to use.
///
/// To manually specify more advanced parameters, use [`ImmutableKdTree`](`immutable::float::kdtree::ImmutableKdTree`) directly.
Expand Down

0 comments on commit cc5f942

Please sign in to comment.