bixverse-rs 0.2.0

//! Various helpers to transfer between Rust and R via the rextendr interface

use extendr_api::prelude::*;
use faer::{Mat, MatRef};
use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
use std::collections::{BTreeMap, HashMap};
use std::ops::{Add, Mul};

use crate::prelude::*;
use crate::utils::traits::FaerRType;

//////////////////
// Type aliases //
//////////////////

/// Type alias for a double nested HashMap
pub type NestedHashMap = FxHashMap<String, FxHashMap<String, FxHashSet<String>>>;

/// Type alias for double nested BtreeMap
pub type NestedBtreeMap = BTreeMap<String, BTreeMap<String, FxHashSet<String>>>;

////////////
// Errors //
////////////

/// Error handling for named numeric conversion
#[derive(Debug)]
pub enum NamedVecError {
    /// Not numeric error
    NotNumeric,
    /// No Names provided error
    NoNames,
    /// Missing values in the data error
    MissingValues,
}

impl std::fmt::Display for NamedVecError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            NamedVecError::NotNumeric => write!(f, "Input is not a numeric vector"),
            NamedVecError::NoNames => write!(f, "Vector has no names attribute"),
            NamedVecError::MissingValues => write!(f, "Vector contains missing values"),
        }
    }
}

impl std::error::Error for NamedVecError {}

impl From<NamedVecError> for extendr_api::Error {
    fn from(err: NamedVecError) -> Self {
        extendr_api::Error::Other(err.to_string())
    }
}

///////////
// Lists //
///////////

/// Transforms a Robj List into a Hashmap
///
/// This function assumes that the R list contains string vector!
///
/// ### Params
///
/// * `r_list` - R list that has names and contains string vectors.
///
/// ### Returns
///
/// A HashMap with as keys the names of the list and values the string vectors.
pub fn r_list_to_hashmap(r_list: List) -> extendr_api::Result<FxHashMap<String, Vec<String>>> {
    let mut result = FxHashMap::with_capacity_and_hasher(r_list.len(), FxBuildHasher);

    for (n, s) in r_list {
        let s_vec = s.as_string_vector().ok_or_else(|| {
            Error::Other(format!(
                "Failed to convert value for key '{}' to string vector",
                n
            ))
        })?;
        result.insert(n.to_string(), s_vec);
    }

    Ok(result)
}

/// Transforms a Robj List into a Hashmap with HashSet values
///
/// This function assumes that the R list contains string vector!
///
/// ### Params
///
/// * `r_list` - R list that has names and contains string vectors.
///
/// ### Returns
///
/// A HashMap with as keys the names of the list and values as HashSets.
pub fn r_list_to_hashmap_set(
    r_list: List,
) -> extendr_api::Result<FxHashMap<String, FxHashSet<String>>> {
    let mut result = FxHashMap::with_capacity_and_hasher(r_list.len(), FxBuildHasher);

    for (n, s) in r_list {
        let s_vec = s.as_string_vector().ok_or_else(|| {
            Error::Other(format!(
                "Failed to convert value for key '{}' to string vector",
                n
            ))
        })?;
        let mut s_hash = FxHashSet::with_capacity_and_hasher(s_vec.len(), FxBuildHasher);
        for item in s_vec {
            s_hash.insert(item);
        }
        result.insert(n.to_string(), s_hash);
    }

    Ok(result)
}

/// Transforms an Robj nested list into a nested HashMap containing further HashMap
///
/// A helper that generates a nested HashMap from a nested R list.
///
/// ### Params
///
/// * `r_nested_list` - A named R list that contains named lists with String vectors.
///
/// ### Returns
///
/// Returns a `NestedHashMap`
#[allow(dead_code)]
pub fn r_nested_list_to_nested_hashmap(r_nested_list: List) -> extendr_api::Result<NestedHashMap> {
    let mut result = FxHashMap::with_capacity_and_hasher(r_nested_list.len(), FxBuildHasher);
    for (n, obj) in r_nested_list {
        let inner_list = obj.as_list().ok_or_else(|| {
            Error::Other(format!("Failed to convert value for key '{}' to list", n))
        })?;
        let inner_hashmap = r_list_to_hashmap_set(inner_list)?;
        result.insert(n.to_string(), inner_hashmap);
    }
    Ok(result)
}

/// Transforms an R list to a vector of HashSets
///
/// ### Params
///
/// * `r_list` - A named R list that contains named lists with String vectors.
///
/// ### Returns
///
/// Returns a Vector of FxHashSets
pub fn r_list_to_hash_vec(r_list: List) -> extendr_api::Result<Vec<FxHashSet<String>>> {
    let mut res = Vec::with_capacity(r_list.len());
    for (n, s) in r_list {
        let s_vec = s.as_string_vector().ok_or_else(|| {
            Error::Other(format!(
                "Failed to convert value for key '{}' to string vector",
                n
            ))
        })?;
        let mut s_hash = FxHashSet::with_capacity_and_hasher(s_vec.len(), FxBuildHasher);
        for item in s_vec {
            s_hash.insert(item);
        }
        res.push(s_hash)
    }

    Ok(res)
}

