pub trait Shapes: Sealed + StorageElement Errors
Error::KernelFailed when the kernel refuses to run, like on a misaligned workspace.
ShapesPackedMatrix packs point clouds of any sizes once each, centered and rotated onto their principal axes.
Its methods score every cloud of one set against every cloud of another into caller-provided slices.
Distances are f32 for f16 clouds and f64 for the others, and the optional rotations are row-major [f32; 9] matrices taking each column cloud onto each row cloud.
On equal sizes, Wasserstein-Procrustes is at most the Kabsch RMSD of a known correspondence, and equals it when that correspondence is an optimal plan.
use numkong::ShapesPackedMatrix;
let tetrahedron = [[0.0_f32, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]];
let triangle = [[1.0_f32, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
let rows = ShapesPackedMatrix::new(&[&tetrahedron[..]]).unwrap();
let columns = ShapesPackedMatrix::new(&[&tetrahedron[..], &triangle[..]]).unwrap();
let mut distances = [0.0_f64; 2];
let mut rotations = [[0.0_f32; 9]; 2];
rows.wassersteinprocrustes(&columns, &mut distances, Some(&mut rotations)).unwrap();
assert!(distances[0] < 1e-6);
rows.chamfers(&columns, &mut distances, None, f64::INFINITY).unwrap();
assert_eq!(columns.shape(), (2, 7)); // clouds and their total points
assert_eq!(columns.cloud_points(1).unwrap(), 3);
let mut unpacked = [[0.0_f32; 3]; 3];
columns.unpack(1, &mut unpacked).unwrap(); // the inverse of packing, within one rounding
assert!(unpacked.iter().flatten().zip(triangle.iter().flatten()).all(|(a, b)| (a - b).abs() < 1e-6));
let grown = tetrahedron.map(|point| point.map(|value| value * 3.7 + 1.0));
let similar = ShapesPackedMatrix::new_normalized(&[&grown[..]]).unwrap();
let normalized = ShapesPackedMatrix::new_normalized(&[&tetrahedron[..]]).unwrap();
normalized.wassersteinprocrustes(&similar, &mut distances[..1], None).unwrap();
assert!(distances[0] < 1e-6 && similar.scale(0).unwrap() > 3.0 * normalized.scale(0).unwrap());ShapesPackedMatrix::new_normalized also divides every cloud by its RMS distance from the centroid, kept as scale(i), so every kind ignores size as well as pose.
The pack layout is private, and the cloud_points(i), centroid(i), frame(i), spreads(i) and max_radius(i) accessors report the rest of each cloud’s pose and bounds. unpack(i, &mut points) writes cloud i back interleaved, in its original pose and scale, within one rounding to the cloud type.
Between normalized packs, the rotation also scales by scale_a / scale_b, and each kind compares the packed coordinates, so mixing absolute and normalized packs in one call is the caller’s choice.
The transport kinds, wassersteinprocrustes and gromovwassersteins, allocate their workspace through the row set’s allocator for the call.
Every kind has a safe _on verb, which runs the capabilities of a Stream::scope and keeps every borrow until the stream joins. new_in with a UnifiedAllocator writes the table and the packs into one block, which a GPU reads as it stands, as the host packs and every capability reads every pack.
No verb allocates, so results and workspace come from the caller, and memory the device cannot reach fails with Status::DeviceMemoryMismatch, whichever allocator made it.
use numkong::{f16, Capabilities, Shapes, ShapesPackedMatrix, Stream, UnifiedAllocator};
let stream = Stream::new(Capabilities::cuda_enabled(0)?, 0)?;
let unified = UnifiedAllocator::new(&stream);
let packs = ShapesPackedMatrix::<f16, _>::new_in(&clouds, unified)?;
let (mut results, mut workspace) = (Vec::new_in(unified), Vec::new_in(unified));
results.resize(clouds.len() * clouds.len(), 0.0_f32);
workspace.resize(f16::chamfers_workspace_size(largest, largest, stream.capabilities())?, 0_u8);
stream.scope(|scope| packs.chamfers_on(scope, &packs, &mut results, None, limit, &mut workspace))?;ShapesCloud packing and the many-to-many shape kernels, for f16, f32 and f64 clouds.15pub trait Shapes: Sealed + StorageElement Error::KernelFailed when the kernel refuses to run, like on a misaligned workspace.
DistanceDistance type and pose type: f32 for f16 clouds, f64 for the others.type Distance: Copy + Default + Send + Sync; shapes_pack_sizeReturns the bytes a pack of points points takes.fn shapes_pack_size(points: usize) -> Result<usize, Error> shapes_packPacks one cloud: centers it and rotates it onto its principal axes.unsafe fn shapes_pack(points: &[[Self; 3]], packed: *mut u8) -> Result<(), Error> packed must have shapes_pack_size(points.len()) writable bytes on an
nk_default_alignment_k boundary.
shapes_pack_normalizedPacks one cloud like Shapes::shapes_pack, also dividing it by its RMS distance from the centroid, which the pack keeps as its scale, or by 1 when the cloud has no spread.unsafe fn shapes_pack_normalized(points: &[[Self; 3]], packed: *mut u8) -> Result<(), Error> The contract of Shapes::shapes_pack.
shapes_packed_shapeReads the pose and bounds of a pack into every non-null output: its number of points, its mean input point, the unit quaternion (x, y, z, w) rotating the centered cloud onto its principal axes, what the rotated cloud was divided by, the descending eigenvalues of the packed points' second moments, and the largest packed point norm.unsafe fn shapes_packed_shape(
packed: *const u8,
points: *mut usize,
centroid: *mut Self::Distance,
frame: *mut Self::Distance,
scale: *mut Self::Distance,
spreads: *mut Self::Distance,
max_radius: *mut Self::Distance,
) -> Result<(), Error> packed must point to a pack of Self, and the non-null outputs must hold 1, 3, 4, 1, 3
and 1 writable values in that order.
shapes_unpackWrites the cloud of a pack back interleaved, in its original pose and scale, within one rounding to Self.unsafe fn shapes_unpack(packed: *const u8, points: *mut [Self; 3]) -> Result<(), Error> UNBOUNDEDThe limit that scores every pair exactly.const UNBOUNDED: Self::Distance; rmsds_packedSuperposed RMSDs between every row cloud and every column cloud of equally many points.unsafe fn rmsds_packed(
a_packed: *const *const u8,
rows: usize,
b_packed: *const *const u8,
columns: usize,
results: *mut Self::Distance,
results_stride: usize,
capabilities: Capabilities,
stream: *mut c_void,
) -> Result<(), Error> a_packed and b_packed must hold rows and columns pointers to packs of Selfresults must hold rows rows of columns distances, results_stride bytes apartstream must be null or a live stream of the group capabilities name, null on the CPU,
and on a GPU everything above must stay alive and untouched until that stream joinshausdorffs_packedHausdorff distances between every row cloud and every column cloud.unsafe fn hausdorffs_packed(
a_packed: *const *const u8,
rows: usize,
b_packed: *const *const u8,
columns: usize,
results: *mut Self::Distance,
results_stride: usize,
rotations: *mut f32,
limit: Self::Distance,
capabilities: Capabilities,
stream: *mut c_void,
) -> Result<(), Error> The contract of Shapes::rmsds_packed, plus rotations must be null or hold
rows × columns × 9 writable values.
chamfers_packedChamfer distances between every row cloud and every column cloud.unsafe fn chamfers_packed(
a_packed: *const *const u8,
rows: usize,
b_packed: *const *const u8,
columns: usize,
results: *mut Self::Distance,
results_stride: usize,
rotations: *mut f32,
limit: Self::Distance,
workspace: *mut u8,
workspace_bytes: usize,
capabilities: Capabilities,
stream: *mut c_void,
) -> Result<(), Error> The contract of Shapes::hausdorffs_packed, plus workspace must have workspace_bytes
writable bytes, on the default alignment, as Shapes::chamfers_workspace_size reports for
the largest clouds.
chamfers_workspace_sizeReturns the workspace bytes Chamfer needs between clouds of up to row_points and column_points points.fn chamfers_workspace_size(
row_points: usize,
column_points: usize,
capabilities: Capabilities,
) -> Result<usize, Error> wassersteinprocrustes_workspace_sizeReturns the workspace bytes Wasserstein-Procrustes needs between clouds of up to row_points and column_points points.fn wassersteinprocrustes_workspace_size(
row_points: usize,
column_points: usize,
capabilities: Capabilities,
) -> Result<usize, Error> wassersteinprocrustes_packedWasserstein-Procrustes distances between every row cloud and every column cloud.unsafe fn wassersteinprocrustes_packed(
a_packed: *const *const u8,
rows: usize,
b_packed: *const *const u8,
columns: usize,
results: *mut Self::Distance,
results_stride: usize,
rotations: *mut f32,
workspace: *mut u8,
workspace_bytes: usize,
capabilities: Capabilities,
stream: *mut c_void,
) -> Result<(), Error> The contract of Shapes::chamfers_packed, plus workspace must have workspace_bytes
writable bytes on an nk_default_alignment_k boundary.
gromovwassersteins_workspace_sizeReturns the workspace bytes Gromov-Wasserstein needs between clouds of up to row_points and column_points points.fn gromovwassersteins_workspace_size(
row_points: usize,
column_points: usize,
capabilities: Capabilities,
) -> Result<usize, Error> gromovwassersteins_packedGromov-Wasserstein distances between every row cloud and every column cloud.unsafe fn gromovwassersteins_packed(
a_packed: *const *const u8,
rows: usize,
b_packed: *const *const u8,
columns: usize,
results: *mut Self::Distance,
results_stride: usize,
rotations: *mut f32,
workspace: *mut u8,
workspace_bytes: usize,
capabilities: Capabilities,
stream: *mut c_void,
) -> Result<(), Error> The contract of Shapes::wassersteinprocrustes_packed.
ShapesPackedMatrixPoint clouds packed one by one onto their principal axes, scored as rows against the columns of another set.27pub struct ShapesPackedMatrix<Scalar: Shapes, Alloc: Allocator = Global> One allocation holds the pointer to every pack, as the kernels take them, then the packs.
Debug, Send, Syncnew_inPacks every cloud of clouds, of any sizes, using a custom allocator.pub fn new_in(clouds: &[&[[Scalar; 3]]], alloc: Alloc) -> Result<Self, Error> new_normalized_inPacks every cloud of clouds like new_in, each divided by its RMS distance from the centroid, so every kind compares shapes regardless of size.pub fn new_normalized_in(clouds: &[&[[Scalar; 3]]], alloc: Alloc) -> Result<Self, Error> allocatorReturns a reference to the allocator.pub fn allocator(&self) -> &Alloc cloudsNumber of packed clouds.pub fn clouds(&self) -> usize pointsTotal points across the packed clouds.pub fn points(&self) -> usize shapeReturns the shape [clouds,points], the clouds and their total points.pub fn shape(&self) -> (usize, usize) pack_sizeBytes one packed cloud of points points takes, to size a buffer without packing.pub fn pack_size(points: usize) -> Result<usize, Error> packsThe pointer to every pack, as the nk_*_packed_* shape kernels take them.pub fn packs(&self) -> &[*const u8] cloud_pointsNumber of points in cloud index.pub fn cloud_points(&self, index: usize) -> Result<usize, Error> centroidMean input point of cloud index.pub fn centroid(&self, index: usize) -> Result<[Scalar::Distance; 3], Error> frameUnit quaternion (x, y, z, w) rotating centered cloud index onto its principal axes.pub fn frame(&self, index: usize) -> Result<[Scalar::Distance; 4], Error> scaleWhat rotated cloud index was divided by: its RMS distance from the centroid for normalized packs, and 1 otherwise.pub fn scale(&self, index: usize) -> Result<Scalar::Distance, Error> spreadsDescending eigenvalues of the second moments of packed cloud index.pub fn spreads(&self, index: usize) -> Result<[Scalar::Distance; 3], Error> max_radiusLargest packed point norm of cloud index.pub fn max_radius(&self, index: usize) -> Result<Scalar::Distance, Error> unpackWrites cloud index into out as it was packed: interleaved, in its original pose and scale, within one rounding to Scalar.pub fn unpack(&self, index: usize, out: &mut [[Scalar; 3]]) -> Result<(), Error> Returns Err if index is out of range or out does not hold
cloud_points points.
rmsdsSuperposed RMSDs from every cloud of self to every cloud of columns, each holding the same number of points, point i matching point i under the best rotation.pub fn rmsds<OtherAlloc: Allocator>(
&self,
columns: &ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &mut [Scalar::Distance],
) -> Result<(), Error> hausdorffsHausdorff distances from every cloud of self to every cloud of columns, the best of the 8 sign flips of the principal axes, zero between two empty clouds and infinite between an empty and a non-empty one.pub fn hausdorffs<OtherAlloc: Allocator>(
&self,
columns: &ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &mut [Scalar::Distance],
rotations: Option<&mut [[f32; 9]]>,
limit: Scalar::Distance,
) -> Result<(), Error> Fills results with [rows,columns] distances and, when given, rotations with
as many matrices taking each column cloud onto each row cloud.
Distances up to limit are exact; past it, a pair may stop early and report a lower bound
that also lies past it, with that bound's flip in rotations.
Shapes::UNBOUNDED scores every pair exactly.
Returns Err if results or rotations does not hold rows × columns entries.
chamfersChamfer distances from every cloud of self to every cloud of columns, the best of the 8 sign flips of the principal axes, with the outputs and conventions of hausdorffs.pub fn chamfers<OtherAlloc: Allocator>(
&self,
columns: &ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &mut [Scalar::Distance],
rotations: Option<&mut [[f32; 9]]>,
limit: Scalar::Distance,
) -> Result<(), Error> wassersteinprocrustesWasserstein-Procrustes distances from every cloud of self to every cloud of columns: the RMSD of the exact optimal transport plan between uniform weights, after the rotation or reflection minimizing it.pub fn wassersteinprocrustes<OtherAlloc: Allocator>(
&self,
columns: &ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &mut [Scalar::Distance],
rotations: Option<&mut [[f32; 9]]>,
) -> Result<(), Error> The workspace comes from self's allocator for the call. Outputs and conventions match
hausdorffs.
gromovwassersteinsGromov-Wasserstein distances from every cloud of self to every cloud of columns: the transport distance between the clouds' own distance matrices, under the best plan found from the Wasserstein-Procrustes plans and the product plan.pub fn gromovwassersteins<OtherAlloc: Allocator>(
&self,
columns: &ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &mut [Scalar::Distance],
rotations: Option<&mut [[f32; 9]]>,
) -> Result<(), Error> The workspace comes from self's allocator for the call. Outputs and conventions match
hausdorffs.
rmsds_onrmsds queued on the stream of scope, with its capabilities.pub fn rmsds_on<'scope, OtherAlloc: Allocator>(
&'scope self,
scope: &'scope Scope<'scope, '_>,
columns: &'scope ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &'scope mut [Scalar::Distance],
) -> Result<(), Error> On a GPU the packs of both sets, results, rotations and workspace must lie in memory
the device reaches, like blocks of a UnifiedAllocator, or the
runtime's check refuses them with Status::DeviceMemoryMismatch.
hausdorffs_onhausdorffs queued on the stream of scope, as in rmsds_on.pub fn hausdorffs_on<'scope, OtherAlloc: Allocator>(
&'scope self,
scope: &'scope Scope<'scope, '_>,
columns: &'scope ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &'scope mut [Scalar::Distance],
rotations: Option<&'scope mut [[f32; 9]]>,
limit: Scalar::Distance,
) -> Result<(), Error> chamfers_onchamfers queued on the stream of scope, as in rmsds_on, in the caller's workspace of at least Shapes::chamfers_workspace_size bytes.pub fn chamfers_on<'scope, OtherAlloc: Allocator>(
&'scope self,
scope: &'scope Scope<'scope, '_>,
columns: &'scope ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &'scope mut [Scalar::Distance],
rotations: Option<&'scope mut [[f32; 9]]>,
limit: Scalar::Distance,
workspace: &'scope mut [u8],
) -> Result<(), Error> wassersteinprocrustes_onwassersteinprocrustes queued on the stream of scope, as in chamfers_on, with Shapes::wassersteinprocrustes_workspace_size bytes of workspace on an nk_default_alignment_k boundary.pub fn wassersteinprocrustes_on<'scope, OtherAlloc: Allocator>(
&'scope self,
scope: &'scope Scope<'scope, '_>,
columns: &'scope ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &'scope mut [Scalar::Distance],
rotations: Option<&'scope mut [[f32; 9]]>,
workspace: &'scope mut [u8],
) -> Result<(), Error> gromovwassersteins_ongromovwassersteins queued on the stream of scope, as in wassersteinprocrustes_on, with Shapes::gromovwassersteins_workspace_size bytes of workspace.pub fn gromovwassersteins_on<'scope, OtherAlloc: Allocator>(
&'scope self,
scope: &'scope Scope<'scope, '_>,
columns: &'scope ShapesPackedMatrix<Scalar, OtherAlloc>,
results: &'scope mut [Scalar::Distance],
rotations: Option<&'scope mut [[f32; 9]]>,
workspace: &'scope mut [u8],
) -> Result<(), Error> newPacks every cloud of clouds, of any sizes, using the global allocator.pub fn new(clouds: &[&[[Scalar; 3]]]) -> Result<Self, Error> new_normalizedPacks every cloud of clouds normalized, using the global allocator.pub fn new_normalized(clouds: &[&[[Scalar; 3]]]) -> Result<Self, Error>