proseg 3.0.12

Probabilistic cell segmentation for in situ spatial transcriptomics
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

use rand::Rng;
use rand_distr::{Binomial, Distribution};

#[derive(Copy, Clone, Debug)]
pub struct MultinomialOutcome {
    pub prob: f32,
    pub index: u32,
}

pub struct Multinomial {
    pub outcomes: Vec<MultinomialOutcome>,
}

// Idea here is to speed up sampling sparse multinomial vectors by sorting the outcome
// in descending order by probability, so we can hopefully bail out early.
impl Multinomial {
    pub fn with_k(k: usize) -> Multinomial {
        Multinomial {
            outcomes: vec![
                MultinomialOutcome {
                    prob: 0.0,
                    index: 0
                };
                k
            ],
        }
    }

    pub fn sample<R, F>(&mut self, rng: &mut R, n: u32, mut nonzero_sample: F)
    where
        R: Rng,
        F: FnMut(usize, u32),
    {
        // sort in descending order by probability
        self.outcomes
            .sort_unstable_by(|a, b| b.prob.partial_cmp(&a.prob).unwrap());

        let norm = self
            .outcomes
            .iter()
            .fold(0.0, |accum, outcome| accum + outcome.prob as f64);

        let mut ρ = 1.0;
        let mut s = n;
        for outcome in self.outcomes.iter() {
            let p = outcome.prob as f64 / norm;
            let mut x = 0;
            if ρ > 0.0 {
                x = Binomial::new(s as u64, (p / ρ).min(1.0))
                    .unwrap()
                    .sample(rng) as u32;
                if x > 0 {
                    nonzero_sample(outcome.index as usize, x);
                }
            }

            s -= x;
            ρ -= p;

            if s == 0 {
                break;
            }
        }
        assert!(s == 0);
    }
}