Skip to main content

conspire/math/tensor/rank_2/sparse_vec/
mod.rs

1#[cfg(test)]
2mod test;
3
4use super::TensorRank2;
5use crate::math::{Tensor, TensorArray, TensorRank0};
6use crate::units::Dimensionless;
7use std::{
8    fmt::{self, Debug, Display, Formatter, Result},
9    iter::Sum,
10    ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign},
11};
12
13/// A sparse vector of rank-2 tensors, storing only inserted entries.
14pub struct TensorRank2SparseVec<const D: usize, I, J, U = Dimensionless>(
15    pub(super) Vec<(usize, TensorRank2<D, I, J, U>)>,
16);
17
18impl<const D: usize, I, J, U> Clone for TensorRank2SparseVec<D, I, J, U> {
19    fn clone(&self) -> Self {
20        Self(self.0.clone())
21    }
22}
23
24impl<const D: usize, I, J, U> Debug for TensorRank2SparseVec<D, I, J, U> {
25    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
26        Debug::fmt(&self.0, f)
27    }
28}
29
30impl<const D: usize, I, J, U> Default for TensorRank2SparseVec<D, I, J, U> {
31    fn default() -> Self {
32        Self(Default::default())
33    }
34}
35
36impl<const D: usize, I, J, U> PartialEq for TensorRank2SparseVec<D, I, J, U> {
37    fn eq(&self, other: &Self) -> bool {
38        self.0 == other.0
39    }
40}
41
42impl<const D: usize, I, J, U> TensorRank2SparseVec<D, I, J, U> {
43    pub fn entries(&self) -> impl Iterator<Item = (usize, &TensorRank2<D, I, J, U>)> {
44        self.0.iter().map(|(column, entry)| (*column, entry))
45    }
46}
47
48impl<const D: usize, I, J, U> FromIterator<TensorRank2<D, I, J, U>>
49    for TensorRank2SparseVec<D, I, J, U>
50{
51    fn from_iter<T>(into_iterator: T) -> Self
52    where
53        T: IntoIterator<Item = TensorRank2<D, I, J, U>>,
54    {
55        Self(into_iterator.into_iter().enumerate().collect())
56    }
57}
58
59impl<const D: usize, I, J, U> Index<usize> for TensorRank2SparseVec<D, I, J, U> {
60    type Output = TensorRank2<D, I, J, U>;
61    fn index(&self, index: usize) -> &Self::Output {
62        match self.0.binary_search_by_key(&index, |&(column, _)| column) {
63            Ok(k) => &self.0[k].1,
64            Err(_) => panic!("Entry ({index}) not present."),
65        }
66    }
67}
68
69impl<const D: usize, I, J, U> IndexMut<usize> for TensorRank2SparseVec<D, I, J, U> {
70    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
71        let k = match self.0.binary_search_by_key(&index, |&(column, _)| column) {
72            Ok(k) => k,
73            Err(k) => {
74                self.0.insert(k, (index, TensorRank2::zero()));
75                k
76            }
77        };
78        &mut self.0[k].1
79    }
80}
81
82impl<const D: usize, I, J, U> Display for TensorRank2SparseVec<D, I, J, U> {
83    fn fmt(&self, f: &mut Formatter) -> Result {
84        write!(f, "Need to implement Display")
85    }
86}
87
88impl<const D: usize, I, J, U> Tensor for TensorRank2SparseVec<D, I, J, U> {
89    type Item = TensorRank2<D, I, J, U>;
90    type Unit = U;
91    fn iter(&self) -> impl Iterator<Item = &Self::Item> {
92        self.0.iter().map(|(_, entry)| entry)
93    }
94    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Self::Item> {
95        self.0.iter_mut().map(|(_, entry)| entry)
96    }
97    fn len(&self) -> usize {
98        self.0.len()
99    }
100    fn size(&self) -> usize {
101        self.0.len() * D * D
102    }
103}
104
105fn merge<const D: usize, I, J, U>(
106    a: TensorRank2SparseVec<D, I, J, U>,
107    b: &TensorRank2SparseVec<D, I, J, U>,
108    sign: TensorRank0,
109) -> TensorRank2SparseVec<D, I, J, U> {
110    let mut merged = a;
111    b.0.iter()
112        .for_each(|(column, entry)| merged[*column] += entry * sign);
113    merged
114}
115
116impl<const D: usize, I, J, U> Add for TensorRank2SparseVec<D, I, J, U> {
117    type Output = Self;
118    fn add(self, other: Self) -> Self {
119        merge(self, &other, 1.0)
120    }
121}
122
123impl<const D: usize, I, J, U> Add<&Self> for TensorRank2SparseVec<D, I, J, U> {
124    type Output = Self;
125    fn add(self, other: &Self) -> Self {
126        merge(self, other, 1.0)
127    }
128}
129
130impl<const D: usize, I, J, U> AddAssign for TensorRank2SparseVec<D, I, J, U> {
131    fn add_assign(&mut self, other: Self) {
132        other
133            .0
134            .into_iter()
135            .for_each(|(column, entry)| self[column] += entry);
136    }
137}
138
139impl<const D: usize, I, J, U> AddAssign<&Self> for TensorRank2SparseVec<D, I, J, U> {
140    fn add_assign(&mut self, other: &Self) {
141        other
142            .0
143            .iter()
144            .for_each(|(column, entry)| self[*column] += entry);
145    }
146}
147
148impl<const D: usize, I, J, U> Sub for TensorRank2SparseVec<D, I, J, U> {
149    type Output = Self;
150    fn sub(self, other: Self) -> Self {
151        merge(self, &other, -1.0)
152    }
153}
154
155impl<const D: usize, I, J, U> Sub<&Self> for TensorRank2SparseVec<D, I, J, U> {
156    type Output = Self;
157    fn sub(self, other: &Self) -> Self {
158        merge(self, other, -1.0)
159    }
160}
161
162impl<const D: usize, I, J, U> SubAssign for TensorRank2SparseVec<D, I, J, U> {
163    fn sub_assign(&mut self, other: Self) {
164        other
165            .0
166            .into_iter()
167            .for_each(|(column, entry)| self[column] -= entry);
168    }
169}
170
171impl<const D: usize, I, J, U> SubAssign<&Self> for TensorRank2SparseVec<D, I, J, U> {
172    fn sub_assign(&mut self, other: &Self) {
173        other
174            .0
175            .iter()
176            .for_each(|(column, entry)| self[*column] -= entry);
177    }
178}
179
180impl<const D: usize, I, J, U> Mul<TensorRank0> for TensorRank2SparseVec<D, I, J, U> {
181    type Output = Self;
182    fn mul(mut self, scalar: TensorRank0) -> Self {
183        self *= &scalar;
184        self
185    }
186}
187
188impl<const D: usize, I, J, U> MulAssign<TensorRank0> for TensorRank2SparseVec<D, I, J, U> {
189    fn mul_assign(&mut self, scalar: TensorRank0) {
190        self.0.iter_mut().for_each(|(_, entry)| *entry *= &scalar);
191    }
192}
193
194impl<const D: usize, I, J, U> MulAssign<&TensorRank0> for TensorRank2SparseVec<D, I, J, U> {
195    fn mul_assign(&mut self, scalar: &TensorRank0) {
196        self.0.iter_mut().for_each(|(_, entry)| *entry *= scalar);
197    }
198}
199
200impl<const D: usize, I, J, U> Div<TensorRank0> for TensorRank2SparseVec<D, I, J, U> {
201    type Output = Self;
202    fn div(mut self, scalar: TensorRank0) -> Self {
203        self /= &scalar;
204        self
205    }
206}
207
208impl<const D: usize, I, J, U> DivAssign<TensorRank0> for TensorRank2SparseVec<D, I, J, U> {
209    fn div_assign(&mut self, scalar: TensorRank0) {
210        self.0.iter_mut().for_each(|(_, entry)| *entry /= &scalar);
211    }
212}
213
214impl<const D: usize, I, J, U> DivAssign<&TensorRank0> for TensorRank2SparseVec<D, I, J, U> {
215    fn div_assign(&mut self, scalar: &TensorRank0) {
216        self.0.iter_mut().for_each(|(_, entry)| *entry /= scalar);
217    }
218}
219
220impl<const D: usize, I, J, U> Sum for TensorRank2SparseVec<D, I, J, U> {
221    fn sum<T>(iter: T) -> Self
222    where
223        T: Iterator<Item = Self>,
224    {
225        iter.fold(Self::default(), |sum, entry| sum + entry)
226    }
227}