Skip to main content

ndstruct/
coo.rs

1//! COO (Coordinate) format for N-dimensions.
2
3mod coo_error;
4mod coo_utils;
5
6#[cfg(feature = "alloc")]
7use alloc::vec::Vec;
8use cl_aux::{ArrayWrapper, SingleTypeStorage};
9pub use coo_error::*;
10use coo_utils::{does_not_have_duplicates_sorted, value, value_mut};
11
12/// COO backed by a static array.
13pub type CooArray<DATA, const D: usize, const DN: usize> = Coo<[([usize; D], DATA); DN], D>;
14/// COO backed by a mutable slice
15pub type CooMut<'data, DATA, const D: usize> = Coo<&'data mut [([usize; D], DATA)], D>;
16/// COO backed by a slice
17pub type CooRef<'data, DATA, const D: usize> = Coo<&'data [([usize; D], DATA)], D>;
18#[cfg(feature = "alloc")]
19/// COO backed by a dynamic vector.
20pub type CooVec<DATA, const D: usize> = Coo<Vec<([usize; D], DATA)>, D>;
21
22/// Base structure for all [Coo] variants.
23///
24/// # Types
25///
26/// * `D`: Number of dimensions
27/// * `DS`: Data Storage
28#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
29#[derive(Clone, Debug, Default, Eq, Ord, PartialEq, PartialOrd)]
30pub struct Coo<DS, const D: usize> {
31  pub(crate) data: DS,
32  pub(crate) dims: ArrayWrapper<usize, D>,
33}
34
35impl<DS, const D: usize> Coo<DS, D> {
36  /// The definitions of all dimensions.
37  ///
38  /// # Example
39  ///
40  /// ```rust
41  /// use ndstruct::doc_tests::coo_array_5;
42  /// assert_eq!(coo_array_5().dims(), &[2, 3, 4, 3, 3]);
43  /// ```
44  #[inline]
45  pub fn dims(&self) -> &[usize; D] {
46    &self.dims
47  }
48}
49
50impl<DATA, DS, const D: usize> Coo<DS, D>
51where
52  DS: AsRef<[<DS as SingleTypeStorage>::Item]> + SingleTypeStorage<Item = ([usize; D], DATA)>,
53{
54  /// Creates a valid COO instance.
55  ///
56  /// # Arguments
57  ///
58  /// * `dims`: Array of dimensions
59  /// * `data`: Data collection
60  ///
61  /// # Example
62  #[cfg_attr(feature = "alloc", doc = "```rust")]
63  #[cfg_attr(not(feature = "alloc"), doc = "```ignore")]
64  /// use ndstruct::coo::{CooArray, CooVec};
65  /// // Sparse array ([8, _, _, _, _, 9, _, _, _, _])
66  /// let mut _sparse_array = CooArray::new([10], [([0], 8.0), ([5], 9.0)]);
67  /// // A bunch of nothing for your overflow needs
68  /// let mut _over_nine: ndstruct::Result<CooVec<(), 9001>>;
69  /// _over_nine = CooVec::new([0; 9001], vec![]);
70  /// ```
71  #[inline]
72  pub fn new(dims: [usize; D], data: DS) -> crate::Result<Self> {
73    if !crate::utils::are_in_ascending_order(data.as_ref(), |a, b| [&a.0, &b.0]) {
74      return Err(CooError::InvalidIndcsOrder.into());
75    }
76    let has_invalid_indcs = !data.as_ref().iter().all(|&(indcs, _)| {
77      indcs.iter().zip(dims.iter()).all(
78        |(data_idx, dim)| {
79          if dim == &0 {
80            true
81          } else {
82            data_idx < dim
83          }
84        },
85      )
86    });
87    if has_invalid_indcs {
88      return Err(CooError::InvalidIndcs.into());
89    }
90    if !does_not_have_duplicates_sorted(data.as_ref(), |a, b| a.0[..] != b.0[..]) {
91      return Err(CooError::DuplicatedIndices.into());
92    }
93    Ok(Self { data, dims: dims.into() })
94  }
95
96  /// The data that is being stored.
97  ///
98  /// # Example
99  ///
100  /// ```rust
101  /// use ndstruct::doc_tests::coo_array_5;
102  /// assert_eq!(coo_array_5().data().first(), Some(&([0, 0, 1, 1, 2].into(), 1)));
103  /// ```
104  #[inline]
105  pub fn data(&self) -> &[([usize; D], DATA)] {
106    self.data.as_ref()
107  }
108
109  /// If any, retrieves an immutable data reference of a given set of indices.
110  ///
111  /// # Arguments
112  ///
113  /// * `indcs`: Indices of the desired data location
114  ///
115  /// # Example
116  ///
117  /// ```rust
118  /// use ndstruct::doc_tests::coo_array_5;
119  /// let coo = coo_array_5();
120  /// assert_eq!(coo.value([0, 0, 0, 0, 0]), None);
121  /// assert_eq!(coo.value([0, 2, 2, 0, 1]), Some(&4));
122  /// ```
123  #[inline]
124  pub fn value(&self, indcs: [usize; D]) -> Option<&DATA> {
125    value(indcs, self.data.as_ref())
126  }
127}
128
129impl<DATA, DS, const D: usize> Coo<DS, D>
130where
131  DS: AsMut<[<DS as SingleTypeStorage>::Item]> + SingleTypeStorage<Item = ([usize; D], DATA)>,
132{
133  /// Mutable version of [`value`](#method.value).
134  #[inline]
135  pub fn value_mut(&mut self, indcs: [usize; D]) -> Option<&mut DATA> {
136    value_mut(indcs, self.data.as_mut())
137  }
138}
139
140#[cfg(feature = "rand")]
141impl<DATA, DS, const D: usize> Coo<DS, D>
142where
143  DS: AsMut<[<DS as SingleTypeStorage>::Item]>
144    + AsRef<[<DS as SingleTypeStorage>::Item]>
145    + Default
146    + SingleTypeStorage<Item = ([usize; D], DATA)>
147    + cl_aux::CapacityUpperBound
148    + cl_aux::Push<<DS as SingleTypeStorage>::Item>,
149{
150  /// Creates a new random and valid instance delimited by the passed arguments.
151  ///
152  /// # Arguments
153  ///
154  /// * `dims`: Array of dimensions
155  /// * `nnz`: Number of Non-Zero elements
156  /// * `rng`: `rand::Rng` trait
157  /// * `cb`: Callback to control data creation
158  ///
159  /// # Example
160  #[cfg_attr(feature = "alloc", doc = "```rust")]
161  #[cfg_attr(not(feature = "alloc"), doc = "```ignore")]
162  /// use ndstruct::coo::CooVec;
163  /// use rand::{Rng, rngs::mock::StepRng};
164  /// let mut rng = StepRng::new(0, 1);
165  /// let dims = [1, 2, 3];
166  /// let mut _random: ndstruct::Result<CooVec<u8, 3>>;
167  /// _random = CooVec::new_controlled_random_rand(dims, 3, &mut rng, |r, _| r.gen());
168  /// ```
169  #[inline]
170  pub fn new_controlled_random_rand<R>(
171    dims: [usize; D],
172    nnz: usize,
173    rng: &mut R,
174    mut cb: impl FnMut(&mut R, &[usize; D]) -> DATA,
175  ) -> crate::Result<Self>
176  where
177    R: rand::Rng,
178  {
179    use rand::distr::Distribution as _;
180    if nnz > crate::utils::max_nnz(&dims) {
181      return Err(CooError::NnzGreaterThanMaximumNnz.into());
182    }
183    let mut data: DS = Default::default();
184    if nnz > data.as_ref().len() {
185      return Err(crate::Error::InsufficientCapacity);
186    }
187    for _ in 0..nnz {
188      let indcs: [usize; D] = ArrayWrapper::from_fn(|idx| {
189        let dim = *dims.get(idx).unwrap_or(&0);
190        if dim == 0 {
191          0
192        } else {
193          rand::distr::Uniform::new(0, dim).ok().map(|el| el.sample(rng)).unwrap_or_default()
194        }
195      })
196      .0;
197      if data.as_ref().iter().all(|value| value.0 != indcs) {
198        data
199          .push({
200            let element = cb(rng, &indcs);
201            (indcs, element)
202          })
203          .map_err(|_err| crate::Error::InsufficientCapacity)?;
204      }
205    }
206    data.as_mut().sort_unstable_by(|a, b| a.0.cmp(&b.0));
207    Coo::new(dims, data)
208  }
209
210  /// Creates a new random and valid instance.
211  ///
212  /// # Arguments
213  ///
214  /// * `rng`: `rand::Rng` trait
215  /// * `upper_bound`: The maximum allowed exclusive dimension
216  #[inline]
217  pub fn new_random_rand<R>(rng: &mut R, upper_bound: usize) -> crate::Result<Self>
218  where
219    R: rand::Rng,
220    rand::distr::StandardUniform: rand::distr::Distribution<DATA>,
221  {
222    let dims = crate::utils::valid_random_dims(rng, upper_bound);
223    let max_nnz = crate::utils::max_nnz(&dims);
224    let nnz = if max_nnz == 0 { 0 } else { rng.random_range(0..max_nnz) };
225    Self::new_controlled_random_rand(dims, nnz, rng, |r, _| r.random())
226  }
227}