sfs_core/spectrum/
count.rs

1use std::ops::{Deref, Index, IndexMut};
2
3use crate::array::Shape;
4
5/// An allele count.
6///
7/// This corresponds to an index in a [`Spectrum`](crate::Spectrum).
8#[derive(Clone, Debug, Eq, Hash, PartialEq)]
9pub struct Count(pub Vec<usize>);
10
11impl Count {
12    /// The number of dimensions of the count.
13    pub fn dimensions(&self) -> usize {
14        self.0.len()
15    }
16
17    pub(crate) fn try_from_shape(shape: Shape) -> Option<Self> {
18        let mut vec = shape.0;
19        for x in vec.iter_mut() {
20            *x = x.checked_sub(1)?;
21        }
22        Some(Self(vec))
23    }
24
25    /// Creates a new count from zeros.
26    pub fn from_zeros(dimensions: usize) -> Self {
27        Self(vec![0; dimensions])
28    }
29
30    pub(crate) fn into_shape(self) -> Shape {
31        let mut vec = self.0;
32        vec.iter_mut().for_each(|x| *x += 1);
33        Shape(vec)
34    }
35
36    /// Set all elements to zero.
37    pub fn set_zero(&mut self) {
38        self.0.iter_mut().for_each(|x| *x = 0);
39    }
40}
41
42impl AsRef<[usize]> for Count {
43    fn as_ref(&self) -> &[usize] {
44        self
45    }
46}
47
48impl Deref for Count {
49    type Target = [usize];
50
51    fn deref(&self) -> &Self::Target {
52        &self.0
53    }
54}
55
56impl From<Vec<usize>> for Count {
57    fn from(shape: Vec<usize>) -> Self {
58        Self(shape)
59    }
60}
61
62impl<const N: usize> From<[usize; N]> for Count {
63    fn from(shape: [usize; N]) -> Self {
64        Self(shape.to_vec())
65    }
66}
67
68impl From<usize> for Count {
69    fn from(shape: usize) -> Self {
70        Self(vec![shape])
71    }
72}
73
74impl Index<usize> for Count {
75    type Output = usize;
76
77    fn index(&self, index: usize) -> &Self::Output {
78        self.0.index(index)
79    }
80}
81
82impl IndexMut<usize> for Count {
83    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
84        self.0.index_mut(index)
85    }
86}
sfs_core/spectrum/count.rs

sfs_core/spectrum/
count.rs
