winsfs_core/sfs/
generics.rs

1//! Types for parameterising an SFS.
2//!
3//! This module contains the traits and structs used for the generics on the
4//! [`SfsBase`](super::SfsBase) struct. In most situations, it should be possible to use the
5//! provided type aliases in the parent module to avoid having to interact with these types
6//! directly.
7
8use std::{fmt, slice};
9
10/// A type that can act as a shape for an SFS.
11///
12/// This is simply a trait alias for things like arrays and slices of usize.
13/// Users should not need to implement this trait or care about its internals.
14pub trait Shape: AsRef<[usize]> + Clone + fmt::Debug + Eq + PartialEq {
15    /// Returns the strides of the shape.
16    fn strides(&self) -> Self;
17
18    /// Returns an iterator of the values of the shape.
19    fn iter(&self) -> slice::Iter<'_, usize> {
20        self.as_ref().iter()
21    }
22
23    /// Returns the length of the shape.
24    fn len(&self) -> usize {
25        self.as_ref().len()
26    }
27
28    /// Returns `true` if the shape is empty, `false` otherwise.
29    fn is_empty(&self) -> bool {
30        // Mainly to appease clippy
31        self.as_ref().is_empty()
32    }
33}
34
35/// An SFS shape that is known at compile-time.
36pub type ConstShape<const D: usize> = [usize; D];
37
38impl<const D: usize> Shape for ConstShape<D> {
39    fn strides(&self) -> Self {
40        let mut strides = [1; D];
41        compute_strides(&self[..], &mut strides);
42        strides
43    }
44
45    fn len(&self) -> usize {
46        D
47    }
48}
49
50/// An SFS shape that is known at run-time.
51pub type DynShape = Box<[usize]>;
52
53impl Shape for DynShape {
54    fn strides(&self) -> Self {
55        let mut strides = vec![1; self.len()];
56        compute_strides(self, &mut strides);
57        strides.into_boxed_slice()
58    }
59}
60
61/// A marker trait for SFS normalisation.
62///
63/// An SFS can either be in a normalised or unnormalised state, i.e. the values of the SFS can
64/// either be in probability space or not. When not in probability space, the values are
65/// typically in count space, in which case they sum to the number of sites in the input from
66/// which the SFS was constructed.
67pub trait Normalisation {
68    /// A boolean indicating whether the spectrum is normalised.
69    const NORM: bool;
70}
71
72/// A marker for a normalised SFS.
73#[derive(Debug, Clone, Copy, Eq, PartialEq)]
74pub struct Norm {}
75impl Normalisation for Norm {
76    const NORM: bool = true;
77}
78
79/// A marker for a (potentially) unnormalised SFS.
80#[derive(Debug, Clone, Copy, Eq, PartialEq)]
81pub struct Unnorm {}
82impl Normalisation for Unnorm {
83    const NORM: bool = false;
84}
85
86/// Compute strides of `shape` into 1-initialised `strides`
87fn compute_strides(shape: &[usize], strides: &mut [usize]) {
88    debug_assert_eq!(shape.len(), strides.len());
89
90    for (i, v) in shape.iter().enumerate().skip(1).rev() {
91        strides.iter_mut().take(i).for_each(|stride| *stride *= v)
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn test_array_strides() {
101        assert_eq!([7].strides(), [1]);
102        assert_eq!([9, 3].strides(), [3, 1]);
103        assert_eq!([3, 7, 5].strides(), [35, 5, 1]);
104        assert_eq!([9, 3, 5, 7].strides(), [105, 35, 7, 1]);
105    }
106}