bixverse-rs 0.6.0

//! Various traits in bixverse

use extendr_api::*;
use faer::traits::{ComplexField, RealField};
use faer_entity::SimpleEntity;
use num_traits::float::TotalOrder;
use num_traits::{Float, FromPrimitive, ToPrimitive};

use std::fmt::Display;

use crate::utils::r_rust_interface::fill_column_major;
use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};

///////////////////
// Numeric trait //
///////////////////

/// Trait for floating-point types used in Bixverse. Has all of the common
/// floating-point operations and traits.
pub trait BixverseFloat:
    Float
    + FromPrimitive
    + ToPrimitive
    + ComplexField
    + Copy
    + 'static
    + RealField
    + AddAssign
    + DivAssign
    + SubAssign
    + MulAssign
    + Div
    + Display
    + Sub
    + Mul
    + Div
    + TotalOrder
    + Clone
{
}

impl<T> BixverseFloat for T where
    T: Float
        + FromPrimitive
        + ToPrimitive
        + ComplexField
        + Copy
        + 'static
        + RealField
        + AddAssign
        + DivAssign
        + SubAssign
        + MulAssign
        + Div
        + Display
        + Sub
        + Mul
        + Div
        + TotalOrder
        + Clone
{
}

/// More general numeric trait that includes also integers
pub trait BixverseNumeric:
    Clone
    + Default
    + Copy
    + Sync
    + Send
    + Add<Output = Self>
    + Mul<Output = Self>
    + Sub<Output = Self>
    + AddAssign
    + PartialEq
    + Sub
    + DivAssign
{
}

impl<T> BixverseNumeric for T where
    T: Clone
        + Default
        + Copy
        + Sync
        + Send
        + Add<Output = Self>
        + Mul<Output = Self>
        + Sub<Output = Self>
        + AddAssign
        + PartialEq
        + Sub
        + DivAssign
{
}

///////////////////
// Index casting //
///////////////////

/// Casts index vectors between `usize` and `u32` (memory-width conversion).
pub trait VecIndexCast<U> {
    /// Casts the index vector element-wise to the target integer type.
    fn index_cast(self) -> Vec<U>;
}

impl VecIndexCast<u32> for Vec<usize> {
    fn index_cast(self) -> Vec<u32> {
        self.into_iter().map(|x| x as u32).collect()
    }
}

impl VecIndexCast<u32> for &[usize] {
    fn index_cast(self) -> Vec<u32> {
        self.iter().map(|&x| x as u32).collect()
    }
}

impl VecIndexCast<usize> for Vec<u32> {
    fn index_cast(self) -> Vec<usize> {
        self.into_iter().map(|x| x as usize).collect()
    }
}

impl VecIndexCast<usize> for &[u32] {
    fn index_cast(self) -> Vec<usize> {
        self.iter().map(|&x| x as usize).collect()
    }
}

////////////////////
// R vector stuff //
////////////////////

/// Converts vectors between integer types for R compatibility.
pub trait VecConvert<U> {
    /// Converts usize vectors into i32 for fast R transfer
    fn r_int_convert(self) -> Vec<U>;

    /// Converts vectors from usize to i32, while adding one (shifting from
    /// 0-indexed data to 1-indexed data) or vice versa from i32 to usize while
    /// subtracting 1 (shifting from 1-indexed data to 0-indexed data).
    fn r_int_convert_shift(self) -> Vec<U>;
}

impl VecConvert<i32> for Vec<usize> {
    fn r_int_convert(self) -> Vec<i32> {
        self.into_iter().map(|x| x as i32).collect()
    }
    fn r_int_convert_shift(self) -> Vec<i32> {
        self.into_iter().map(|x| (x as i32) + 1).collect()
    }
}

impl VecConvert<i32> for &[&usize] {
    fn r_int_convert(self) -> Vec<i32> {
        self.iter().map(|&&x| x as i32).collect()
    }
    fn r_int_convert_shift(self) -> Vec<i32> {
        self.iter().map(|&&x| (x as i32) + 1).collect()
    }
}

