winsfs_core/saf/
site.rs

1use super::{Lifetime, ShapeError};
2
3/// A type that can be cheaply converted to a SAF site.
4///
5/// This is akin to GATified [`AsRef`] for SAF sites.
6pub trait AsSiteView<const N: usize>: for<'a> Lifetime<'a, Item = SiteView<'a, N>> {
7    /// Returns a SAF site view of `self`.
8    fn as_site_view(&self) -> <Self as Lifetime<'_>>::Item;
9}
10
11macro_rules! impl_shared_site_methods {
12    () => {
13        /// Returns the values of the site as a flat slice.
14        ///
15        /// See the [`Site`] documentation for details on the storage order.
16        pub fn as_slice(&self) -> &[f32] {
17            &self.values
18        }
19
20        /// Returns an iterator over all values in the site.
21        ///
22        /// See the [`Site`] documentation for details on the storage order.
23        pub fn iter(&self) -> ::std::slice::Iter<f32> {
24            self.values.iter()
25        }
26
27        /// Returns the shape of the site.
28        #[inline]
29        pub fn shape(&self) -> [usize; N] {
30            self.shape
31        }
32
33        /// Returns an array of slices corresponding to the sites in each population.
34        #[inline]
35        pub fn split(&self) -> [&[f32]; N] {
36            let mut buf = &self.values[..];
37
38            self.shape.map(|i| {
39                let (hd, tl) = buf.split_at(i);
40                buf = tl;
41                hd
42            })
43        }
44    };
45}
46
47/// A single site of SAF likelihoods for `N` populations.
48///
49/// Internally, the site is stored as a contiguous block of memory, with all values from the first
50/// population first, then the second, and so on. [`Site::shape`] gives the number of values for
51/// each population.
52#[derive(Clone, Debug, PartialEq)]
53pub struct Site<const N: usize> {
54    values: Vec<f32>,
55    shape: [usize; N],
56}
57
58impl<const N: usize> Site<N> {
59    /// Returns a mutable reference to the values of the SAF site as a flat slice.
60    ///
61    /// See the [`Site`] documentation for details on the storage order.
62    pub fn as_mut_slice(&mut self) -> &mut [f32] {
63        &mut self.values
64    }
65
66    /// Returns an iterator over all values in the site.
67    ///
68    /// See the [`Site`] documentation for details on the storage order.
69    pub fn iter_mut(&mut self) -> ::std::slice::IterMut<f32> {
70        self.values.iter_mut()
71    }
72
73    /// Returns a new SAF site.
74    ///
75    /// The number of provided values must be equal to the sum of shapes.
76    /// See the [`Site`] documentation for details on the storage order.
77    ///
78    /// # Examples
79    ///
80    /// ```
81    /// use winsfs_core::saf::Site;
82    /// let vec = vec![0.0, 0.1, 0.2, 1.0, 1.2];
83    /// let shape = [3, 2];
84    /// let site = Site::new(vec, shape).unwrap();
85    /// assert_eq!(site.split(), [&[0.0, 0.1, 0.2][..], &[1.0, 1.2][..]]);
86    /// ```
87    ///
88    /// A [`ShapeError`] is thrown if the shape does not fit the number of values:
89    ///
90    /// ```
91    /// use winsfs_core::saf::Site;
92    /// let vec = vec![0.0, 0.1, 0.2, 1.0, 1.2];
93    /// let wrong_shape = [6, 2];
94    /// assert!(Site::new(vec, wrong_shape).is_err());
95    /// ```
96    pub fn new(values: Vec<f32>, shape: [usize; N]) -> Result<Self, ShapeError<N>> {
97        let len = values.len();
98        let width: usize = shape.iter().sum();
99
100        if len == width {
101            Ok(Self::new_unchecked(values, shape))
102        } else {
103            Err(ShapeError { len, shape })
104        }
105    }
106
107    /// Returns a new SAF site without checking that the shape fits the number of values.
108    pub(crate) fn new_unchecked(values: Vec<f32>, shape: [usize; N]) -> Self {
109        Self { values, shape }
110    }
111
112    /// Returns a view of the entire site.
113    #[inline]
114    pub fn view(&self) -> SiteView<N> {
115        SiteView::new_unchecked(self.values.as_slice(), self.shape)
116    }
117
118    impl_shared_site_methods! {}
119}
120
121impl<'a, const N: usize> Lifetime<'a> for Site<N> {
122    type Item = SiteView<'a, N>;
123}
124
125impl<const N: usize> AsSiteView<N> for Site<N> {
126    #[inline]
127    fn as_site_view(&self) -> <Self as Lifetime<'_>>::Item {
128        self.view()
129    }
130}
131
132impl<'a, 'b, const N: usize> Lifetime<'a> for &'b Site<N> {
133    type Item = SiteView<'a, N>;
134}
135
136impl<'a, const N: usize> AsSiteView<N> for &'a Site<N> {
137    #[inline]
138    fn as_site_view(&self) -> <Self as Lifetime<'_>>::Item {
139        self.view()
140    }
141}
142
143/// A view of a single site of SAF likelihoods for `N` populations.
144#[derive(Clone, Copy, Debug, PartialEq)]
145pub struct SiteView<'a, const N: usize> {
146    values: &'a [f32],
147    shape: [usize; N],
148}
149
150impl<'a, const N: usize> SiteView<'a, N> {
151    /// Returns a new SAF site view;
152    ///
153    /// The number of provided values must be equal to the sum of shapes.
154    /// See the [`Site`] documentation for details on the storage order.
155    ///
156    /// # Examples
157    ///
158    /// ```
159    /// use winsfs_core::saf::SiteView;
160    /// let slice = &[0.0, 0.1, 0.2, 1.0, 1.2];
161    /// let shape = [3, 2];
162    /// let site = SiteView::new(slice, shape).unwrap();
163    /// assert_eq!(site.split(), [&[0.0, 0.1, 0.2][..], &[1.0, 1.2][..]]);
164    /// ```
165    ///
166    /// A [`ShapeError`] is thrown if the shape does not fit the number of values:
167    ///
168    /// ```
169    /// use winsfs_core::saf::SiteView;
170    /// let slice = &[0.0, 0.1, 0.2, 1.0, 1.2];
171    /// let wrong_shape = [6, 2];
172    /// assert!(SiteView::new(slice, wrong_shape).is_err());
173    /// ```
174    pub fn new(values: &'a [f32], shape: [usize; N]) -> Result<Self, ShapeError<N>> {
175        let len = values.len();
176        let width: usize = shape.iter().sum();
177
178        if len == width {
179            Ok(Self::new_unchecked(values, shape))
180        } else {
181            Err(ShapeError { len, shape })
182        }
183    }
184
185    /// Returns a new SAF site view without checking that the shape fits the number of values.
186    pub(crate) fn new_unchecked(values: &'a [f32], shape: [usize; N]) -> Self {
187        Self { values, shape }
188    }
189
190    impl_shared_site_methods! {}
191}
192
193impl<'a, 'b, const N: usize> Lifetime<'a> for SiteView<'b, N> {
194    type Item = SiteView<'a, N>;
195}
196
197impl<'a, const N: usize> AsSiteView<N> for SiteView<'a, N> {
198    #[inline]
199    fn as_site_view(&self) -> <Self as Lifetime<'_>>::Item {
200        *self
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    #[test]
209    fn test_split_1d() {
210        let vec = vec![0., 1., 2.];
211        let site = Site::new(vec, [3]).unwrap();
212        assert_eq!(site.split(), [&[0., 1., 2.][..]]);
213        assert_eq!(site.split(), site.view().split());
214    }
215
216    #[test]
217    fn test_split_2d() {
218        let vec = vec![0., 1., 10., 11., 12., 13.];
219        let site = Site::new(vec, [2, 4]).unwrap();
220        assert_eq!(site.split(), [&[0., 1.][..], &[10., 11., 12., 13.][..]]);
221        assert_eq!(site.split(), site.view().split());
222    }
223
224    #[test]
225    fn test_split_3d() {
226        let vec = vec![0., 10., 11., 12., 20., 21.];
227        let site = Site::new(vec, [1, 3, 2]).unwrap();
228        assert_eq!(
229            site.split(),
230            [&[0.,][..], &[10., 11., 12.][..], &[20., 21.][..]]
231        );
232        assert_eq!(site.split(), site.view().split());
233    }
234}