hsc/
subclone.rs

1use crate::genotype::NeutralMutationPoisson;
2use crate::Probs;
3use crate::{stemcell::StemCell, write2file, MAX_SUBCLONES};
4use anyhow::{ensure, Context};
5use rand::seq::IteratorRandom;
6use rand::Rng;
7use rand_distr::{Bernoulli, Distribution, Gamma};
8use sosa::ReactionRates;
9use std::num::NonZeroUsize;
10use std::path::Path;
11
12/// Distribution probabilities for the simulations upon cell division.
13#[derive(Debug, Default, Clone)]
14pub struct Distributions {
15    /// arrival rate of fit mutations per division per cell
16    pub u: f32,
17    /// neutral mutations
18    pub neutral_poisson: NeutralMutationPoisson,
19}
20
21impl Distributions {
22    pub fn new(probs: Probs, verbosity: u8) -> Self {
23        let (u, background, division) = match probs {
24            Probs::Symmetric { u, probs_per_year } => {
25                (u, probs_per_year.mu_background, probs_per_year.mu_division)
26            }
27            Probs::Asymmetric {
28                u, probs_per_year, ..
29            } => (u, probs_per_year.mu_background, probs_per_year.mu_division),
30        };
31        if verbosity > 1 {
32            println!(
33                "creating distributions with u: {u}, lambda_background: {background}, lambda_division: {division}"
34            );
35        }
36
37        Self {
38            u,
39            neutral_poisson: NeutralMutationPoisson::new(division, background).unwrap(),
40        }
41    }
42}
43
44pub fn from_mean_std_to_shape_scale(mean: f32, std: f32) -> (f32, f32) {
45    (mean.powf(2.) / std.powf(2.), std.powf(2.) / mean)
46}
47
48/// Fitness models implemented so far.
49#[derive(Clone, Copy, Debug)]
50pub enum Fitness {
51    /// According to the `Neutral` fitness model, all subclones have the same
52    /// birth-rates `b0`, where `b0` is the birth-rate of the wild-type neutral
53    /// clone.
54    Neutral,
55    /// According to the `Fixed` fitness model, the birth-rates for the
56    /// non-wildtype clones are `b0(1+s)`, where `b0` is the birth-rate of the
57    /// wild-type neutral clone.
58    /// That is, all the subclones have the same proliferate rates except the
59    /// wild-type one.
60    Fixed { s: f32 },
61    /// According to the `GammaSampled` fitness model, the birth-rates for the
62    /// non-wildtype clones are `b0(1+s_i)`, where `b0` is the birth-rate of the
63    /// wild-type neutral clone and `s_i` are sampled from the Gamma
64    /// distribution.
65    /// In this case, all subclones can have different proliferate rates.
66    GammaSampled { shape: f32, scale: f32 },
67}
68
69impl Fitness {
70    pub fn get_mean_std(&self) -> (f32, f32) {
71        match self {
72            Fitness::Neutral => (0., 0.),
73            Fitness::Fixed { s } => (*s, 0.),
74            Fitness::GammaSampled { shape, scale } => {
75                (shape * scale, f32::sqrt(shape * scale.powi(2)))
76            }
77        }
78    }
79}
80
81/// Id of the [`SubClone`]s.
82pub type CloneId = usize;
83
84#[derive(Debug, Clone)]
85/// A group of cells sharing the same genetic background with a specific
86/// proliferation rate.
87///
88/// The main loop of the simulation delegates the proliferation of cells to
89/// this structure, meaning that the `SubClone` will randomly pick one of its
90/// cells and make it proliferate.
91/// Upon proliferation, the cell can be assigned to a new clone with
92/// probability `p` (see [`crate::process::CellDivisionProbabilities`]).
93pub struct SubClone {
94    cells: Vec<StemCell>,
95    pub id: CloneId,
96}
97
98impl Iterator for SubClone {
99    type Item = StemCell;
100
101    fn next(&mut self) -> Option<Self::Item> {
102        self.cells.pop()
103    }
104}
105
106impl SubClone {
107    pub fn new(id: CloneId, cell_capacity: usize) -> SubClone {
108        SubClone {
109            cells: Vec::with_capacity(cell_capacity),
110            id,
111        }
112    }
113
114    pub fn with_capacity(id: usize, capacity: usize) -> SubClone {
115        SubClone {
116            cells: Vec::with_capacity(capacity),
117            id,
118        }
119    }
120
121    pub fn empty_with_id(id: usize) -> SubClone {
122        SubClone { cells: vec![], id }
123    }
124
125    pub fn get_mut_cells(&mut self) -> &mut [StemCell] {
126        &mut self.cells
127    }
128
129    pub fn get_cells(&self) -> &[StemCell] {
130        &self.cells
131    }
132
133    pub fn get_cells_subclones_idx(&self) -> Vec<(&StemCell, usize)> {
134        self.cells.iter().map(|cell| (cell, self.id)).collect()
135    }
136
137    pub fn is_empty(&self) -> bool {
138        self.get_cells().is_empty()
139    }
140
141    pub fn assign_cell(&mut self, cell: StemCell) {
142        self.cells.push(cell);
143    }
144
145    pub fn random_cell(&mut self, rng: &mut impl Rng) -> anyhow::Result<StemCell> {
146        ensure!(!self.cells.is_empty());
147        Ok(self
148            .cells
149            .swap_remove(rng.random_range(0..self.cells.len())))
150    }
151
152    pub fn cell_count(&self) -> u64 {
153        self.cells.len() as u64
154    }
155}
156
157pub fn next_clone(
158    subclones: &SubClones,
159    old_subclone_id: CloneId,
160    p: f64,
161    rng: &mut impl Rng,
162    verbosity: u8,
163) -> CloneId {
164    //! Returns a [`CloneId`] to which the stem cell must be assigned by
165    //! sampling a fit variant.
166    //!
167    //! This id can be either a new id if a new variant has been arisen, else
168    //! it will be the old clone.
169    // this will panic when u is not small and the cell hasn't divided much,
170    // i.e. when p > 1
171    if Bernoulli::new(p).unwrap().sample(rng) {
172        let mut rnd_clone_id = rng.random_range(0..MAX_SUBCLONES);
173        let mut counter = 0;
174        // the new random clone cannot have `subclone_id` id and must be empty
175        while (rnd_clone_id == old_subclone_id
176            || !subclones.get_clone(rnd_clone_id).unwrap().is_empty())
177            && counter <= MAX_SUBCLONES
178        {
179            rnd_clone_id = rng.random_range(0..MAX_SUBCLONES);
180            counter += 1;
181        }
182        assert!(counter <= MAX_SUBCLONES, "max number of clones reached");
183        if verbosity > 1 {
184            println!("new fit variant: assign cell to clone {rnd_clone_id}");
185        }
186        return rnd_clone_id;
187    }
188    if verbosity > 1 {
189        println!("no new fit variants with p {p}",);
190    }
191    old_subclone_id
192}
193
194/// A collection of subclones each having their proliferative advantage.
195/// This collection contains the neutral clone as well.
196#[derive(Clone, Debug)]
197pub struct SubClones([SubClone; MAX_SUBCLONES]);
198
199impl SubClones {
200    pub fn new(cells: Vec<StemCell>, capacity: usize, verbosity: u8) -> Self {
201        //! Returns all the newly initiated subclones by assigning all `cells`
202        //! to the neutral clone.
203        //!
204        //! All clones have their own `capacity`.
205        // initial state
206        let mut subclones: [SubClone; MAX_SUBCLONES] =
207            std::array::from_fn(|i| SubClone::new(i, capacity));
208        let tot_cells = cells.len();
209        if verbosity > 1 {
210            println!("assigning {tot_cells} cells to the wild-type clone");
211        }
212        for cell in cells {
213            subclones[0].assign_cell(cell);
214        }
215        assert_eq!(subclones[0].cells.len(), tot_cells);
216
217        Self(subclones)
218    }
219
220    pub fn new_empty() -> Self {
221        SubClones(std::array::from_fn(SubClone::empty_with_id))
222    }
223
224    pub fn with_capacity(capacity: usize) -> Self {
225        SubClones(std::array::from_fn(|id| {
226            SubClone::with_capacity(id, capacity)
227        }))
228    }
229
230    pub fn compute_tot_cells(&self) -> u64 {
231        self.0.iter().map(|subclone| subclone.cell_count()).sum()
232    }
233
234    pub(crate) fn get_clone(&self, id: usize) -> Option<&SubClone> {
235        self.0.get(id)
236    }
237
238    pub(crate) fn get_mut_clone_unchecked(&mut self, id: usize) -> &mut SubClone {
239        &mut self.0[id]
240    }
241
242    pub fn get_neutral_clone(&self) -> &SubClone {
243        &self.0[0]
244    }
245
246    pub fn array_of_gillespie_reactions(&self) -> [CloneId; MAX_SUBCLONES] {
247        core::array::from_fn(|i| i)
248    }
249
250    pub(crate) fn into_subsampled(self, nb_cells: NonZeroUsize, rng: &mut impl Rng) -> Self {
251        //! Take a subsample of the population that is stored in these subclones.
252        let mut subclones = Self::new_empty();
253
254        for (cell, clone_id) in self
255            .0
256            .into_iter()
257            .flat_map(|subclone| {
258                // for all subclones get all the cells with the its subclone_id.
259                // This will return all the population of cells, as a vec of
260                // tuple with the first entry being a stem cell and the second
261                // entry the id of the clone to which the cell belongs.
262                // Then subsample this population, by getting a `nb_cells`
263                // ammount of tuples (cell, id).
264                subclone
265                    .cells
266                    .into_iter()
267                    .map(|cell| (cell, subclone.id))
268                    .collect::<Vec<(StemCell, usize)>>()
269            })
270            .choose_multiple(rng, nb_cells.get())
271        {
272            subclones.0[clone_id].cells.push(cell);
273        }
274        subclones
275    }
276
277    pub fn get_mut_cells(&mut self) -> Vec<&mut StemCell> {
278        self.0
279            .iter_mut()
280            .flat_map(|subclone| subclone.get_mut_cells())
281            .collect()
282    }
283
284    pub fn get_cells(&self) -> Vec<&StemCell> {
285        self.0
286            .iter()
287            .flat_map(|subclone| subclone.get_cells())
288            .collect()
289    }
290
291    pub(crate) fn get_cells_with_clones_idx(&self) -> Vec<(&StemCell, usize)> {
292        self.0
293            .iter()
294            .flat_map(|subclone| subclone.get_cells_subclones_idx())
295            .collect()
296    }
297
298    pub(crate) fn get_cells_subsampled_with_clones_idx(
299        &self,
300        nb_cells: usize,
301        rng: &mut impl Rng,
302    ) -> Vec<(&StemCell, usize)> {
303        self.0
304            .iter()
305            .flat_map(|subclone| subclone.get_cells_subclones_idx())
306            .choose_multiple(rng, nb_cells)
307    }
308
309    pub fn gillespie_rates(
310        &self,
311        fitness: &Fitness,
312        b0: f32,
313        rng: &mut impl Rng,
314    ) -> ReactionRates<MAX_SUBCLONES> {
315        //! Create the Gillespie reaction rates according to the `fitness`
316        //! model with `b0` being the proliferative rate of the wild-type clone.
317        match fitness {
318            Fitness::Fixed { s } => ReactionRates(core::array::from_fn(|i| {
319                if i == 0 {
320                    b0
321                } else {
322                    b0 * (1. + s)
323                }
324            })),
325            &Fitness::GammaSampled { shape, scale } => {
326                let gamma = Gamma::new(shape, scale).unwrap();
327                ReactionRates(core::array::from_fn(|i| {
328                    if i == 0 {
329                        b0
330                    } else {
331                        b0 * (1. + gamma.sample(rng))
332                    }
333                }))
334            }
335            &Fitness::Neutral => ReactionRates(core::array::from_fn(|_| b0)),
336        }
337    }
338}
339
340impl Default for SubClones {
341    fn default() -> Self {
342        //! Create new subclones each having one cell (this creates also the cells).
343        let subclones = std::array::from_fn(|i| {
344            let mut subclone = SubClone::new(i, 2);
345            subclone.assign_cell(StemCell::new());
346            subclone
347        });
348        Self(subclones)
349    }
350}
351
352impl From<Vec<(StemCell, usize)>> for SubClones {
353    fn from(cells: Vec<(StemCell, usize)>) -> Self {
354        let mut subclones = SubClones::new_empty();
355        for (cell, id) in cells.into_iter() {
356            subclones.get_mut_clone_unchecked(id).assign_cell(cell);
357        }
358        subclones
359    }
360}
361
362/// Number of cells in subclones.
363pub struct Variants {}
364
365impl Variants {
366    pub fn variant_counts(subclones: &SubClones) -> [u64; MAX_SUBCLONES] {
367        //! The total variant count is the number of cells in all subclones.
368        //! ```
369        //! use hsc::MAX_SUBCLONES;
370        //! # use hsc::{stemcell::StemCell, process::{Moran}, subclone::Variants};
371        //! // create a process with one cell in each `MAX_SUBCLONES` subclones
372        //! let mut hsc = Moran::default();
373        //!
374        //! assert_eq!(Variants::variant_counts(&hsc.subclones), [1; MAX_SUBCLONES]);
375        //! ```
376        std::array::from_fn(|i| subclones.0[i].cell_count())
377    }
378
379    pub fn variant_fractions(subclones: &SubClones) -> Vec<f32> {
380        //! The proportion of cells in all subclones.
381        //!
382        //! ```
383        //! use hsc::MAX_SUBCLONES;
384        //! use hsc::subclone::SubClones;
385        //! # use hsc::{stemcell::StemCell, process::{Moran}, subclone::Variants};
386        //! // create a process with one cell in each `MAX_SUBCLONES` subclones
387        //! let mut hsc = Moran::default();
388        //!
389        //! assert_eq!(
390        //!     Variants::variant_fractions(&hsc.subclones),
391        //!     [1. / MAX_SUBCLONES as f32; MAX_SUBCLONES]
392        //! );
393        //!
394        //! let subclones = SubClones::new(vec![StemCell::new()], MAX_SUBCLONES, 0);
395        //! let mut variant_fraction = [0.; MAX_SUBCLONES];
396        //! variant_fraction[0] = 1.;
397        //!
398        //! assert_eq!(
399        //!     Variants::variant_fractions(&subclones),
400        //!     variant_fraction
401        //! );
402        //! ```
403        let tot_cells = subclones.compute_tot_cells();
404        Variants::variant_counts(subclones)
405            .into_iter()
406            .map(|frac| (frac as f32 / tot_cells as f32))
407            .collect()
408    }
409}
410
411pub fn save_variant_fraction(
412    subclones: &SubClones,
413    path2file: &Path,
414    verbosity: u8,
415) -> anyhow::Result<()> {
416    let path2file = path2file.with_extension("csv");
417    let total_variant_frac = Variants::variant_fractions(subclones);
418    if verbosity > 0 {
419        println!("total variant fraction in {path2file:#?}")
420    }
421    write2file(&total_variant_frac, &path2file, None, false)?;
422    Ok(())
423}
424
425pub fn proliferating_cell(
426    subclones: &mut SubClones,
427    subclone_id: CloneId,
428    verbosity: u8,
429    rng: &mut impl Rng,
430) -> StemCell {
431    //! Determine which cells will proliferate by randomly selecting a cell
432    //! from the subclone with id `subclone_id`.
433    //! This will also remove the random selected cell from the subclone,
434    //! hence subclones are borrowed mutably.
435    if verbosity > 2 {
436        println!(
437            "a cell from clone {:#?} will divide",
438            subclones.0[subclone_id]
439        );
440    }
441    subclones.0[subclone_id]
442        .random_cell(rng)
443        .with_context(|| "found empty subclone")
444        .unwrap()
445}
446
447#[cfg(test)]
448mod tests {
449    use std::num::NonZeroU8;
450
451    use crate::tests::LambdaFromNonZeroU8;
452
453    use super::*;
454    use quickcheck_macros::quickcheck;
455    use rand::{rngs::SmallRng, SeedableRng};
456
457    #[quickcheck]
458    fn assign_cell_test(id: usize) -> bool {
459        let mut neutral_clone = SubClone { cells: vec![], id };
460        let cell = StemCell::new();
461        assert!(neutral_clone.cells.is_empty());
462
463        neutral_clone.assign_cell(cell);
464        !neutral_clone.cells.is_empty()
465    }
466
467    #[quickcheck]
468    fn proliferating_cell_test(seed: u64, mut subclone_id: u8) -> bool {
469        let mut rng = SmallRng::seed_from_u64(seed);
470        let mut subclones = SubClones::default();
471        if subclone_id >= subclones.0.len() as u8 {
472            subclone_id = 1;
473        }
474        let number_of_cells = subclones.0[subclone_id as usize].get_cells().len();
475        let tot_cells = subclones.compute_tot_cells();
476
477        let _ = proliferating_cell(&mut subclones, subclone_id as usize, 1, &mut rng);
478
479        let subclones_have_lost_cell =
480            Variants::variant_counts(&subclones).iter().sum::<u64>() < tot_cells;
481        let subclone_has_lost_cell =
482            subclones.0[subclone_id as usize].get_cells().len() == number_of_cells - 1;
483        subclones_have_lost_cell && subclone_has_lost_cell
484    }
485
486    #[test]
487    fn variant_fraction() {
488        let subclones = SubClones::default();
489        assert_eq!(
490            Variants::variant_fractions(&subclones),
491            &[1. / MAX_SUBCLONES as f32; MAX_SUBCLONES]
492        );
493    }
494
495    #[quickcheck]
496    fn division_no_new_clone(seed: u64, cells_present: NonZeroU8) -> bool {
497        let mut rng = SmallRng::seed_from_u64(seed);
498        let mut cell = StemCell::new();
499        cell.set_last_division_time(1.1).unwrap();
500
501        let cells = vec![cell; cells_present.get() as usize];
502        let subclones = SubClones::new(cells, cells_present.get() as usize + 1, 0);
503
504        let old_id = 0;
505        let id = next_clone(&subclones, old_id, 0., &mut rng, 0);
506        id == old_id
507    }
508
509    #[quickcheck]
510    fn division_new_clone(seed: u64, cells_present: NonZeroU8) -> bool {
511        let mut rng = SmallRng::seed_from_u64(seed);
512        let mut cell = StemCell::new();
513        cell.set_last_division_time(1.1).unwrap();
514
515        let cells = vec![cell; cells_present.get() as usize];
516        let before_assignment = cells.len();
517        let subclones = SubClones::new(cells, cells_present.get() as usize + 1, 0);
518
519        next_clone(&subclones, 0, 1., &mut rng, 0);
520        subclones.0[0].cell_count() as usize == before_assignment
521    }
522
523    #[test]
524    #[should_panic]
525    fn assign_all_clones_occupied() {
526        let mut rng = SmallRng::seed_from_u64(26);
527
528        let subclones = SubClones::default();
529        let mut cell = StemCell::new();
530        cell.set_last_division_time(1.).unwrap();
531
532        next_clone(&subclones, 0, 1., &mut rng, 0);
533    }
534
535    #[quickcheck]
536    fn new_distribution_test(
537        lambda_division: LambdaFromNonZeroU8,
538        lambda_background: LambdaFromNonZeroU8,
539    ) -> bool {
540        let probs = Probs::new(lambda_background.0, lambda_division.0, 0.01, 0., 10, 0);
541        let distrs = Distributions::new(probs, 0);
542        distrs.neutral_poisson.eq(&NeutralMutationPoisson::new(
543            lambda_division.0,
544            lambda_background.0,
545        )
546        .unwrap())
547    }
548}