winsfs_core/saf.rs
1//! Site allelele frequency ("SAF") likelihoods.
2//!
3//! SAF likelihoods represent a generalisation of genotype likelihoods from individuals to
4//! populations. To estimate the N-dimensional SFS, we need the SAF likelihoods from intersecting
5//! sites for those N populations. We represent those by the owned type [`Saf`] or the borrowed type
6//! [`SafView`], which may represent the full data or only a smaller block of sites.
7
8use std::{cmp::Ordering, error::Error, fmt, io};
9
10use angsd_saf as saf;
11
12use rand::Rng;
13
14use crate::ArrayExt;
15
16pub mod iter;
17use iter::{BlockIter, ParBlockIter, ParSiteIter, SiteIter};
18
19mod site;
20pub use site::{AsSiteView, Site, SiteView};
21
22// This exists to ensure that the lifetime bound on `Lifetime` cannot be changed by an external
23// implementor.
24// See: sabrinajewson.org/blog/the-better-alternative-to-lifetime-gats#the-better-gats
25mod sealed {
26 pub trait Sealed: Sized {}
27 pub struct Bounds<T>(T);
28 impl<T> Sealed for Bounds<T> {}
29}
30use sealed::{Bounds, Sealed};
31
32/// Stable workaround for lifetime GATs.
33///
34/// See the [GAT tracking issue][gat_tracking_issue] and [stabilisation PR][gat_stabilisation],
35/// and in particular [this blog post][sabrina_jewson] for details on this workaround.
36///
37/// [gat_tracking_issue]: https://github.com/rust-lang/rust/issues/44265
38/// [gat_stabilisation]: https://github.com/rust-lang/rust/pull/96709
39/// [sabrina_jewson]: https://sabrinajewson.org/blog/the-better-alternative-to-lifetime-gats
40pub trait Lifetime<'a, SELF: Sealed = Bounds<&'a Self>> {
41 // TODO: Replace with GAT when stable.
42
43 /// The inner item, the lifetime of which should be tied to `Self`.
44 type Item;
45}
46
47/// A type that can be cheaply converted to a SAF view.
48///
49/// This is akin to GATified [`AsRef`] for SAF views.
50pub trait AsSafView<const N: usize>: for<'a> Lifetime<'a, Item = SafView<'a, N>> {
51 /// Returns a SAF view of `self`.
52 fn as_saf_view(&self) -> <Self as Lifetime<'_>>::Item;
53}
54
55/// Creates a SAF matrix of a single population.
56///
57/// This is mainly intended for readability in doc-tests.
58///
59/// # Examples
60///
61/// ```
62/// use winsfs_core::saf1d;
63/// let saf = saf1d![
64/// [0.0, 0.1, 0.2],
65/// [0.3, 0.4, 0.5],
66/// [0.6, 0.7, 0.8],
67/// [0.9, 0.10, 0.11],
68/// [0.12, 0.13, 0.14],
69/// ];
70/// assert_eq!(saf.sites(), 5);
71/// assert_eq!(saf.shape(), [3]);
72/// assert_eq!(saf.get_site(0).as_slice(), &[0.0, 0.1, 0.2]);
73/// assert_eq!(saf.get_site(2).as_slice(), &[0.6, 0.7, 0.8]);
74/// ```
75#[macro_export]
76macro_rules! saf1d {
77 ($([$($x:literal),+ $(,)?]),+ $(,)?) => {{
78 let (shape, vec) = $crate::matrix!($([$($x),+]),+);
79 $crate::saf::Saf::new(vec, [shape[0]]).unwrap()
80 }};
81}
82
83/// Creates a joint SAF matrix of two populations.
84///
85/// This is mainly intended for readability in doc-tests.
86///
87/// # Examples
88///
89/// ```
90/// use winsfs_core::saf2d;
91/// let saf = saf2d![
92/// [0.0, 0.1, 0.2 ; 1.0, 1.1],
93/// [0.3, 0.4, 0.5 ; 1.2, 1.3],
94/// [0.6, 0.7, 0.8 ; 1.4, 1.5],
95/// [0.9, 0.10, 0.11 ; 1.6, 1.7],
96/// [0.12, 0.13, 0.14 ; 1.8, 1.9],
97/// ];
98/// assert_eq!(saf.sites(), 5);
99/// assert_eq!(saf.shape(), [3, 2]);
100/// assert_eq!(saf.get_site(0).as_slice(), &[0.0, 0.1, 0.2, 1.0, 1.1]);
101/// assert_eq!(saf.get_site(2).as_slice(), &[0.6, 0.7, 0.8, 1.4, 1.5]);
102/// ```
103#[macro_export]
104macro_rules! saf2d {
105 ($([$($x:literal),+ $(,)?; $($y:literal),+ $(,)?]),+ $(,)?) => {{
106 let x_cols = vec![$($crate::matrix!(count: $($x),+)),+];
107 let y_cols = vec![$($crate::matrix!(count: $($y),+)),+];
108 for cols in [&x_cols, &y_cols] {
109 assert!(cols.windows(2).all(|w| w[0] == w[1]));
110 }
111 let vec = vec![$($($x),+, $($y),+),+];
112 $crate::saf::Saf::new(vec, [x_cols[0], y_cols[0]]).unwrap()
113 }};
114}
115
116macro_rules! impl_shared_saf_methods {
117 () => {
118 /// Returns the values of the SAF as a flat slice.
119 ///
120 /// See the [`Saf`] documentation for details on the storage order.
121 pub fn as_slice(&self) -> &[f32] {
122 &self.values
123 }
124
125 /// Returns an iterator over all values in the SAF.
126 ///
127 /// See the [`Saf`] documentation for details on the storage order.
128 pub fn iter(&self) -> ::std::slice::Iter<f32> {
129 self.values.iter()
130 }
131
132 /// Returns a single site in the SAF.
133 pub fn get_site(&self, index: usize) -> SiteView<N> {
134 let width = self.width();
135
136 SiteView::new_unchecked(&self.values[index * width..][..width], self.shape)
137 }
138
139 /// Returns the number of sites in the SAF.
140 #[inline]
141 pub fn sites(&self) -> usize {
142 self.values.len() / self.width()
143 }
144
145 /// Returns the shape of the SAF.
146 #[inline]
147 pub fn shape(&self) -> [usize; N] {
148 self.shape
149 }
150
151 /// Returns the width of the SAF, i.e. the number of elements per site.
152 #[inline]
153 pub(self) fn width(&self) -> usize {
154 self.shape.iter().sum()
155 }
156 };
157}
158
159/// Joint SAF likelihood matrix for `N` populations.
160///
161/// Internally, the matrix is represented with each site in continuous memory.
162/// That is, the first values are those from the first site of all populations,
163/// then comes the next site, and so on. [`Saf::shape`] gives the number of values
164/// per site per population. This should only be important when operating directly
165/// on the underlying storage, e.g. using [`Saf::as_slice`] or [`Saf::as_mut_slice`].
166#[derive(Clone, Debug, PartialEq)]
167pub struct Saf<const N: usize> {
168 values: Vec<f32>,
169 shape: [usize; N],
170}
171
172impl<const N: usize> Saf<N> {
173 /// Returns a mutable reference to the values of the SAF as a flat slice.
174 ///
175 /// See the [`Saf`] documentation for details on the storage order.
176 pub fn as_mut_slice(&mut self) -> &mut [f32] {
177 &mut self.values
178 }
179
180 /// Returns an iterator over mutable references to all values in the SAF.
181 ///
182 /// See the [`Saf`] documentation for details on the storage order.
183 pub fn iter_mut(&mut self) -> ::std::slice::IterMut<f32> {
184 self.values.iter_mut()
185 }
186
187 /// Returns an iterator over blocks of sites in the SAF.
188 ///
189 /// If the number of sites in the SAF is not evenly divided by `block_size`,
190 /// the last block will be smaller than the others.
191 ///
192 /// # Examples
193 ///
194 /// ```
195 /// use winsfs_core::saf1d;
196 /// let saf = saf1d![
197 /// [0.0, 0.1, 0.2],
198 /// [0.3, 0.4, 0.5],
199 /// [0.6, 0.7, 0.8],
200 /// [0.9, 0.10, 0.11],
201 /// [0.12, 0.13, 0.14],
202 /// ];
203 /// let mut iter = saf.iter_blocks(2);
204 /// assert_eq!(
205 /// iter.next().unwrap(),
206 /// saf1d![[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]].view()
207 /// );
208 /// assert_eq!(
209 /// iter.next().unwrap(),
210 /// saf1d![[0.6, 0.7, 0.8], [0.9, 0.10, 0.11]].view()
211 /// );
212 /// assert_eq!(iter.next().unwrap(), saf1d![[0.12, 0.13, 0.14]].view());
213 /// assert!(iter.next().is_none());
214 /// ```
215 pub fn iter_blocks(&self, block_size: usize) -> BlockIter<N> {
216 BlockIter::new(self.view(), block_size)
217 }
218
219 /// Returns an iterator over the sites in the SAF.
220 ///
221 /// # Examples
222 ///
223 /// ```
224 /// use winsfs_core::saf1d;
225 /// let saf = saf1d![
226 /// [0.0, 0.1, 0.2],
227 /// [0.3, 0.4, 0.5],
228 /// [0.6, 0.7, 0.8],
229 /// ];
230 /// let mut iter = saf.iter_sites();
231 /// assert_eq!(iter.next().unwrap().as_slice(), [0.0, 0.1, 0.2]);
232 /// assert_eq!(iter.next().unwrap().as_slice(), [0.3, 0.4, 0.5]);
233 /// assert_eq!(iter.next().unwrap().as_slice(), [0.6, 0.7, 0.8]);
234 /// assert!(iter.next().is_none());
235 /// ```
236 pub fn iter_sites(&self) -> SiteIter<N> {
237 SiteIter::new(self.view())
238 }
239
240 /// Returns a new SAF.
241 ///
242 /// The number of provided values must be a multiple of the sum of shapes.
243 /// See the [`Saf`] documentation for details on the storage order.
244 ///
245 /// # Examples
246 ///
247 /// ```
248 /// use winsfs_core::{saf::Saf, saf2d};
249 /// let vec = vec![0.0, 0.1, 0.2, 1.0, 1.1, 0.3, 0.4, 0.5, 1.2, 1.3];
250 /// let shape = [3, 2];
251 /// assert_eq!(
252 /// Saf::new(vec, shape).unwrap(),
253 /// saf2d![
254 /// [0.0, 0.1, 0.2 ; 1.0, 1.1],
255 /// [0.3, 0.4, 0.5 ; 1.2, 1.3],
256 /// ],
257 /// );
258 /// ```
259 /// A [`ShapeError`] is thrown if the shape does not fit the number of values:
260 ///
261 /// ```
262 /// use winsfs_core::saf::Saf;
263 /// let vec = vec![0.0, 0.1, 0.2, 1.0, 1.1, 0.3, 0.4, 0.5, 1.2, 1.3];
264 /// let wrong_shape = [4, 2];
265 /// assert!(Saf::new(vec, wrong_shape).is_err());
266 /// ```
267 pub fn new(values: Vec<f32>, shape: [usize; N]) -> Result<Self, ShapeError<N>> {
268 let len = values.len();
269 let width: usize = shape.iter().sum();
270
271 if len % width == 0 {
272 Ok(Self::new_unchecked(values, shape))
273 } else {
274 Err(ShapeError { len, shape })
275 }
276 }
277
278 /// Returns a new SAF without checking that the shape fits the number of values.
279 pub(crate) fn new_unchecked(values: Vec<f32>, shape: [usize; N]) -> Self {
280 Self { values, shape }
281 }
282
283 /// Returns a parallel iterator over the blocks in the SAF.
284 ///
285 /// This is the parallel version of [`Saf::iter_blocks`].
286 /// If the number of sites in the SAF is not evenly divided by `block_size`,
287 /// the last block will be smaller than the others.
288 ///
289 /// # Examples
290 ///
291 /// ```
292 /// use winsfs_core::{saf::SafView, saf1d};
293 /// use rayon::iter::ParallelIterator;
294 /// let saf = saf1d![
295 /// [0.0, 0.1, 0.2],
296 /// [0.3, 0.4, 0.5],
297 /// [0.6, 0.7, 0.8],
298 /// [0.9, 0.10, 0.11],
299 /// [0.12, 0.13, 0.14],
300 /// ];
301 /// let blocks: Vec<SafView<1>> = saf.par_iter_blocks(2).collect();
302 /// assert_eq!(blocks.len(), 3);
303 /// assert_eq!(
304 /// blocks[0],
305 /// saf1d![[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]].view()
306 /// );
307 /// assert_eq!(
308 /// blocks[1],
309 /// saf1d![[0.6, 0.7, 0.8], [0.9, 0.10, 0.11]].view()
310 /// );
311 /// assert_eq!(blocks[2], saf1d![[0.12, 0.13, 0.14]].view());
312 pub fn par_iter_blocks(&self, block_size: usize) -> ParBlockIter<N> {
313 ParBlockIter::new(self.view(), block_size)
314 }
315
316 /// Returns a parallel iterator over the sites in the SAF.
317 ///
318 /// This is the parallel version of [`Saf::iter_sites`].
319 ///
320 /// # Examples
321 ///
322 /// ```
323 /// use winsfs_core::saf1d;
324 /// use rayon::iter::ParallelIterator;
325 /// let saf = saf1d![
326 /// [1., 1., 1.],
327 /// [1., 1., 1.],
328 /// [1., 1., 1.],
329 /// ];
330 /// saf.par_iter_sites().all(|site| site.as_slice() == &[1., 1., 1.]);
331 pub fn par_iter_sites(&self) -> ParSiteIter<N> {
332 ParSiteIter::new(self.view())
333 }
334
335 /// Creates a new SAF by reading intersecting sites among SAF readers.
336 ///
337 /// SAF files contain values in log-space. The returned values will be exponentiated
338 /// to get out of log-space.
339 ///
340 /// # Panics
341 ///
342 /// Panics if `N == 0`.
343 pub fn read<R>(readers: [saf::ReaderV3<R>; N]) -> io::Result<Self>
344 where
345 R: io::BufRead + io::Seek,
346 {
347 Self::read_inner_impl(readers, |values, item, _| {
348 values.extend_from_slice(item);
349 })
350 }
351
352 /// Creates a new SAF by reading intersecting sites among banded SAF readers.
353 ///
354 /// SAF files contain values in log-space. The returned values will be exponentiated
355 /// to get out of log-space.
356 ///
357 /// Note that this simply fills all non-explicitly represented values in the banded SAF
358 /// with zeros (after getting out of log-space). Hence, this amounts to in some sense "undoing"
359 /// the banding.
360 ///
361 /// # Panics
362 ///
363 /// Panics if `N == 0`.
364 pub fn read_from_banded<R>(readers: [saf::ReaderV4<R>; N]) -> io::Result<Self>
365 where
366 R: io::BufRead + io::Seek,
367 {
368 Self::read_inner_impl(readers, |values, item, alleles| {
369 let full_likelihoods = &item.clone().into_full(alleles, f32::NEG_INFINITY);
370 values.extend_from_slice(full_likelihoods);
371 })
372 }
373
374 /// The inner implementor of readers from full SAF and banded SAF.
375 fn read_inner_impl<R, V, F>(readers: [saf::Reader<R, V>; N], f: F) -> io::Result<Self>
376 where
377 R: io::BufRead + io::Seek,
378 V: saf::version::Version,
379 F: Fn(&mut Vec<f32>, &V::Item, usize),
380 {
381 assert!(N > 0);
382
383 let max_sites = readers
384 .iter()
385 .map(|reader| reader.index().total_sites())
386 .min()
387 .unwrap();
388
389 let shape = readers.by_ref().map(|reader| reader.index().alleles() + 1);
390
391 // The number of intersecting sites is as most the smallest number of sites,
392 // so we preallocate this number and free excess capacity at the end.
393 let capacity = shape.iter().map(|shape| shape * max_sites).sum();
394 let mut values = Vec::with_capacity(capacity);
395
396 let mut intersect = saf::Intersect::new(Vec::from(readers));
397 let mut bufs = intersect.create_record_bufs();
398
399 while intersect.read_records(&mut bufs)?.is_not_done() {
400 for (buf, alleles) in bufs.iter().zip(shape.iter().map(|x| x - 1)) {
401 f(&mut values, buf.item(), alleles)
402 }
403 }
404 // The allocated capacity is an overestimate unless all sites in smallest file intersected.
405 values.shrink_to_fit();
406
407 // Representation in SAF file is in log-space.
408 values.iter_mut().for_each(|x| *x = x.exp());
409
410 Ok(Self::new_unchecked(values, shape))
411 }
412
413 /// Shuffles the SAF sitewise according to a random permutation.
414 pub fn shuffle<R>(&mut self, rng: &mut R)
415 where
416 R: Rng,
417 {
418 let width = self.width();
419
420 // Modified from rand::seq::SliceRandom::shuffle
421 for i in (1..self.sites()).rev() {
422 let j = rng.gen_range(0..i + 1);
423
424 self.swap_sites(i, j, width);
425 }
426 }
427
428 /// Swap sites `i` and `j` in SAF.
429 ///
430 /// `width` is passed in to avoid recalculating for each swap.
431 ///
432 /// # Panics
433 ///
434 /// If `i` or `j` are greater than the number of sites according to `width`.
435 fn swap_sites(&mut self, mut i: usize, mut j: usize, width: usize) {
436 debug_assert_eq!(width, self.width());
437
438 match i.cmp(&j) {
439 Ordering::Less => (i, j) = (j, i),
440 Ordering::Equal => {
441 if i >= self.sites() || j >= self.sites() {
442 panic!("index out of bounds for swapping sites")
443 } else {
444 return;
445 }
446 }
447 Ordering::Greater => (),
448 }
449
450 let (hd, tl) = self.as_mut_slice().split_at_mut(i * width);
451
452 let left = &mut hd[j * width..][..width];
453 let right = &mut tl[..width];
454
455 left.swap_with_slice(right)
456 }
457
458 /// Returns a view of the entire SAF.
459 pub fn view(&self) -> SafView<N> {
460 SafView {
461 values: self.values.as_slice(),
462 shape: self.shape,
463 }
464 }
465
466 impl_shared_saf_methods! {}
467}
468
469impl<'a, const N: usize> Lifetime<'a> for Saf<N> {
470 type Item = SafView<'a, N>;
471}
472
473impl<const N: usize> AsSafView<N> for Saf<N> {
474 #[inline]
475 fn as_saf_view(&self) -> <Self as Lifetime<'_>>::Item {
476 self.view()
477 }
478}
479
480impl<'a, 'b, const N: usize> Lifetime<'a> for &'b Saf<N> {
481 type Item = SafView<'a, N>;
482}
483
484impl<'a, const N: usize> AsSafView<N> for &'a Saf<N> {
485 #[inline]
486 fn as_saf_view(&self) -> <Self as Lifetime<'_>>::Item {
487 self.view()
488 }
489}
490
491/// A view of a joint SAF likelihood matrix for `N` populations.
492///
493/// This may or may not be the entire matrix, but it always represents a contiguous block of sites.
494#[derive(Clone, Copy, Debug, PartialEq)]
495pub struct SafView<'a, const N: usize> {
496 values: &'a [f32],
497 shape: [usize; N],
498}
499
500impl<'a, const N: usize> SafView<'a, N> {
501 /// Returns an iterator over blocks of sites in the SAF.
502 ///
503 /// If the number of sites in the SAF is not evenly divided by `block_size`,
504 /// the last block will be smaller than the others.
505 ///
506 /// # Examples
507 ///
508 /// ```
509 /// use winsfs_core::saf1d;
510 /// let saf = saf1d![
511 /// [0.0, 0.1, 0.2],
512 /// [0.3, 0.4, 0.5],
513 /// [0.6, 0.7, 0.8],
514 /// [0.9, 0.10, 0.11],
515 /// [0.12, 0.13, 0.14],
516 /// ];
517 /// let mut iter = saf.view().iter_blocks(2);
518 /// assert_eq!(
519 /// iter.next().unwrap(),
520 /// saf1d![[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]].view()
521 /// );
522 /// assert_eq!(
523 /// iter.next().unwrap(),
524 /// saf1d![[0.6, 0.7, 0.8], [0.9, 0.10, 0.11]].view()
525 /// );
526 /// assert_eq!(iter.next().unwrap(), saf1d![[0.12, 0.13, 0.14]].view());
527 /// assert!(iter.next().is_none());
528 /// ```
529 pub fn iter_blocks(&self, block_size: usize) -> BlockIter<'a, N> {
530 BlockIter::new(*self, block_size)
531 }
532
533 /// Returns an iterator over the sites in the SAF.
534 ///
535 /// # Examples
536 ///
537 /// ```
538 /// use winsfs_core::saf1d;
539 /// let saf = saf1d![
540 /// [0.0, 0.1, 0.2],
541 /// [0.3, 0.4, 0.5],
542 /// [0.6, 0.7, 0.8],
543 /// ];
544 /// let mut iter = saf.view().iter_sites();
545 /// assert_eq!(iter.next().unwrap().as_slice(), [0.0, 0.1, 0.2]);
546 /// assert_eq!(iter.next().unwrap().as_slice(), [0.3, 0.4, 0.5]);
547 /// assert_eq!(iter.next().unwrap().as_slice(), [0.6, 0.7, 0.8]);
548 /// assert!(iter.next().is_none());
549 /// ```
550 pub fn iter_sites(&self) -> SiteIter<'a, N> {
551 SiteIter::new(*self)
552 }
553
554 /// Returns a new SAF view.
555 ///
556 /// The number of provided values must be a multiple of the sum of shapes.
557 /// See the [`Saf`] documentation for details on the storage order.
558 ///
559 /// To create an owned SAF matrix, see [`Saf::new`] for the equivalent method.
560 ///
561 /// # Examples
562 ///
563 /// ```
564 /// use winsfs_core::{saf::SafView, saf2d};
565 /// let slice = &[0.0, 0.1, 0.2, 1.0, 1.1, 0.3, 0.4, 0.5, 1.2, 1.3];
566 /// let shape = [3, 2];
567 /// assert_eq!(
568 /// SafView::new(slice, shape).unwrap(),
569 /// saf2d![
570 /// [0.0, 0.1, 0.2 ; 1.0, 1.1],
571 /// [0.3, 0.4, 0.5 ; 1.2, 1.3],
572 /// ].view(),
573 /// );
574 /// ```
575 /// A [`ShapeError`] is thrown if the shape does not fit the number of values:
576 ///
577 /// ```
578 /// use winsfs_core::saf::SafView;
579 /// let slice = &[0.0, 0.1, 0.2, 1.0, 1.1, 0.3, 0.4, 0.5, 1.2, 1.3];
580 /// let wrong_shape = [4, 2];
581 /// assert!(SafView::new(slice, wrong_shape).is_err());
582 /// ```
583 pub fn new(values: &'a [f32], shape: [usize; N]) -> Result<Self, ShapeError<N>> {
584 let len = values.len();
585 let width: usize = shape.iter().sum();
586
587 if len % width == 0 {
588 Ok(Self::new_unchecked(values, shape))
589 } else {
590 Err(ShapeError { len, shape })
591 }
592 }
593
594 /// Returns a new SAF view without checking that the shape fits the number of values.
595 pub(crate) fn new_unchecked(values: &'a [f32], shape: [usize; N]) -> Self {
596 Self { values, shape }
597 }
598
599 /// Returns a parallel iterator over the blocks in the SAF.
600 ///
601 /// This is the parallel version of [`SafView::iter_blocks`].
602 /// If the number of sites in the SAF is not evenly divided by `block_size`,
603 /// the last block will be smaller than the others.
604 ///
605 /// # Examples
606 ///
607 /// ```
608 /// use winsfs_core::{saf::SafView, saf1d};
609 /// use rayon::iter::ParallelIterator;
610 /// let saf = saf1d![
611 /// [0.0, 0.1, 0.2],
612 /// [0.3, 0.4, 0.5],
613 /// [0.6, 0.7, 0.8],
614 /// [0.9, 0.10, 0.11],
615 /// [0.12, 0.13, 0.14],
616 /// ];
617 /// let view = saf.view();
618 /// let blocks: Vec<SafView<1>> = view.par_iter_blocks(2).collect();
619 /// assert_eq!(blocks.len(), 3);
620 /// assert_eq!(
621 /// blocks[0],
622 /// saf1d![[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]].view()
623 /// );
624 /// assert_eq!(
625 /// blocks[1],
626 /// saf1d![[0.6, 0.7, 0.8], [0.9, 0.10, 0.11]].view()
627 /// );
628 /// assert_eq!(blocks[2], saf1d![[0.12, 0.13, 0.14]].view());
629 pub fn par_iter_blocks(&self, block_size: usize) -> ParBlockIter<N> {
630 ParBlockIter::new(*self, block_size)
631 }
632
633 /// Returns a parallel iterator over the sites in the SAF.
634 ///
635 /// This is the parallel version of [`SafView::iter_sites`].
636 ///
637 /// # Examples
638 ///
639 /// ```
640 /// use winsfs_core::saf1d;
641 /// use rayon::iter::ParallelIterator;
642 /// let saf = saf1d![
643 /// [1., 1., 1.],
644 /// [1., 1., 1.],
645 /// [1., 1., 1.],
646 /// ];
647 /// saf.view().par_iter_sites().all(|site| site.as_slice() == &[1., 1., 1.]);
648 pub fn par_iter_sites(&self) -> ParSiteIter<N> {
649 ParSiteIter::new(*self)
650 }
651
652 impl_shared_saf_methods! {}
653}
654
655impl<'a, 'b, const N: usize> Lifetime<'a> for SafView<'b, N> {
656 type Item = SafView<'a, N>;
657}
658
659impl<'a, const N: usize> AsSafView<N> for SafView<'a, N> {
660 #[inline]
661 fn as_saf_view(&self) -> <Self as Lifetime<'_>>::Item {
662 *self
663 }
664}
665
666/// An error associated with SAF or SAF site construction using invalid shape.
667#[derive(Clone, Debug)]
668pub struct ShapeError<const N: usize> {
669 shape: [usize; N],
670 len: usize,
671}
672
673impl<const N: usize> fmt::Display for ShapeError<N> {
674 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
675 write!(
676 f,
677 "cannot construct shape {} from {} values",
678 self.shape.map(|x| x.to_string()).join("/"),
679 self.len,
680 )
681 }
682}
683
684impl<const N: usize> Error for ShapeError<N> {}
685
686#[cfg(test)]
687mod tests {
688 #[test]
689 fn test_swap_1d() {
690 let mut saf = saf1d![
691 [0., 0., 0.],
692 [1., 1., 1.],
693 [2., 2., 2.],
694 [3., 3., 3.],
695 [4., 4., 4.],
696 [5., 5., 5.],
697 ];
698
699 let width = saf.width();
700 assert_eq!(width, 3);
701
702 saf.swap_sites(3, 3, width);
703 assert_eq!(saf.get_site(3).as_slice(), &[3., 3., 3.]);
704
705 saf.swap_sites(0, 1, width);
706 assert_eq!(saf.get_site(0).as_slice(), &[1., 1., 1.]);
707 assert_eq!(saf.get_site(1).as_slice(), &[0., 0., 0.]);
708
709 saf.swap_sites(5, 0, width);
710 assert_eq!(saf.get_site(0).as_slice(), &[5., 5., 5.]);
711 assert_eq!(saf.get_site(5).as_slice(), &[1., 1., 1.]);
712 }
713
714 #[test]
715 fn test_swap_2d() {
716 #[rustfmt::skip]
717 let mut saf = saf2d![
718 [0., 0., 0.; 10., 10.],
719 [1., 1., 1.; 11., 11.],
720 [2., 2., 2.; 12., 12.],
721 [3., 3., 3.; 13., 13.],
722 [4., 4., 4.; 14., 14.],
723 [5., 5., 5.; 15., 15.],
724 ];
725
726 let width = saf.width();
727 assert_eq!(width, 5);
728
729 saf.swap_sites(0, 5, width);
730 assert_eq!(saf.get_site(0).as_slice(), &[5., 5., 5., 15., 15.,]);
731 saf.swap_sites(5, 0, width);
732 assert_eq!(saf.get_site(0).as_slice(), &[0., 0., 0., 10., 10.,]);
733 }
734
735 #[test]
736 #[should_panic]
737 fn test_swap_panics_out_of_bounds() {
738 let mut saf = saf1d![
739 [0., 0., 0.],
740 [1., 1., 1.],
741 [2., 2., 2.],
742 [3., 3., 3.],
743 [4., 4., 4.],
744 [5., 5., 5.],
745 ];
746
747 saf.swap_sites(6, 5, saf.width());
748 }
749}