conspire/math/tensor/rank_0/list/
mod.rs

1#[cfg(test)]
2mod test;
3
4use super::{super::super::write_tensor_rank_0, Tensor, TensorArray, TensorRank0};
5use std::{
6    fmt::{Display, Formatter, Result},
7    ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign},
8};
9
10/// A list of tensors of rank 0 (a list of scalars).
11///
12/// `W` is the list length.
13#[derive(Clone, Debug)]
14pub struct TensorRank0List<const W: usize>([TensorRank0; W]);
15
16pub const fn tensor_rank_0_list<const W: usize>(array: [TensorRank0; W]) -> TensorRank0List<W> {
17    TensorRank0List(array)
18}
19
20/// Display implementation for rank-0 lists.
21impl<const W: usize> Display for TensorRank0List<W> {
22    fn fmt(&self, f: &mut Formatter) -> Result {
23        write!(f, "\x1B[s")?;
24        write!(f, "[")?;
25        self.0.chunks(5).enumerate().try_for_each(|(i, chunk)| {
26            chunk
27                .iter()
28                .try_for_each(|entry| write_tensor_rank_0(f, entry))?;
29            if (i + 1) * 5 < W {
30                writeln!(f, "\x1B[2D,")?;
31                write!(f, "\x1B[u")?;
32                write!(f, "\x1B[{}B ", i + 1)?;
33            }
34            Ok(())
35        })?;
36        write!(f, "\x1B[2D]")?;
37        Ok(())
38    }
39}
40
41impl<const W: usize> Tensor for TensorRank0List<W> {
42    type Item = TensorRank0;
43    fn full_contraction(&self, tensor_rank_0_list: &Self) -> TensorRank0 {
44        self.iter()
45            .zip(tensor_rank_0_list.iter())
46            .map(|(self_entry, tensor_rank_0)| self_entry * tensor_rank_0)
47            .sum()
48    }
49    fn iter(&self) -> impl Iterator<Item = &TensorRank0> {
50        self.0.iter()
51    }
52    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Self::Item> {
53        self.0.iter_mut()
54    }
55}
56
57impl<const W: usize> TensorArray for TensorRank0List<W> {
58    type Array = [TensorRank0; W];
59    type Item = TensorRank0;
60    fn as_array(&self) -> Self::Array {
61        self.0
62    }
63    fn identity() -> Self {
64        Self([1.0; W])
65    }
66    fn new(array: Self::Array) -> Self {
67        Self(array)
68    }
69    fn zero() -> Self {
70        Self([0.0; W])
71    }
72}
73
74impl<const W: usize> FromIterator<TensorRank0> for TensorRank0List<W> {
75    fn from_iter<Ii: IntoIterator<Item = TensorRank0>>(into_iterator: Ii) -> Self {
76        let mut tensor_rank_0_list = Self::zero();
77        tensor_rank_0_list
78            .iter_mut()
79            .zip(into_iterator)
80            .for_each(|(tensor_rank_0, entry)| *tensor_rank_0 = entry);
81        tensor_rank_0_list
82    }
83}
84
85impl<const W: usize> Index<usize> for TensorRank0List<W> {
86    type Output = TensorRank0;
87    fn index(&self, index: usize) -> &Self::Output {
88        &self.0[index]
89    }
90}
91
92impl<const W: usize> IndexMut<usize> for TensorRank0List<W> {
93    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
94        &mut self.0[index]
95    }
96}
97
98impl<const W: usize> Add for TensorRank0List<W> {
99    type Output = Self;
100    fn add(mut self, tensor_rank_0_list: Self) -> Self::Output {
101        self += tensor_rank_0_list;
102        self
103    }
104}
105
106impl<const W: usize> Add<&Self> for TensorRank0List<W> {
107    type Output = Self;
108    fn add(mut self, tensor_rank_0_list: &Self) -> Self::Output {
109        self += tensor_rank_0_list;
110        self
111    }
112}
113
114impl<const W: usize> AddAssign for TensorRank0List<W> {
115    fn add_assign(&mut self, tensor_rank_0_list: Self) {
116        self.iter_mut()
117            .zip(tensor_rank_0_list.iter())
118            .for_each(|(self_entry, tensor_rank_0)| *self_entry += tensor_rank_0);
119    }
120}
121
122impl<const W: usize> AddAssign<&Self> for TensorRank0List<W> {
123    fn add_assign(&mut self, tensor_rank_0_list: &Self) {
124        self.iter_mut()
125            .zip(tensor_rank_0_list.iter())
126            .for_each(|(self_entry, tensor_rank_0)| *self_entry += tensor_rank_0);
127    }
128}
129
130impl<const W: usize> Div<TensorRank0> for TensorRank0List<W> {
131    type Output = Self;
132    fn div(mut self, tensor_rank_0: TensorRank0) -> Self::Output {
133        self /= tensor_rank_0;
134        self
135    }
136}
137
138impl<const W: usize> Div<&TensorRank0> for TensorRank0List<W> {
139    type Output = Self;
140    fn div(mut self, tensor_rank_0: &TensorRank0) -> Self::Output {
141        self /= tensor_rank_0;
142        self
143    }
144}
145
146impl<const W: usize> DivAssign<TensorRank0> for TensorRank0List<W> {
147    fn div_assign(&mut self, tensor_rank_0: TensorRank0) {
148        self.iter_mut().for_each(|self_i| *self_i /= &tensor_rank_0);
149    }
150}
151
152impl<const W: usize> DivAssign<&TensorRank0> for TensorRank0List<W> {
153    fn div_assign(&mut self, tensor_rank_0: &TensorRank0) {
154        self.iter_mut().for_each(|self_i| *self_i /= tensor_rank_0);
155    }
156}
157
158impl<const W: usize> Mul<TensorRank0> for TensorRank0List<W> {
159    type Output = Self;
160    fn mul(mut self, tensor_rank_0: TensorRank0) -> Self::Output {
161        self *= tensor_rank_0;
162        self
163    }
164}
165
166// impl<const W: usize> Mul<TensorRank0> for &TensorRank0List<W> {
167//     type Output = TensorRank0List<W>;
168//     fn mul(self, tensor_rank_0: TensorRank0) -> Self::Output {
169//         self.iter().map(|self_i| self_i * tensor_rank_0).collect()
170//     }
171// }
172
173impl<const W: usize> Mul<&TensorRank0> for TensorRank0List<W> {
174    type Output = Self;
175    fn mul(mut self, tensor_rank_0: &TensorRank0) -> Self::Output {
176        self *= tensor_rank_0;
177        self
178    }
179}
180
181impl<const W: usize> Mul<&TensorRank0> for &TensorRank0List<W> {
182    type Output = TensorRank0List<W>;
183    fn mul(self, tensor_rank_0: &TensorRank0) -> Self::Output {
184        self.iter().map(|self_i| self_i * tensor_rank_0).collect()
185    }
186}
187
188impl<const W: usize> MulAssign<TensorRank0> for TensorRank0List<W> {
189    fn mul_assign(&mut self, tensor_rank_0: TensorRank0) {
190        self.iter_mut().for_each(|self_i| *self_i *= &tensor_rank_0);
191    }
192}
193
194impl<const W: usize> MulAssign<&TensorRank0> for TensorRank0List<W> {
195    fn mul_assign(&mut self, tensor_rank_0: &TensorRank0) {
196        self.iter_mut().for_each(|self_i| *self_i *= tensor_rank_0);
197    }
198}
199
200impl<const W: usize> Mul for TensorRank0List<W> {
201    type Output = TensorRank0;
202    fn mul(self, tensor_rank_0_list: Self) -> Self::Output {
203        self.iter()
204            .zip(tensor_rank_0_list.iter())
205            .map(|(self_entry, tensor_rank_0_list_entry)| self_entry * tensor_rank_0_list_entry)
206            .sum()
207    }
208}
209
210impl<const W: usize> Mul<&Self> for TensorRank0List<W> {
211    type Output = TensorRank0;
212    fn mul(self, tensor_rank_0_list: &Self) -> Self::Output {
213        self.iter()
214            .zip(tensor_rank_0_list.iter())
215            .map(|(self_entry, tensor_rank_0_list_entry)| self_entry * tensor_rank_0_list_entry)
216            .sum()
217    }
218}
219
220impl<const W: usize> Mul<TensorRank0List<W>> for &TensorRank0List<W> {
221    type Output = TensorRank0;
222    fn mul(self, tensor_rank_0_list: TensorRank0List<W>) -> Self::Output {
223        self.iter()
224            .zip(tensor_rank_0_list.iter())
225            .map(|(self_entry, tensor_rank_0_list_entry)| self_entry * tensor_rank_0_list_entry)
226            .sum()
227    }
228}
229
230impl<const W: usize> Mul for &TensorRank0List<W> {
231    type Output = TensorRank0;
232    fn mul(self, tensor_rank_0_list: Self) -> Self::Output {
233        self.iter()
234            .zip(tensor_rank_0_list.iter())
235            .map(|(self_entry, tensor_rank_0_list_entry)| self_entry * tensor_rank_0_list_entry)
236            .sum()
237    }
238}
239
240impl<const W: usize> Sub for TensorRank0List<W> {
241    type Output = Self;
242    fn sub(mut self, tensor_rank_0_list: Self) -> Self::Output {
243        self -= tensor_rank_0_list;
244        self
245    }
246}
247
248impl<const W: usize> Sub<&Self> for TensorRank0List<W> {
249    type Output = Self;
250    fn sub(mut self, tensor_rank_0_list: &Self) -> Self::Output {
251        self -= tensor_rank_0_list;
252        self
253    }
254}
255
256impl<const W: usize> SubAssign for TensorRank0List<W> {
257    fn sub_assign(&mut self, tensor_rank_0_list: Self) {
258        self.iter_mut()
259            .zip(tensor_rank_0_list.iter())
260            .for_each(|(self_entry, tensor_rank_0)| *self_entry -= tensor_rank_0);
261    }
262}
263
264impl<const W: usize> SubAssign<&Self> for TensorRank0List<W> {
265    fn sub_assign(&mut self, tensor_rank_0_list: &Self) {
266        self.iter_mut()
267            .zip(tensor_rank_0_list.iter())
268            .for_each(|(self_entry, tensor_rank_0)| *self_entry -= tensor_rank_0);
269    }
270}