bixverse-rs 0.4.2

//! Contains traits and their implementations, specifically designed for the
//! single cell workflows in this crate.

use bincode::{Decode, Encode};
use half::f16;
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::iter::Sum;

/////////
// F16 //
/////////

/// A serialisable wrapper for IEEE 754 half-precision floats.
///
/// Stores the raw bits as u16 for compatibility with bincode and serde.
#[derive(Encode, Decode, Serialize, Deserialize, Debug, Clone, Copy, Default)]
pub struct F16(u16);

/////////////////
// Conversions //
/////////////////

impl From<f16> for F16 {
    fn from(f: f16) -> Self {
        F16(f.to_bits())
    }
}

impl From<F16> for f16 {
    fn from(f: F16) -> Self {
        f16::from_bits(f.0)
    }
}

//////////////////
// Iterator sum //
//////////////////

impl Sum for F16 {
    fn sum<I: Iterator<Item = F16>>(iter: I) -> Self {
        let sum: f16 = iter.map(f16::from).sum();
        F16::from(sum)
    }
}

impl<'a> Sum<&'a F16> for F16 {
    fn sum<I: Iterator<Item = &'a F16>>(iter: I) -> Self {
        let sum: f16 = iter.map(|&f| f16::from(f)).sum();
        F16::from(sum)
    }
}

//////////////
// Math ops //
//////////////

impl std::ops::Add for F16 {
    type Output = F16;

    fn add(self, other: F16) -> F16 {
        let a = f16::from(self);
        let b = f16::from(other);
        F16::from(a + b)
    }
}

impl<'a> std::ops::Add<&'a F16> for F16 {
    type Output = F16;

    fn add(self, other: &'a F16) -> F16 {
        let a = f16::from(self);
        let b = f16::from(*other);
        F16::from(a + b)
    }
}

impl std::ops::Mul for F16 {
    type Output = F16;

    fn mul(self, other: F16) -> F16 {
        let a = f16::from(self);
        let b = f16::from(other);
        F16::from(a * b)
    }
}

impl<'a> std::ops::Mul<&'a F16> for F16 {
    type Output = F16;

    fn mul(self, other: &'a F16) -> F16 {
        let a = f16::from(self);
        let b = f16::from(*other);
        F16::from(a * b)
    }
}

impl std::ops::Sub for F16 {
    type Output = F16;
    fn sub(self, other: F16) -> F16 {
        let a = f16::from(self);
        let b = f16::from(other);
        F16::from(a - b)
    }
}

impl<'a> std::ops::Sub<&'a F16> for F16 {
    type Output = F16;
    fn sub(self, other: &'a F16) -> F16 {
        let a = f16::from(self);
        let b = f16::from(*other);
        F16::from(a - b)
    }
}

impl std::ops::AddAssign for F16 {
    fn add_assign(&mut self, other: F16) {
        let a = f16::from(*self);
        let b = f16::from(other);
        *self = F16::from(a + b);
    }
}

impl<'a> std::ops::AddAssign<&'a F16> for F16 {
    fn add_assign(&mut self, other: &'a F16) {
        let a = f16::from(*self);
        let b = f16::from(*other);
        *self = F16::from(a + b);
    }
}

impl std::ops::DivAssign for F16 {
    fn div_assign(&mut self, other: F16) {
        let a = f16::from(*self);
        let b = f16::from(other);
        *self = F16::from(a / b);
    }
}

impl<'a> std::ops::DivAssign<&'a F16> for F16 {
    fn div_assign(&mut self, other: &'a F16) {
        let a = f16::from(*self);
        let b = f16::from(*other);
        *self = F16::from(a / b);
    }
}

impl std::ops::Div for F16 {
    type Output = F16;
    fn div(self, other: F16) -> F16 {
        let a = f16::from(self);
        let b = f16::from(other);
        F16::from(a / b)
    }
}

impl<'a> std::ops::Div<&'a F16> for F16 {
    type Output = F16;
    fn div(self, other: &'a F16) -> F16 {
        let a = f16::from(self);
        let b = f16::from(*other);
        F16::from(a / b)
    }
}

//////////////
// Equality //
//////////////

impl PartialEq for F16 {
    fn eq(&self, other: &Self) -> bool {
        let a = f16::from(*self);
        let b = f16::from(*other);

        if a.is_nan() && b.is_nan() {
            false
        } else {
            a == b
        }
    }
}

impl Eq for F16 {}

impl PartialOrd for F16 {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for F16 {
    fn cmp(&self, other: &Self) -> Ordering {
        let a = half::f16::from(*self);
        let b = half::f16::from(*other);

        match (a.is_nan(), b.is_nan()) {
            (true, true) => Ordering::Equal,
            (true, false) => Ordering::Greater,
            (false, true) => Ordering::Less,
            (false, false) => a.partial_cmp(&b).unwrap(),
        }
    }
}

impl F16 {
    /// Get the raw bits representation
    pub fn to_bits(self) -> u16 {
        self.0
    }

    /// Create from raw bits
    pub fn from_bits(bits: u16) -> Self {
        F16(bits)
    }

    /// From f32
    pub fn from_f32(f: f32) -> Self {
        F16::from(f16::from_f32(f))
    }

    /// Convert to little-endian bytes
    pub fn to_le_bytes(self) -> [u8; 2] {
        self.0.to_le_bytes()
    }

    /// Create from little-endian bytes
    pub fn from_le_bytes(bytes: [u8; 2]) -> Self {
        F16(u16::from_le_bytes(bytes))
    }

    /// From F16 to f32
    pub fn to_f32(self) -> f32 {
        f16::from_bits(self.0).to_f32()
    }

    /// From F16 fo f64
    pub fn to_f64(self) -> f64 {
        f16::from_bits(self.0).to_f32() as f64
    }
}

///////////////////////
// f32 and u16 stuff //
///////////////////////

/// Trait for types that can be converted to f32 and u16.
pub trait FloatAndUInt: Copy {
    /// Transform u16 to f32 for fast conversions
    fn to_f32(self) -> f32;

    /// Transform u32/f32 to u16, **saturating** at `u16::MAX`.
    ///
    /// Prefer [`Self::to_u32`] plus `RawCounts::from_u32_auto` on any path that
    /// reaches the binary format: that narrows to u16 only when every value
    /// fits, instead of silently clipping high-expression genes.
    fn to_u16(self) -> u16;

    /// Transform f32/u16 to u32
    fn to_u32(self) -> u32;
}

impl FloatAndUInt for i32 {
    fn to_f32(self) -> f32 {
        self as f32
    }
    fn to_u16(self) -> u16 {
        self.min(u16::MAX as i32).max(0) as u16
    }
    fn to_u32(self) -> u32 {
        self.max(0) as u32
    }
}

impl FloatAndUInt for u32 {
    fn to_f32(self) -> f32 {
        self as f32
    }
    fn to_u16(self) -> u16 {
        self.min(u16::MAX as u32) as u16
    }
    fn to_u32(self) -> u32 {
        self
    }
}

impl FloatAndUInt for usize {
    fn to_f32(self) -> f32 {
        self as f32
    }
    fn to_u16(self) -> u16 {
        self.min(u16::MAX as usize) as u16
    }
    fn to_u32(self) -> u32 {
        self as u32
    }
}

impl FloatAndUInt for u16 {
    fn to_f32(self) -> f32 {
        self as f32
    }
    fn to_u16(self) -> u16 {
        self
    }
    fn to_u32(self) -> u32 {
        self as u32
    }
}