data_beans/alg/collapse_data/
reassign_cells.rs

1//! Reassign cells: the first pass over the initial hash. Cells move
2//! between the top nodes under the DC-Poisson likelihood of the node
3//! profiles, so a node the marginal sketch left mixed across lineages
4//! becomes one lineage before the tree grows below it. Cells only move within their batch: the entities are cells,
5//! the groups are (node, batch) pairs, and a cell's candidates are the
6//! groups of its own batch.
7
8use super::*;
9use crate::alg::dc_poisson::{
10    refine_with_candidates, FeatureWeighting, Profiles, RefineContext, RefineParams,
11};
12use nalgebra_sparse::CscMatrix;
13use rand::rngs::SmallRng;
14use rand::SeedableRng;
15
16#[derive(Clone, Debug)]
17pub struct ReassignCellsParams {
18    pub num_gibbs: usize,
19    pub num_greedy: usize,
20    pub feature_weighting: FeatureWeighting,
21    pub seed: u64,
22}
23
24impl Default for ReassignCellsParams {
25    fn default() -> Self {
26        Self {
27            num_gibbs: 3,
28            num_greedy: 10,
29            feature_weighting: FeatureWeighting::FisherInfoNb,
30            seed: 42,
31        }
32    }
33}
34
35/// Move active cells between nodes in place; returns the number of moves.
36/// `csc` holds every cell's counts (genes x cells, global column order);
37/// `node_of_cell` holds compact node ids and is rewritten for active cells.
38pub(crate) fn reassign_cells_to_nodes(
39    csc: &CscMatrix<f32>,
40    col_to_batch: &[usize],
41    num_batches: usize,
42    active: &[bool],
43    node_of_cell: &mut [usize],
44    params: &ReassignCellsParams,
45) -> usize {
46    let n = csc.ncols();
47    assert_eq!(col_to_batch.len(), n);
48    assert_eq!(active.len(), n);
49    assert_eq!(node_of_cell.len(), n);
50    let num_batches = num_batches.max(1);
51    let num_nodes = node_of_cell.iter().copied().max().map_or(0, |m| m + 1);
52    if num_nodes < 2 || params.num_gibbs + params.num_greedy == 0 {
53        return 0;
54    }
55    let cells: Vec<usize> = (0..n).filter(|&c| active[c]).collect();
56    if cells.is_empty() {
57        return 0;
58    }
59
60    ////////////////////////////////////
61    // Groups = occupied (node, batch) //
62    ////////////////////////////////////
63
64    let mut group_of_pair = vec![usize::MAX; num_nodes * num_batches];
65    let mut group_pair: Vec<(usize, usize)> = Vec::new();
66    let mut labels: Vec<usize> = Vec::with_capacity(cells.len());
67    for &c in &cells {
68        let key = node_of_cell[c] * num_batches + col_to_batch[c];
69        if group_of_pair[key] == usize::MAX {
70            group_of_pair[key] = group_pair.len();
71            group_pair.push((node_of_cell[c], col_to_batch[c]));
72        }
73        labels.push(group_of_pair[key]);
74    }
75    let k = group_pair.len();
76    let mut groups_of_batch: Vec<Vec<usize>> = vec![Vec::new(); num_batches];
77    for (g, &(_, b)) in group_pair.iter().enumerate() {
78        groups_of_batch[b].push(g);
79    }
80    let candidates: Vec<Vec<usize>> = cells
81        .iter()
82        .map(|&c| groups_of_batch[col_to_batch[c]].clone())
83        .collect();
84
85    //////////////////////////////
86    // Cell profiles and sweeps //
87    //////////////////////////////
88
89    let rows: Vec<Vec<(usize, f32)>> = cells
90        .iter()
91        .map(|&c| {
92            let col = csc.col(c);
93            col.row_indices()
94                .iter()
95                .zip(col.values())
96                .map(|(&g, &v)| (g, v))
97                .collect()
98        })
99        .collect();
100    let mut profiles = Profiles::from_gene_sums(&rows, csc.nrows());
101    profiles.apply_feature_weighting(params.feature_weighting);
102    let refine = RefineParams {
103        num_gibbs: params.num_gibbs,
104        num_greedy: params.num_greedy,
105        feature_weighting: params.feature_weighting,
106        seed: params.seed,
107        ..RefineParams::default()
108    };
109    let mut rng = SmallRng::seed_from_u64(params.seed);
110    let moves = refine_with_candidates(
111        &mut labels,
112        &candidates,
113        &mut rng,
114        &RefineContext {
115            profiles: &profiles,
116            k,
117            params: &refine,
118            level_label: "reassign cells",
119        },
120    );
121    for (i, &c) in cells.iter().enumerate() {
122        node_of_cell[c] = group_pair[labels[i]].0;
123    }
124    info!(
125        "reassign cells: {} of {} cells moved between {} top nodes ({} node x batch groups)",
126        moves,
127        cells.len(),
128        num_nodes,
129        k
130    );
131    moves
132}
133
134#[cfg(test)]
135#[path = "reassign_cells_tests.rs"]
136mod tests;
data_beans/alg/collapse_data/reassign_cells.rs

data_beans/alg/collapse_data/
reassign_cells.rs
