sfs_core/spectrum/
stat.rs

1use std::fmt;
2
3pub mod theta;
4pub use theta::Theta;
5
6pub mod d;
7pub use d::D;
8
9use super::{Sfs, Shape, Spectrum, State};
10
11pub type Pi = Theta<theta::Tajima>;
12
13#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
14pub struct PiXY(pub f64);
15
16impl PiXY {
17    pub fn from_spectrum<S: State>(spectrum: &Spectrum<S>) -> Result<Self, DimensionError> {
18        if spectrum.dimensions() == 2 {
19            Ok(Self::from_spectrum_unchecked(spectrum))
20        } else {
21            Err(DimensionError {
22                expected: 2,
23                actual: spectrum.dimensions(),
24            })
25        }
26    }
27
28    fn from_spectrum_unchecked<S: State>(spectrum: &Spectrum<S>) -> Self {
29        let (n1, n2) = if let &[n1, n2] = spectrum.shape().as_ref() {
30            (n1 - 1, n2 - 1)
31        } else {
32            panic!("dimensions do not fit");
33        };
34
35        let num = (0..=n1)
36            .flat_map(|m1| (0..=n2).map(move |m2| (m1, m2)))
37            .take(spectrum.elements() - 1)
38            .skip(1)
39            .map(|(m1, m2)| {
40                let p1 = m1 * (n2 - m2);
41                let p2 = m2 * (n1 - m1);
42                spectrum[[m1, m2]] * ((p1 + p2) as f64)
43            })
44            .sum::<f64>();
45
46        let denom = (n1 * n2) as f64;
47
48        Self(num / denom)
49    }
50}
51
52#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
53pub struct F2(pub f64);
54
55impl F2 {
56    pub fn from_sfs(sfs: &Sfs) -> Result<Self, DimensionError> {
57        if sfs.dimensions() == 2 {
58            Ok(Self::from_sfs_unchecked(sfs))
59        } else {
60            Err(DimensionError {
61                expected: 2,
62                actual: sfs.dimensions(),
63            })
64        }
65    }
66
67    fn from_sfs_unchecked(sfs: &Sfs) -> Self {
68        Self(
69            sfs.array
70                .iter()
71                .zip(sfs.iter_frequencies())
72                .map(|(v, fs)| v * (fs[0] - fs[1]).powi(2))
73                .sum(),
74        )
75    }
76}
77
78#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
79pub struct F3(pub f64);
80
81impl F3 {
82    pub fn from_sfs(sfs: &Sfs) -> Result<Self, DimensionError> {
83        if sfs.dimensions() == 3 {
84            Ok(Self::from_sfs_unchecked(sfs))
85        } else {
86            Err(DimensionError {
87                expected: 3,
88                actual: sfs.dimensions(),
89            })
90        }
91    }
92
93    fn from_sfs_unchecked(sfs: &Sfs) -> Self {
94        Self(
95            sfs.array
96                .iter()
97                .zip(sfs.iter_frequencies())
98                .map(|(v, fs)| v * (fs[0] - fs[1]) * (fs[0] - fs[2]))
99                .sum(),
100        )
101    }
102}
103
104#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
105pub struct F4(pub f64);
106
107impl F4 {
108    pub fn from_sfs(sfs: &Sfs) -> Result<Self, DimensionError> {
109        if sfs.dimensions() == 4 {
110            Ok(Self::from_sfs_unchecked(sfs))
111        } else {
112            Err(DimensionError {
113                expected: 4,
114                actual: sfs.dimensions(),
115            })
116        }
117    }
118
119    fn from_sfs_unchecked(sfs: &Sfs) -> Self {
120        Self(
121            sfs.array
122                .iter()
123                .zip(sfs.iter_frequencies())
124                .map(|(v, fs)| v * (fs[0] - fs[1]) * (fs[2] - fs[3]))
125                .sum(),
126        )
127    }
128}
129
130#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
131pub struct Fst(pub f64);
132
133impl Fst {
134    pub fn from_sfs(sfs: &Sfs) -> Result<Self, DimensionError> {
135        if sfs.dimensions() == 2 {
136            Ok(Self::from_sfs_unchecked(sfs))
137        } else {
138            Err(DimensionError {
139                expected: 2,
140                actual: sfs.dimensions(),
141            })
142        }
143    }
144
145    fn from_sfs_unchecked(sfs: &Sfs) -> Self {
146        // We only want the polymorphic parts of the spectrum and corresponding frequencies,
147        // so we drop the first and last values
148        let polymorphic_iter = sfs
149            .array
150            .iter()
151            .zip(sfs.iter_frequencies())
152            .take(sfs.elements() - 1)
153            .skip(1);
154
155        let shape = sfs.shape();
156        let n_i_sub = (shape[0] - 2) as f64;
157        let n_j_sub = (shape[1] - 2) as f64;
158
159        let (num, denom) = polymorphic_iter
160            .map(|(v, fs)| {
161                let f_i = fs[0];
162                let f_j = fs[1];
163                let g_i = 1. - f_i;
164                let g_j = 1. - f_j;
165
166                let num = (f_i - f_j).powi(2) - f_i * g_i / n_i_sub - f_j * g_j / n_j_sub;
167                let denom = f_i * g_j + f_j * g_i;
168                (v * num, v * denom)
169            })
170            .fold((0., 0.), |(n_sum, d_sum), (n, d)| (n_sum + n, d_sum + d));
171
172        Self(num / denom)
173    }
174}
175
176#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
177pub struct King(pub f64);
178
179impl King {
180    pub fn from_spectrum<S: State>(spectrum: &Spectrum<S>) -> Result<Self, ShapeError> {
181        if spectrum.shape().0 == [3, 3] {
182            Ok(Self::from_spectrum_unchecked(spectrum))
183        } else {
184            Err(ShapeError {
185                expected: Shape(vec![3, 3]),
186                actual: spectrum.shape().clone(),
187            })
188        }
189    }
190
191    fn from_spectrum_unchecked<S: State>(spectrum: &Spectrum<S>) -> Self {
192        let s = spectrum;
193
194        let numer = s[[1, 1]] - 2. * (s[[0, 2]] + s[[2, 0]]);
195        let denom = s[[0, 1]] + s[[1, 0]] + 2. * s[[1, 1]] + s[[1, 2]] + s[[2, 1]];
196
197        Self(numer / denom)
198    }
199}
200
201#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
202pub struct R0(pub f64);
203
204impl R0 {
205    pub fn from_spectrum<S: State>(spectrum: &Spectrum<S>) -> Result<Self, ShapeError> {
206        if spectrum.shape().0 == [3, 3] {
207            Ok(Self::from_spectrum_unchecked(spectrum))
208        } else {
209            Err(ShapeError {
210                expected: Shape(vec![3, 3]),
211                actual: spectrum.shape().clone(),
212            })
213        }
214    }
215
216    fn from_spectrum_unchecked<S: State>(spectrum: &Spectrum<S>) -> Self {
217        let s = spectrum;
218
219        Self((s[[0, 2]] + s[[2, 0]]) / s[[1, 1]])
220    }
221}
222
223#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
224pub struct R1(pub f64);
225
226impl R1 {
227    pub fn from_spectrum<S: State>(spectrum: &Spectrum<S>) -> Result<Self, ShapeError> {
228        if spectrum.shape().0 == [3, 3] {
229            Ok(Self::from_spectrum_unchecked(spectrum))
230        } else {
231            Err(ShapeError {
232                expected: Shape(vec![3, 3]),
233                actual: spectrum.shape().clone(),
234            })
235        }
236    }
237
238    fn from_spectrum_unchecked<S: State>(spectrum: &Spectrum<S>) -> Self {
239        let denom = [[0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1]]
240            .iter()
241            .map(|&i| spectrum[i])
242            .sum::<f64>();
243        Self(spectrum[[1, 1]] / denom)
244    }
245}
246
247/// An error associated with the calculation of a statistic from a spectrum.
248#[derive(Debug)]
249pub enum StatisticError {
250    /// The statistic is not defined for an array of the provided dimensionality.
251    DimensionError(DimensionError),
252    /// The statistic is not defined for an array of the provided shape.
253    ShapeError(ShapeError),
254}
255
256impl fmt::Display for StatisticError {
257    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
258        match self {
259            StatisticError::DimensionError(e) => write!(f, "{e}"),
260            StatisticError::ShapeError(e) => write!(f, "{e}"),
261        }
262    }
263}
264
265impl std::error::Error for StatisticError {}
266
267impl From<ShapeError> for StatisticError {
268    fn from(e: ShapeError) -> Self {
269        Self::ShapeError(e)
270    }
271}
272
273impl From<DimensionError> for StatisticError {
274    fn from(e: DimensionError) -> Self {
275        Self::DimensionError(e)
276    }
277}
278
279#[derive(Debug)]
280pub struct DimensionError {
281    expected: usize,
282    actual: usize,
283}
284
285impl fmt::Display for DimensionError {
286    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
287        let &DimensionError { expected, actual } = self;
288        write!(
289            f,
290            "expected SFS with dimension {expected}, found SFS with dimension {actual}"
291        )
292    }
293}
294
295impl std::error::Error for DimensionError {}
296
297#[derive(Debug)]
298pub struct ShapeError {
299    expected: Shape,
300    actual: Shape,
301}
302
303impl fmt::Display for ShapeError {
304    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
305        let &ShapeError { expected, actual } = &self;
306        write!(
307            f,
308            "expected SFS with shape {expected}, found SFS with shape {actual}"
309        )
310    }
311}
312
313impl std::error::Error for ShapeError {}
314
315#[cfg(test)]
316mod tests {
317    use crate::Scs;
318
319    #[test]
320    fn test_pi_xy() {
321        let data = vec![
322            0.983094, 0.000819, 0.000298, 0.001926, 0.000439, 0.000266, 0.000753, 0.000380,
323            0.000330, 0.000256, 0.000217, 0.000398, 0.000102, 0.000166, 0.010558,
324        ];
325
326        let sfs = Scs::new(data, [5, 3]).unwrap();
327
328        assert_approx_eq!(sfs.pi_xy().unwrap(), 0.002925);
329    }
330}
sfs_core/spectrum/stat.rs

sfs_core/spectrum/
stat.rs
