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
13pub 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}