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//!
8//! The pass runs without the member-count limits of [`RefineParams`]. The
9//! top nodes of the marginal sketch are mixes of lineages, and making them
10//! one lineage each takes more than a trade of members: a mixed node gives
11//! its cells to the nodes of their lineages and may empty, and a node that
12//! gathers a lineage grows. Under a floor the leftovers of every mixed
13//! node stay where they are, and the nodes end up far less pure. The
14//! number of leaves is not lost with the emptied nodes: the tree grown
15//! below the remaining ones splits them to the same targets.
16//!
17//! So nodes do empty and grow here, by the profiles of their cells. What
18//! the score does not do is move cells between two nodes of one profile
19//! for their size alone (see
20//! [`DcPoissonStats`](crate::alg::dc_poisson::DcPoissonStats)).
21
22use super::*;
23use crate::alg::dc_poisson::{
24    refine_with_candidates, FeatureWeighting, Profiles, RefineContext, RefineParams,
25};
26use nalgebra_sparse::CscMatrix;
27use rand::rngs::SmallRng;
28use rand::SeedableRng;
29
30#[derive(Clone, Debug)]
31pub struct ReassignCellsParams {
32    pub num_gibbs: usize,
33    pub num_greedy: usize,
34    pub feature_weighting: FeatureWeighting,
35    pub seed: u64,
36}
37
38impl Default for ReassignCellsParams {
39    fn default() -> Self {
40        Self {
41            num_gibbs: 3,
42            num_greedy: 10,
43            feature_weighting: FeatureWeighting::FisherInfoNb,
44            seed: 42,
45        }
46    }
47}
48
49/// Move active cells between nodes in place; returns the number of moves.
50/// `csc` holds every cell's counts (genes x cells, global column order);
51/// `node_of_cell` holds compact node ids and is rewritten for active cells.
52pub(crate) fn reassign_cells_to_nodes(
53    csc: &CscMatrix<f32>,
54    col_to_batch: &[usize],
55    num_batches: usize,
56    active: &[bool],
57    node_of_cell: &mut [usize],
58    params: &ReassignCellsParams,
59) -> usize {
60    let n = csc.ncols();
61    assert_eq!(col_to_batch.len(), n);
62    assert_eq!(active.len(), n);
63    assert_eq!(node_of_cell.len(), n);
64    let num_batches = num_batches.max(1);
65    let num_nodes = node_of_cell.iter().copied().max().map_or(0, |m| m + 1);
66    if num_nodes < 2 || params.num_gibbs + params.num_greedy == 0 {
67        return 0;
68    }
69    let cells: Vec<usize> = (0..n).filter(|&c| active[c]).collect();
70    if cells.is_empty() {
71        return 0;
72    }
73
74    ////////////////////////////////////
75    // Groups = occupied (node, batch) //
76    ////////////////////////////////////
77
78    let mut group_of_pair = vec![usize::MAX; num_nodes * num_batches];
79    let mut group_pair: Vec<(usize, usize)> = Vec::new();
80    let mut labels: Vec<usize> = Vec::with_capacity(cells.len());
81    for &c in &cells {
82        let key = node_of_cell[c] * num_batches + col_to_batch[c];
83        if group_of_pair[key] == usize::MAX {
84            group_of_pair[key] = group_pair.len();
85            group_pair.push((node_of_cell[c], col_to_batch[c]));
86        }
87        labels.push(group_of_pair[key]);
88    }
89    let k = group_pair.len();
90    let mut groups_of_batch: Vec<Vec<usize>> = vec![Vec::new(); num_batches];
91    for (g, &(_, b)) in group_pair.iter().enumerate() {
92        groups_of_batch[b].push(g);
93    }
94    let candidates: Vec<Vec<usize>> = cells
95        .iter()
96        .map(|&c| groups_of_batch[col_to_batch[c]].clone())
97        .collect();
98
99    //////////////////////////////
100    // Cell profiles and sweeps //
101    //////////////////////////////
102
103    let rows: Vec<Vec<(usize, f32)>> = cells
104        .iter()
105        .map(|&c| {
106            let col = csc.col(c);
107            col.row_indices()
108                .iter()
109                .zip(col.values())
110                .map(|(&g, &v)| (g, v))
111                .collect()
112        })
113        .collect();
114    let mut profiles = Profiles::from_gene_sums(&rows, csc.nrows());
115    profiles.apply_feature_weighting(params.feature_weighting);
116    let refine = RefineParams {
117        num_gibbs: params.num_gibbs,
118        num_greedy: params.num_greedy,
119        feature_weighting: params.feature_weighting,
120        seed: params.seed,
121        size_floor: 0.0,
122        size_cap: 0.0,
123        ..RefineParams::default()
124    };
125    let mut rng = SmallRng::seed_from_u64(params.seed);
126    let moves = refine_with_candidates(
127        &mut labels,
128        &candidates,
129        &mut rng,
130        &RefineContext {
131            profiles: &profiles,
132            k,
133            params: &refine,
134            level_label: "reassign cells",
135        },
136    );
137    for (i, &c) in cells.iter().enumerate() {
138        node_of_cell[c] = group_pair[labels[i]].0;
139    }
140    info!(
141        "reassign cells: {} of {} cells moved between {} top nodes ({} node x batch groups)",
142        moves,
143        cells.len(),
144        num_nodes,
145        k
146    );
147    moves
148}
149
150#[cfg(test)]
151#[path = "reassign_cells_tests.rs"]
152mod tests;
data_beans/alg/collapse_data/reassign_cells.rs

data_beans/alg/collapse_data/
reassign_cells.rs
