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#[derive(Debug, Default, Clone)]
14pub struct Distributions {
15 pub u: f32,
17 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#[derive(Clone, Copy, Debug)]
50pub enum Fitness {
51 Neutral,
55 Fixed { s: f32 },
61 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
81pub type CloneId = usize;
83
84#[derive(Debug, Clone)]
85pub 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 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 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#[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 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 let mut subclones = Self::new_empty();
253
254 for (cell, clone_id) in self
255 .0
256 .into_iter()
257 .flat_map(|subclone| {
258 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 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 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
362pub struct Variants {}
364
365impl Variants {
366 pub fn variant_counts(subclones: &SubClones) -> [u64; MAX_SUBCLONES] {
367 std::array::from_fn(|i| subclones.0[i].cell_count())
377 }
378
379 pub fn variant_fractions(subclones: &SubClones) -> Vec<f32> {
380 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 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}