/// Transform a Robj List into a BTreeMap with the values as HashSet
///
/// Use where ordering of the values matters as the HashMaps have non-deterministic
/// ordering
///
/// ### Params
///
/// * `r_list` - R list that has names and contains string vectors.
///
/// ### Returns
///
/// A BTreeMap with as keys the names of the list and values as HashSets.
pub fn r_list_to_btree_set(
    r_list: List,
) -> extendr_api::Result<BTreeMap<String, FxHashSet<String>>> {
    let mut result = BTreeMap::new();
    for (n, s) in r_list {
        let s_vec = s.as_string_vector().ok_or_else(|| {
            Error::Other(format!(
                "Failed to convert value for key '{}' to string vector",
                n
            ))
        })?;
        let mut s_hash = FxHashSet::with_capacity_and_hasher(s_vec.len(), FxBuildHasher);
        for item in s_vec {
            s_hash.insert(item);
        }
        result.insert(n.to_string(), s_hash);
    }
    Ok(result)
}

/// Transform an Robj nested list into a nested BtreeMap
///
/// A helper that generates a nested BTreeMap from a nested R list.
///
/// ### Params
///
/// * `r_nested_list` - A named R list that contains named lists with String vectors.
///
/// ### Returns
///
/// Returns a `NestedBtreeMap`
pub fn r_nested_list_to_btree_nest(r_nested_list: List) -> extendr_api::Result<NestedBtreeMap> {
    let mut result = BTreeMap::new();

    for (n, obj) in r_nested_list {
        let inner_list = obj.as_list().ok_or_else(|| {
            Error::Other(format!("Failed to convert value for key '{}' to list", n))
        })?;
        let inner_tree = r_list_to_btree_set(inner_list)?;
        result.insert(n.to_string(), inner_tree);
    }

    Ok(result)
}

/////////////
// Vectors //
/////////////

/// Type alias for named numeric vectors
///
/// ### Fields
///
/// * `0` The names of the vector
/// * `1` The values of the vector
pub type NamedNumericVec = (Vec<String>, Vec<f64>);

/// Transforms a Robj List into an array of String arrays.
///
/// ### Params
///
/// * `r_list` - R list that has names and contains string vectors.
///
/// ### Returns
///
/// A vector of vectors with Strings
pub fn r_list_to_str_vec(r_list: List) -> extendr_api::Result<Vec<Vec<String>>> {
    let mut result = Vec::with_capacity(r_list.len());

    for (n, s) in r_list.into_iter() {
        let s_vec = s.as_string_vector().ok_or_else(|| {
            Error::Other(format!(
                "Failed to convert value to string vector at key '{}'",
                n
            ))
        })?;
        result.push(s_vec);
    }

    Ok(result)
}

/// Get the names and numeric values from a named R vector
///
/// ### Params
///
/// * `named_vec` - Robj that represents a named numeric in R
///
/// ### Returns
///
/// The `NamedNumericVec` type alias.
pub fn r_named_vec_data(named_vec: Robj) -> extendr_api::Result<NamedNumericVec> {
    let values = named_vec
        .as_real_vector()
        .ok_or(NamedVecError::NotNumeric)?;

    let names_attr = named_vec.names().ok_or(NamedVecError::NoNames)?;

    let names: Vec<String> = names_attr.into_iter().map(|s| s.to_string()).collect();

    Ok((names, values))
}

//////////////
// Matrices //
//////////////

/// Transform an R matrix to a Faer one
///
/// ### Params
///
/// * `x` - The R matrix to transform into a faer MatRef (with `f64`)
///
/// ### Returns
///
/// The faer `MatRef` from the original R matrix.
pub fn r_matrix_to_faer<T>(x: &RMatrix<T>) -> MatRef<'_, T>
where
    T: Copy + Clone,
    extendr_api::Robj: for<'a> extendr_api::AsTypedSlice<'a, T>,
{
    let ncol = x.ncols();
    let nrow = x.nrows();
    let data = x.data();
    MatRef::from_column_major_slice(data, nrow, ncol)
}

/// Transform an R matrix into a nested vector of booleans
///
/// ### Params
///
/// * `x` - The R matrix to transform into a vector of vectors with booleans
///
/// ### Returns
///
/// The nested vector with the outer vector representing the columns.
pub fn r_matrix_to_vec_bool(x: &RMatrix<Rbool>) -> Vec<Vec<bool>> {
    let ncol = x.ncols();
    let nrow = x.nrows();
    let data = x.data();

    (0..ncol)
        .map(|j| (0..nrow).map(|i| data[i + j * nrow].to_bool()).collect())
        .collect()
}

