conspire/math/tensor/list/
mod.rs

1use crate::math::{ContractWith, Differentiate, Erase, Quantity, Tensor, TensorArray, TensorRank0};
2use std::{
3    array::{self, from_fn},
4    fmt::{Display, Formatter, Result},
5    iter::Sum,
6    ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign},
7    slice,
8};
9
10/// A fixed-size collection of tensors.
11#[derive(Clone, Debug, PartialEq)]
12#[repr(transparent)]
13pub struct TensorList<T, const N: usize>([T; N])
14where
15    T: Tensor;
16
17impl<T, const N: usize> Erase for TensorList<T, N>
18where
19    T: Erase + Tensor,
20    <T as Erase>::Erased: Tensor,
21{
22    type Erased = TensorList<<T as Erase>::Erased, N>;
23    fn erase(&self) -> &Self::Erased {
24        unsafe { &*(self as *const Self as *const Self::Erased) }
25    }
26}
27
28impl<T, const N: usize> TensorList<T, N>
29where
30    T: Tensor,
31{
32    /// Associated function for const type conversion.
33    pub const fn const_from(array: [T; N]) -> Self {
34        Self(array)
35    }
36}
37
38impl<T, const N: usize> Default for TensorList<T, N>
39where
40    T: Tensor,
41{
42    fn default() -> Self {
43        Self(from_fn(|_| T::default()))
44    }
45}
46
47impl<T, const N: usize> From<[T; N]> for TensorList<T, N>
48where
49    T: Tensor,
50{
51    fn from(tensor_array: [T; N]) -> Self {
52        Self(tensor_array)
53    }
54}
55
56impl<T, const N: usize> From<TensorList<T, N>> for [T; N]
57where
58    T: Tensor,
59{
60    fn from(tensor_list: TensorList<T, N>) -> Self {
61        tensor_list.0
62    }
63}
64
65impl<T, const N: usize> Display for TensorList<T, N>
66where
67    T: Tensor,
68{
69    fn fmt(&self, f: &mut Formatter) -> Result {
70        write!(f, "Need to implement Display")
71        // write!(f, "\x1B[s")?;
72        // write!(f, "[[")?;
73        // self.iter().enumerate().try_for_each(|(i, tensor_rank_1)| {
74        //     tensor_rank_1
75        //         .iter()
76        //         .try_for_each(|entry| write_tensor_rank_0(f, entry))?;
77        //     if i + 1 < W {
78        //         writeln!(f, "\x1B[2D],")?;
79        //         write!(f, "\x1B[u")?;
80        //         write!(f, "\x1B[{}B [", i + 1)?;
81        //     }
82        //     Ok(())
83        // })?;
84        // write!(f, "\x1B[2D]]")
85    }
86}
87
88impl<T, const N: usize> Index<usize> for TensorList<T, N>
89where
90    T: Tensor,
91{
92    type Output = T;
93    fn index(&self, index: usize) -> &Self::Output {
94        &self.0[index]
95    }
96}
97
98impl<T, const N: usize> IndexMut<usize> for TensorList<T, N>
99where
100    T: Tensor,
101{
102    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
103        &mut self.0[index]
104    }
105}
106
107impl<T, const N: usize> Tensor for TensorList<T, N>
108where
109    T: Tensor,
110{
111    type Item = T;
112    type Unit = <T as Tensor>::Unit;
113    fn iter(&self) -> impl Iterator<Item = &Self::Item> {
114        self.0.iter()
115    }
116    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Self::Item> {
117        self.0.iter_mut()
118    }
119    fn len(&self) -> usize {
120        N
121    }
122    fn size(&self) -> usize {
123        N * self[0].size() // fine if T is TensorArray
124    }
125}
126
127impl<T, const N: usize> FromIterator<T> for TensorList<T, N>
128where
129    T: Tensor,
130{
131    fn from_iter<Ii: IntoIterator<Item = T>>(into_iterator: Ii) -> Self {
132        let mut tensor_list = Self::default();
133        tensor_list
134            .iter_mut()
135            .zip(into_iterator)
136            .for_each(|(tensor_list_entry, entry)| *tensor_list_entry = entry);
137        tensor_list
138    }
139}
140
141impl<T, const N: usize> IntoIterator for TensorList<T, N>
142where
143    T: Tensor,
144{
145    type Item = T;
146    type IntoIter = array::IntoIter<Self::Item, N>;
147    fn into_iter(self) -> Self::IntoIter {
148        self.0.into_iter()
149    }
150}
151
152impl<'a, T, const N: usize> IntoIterator for &'a TensorList<T, N>
153where
154    T: Tensor,
155{
156    type Item = &'a T;
157    type IntoIter = slice::Iter<'a, T>;
158    fn into_iter(self) -> Self::IntoIter {
159        self.0.iter()
160    }
161}
162
163impl<T, const N: usize> Sum for TensorList<T, N>
164where
165    T: Tensor,
166{
167    fn sum<Ii>(iter: Ii) -> Self
168    where
169        Ii: Iterator<Item = Self>,
170    {
171        iter.reduce(|mut acc, item| {
172            acc += item;
173            acc
174        })
175        .unwrap_or_else(Self::default)
176    }
177}
178
179impl<T, const N: usize> TensorArray for TensorList<T, N>
180where
181    T: Tensor + TensorArray,
182{
183    type Array = [T::Array; N];
184    type Item = T;
185    fn as_array(&self) -> Self::Array {
186        from_fn(|i| self[i].as_array())
187    }
188    fn identity() -> Self {
189        Self(from_fn(|_| Self::Item::identity()))
190    }
191    fn zero() -> Self {
192        Self(from_fn(|_| Self::Item::zero()))
193    }
194}
195
196impl<T, const N: usize> Div<TensorRank0> for TensorList<T, N>
197where
198    T: Tensor,
199{
200    type Output = Self;
201    fn div(mut self, tensor_rank_0: TensorRank0) -> Self::Output {
202        self /= tensor_rank_0;
203        self
204    }
205}
206
207impl<T, const N: usize> Div<&TensorRank0> for TensorList<T, N>
208where
209    T: Tensor,
210{
211    type Output = Self;
212    fn div(mut self, tensor_rank_0: &TensorRank0) -> Self::Output {
213        self /= tensor_rank_0;
214        self
215    }
216}
217
218impl<T, const N: usize> DivAssign<TensorRank0> for TensorList<T, N>
219where
220    T: Tensor,
221{
222    fn div_assign(&mut self, tensor_rank_0: TensorRank0) {
223        self.iter_mut().for_each(|entry| *entry /= &tensor_rank_0);
224    }
225}
226
227impl<T, const N: usize> DivAssign<&TensorRank0> for TensorList<T, N>
228where
229    T: Tensor,
230{
231    fn div_assign(&mut self, tensor_rank_0: &TensorRank0) {
232        self.iter_mut().for_each(|entry| *entry /= tensor_rank_0);
233    }
234}
235
236impl<T, const N: usize, V> Mul<Quantity<V>> for TensorList<T, N>
237where
238    T: Mul<Quantity<V>> + Tensor,
239    <T as Mul<Quantity<V>>>::Output: Tensor,
240{
241    type Output = TensorList<<T as Mul<Quantity<V>>>::Output, N>;
242    fn mul(self, quantity: Quantity<V>) -> Self::Output {
243        self.into_iter().map(|entry| entry * quantity).collect()
244    }
245}
246
247impl<T, const N: usize, V> Mul<Quantity<V>> for &TensorList<T, N>
248where
249    T: Clone + Mul<Quantity<V>> + Tensor,
250    <T as Mul<Quantity<V>>>::Output: Tensor,
251{
252    type Output = TensorList<<T as Mul<Quantity<V>>>::Output, N>;
253    fn mul(self, quantity: Quantity<V>) -> Self::Output {
254        self.iter().map(|entry| entry.clone() * quantity).collect()
255    }
256}
257
258impl<T, const N: usize, V> Div<Quantity<V>> for TensorList<T, N>
259where
260    T: Div<Quantity<V>> + Tensor,
261    <T as Div<Quantity<V>>>::Output: Tensor,
262{
263    type Output = TensorList<<T as Div<Quantity<V>>>::Output, N>;
264    fn div(self, quantity: Quantity<V>) -> Self::Output {
265        self.into_iter().map(|entry| entry / quantity).collect()
266    }
267}
268
269impl<T, const N: usize> Mul<TensorRank0> for TensorList<T, N>
270where
271    T: Tensor,
272{
273    type Output = Self;
274    fn mul(mut self, tensor_rank_0: TensorRank0) -> Self::Output {
275        self *= tensor_rank_0;
276        self
277    }
278}
279
280impl<T, const N: usize> Mul<&TensorRank0> for TensorList<T, N>
281where
282    T: Tensor,
283{
284    type Output = Self;
285    fn mul(mut self, tensor_rank_0: &TensorRank0) -> Self::Output {
286        self *= tensor_rank_0;
287        self
288    }
289}
290
291impl<T, const N: usize> Mul<TensorRank0> for &TensorList<T, N>
292where
293    T: Tensor,
294    for<'a> &'a T: Mul<&'a TensorRank0, Output = T>,
295{
296    type Output = TensorList<T, N>;
297    fn mul(self, tensor_rank_0: TensorRank0) -> Self::Output {
298        self.iter().map(|self_i| self_i * &tensor_rank_0).collect()
299    }
300}
301
302impl<T, const N: usize> Mul<&TensorRank0> for &TensorList<T, N>
303where
304    T: Tensor,
305{
306    type Output = TensorList<T, N>;
307    fn mul(self, tensor_rank_0: &TensorRank0) -> Self::Output {
308        //
309        // Cloning for now to avoid trait recursion nightmare.
310        //
311        self.clone() * tensor_rank_0
312    }
313}
314
315impl<T, const N: usize> MulAssign<TensorRank0> for TensorList<T, N>
316where
317    T: Tensor,
318{
319    fn mul_assign(&mut self, tensor_rank_0: TensorRank0) {
320        self.iter_mut().for_each(|entry| *entry *= &tensor_rank_0);
321    }
322}
323
324impl<T, const N: usize> MulAssign<&TensorRank0> for TensorList<T, N>
325where
326    T: Tensor,
327{
328    fn mul_assign(&mut self, tensor_rank_0: &TensorRank0) {
329        self.iter_mut().for_each(|entry| *entry *= tensor_rank_0);
330    }
331}
332
333impl<T, const N: usize> Add for TensorList<T, N>
334where
335    T: Tensor,
336{
337    type Output = Self;
338    fn add(mut self, tensor_list: Self) -> Self::Output {
339        self += tensor_list;
340        self
341    }
342}
343
344impl<T, const N: usize> Add<&Self> for TensorList<T, N>
345where
346    T: Tensor,
347{
348    type Output = Self;
349    fn add(mut self, tensor_list: &Self) -> Self::Output {
350        self += tensor_list;
351        self
352    }
353}
354
355impl<T, const N: usize> AddAssign for TensorList<T, N>
356where
357    T: Tensor,
358{
359    fn add_assign(&mut self, tensor_list: Self) {
360        self.iter_mut()
361            .zip(tensor_list)
362            .for_each(|(self_entry, entry)| *self_entry += entry);
363    }
364}
365
366impl<T, const N: usize> AddAssign<&Self> for TensorList<T, N>
367where
368    T: Tensor,
369{
370    fn add_assign(&mut self, tensor_list: &Self) {
371        self.iter_mut()
372            .zip(tensor_list.iter())
373            .for_each(|(self_entry, entry)| *self_entry += entry);
374    }
375}
376
377impl<T, const N: usize> Sub for TensorList<T, N>
378where
379    T: Tensor,
380{
381    type Output = Self;
382    fn sub(mut self, tensor_list: Self) -> Self::Output {
383        self -= tensor_list;
384        self
385    }
386}
387
388impl<T, const N: usize> Sub<&Self> for TensorList<T, N>
389where
390    T: Tensor,
391{
392    type Output = Self;
393    fn sub(mut self, tensor_list: &Self) -> Self::Output {
394        self -= tensor_list;
395        self
396    }
397}
398
399impl<T, const N: usize> SubAssign for TensorList<T, N>
400where
401    T: Tensor,
402{
403    fn sub_assign(&mut self, tensor_list: Self) {
404        self.iter_mut()
405            .zip(tensor_list)
406            .for_each(|(self_entry, entry)| *self_entry -= entry);
407    }
408}
409
410impl<T, const N: usize> SubAssign<&Self> for TensorList<T, N>
411where
412    T: Tensor,
413{
414    fn sub_assign(&mut self, tensor_list: &Self) {
415        self.iter_mut()
416            .zip(tensor_list.iter())
417            .for_each(|(self_entry, entry)| *self_entry -= entry);
418    }
419}
420
421impl<T, V, const N: usize> ContractWith<TensorList<V, N>> for TensorList<T, N>
422where
423    T: ContractWith<V> + Tensor,
424    V: Tensor,
425    <T as ContractWith<V>>::Output: Sum,
426{
427    type Output = <T as ContractWith<V>>::Output;
428    fn contract_with(&self, tensor_list: &TensorList<V, N>) -> Self::Output {
429        self.iter()
430            .zip(tensor_list.iter())
431            .map(|(entry, other)| entry.contract_with(other))
432            .sum()
433    }
434}
435
436impl<E, T, const N: usize> Differentiate<T> for TensorList<E, N>
437where
438    E: Differentiate<T> + Tensor,
439    <E as Differentiate<T>>::Derivative: Tensor,
440{
441    type Derivative = TensorList<<E as Differentiate<T>>::Derivative, N>;
442}
conspire/math/tensor/list/mod.rs

conspire/math/tensor/list/
mod.rs