impl VecConvert<i32> for Vec<u32> {
    fn r_int_convert(self) -> Vec<i32> {
        self.into_iter().map(|x| x as i32).collect()
    }
    fn r_int_convert_shift(self) -> Vec<i32> {
        self.into_iter().map(|x| (x as i32) + 1).collect()
    }
}

impl VecConvert<usize> for Vec<i32> {
    fn r_int_convert(self) -> Vec<usize> {
        self.into_iter().map(|x| x as usize).collect()
    }
    fn r_int_convert_shift(self) -> Vec<usize> {
        self.into_iter().map(|x| (x - 1) as usize).collect()
    }
}

impl VecConvert<i32> for &[usize] {
    fn r_int_convert(self) -> Vec<i32> {
        self.iter().map(|&x| x as i32).collect()
    }
    fn r_int_convert_shift(self) -> Vec<i32> {
        self.iter().map(|&x| (x as i32) + 1).collect()
    }
}

impl VecConvert<usize> for &[i32] {
    fn r_int_convert(self) -> Vec<usize> {
        self.iter().map(|&x| x as usize).collect()
    }
    fn r_int_convert_shift(self) -> Vec<usize> {
        self.iter().map(|&x| (x - 1) as usize).collect()
    }
}

/// Converts vectors between floating-point types for R compatibility.
pub trait VecFloatConvert<U> {
    /// Converts the R double vectors from/to f64 to f32
    fn r_float_convert(self) -> Vec<U>;
}

impl VecFloatConvert<f32> for Vec<f64> {
    fn r_float_convert(self) -> Vec<f32> {
        self.into_iter().map(|x| x as f32).collect()
    }
}

impl VecFloatConvert<f64> for Vec<f32> {
    fn r_float_convert(self) -> Vec<f64> {
        self.into_iter().map(|x| x as f64).collect()
    }
}

impl VecFloatConvert<f64> for &[f32] {
    fn r_float_convert(self) -> Vec<f64> {
        self.iter().map(|x| *x as f64).collect()
    }
}

impl VecFloatConvert<f32> for &[f64] {
    fn r_float_convert(self) -> Vec<f32> {
        self.iter().map(|&x| x as f32).collect()
    }
}

////////////////
// R and faer //
////////////////

/// Bridge between faer matrix types and R matrix types.
///
/// Defines how to convert faer matrices to R-compatible arrays.
pub trait FaerRType: SimpleEntity + Copy + Clone + 'static {
    /// Type definition to allow R conversion
    type RType: Copy + Clone;

    /// Transform an faer matrix (f32/f64) into an R matrix (f64)
    fn to_r_matrix(x: faer::MatRef<Self>) -> extendr_api::RArray<Self::RType, 2>;
}

impl FaerRType for f64 {
    type RType = f64;
    fn to_r_matrix(x: faer::MatRef<Self>) -> extendr_api::RArray<Self, 2> {
        let mut out = RMatrix::new(x.nrows(), x.ncols());
        fill_column_major(out.data_mut(), x, |v| v);
        out
    }
}

impl FaerRType for i32 {
    type RType = i32;
    fn to_r_matrix(x: faer::MatRef<Self>) -> extendr_api::RArray<Self, 2> {
        let mut out = RMatrix::new(x.nrows(), x.ncols());
        fill_column_major(out.data_mut(), x, |v| v);
        out
    }
}

impl FaerRType for f32 {
    type RType = f64;
    fn to_r_matrix(x: faer::MatRef<Self>) -> extendr_api::RArray<f64, 2> {
        let mut out = RMatrix::new(x.nrows(), x.ncols());
        fill_column_major(out.data_mut(), x, |v| v as f64);
        out
    }
}

///////////////////
// R error stuff //
///////////////////

/// Trait to transform errors directly into extendr results
pub trait IntoExtendrErr<T> {
    /// Transforms itself into an extendr_api::Result and transforms any errors
    /// in doing so.
    fn to_extendr(self) -> extendr_api::Result<T>;
}

impl<T, E: std::fmt::Display> IntoExtendrErr<T> for std::result::Result<T, E> {
    fn to_extendr(self) -> extendr_api::Result<T> {
        self.map_err(|e| extendr_api::Error::Other(e.to_string()))
    }
}