data-beans 0.6.15

Sparse genomics data backends, QC, algorithms, and simulation
Documentation
// #![allow(dead_code)]

use crate::sparse_io_vector::*;
use log::info;

/// `sparse_io_stack` is a stack of `sparse_io_vector`
pub struct SparseIoStack {
    pub stack: Vec<SparseIoVec>,
    column_names: Vec<Box<str>>,
}

impl Default for SparseIoStack {
    /// an empty sparse io vector for horizontal data integration
    fn default() -> Self {
        Self::new()
    }
}

impl SparseIoStack {
    pub fn new() -> Self {
        Self {
            stack: vec![],
            column_names: vec![],
        }
    }

    pub fn push(&mut self, data: SparseIoVec) -> anyhow::Result<()> {
        if self.stack.is_empty() {
            self.column_names.extend(data.column_names()?);
            self.stack.push(data);
            return Ok(());
        }

        info!("Checking column names...");
        if self.column_names != data.column_names()? {
            return Err(anyhow::anyhow!("column names don't match"));
        }
        self.stack.push(data);
        Ok(())
    }

    /// number of data types
    pub fn num_types(&self) -> usize {
        self.stack.len()
    }

    /// number of shared columns
    pub fn num_columns(&self) -> anyhow::Result<usize> {
        self.stack
            .iter()
            .map(|x| x.num_columns())
            .max()
            .ok_or(anyhow::anyhow!("can't figure out the max"))
    }

    /// Get the shared column names.
    pub fn column_names(&self) -> anyhow::Result<Vec<Box<str>>> {
        if self.column_names.len() != self.num_columns()? {
            return Err(anyhow::anyhow!("inconsistent columns"));
        }
        Ok(self.column_names.clone())
    }

    /// Exclude columns (cells) across every layer in the stack in
    /// lockstep. `keep[global_col]` is `true` to keep, `false` to drop.
    /// Mirror of [`SparseIoVec::mask_columns`]; the members share one
    /// synchronized cell axis, so the same mask applies to all. MUST be
    /// called before batch/group registration.
    pub fn mask_columns_all(&mut self, keep: &[bool]) -> anyhow::Result<()> {
        let n = self.num_columns()?;
        if keep.len() != n {
            return Err(anyhow::anyhow!(
                "mask_columns_all: keep.len()={} != num_columns={}",
                keep.len(),
                n
            ));
        }
        for layer in self.stack.iter_mut() {
            layer.mask_columns(keep)?;
        }
        // Refresh the cached shared column names from the (now filtered)
        // first member; all members were masked identically.
        self.column_names = match self.stack.first() {
            Some(first) => first.column_names()?,
            None => vec![],
        };
        Ok(())
    }

    /// Register batch membership for all layers in the stack.
    pub fn register_batch_membership<T>(&mut self, batch_membership: &[T])
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        for layer in self.stack.iter_mut() {
            layer.register_batch_membership(batch_membership);
        }
    }

    /// Get the row names combined across all the types. We will
    /// append additional data type index: `format!("{}/{}", x, d)`
    ///
    pub fn row_names(&self) -> anyhow::Result<Vec<Box<str>>> {
        Ok(self
            .stack
            .iter()
            .enumerate()
            .map(|(d, x)| {
                x.row_names().map(|x| {
                    x.into_iter()
                        .map(|y| format!("{}/{}", y, d).into_boxed_str())
                        .collect::<Vec<_>>()
                })
            })
            .collect::<anyhow::Result<Vec<_>>>()?
            .into_iter()
            .flatten()
            .collect())
    }
}