/// Transform a faer into an R matrix
///
/// ### Params
///
/// * `x` - faer `MatRef` matrix to transform into an R matrix
///
/// ###
///
/// The R matrix based on the faer matrix.
pub fn faer_to_r_matrix<T>(x: MatRef<T>) -> extendr_api::RArray<T::RType, 2>
where
    T: FaerRType,
{
    T::to_r_matrix(x)
}

/// Transform an R matrix into a f32 one
///
/// ### Params
///
/// * `x` - R matrix with f64.
///
/// ### Returns
///
/// A faer Mat with f32
pub fn r_matrix_to_faer_fp32(x: &RMatrix<f64>) -> Mat<f32> {
    let ncol = x.ncols();
    let nrow = x.nrows();
    let data = x.data();
    let data_fp32 = data.iter().map(|x| *x as f32).collect::<Vec<f32>>();
    Mat::from_fn(nrow, ncol, |i, j| data_fp32[i + j * nrow])
}

/// Transform a SparseColumnMatrix to an R list
///
/// ### Params
///
/// * `sparse` - SparseColumnMatrix structure
///
/// ### Returns
///
/// R list with the following slots:
/// * `data` - The values
/// * `row_indices` - The row indices
/// * `col_ptr` - The column pointers
/// * `ncol` - Number of columns
/// * `nrow` - Number of rows
/// * `cs_type` - Compressed Sparse Format type
pub fn sparse_data_to_list<T>(sparse: CompressedSparseData2<T>) -> List
where
    T: Into<Robj> + Clone + Default + Into<f64> + Sync + Add + PartialEq + Mul,
{
    let data: Vec<f64> = sparse
        .data
        .into_iter()
        .map(|x| {
            let robj: Robj = x.into();
            robj.as_real().unwrap()
        })
        .collect();
    let indptr = sparse
        .indptr
        .iter()
        .map(|x| *x as i32)
        .collect::<Vec<i32>>();
    let indices = sparse
        .indices
        .iter()
        .map(|x| *x as i32)
        .collect::<Vec<i32>>();
    let nrow = sparse.shape.0;
    let ncol = sparse.shape.1;
    let cs_type = match sparse.cs_type {
        CompressedSparseFormat::Csr => "csr",
        CompressedSparseFormat::Csc => "csc",
    };

    list!(
        data = data,
        indptr = indptr,
        indices = indices,
        nrow = nrow,
        ncol = ncol,
        cs_type = cs_type.to_string()
    )
}

/// Transform an R list storing CSR/C data into CompressedSparseData2
///
/// ### Params
///
/// * `r_list` - R list that has the following elements: `indptr`, `indices`,
///   `data`, `nrow`, `ncol` and `format`
/// * `populate_data_2` - Boolean. If set to `true`, the data will be also
///   copied into data_2 of the `CompressedSparseData2`.
///
/// ### Returns
///
/// The CompressedSparseData2 Rust object with the data
pub fn list_to_sparse_matrix<T>(
    r_list: List,
    populate_data_2: bool,
) -> Result<CompressedSparseData2<T>, BixverseErrors>
where
    T: Clone + Default + TryFrom<Robj>,
{
    let r_data: HashMap<&str, Robj> = r_list
        .try_into()
        .map_err(|_| BixverseErrors::RListParse("not a named list"))?;

    let indptr: Vec<usize> = r_data
        .get("indptr")
        .and_then(|v| v.as_integer_slice())
        .unwrap()
        .iter()
        .map(|&x| x as usize)
        .collect();

    let indices: Vec<usize> = r_data
        .get("indices")
        .and_then(|v| v.as_integer_slice())
        .unwrap()
        .iter()
        .map(|&x| x as usize)
        .collect();

    let data: Vec<T> = r_data
        .get("data")
        .unwrap()
        .as_real_slice()
        .unwrap()
        .iter()
        .map(|&x| T::try_from(Robj::from(x)).ok().unwrap())
        .collect();

    let nrow = r_data.get("nrow").and_then(|v| v.as_integer()).unwrap() as usize;
    let ncol = r_data.get("ncol").and_then(|v| v.as_integer()).unwrap() as usize;

    let format_str = r_data.get("format").and_then(|v| v.as_str()).unwrap();
    let cs_type = match format_str {
        "csr" => CompressedSparseFormat::Csr,
        "csc" => CompressedSparseFormat::Csc,
        _ => panic!("Unknown format"),
    };

    let data_2 = if populate_data_2 {
        Some(data.clone())
    } else {
        None
    };

    Ok(CompressedSparseData2 {
        data,
        indices,
        indptr,
        cs_type,
        data_2,
        shape: (nrow, ncol),
    })
}