data_beans/sparse_io_vector/

batch.rs

#![allow(dead_code)]

use super::*;

impl SparseIoVec {
    /////////////////////////////////////
    // data dictionary related methods //
    /////////////////////////////////////

    /// Register batch membership information along with the feature
    /// matrix for quick look up operations.
    ///
    /// # Arguments
    /// * `feature_matrix` - A feature matrix where each column corresponds to a cell.
    /// * `batch_membership` - A vector of batch membership information for each cell.
    #[cfg(feature = "ndarray")]
    pub fn register_batches_ndarray<T>(
        &mut self,
        feature_matrix: &ndarray::Array2<f32>,
        batch_membership: &[T],
    ) -> anyhow::Result<()>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        anyhow::ensure!(
            feature_matrix.ncols() == batch_membership.len(),
            "the feature matrix has {} columns but there are {} batch labels",
            feature_matrix.ncols(),
            batch_membership.len(),
        );
        self._register_batches(
            feature_matrix,
            batch_membership,
            |feature_matrix, batch_cells| {
                let columns = batch_cells
                    .iter()
                    .map(|&c| feature_matrix.column(c))
                    .collect::<Vec<_>>();
                ColumnDict::<usize>::from_ndarray_views(columns, batch_cells.clone())
            },
        )
    }

    /// Register batch membership information along with the feature
    /// matrix for quick look up operations.
    ///
    /// # Arguments
    /// * `feature_matrix` - A feature matrix where each column corresponds to a cell.
    /// * `batch_membership` - A vector of batch membership information for each cell.
    pub fn register_batches_dmatrix<T>(
        &mut self,
        feature_matrix: &nalgebra::DMatrix<f32>,
        batch_membership: &[T],
    ) -> anyhow::Result<()>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        anyhow::ensure!(
            feature_matrix.ncols() == batch_membership.len(),
            "the feature matrix has {} columns but there are {} batch labels",
            feature_matrix.ncols(),
            batch_membership.len(),
        );

        self._register_batches(
            feature_matrix,
            batch_membership,
            |feature_matrix, batch_cells| {
                let columns = batch_cells
                    .iter()
                    .map(|&c| feature_matrix.column(c))
                    .collect::<Vec<_>>();
                ColumnDict::<usize>::from_dvector_views(columns, batch_cells.clone())
            },
        )
    }

    fn _register_batches<M, F, T>(
        &mut self,
        feature_matrix: &M,
        batch_membership: &[T],
        create_column_dict: F,
    ) -> anyhow::Result<()>
    where
        M: Sync,
        F: Fn(&M, &Vec<usize>) -> ColumnDict<usize> + Sync,
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        let ntot = self.num_columns();
        anyhow::ensure!(
            batch_membership.len() == ntot,
            "batch membership has {} entries but there are {ntot} columns",
            batch_membership.len(),
        );
        let batches = partition_by_membership(batch_membership, None);
        let mut col_to_batch = vec![0; ntot];

        // A ColumnDict build is internally parallel (instant-distance's rayon
        // insert) only ABOVE EXACT_THRESHOLD points; smaller batches use the
        // exact backend and build single-threaded. So when batches are many
        // (>= n_threads) run the outer loop in parallel — that is the only
        // parallelism the small batches get, and large ones just nest into the
        // same work-stealing pool (safe, no true oversubscription). When batches
        // are few, build sequentially so each large build owns the pool and only
        // one HNSW index is under construction at a time (bounds peak memory).
        let n_threads = rayon::current_num_threads();
        let outer_parallel = batches.len() >= n_threads;

        info!(
            "building per-batch kNN indices ({} batches, {} cells, {}) ...",
            batches.len(),
            ntot,
            if outer_parallel {
                "parallel over batches"
            } else {
                "sequential over batches"
            }
        );

        let prog_bar =
            crate::sparse_data_visitors::styled_progress_bar(batches.len() as u64, "batches kNN");

        // Canonical batch id = SORTED LABEL order (consistent with
        // `register_batch_membership`), so per-batch stats / δ columns carry the
        // same batch ids as the rest of the pipeline. Previously this sorted by
        // descending size and used the processing position as the id, which
        // silently permuted batch ids vs the label order — a latent bug for δ /
        // `AdjMethod::Batch` consumers.
        let mut batches_vec: Vec<_> = batches.into_iter().collect();
        batches_vec.sort_by_key(|a| a.0.to_string());
        // Schedule largest batches first (LPT) so the heavy tail overlaps with
        // the small batches. The canonical id rides along; `sort_by_key(idx)`
        // below restores canonical order.
        let mut enumerated: Vec<_> = batches_vec.iter().enumerate().collect();
        enumerated.sort_by_key(|(_, (_, cells))| std::cmp::Reverse(cells.len()));
        let mut idx_name_glob_dict: Vec<_> = if outer_parallel {
            enumerated
                .into_par_iter()
                .progress_with(prog_bar.clone())
                .map(|(batch_index, (batch_name, batch_glob_indices))| {
                    (
                        batch_index,
                        batch_name.to_string().into_boxed_str(),
                        batch_glob_indices.clone(),
                        create_column_dict(feature_matrix, batch_glob_indices),
                    )
                })
                .collect()
        } else {
            enumerated
                .into_iter()
                .progress_with(prog_bar.clone())
                .map(|(batch_index, (batch_name, batch_glob_indices))| {
                    (
                        batch_index,
                        batch_name.to_string().into_boxed_str(),
                        batch_glob_indices.clone(),
                        create_column_dict(feature_matrix, batch_glob_indices),
                    )
                })
                .collect()
        };
        prog_bar.finish_and_clear();

        idx_name_glob_dict.sort_by_key(|&(idx, _, _, _)| idx);

        let mut batch_names = vec![];
        let mut batch_to_cols = vec![];
        let mut dictionaries = vec![];

        for (batch_idx, batch_name, glob_indices, dict) in idx_name_glob_dict.into_iter() {
            dict.names()
                .iter()
                .for_each(|&cell| col_to_batch[cell] = batch_idx);

            batch_names.push(batch_name);
            batch_to_cols.push(glob_indices);
            dictionaries.push(dict);
        }

        self.derived.batch_knn_lookup = Some(dictionaries);
        self.derived.col_to_batch = Some(col_to_batch);
        self.derived.batch_to_cols = Some(batch_to_cols);
        self.derived.batch_idx_to_name = Some(batch_names);

        if self.num_batches() > 2 {
            self.sort_batch_proximity()?;
        }

        Ok(())
    }

    fn sort_batch_proximity(&mut self) -> anyhow::Result<()> {
        let lookups = self
            .derived
            .batch_knn_lookup
            .as_ref()
            .ok_or(anyhow::anyhow!("no knn lookup"))?;

        use nalgebra::DMatrix;

        info!("retrieving batch-specific lookups");
        let batch_data = lookups
            .iter()
            .flat_map(|dict| {
                let data: Vec<f32> = dict.points().flatten().copied().collect();
                let ncols = dict.num_points();
                let nrows = data.len() / ncols;
                DMatrix::from_vec(nrows, ncols, data)
                    .column_mean()
                    .data
                    .as_vec()
                    .clone()
            })
            .collect::<Vec<_>>();

        let ncols = self.num_batches();
        let nrows = batch_data.len() / ncols;
        let batch_features = DMatrix::<f32>::from_vec(nrows, ncols, batch_data);

        info!(
            "built feature matrix across batches: {} x {}",
            batch_features.nrows(),
            batch_features.ncols()
        );

        let nbatches = self.num_batches();
        let batches = (0..nbatches).collect();

        let dict = ColumnDict::<usize>::from_dvector_views(
            batch_features.column_iter().collect(),
            batches,
        );

        let ret: Vec<Vec<usize>> = (0..nbatches)
            .into_par_iter()
            .map(|b| {
                dict.search_by_query_name(&b, nbatches, false)
                    .map(|(others, _)| others)
            })
            .collect::<anyhow::Result<Vec<Vec<usize>>>>()?;
        self.derived.between_batch_proximity = Some(ret);

        Ok(())
    }

    pub fn batch_name_map(&self) -> Option<HashMap<Box<str>, usize>> {
        self.derived.batch_idx_to_name.as_ref().map(|names| {
            names
                .iter()
                .enumerate()
                .map(|(idx, name)| (name.clone(), idx))
                .collect::<HashMap<Box<str>, usize>>()
        })
    }

    pub fn num_batches(&self) -> usize {
        if let Some(v) = &self.derived.batch_to_cols {
            v.len()
        } else if let Some(v) = &self.derived.batch_knn_lookup {
            v.len()
        } else {
            0
        }
    }

    /// Borrow the per-batch HNSW lookups populated by
    /// `build_hnsw_per_batch` / `register_batches_dmatrix`. Returns `None`
    /// before the indices have been built.
    pub fn batch_knn_lookup(&self) -> Option<&Vec<ColumnDict<usize>>> {
        self.derived.batch_knn_lookup.as_ref()
    }

    /// Register batch membership information without building HNSW
    /// indices. This is a lightweight alternative to `register_batches_dmatrix`
    /// for use with pb-sample based batch correction.
    ///
    /// # Panics
    /// If `batch_membership` is not one entry per column;
    /// [`try_register_batch_membership`](Self::try_register_batch_membership)
    /// returns that as an error instead.
    pub fn register_batch_membership<T>(&mut self, batch_membership: &[T])
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        self.try_register_batch_membership(batch_membership)
            .unwrap_or_else(|e| panic!("{e}"));
    }

    /// [`register_batch_membership`](Self::register_batch_membership) that
    /// returns a membership of the wrong length as an error.
    ///
    /// # Errors
    /// If `batch_membership` is not one entry per column: a shorter slice
    /// would leave the unlisted columns in the first batch, absent from that
    /// batch's column list. Nothing is registered then.
    pub fn try_register_batch_membership<T>(&mut self, batch_membership: &[T]) -> anyhow::Result<()>
    where
        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
    {
        let ntot = self.num_columns();
        anyhow::ensure!(
            batch_membership.len() == ntot,
            "batch membership has {} entries but there are {ntot} columns",
            batch_membership.len(),
        );
        let batches = partition_by_membership(batch_membership, None);
        let mut col_to_batch = vec![0; ntot];

        let mut sorted_batches: Vec<_> = batches.into_iter().collect();
        sorted_batches.sort_by_key(|a| a.0.to_string());

        let mut batch_names = Vec::with_capacity(sorted_batches.len());
        let mut batch_to_cols = Vec::with_capacity(sorted_batches.len());

        for (batch_idx, (batch_name, glob_indices)) in sorted_batches.into_iter().enumerate() {
            for &cell in &glob_indices {
                col_to_batch[cell] = batch_idx;
            }
            batch_names.push(batch_name.to_string().into_boxed_str());
            batch_to_cols.push(glob_indices);
        }

        self.derived.col_to_batch = Some(col_to_batch);
        self.derived.batch_to_cols = Some(batch_to_cols);
        self.derived.batch_idx_to_name = Some(batch_names);
        Ok(())
    }

    /// Declare that each column stands for more than one observation.
    ///
    /// A column is normally one cell, so every statistic that divides by a
    /// count adds `1` per column. That breaks when a column is a *summary* of
    /// many cells — a carried pseudobulk, or a bulk sample — because the
    /// per-cell rate `μ = Σy / n` would divide a whole group's counts by one.
    ///
    /// With multiplicities registered, a column holding the **mean** profile of
    /// `m` cells and a weight of `m` contributes exactly what those `m` cells
    /// would have: `m·mean` to the sums and `m` to the count.
    ///
    /// Absent (the default) every column weighs `1`, and every accumulation is
    /// bit-for-bit what it was before this existed.
    ///
    /// # Errors
    /// If `multiplicity` is not one entry per column, or holds a non-finite or
    /// non-positive weight — a zero would silently delete a column from the
    /// denominator while leaving its counts in the numerator.
    pub fn register_column_multiplicity(&mut self, multiplicity: &[f32]) -> anyhow::Result<()> {
        let ntot = self.num_columns();
        anyhow::ensure!(
            multiplicity.len() == ntot,
            "column multiplicity has {} entries but there are {ntot} columns",
            multiplicity.len(),
        );
        if let Some((i, w)) = multiplicity
            .iter()
            .enumerate()
            .find(|(_, w)| !w.is_finite() || **w <= 0.0)
        {
            anyhow::bail!("column {i} has multiplicity {w}; weights must be finite and positive");
        }
        self.col_multiplicity = Some(multiplicity.to_vec());
        Ok(())
    }

    /// Weight of a single column — `1.0` when no multiplicities are registered.
    #[must_use]
    pub fn column_multiplicity(&self, col: usize) -> f32 {
        self.col_multiplicity.as_ref().map_or(1.0, |m| m[col])
    }

    /// True when any column stands for more than one observation.
    #[must_use]
    pub fn has_column_multiplicity(&self) -> bool {
        self.col_multiplicity.is_some()
    }

    /// The whole multiplicity vector, one weight per column — `None` when no
    /// multiplicities are registered (every column is one observation).
    ///
    /// Prefer this over gathering [`Self::column_multiplicity`] in a loop:
    /// callers were rebuilding the vector element-by-element, re-encoding the
    /// "absent means 1.0" default at every site.
    #[must_use]
    pub fn column_multiplicities(&self) -> Option<&[f32]> {
        self.col_multiplicity.as_deref()
    }

    pub fn batch_names(&self) -> Option<Vec<Box<str>>> {
        self.derived.batch_idx_to_name.clone()
    }

    pub fn batch_to_columns(&self, batch: usize) -> Option<&Vec<usize>> {
        if let Some(batch_to_cols) = &self.derived.batch_to_cols {
            Some(&batch_to_cols[batch])
        } else {
            None
        }
    }

    pub fn get_batch_membership<I>(&self, cells: I) -> Vec<usize>
    where
        I: Iterator<Item = usize>,
    {
        let cell_to_batch = self
            .derived
            .col_to_batch
            .as_ref()
            .expect("cell_to_batch not initialized");
        cells.into_iter().map(|c| cell_to_batch[c]).collect()
    }

    pub fn column_names(&self) -> anyhow::Result<Vec<Box<str>>> {
        debug_assert_eq!(self.num_columns(), self.column_names_with_data_tag.len());
        Ok(self.column_names_with_data_tag.clone())
    }
}