Skip to main content

Crate argminmax

Crate argminmax 

Source
Expand description

A crate for finding the index of the minimum and maximum values in an array.

These operations are optimized for speed using SIMD instructions (when available).
The SIMD inner loop is branchless, so its runtime does not depend on the order of the data (no best case / worst case). The masked functions (which skip the nulls) search a part of the data again when all its valid values are the max (or all the min) value of the data type, NaN or infinite, so they can be slower on such data. Furthermore, runtime CPU feature detection is used to choose the fastest implementation for the current CPU (with a scalar fallback).

The SIMD implementation is enabled for the following architectures:

i128 and u128 always use the scalar implementation.

§Description

This crate provides two traits: ArgMinMax and NaNArgMinMax.

These traits are implemented for slice and Vec.

  • For ArgMinMax the supported data types are
    • ints: i8, i16, i32, i64, i128
    • uints: u8, u16, u32, u64, u128
    • floats: f16, bf16, f32, f64 (see Features)
  • For NaNArgMinMax the supported data types are
    • floats: f16, bf16, f32, f64 (see Features)

Both traits differ in how they handle NaNs:

  • ArgMinMax ignores NaNs and returns the index of the minimum and maximum values in an array.
  • NaNArgMinMax returns the index of the first NaN in an array if there is one, otherwise it returns the index of the minimum and maximum values in an array.

Both traits have a counterpart that skips the null elements, which an Arrow validity bitmap marks: ArgMinMaxMasked and NaNArgMinMaxMasked.

§Caution

When dealing with floats and you are sure that there are no NaNs in the array, you should use ArgMinMax instead of NaNArgMinMax for performance reasons. The former is 5%-30% faster than the latter.

§Features

This crate has several features.

  • float (default) - enables the traits for floats (f32 and f64).
  • nightly_simd - enables NEON SIMD instructions on 32-bit ARM (requires a nightly compiler; no effect on other architectures).
  • half - enables the traits for f16 and bf16 (requires the half crate).
  • ndarray - adds the traits to ndarray::ArrayBase (requires the ndarray crate).
  • arrow - adds the traits to arrow::array::PrimitiveArray, skipping the nulls (requires the arrow crate).

For arrow, the functions panic when all values are null (as for an empty array). To get an Option instead, check the null count first, e.g. (array.null_count() < array.len()).then(|| array.argmin()) (arrow keeps the null count, so this check is O(1)), or use the ArgMinMaxMasked & NaNArgMinMaxMasked methods on the values and the validity bitmap, which return None.

§Examples

Three examples are provided below.

§Example with integers

use argminmax::ArgMinMax;

let a: Vec<i32> = vec![0, 1, 2, 3, 4, 5];
let (imin, imax) = a.argminmax();
assert_eq!(imin, 0);
assert_eq!(imax, 5);

§Example with NaNs (default float feature)

ⓘ
use argminmax::ArgMinMax; // argminmax ignores NaNs
use argminmax::NaNArgMinMax; // nanargminmax returns index of first NaN

let a: Vec<f32> = vec![f32::NAN, 1.0, f32::NAN, 3.0, 4.0, 5.0];
let (imin, imax) = a.argminmax(); // ArgMinMax::argminmax
assert_eq!(imin, 1);
assert_eq!(imax, 5);
let (imin, imax) = a.nanargminmax(); // NaNArgMinMax::nanargminmax
assert_eq!(imin, 0);
assert_eq!(imax, 0);

§Example with nulls

use argminmax::ArgMinMaxMasked;

let a: Vec<i32> = vec![-5, 1, 2, 3, 4, 9];
// The validity bitmap: bit 1 + i is set iff element i is valid (thus, elements 0 and 5 are null)
let validity: Vec<u8> = vec![0b0011_1100];
let (imin, imax) = a.argminmax_masked(&validity, 1).unwrap();
assert_eq!(imin, 1);
assert_eq!(imax, 4);
// There are no valid elements
assert_eq!(a.argminmax_masked(&[0], 0), None);

Modules§

dtype_strategy
The strategy that is used to handle the data type.
scalar
Scalar implementation of the argminmax functions.
simd
SIMD implementations of the argminmax functions.

Traits§

ArgMinMax
Trait for finding the minimum and maximum values in an array. For floats, NaNs are ignored.
ArgMinMaxMasked
Trait for finding the minimum and maximum values in an array, skipping the null elements. For floats, NaNs are ignored.
NaNArgMinMax
Trait for finding the minimum and maximum values in an array. For floats, NaNs are propagated - index of the first NaN is returned.
NaNArgMinMaxMasked
Trait for finding the minimum and maximum values in an array, skipping the null elements. For floats, NaNs are propagated - index of the first (valid) NaN is returned.