Skip to main content

conspire/math/tensor/tuple/
mod.rs

1pub(crate) mod list;
2pub(crate) mod vec;
3
4use crate::math::{
5    Hessian, Jacobian, Solution, SquareMatrix, Tensor, TensorRank0, TensorRank2, TensorRank4,
6    Vector,
7};
8use std::{
9    fmt::{Display, Formatter, Result},
10    iter::Sum,
11    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign},
12};
13
14/// A fixed-size nested collection of different tensors.
15#[derive(Clone, Debug, PartialEq)]
16pub struct TensorTuple<T1, T2>(pub T1, pub T2)
17where
18    T1: Tensor,
19    T2: Tensor;
20
21impl<T1, T2> Default for TensorTuple<T1, T2>
22where
23    T1: Tensor,
24    T2: Tensor,
25{
26    fn default() -> Self {
27        Self(T1::default(), T2::default())
28    }
29}
30
31impl<T1, T2> From<(T1, T2)> for TensorTuple<T1, T2>
32where
33    T1: Tensor,
34    T2: Tensor,
35{
36    fn from(tuple: (T1, T2)) -> Self {
37        Self(tuple.0, tuple.1)
38    }
39}
40
41impl<'a, T1, T2> From<&'a TensorTuple<T1, T2>> for (&'a T1, &'a T2)
42where
43    T1: Tensor,
44    T2: Tensor,
45{
46    fn from(tensor_tuple: &'a TensorTuple<T1, T2>) -> Self {
47        (&tensor_tuple.0, &tensor_tuple.1)
48    }
49}
50
51impl<T1, T2> From<TensorTuple<T1, T2>> for (T1, T2)
52where
53    T1: Tensor,
54    T2: Tensor,
55{
56    fn from(tensor_tuple: TensorTuple<T1, T2>) -> Self {
57        (tensor_tuple.0, tensor_tuple.1)
58    }
59}
60
61impl<T1, T2> From<Vector> for TensorTuple<T1, T2>
62where
63    T1: Tensor,
64    T2: Tensor,
65{
66    fn from(_vector: Vector) -> Self {
67        unimplemented!()
68    }
69}
70
71impl<T1, T2> Display for TensorTuple<T1, T2>
72where
73    T1: Tensor,
74    T2: Tensor,
75{
76    fn fmt(&self, f: &mut Formatter) -> Result {
77        write!(f, "Need to implement Display")
78    }
79}
80
81impl<T1, T2> Tensor for TensorTuple<T1, T2>
82where
83    T1: Tensor,
84    T2: Tensor,
85{
86    type Item = T1::Item;
87    fn full_contraction(&self, tensor_tuple: &Self) -> TensorRank0 {
88        self.0.full_contraction(&tensor_tuple.0) + self.1.full_contraction(&tensor_tuple.1)
89    }
90    fn iter(&self) -> impl Iterator<Item = &Self::Item> {
91        if self.size() == 0 {
92            self.0.iter()
93        } else {
94            unimplemented!()
95        }
96    }
97    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Self::Item> {
98        if self.size() == 0 {
99            self.0.iter_mut()
100        } else {
101            unimplemented!()
102        }
103    }
104    fn len(&self) -> usize {
105        unimplemented!()
106    }
107    fn norm_inf(&self) -> TensorRank0 {
108        self.0.norm_inf().max(self.1.norm_inf())
109    }
110    fn norm_l1(&self) -> TensorRank0 {
111        self.0.norm_l1() + self.1.norm_l1()
112    }
113    fn norm_p_sum(&self, p: TensorRank0) -> TensorRank0 {
114        self.0.norm_p_sum(p) + self.1.norm_p_sum(p)
115    }
116    fn size(&self) -> usize {
117        self.0.size() + self.1.size()
118    }
119}
120
121impl<const D: usize, const I: usize, const J: usize, const K: usize, const L: usize> Hessian
122    for TensorTuple<
123        TensorRank4<D, I, J, I, J>,
124        TensorTuple<
125            TensorRank4<D, K, L, I, J>,
126            TensorTuple<TensorRank4<D, I, J, K, L>, TensorRank4<D, K, L, K, L>>,
127        >,
128    >
129{
130    fn entry(&self, row: usize, column: usize) -> TensorRank0 {
131        let offset = D * D;
132        match (row < offset, column < offset) {
133            (true, true) => self.0.entry(row, column),
134            (false, true) => self.1.0.entry(row - offset, column),
135            (true, false) => self.1.1.0.entry(row, column - offset),
136            (false, false) => self.1.1.1.entry(row - offset, column - offset),
137        }
138    }
139    fn fill_into(self, square_matrix: &mut SquareMatrix) {
140        let offset = D * D;
141        let (tangent_0, tangent_123) = self.into();
142        let (tangent_1, tangent_23) = tangent_123.into();
143        let (tangent_2, tangent_3) = tangent_23.into();
144
145        tangent_0.into_iter().zip(tangent_1.into_iter().zip(tangent_2.into_iter().zip(tangent_3))).enumerate()
146            .for_each(|(i, (tangent_0_i, (tangent_1_i, (tangent_2_i, tangent_3_i))))| {
147                tangent_0_i.into_iter().zip(tangent_1_i.into_iter().zip(tangent_2_i.into_iter().zip(tangent_3_i))).enumerate()
148                    .for_each(|(j, (tangent_0_ij, (tangent_1_ij, (tangent_2_ij, tangent_3_ij))))| {
149                        tangent_0_ij.into_iter().zip(tangent_1_ij.into_iter().zip(tangent_2_ij.into_iter().zip(tangent_3_ij))).enumerate()
150                            .for_each(|(k, (tangent_0_ijk, (tangent_1_ijk, (tangent_2_ijk, tangent_3_ijk))))| {
151                                tangent_0_ijk.into_iter().zip(tangent_1_ijk.into_iter().zip(tangent_2_ijk.into_iter().zip(tangent_3_ijk))).enumerate()
152                                    .for_each(|(l, (tangent_0_ijkl, (tangent_1_ijkl, (tangent_2_ijkl, tangent_3_ijkl))))| {
153                                        square_matrix[D * i + j][D * k + l] = tangent_0_ijkl;
154                                        square_matrix[offset + D * i + j][D * k + l] = tangent_1_ijkl;
155                                        square_matrix[D * i + j][offset + D * k + l] = tangent_2_ijkl;
156                                        square_matrix[offset + D * i + j][offset + D * k + l] = tangent_3_ijkl;
157                                    })
158                            })
159                    })
160            })
161    }
162}
163
164impl<const D: usize, const I: usize, const J: usize, const K: usize, const L: usize> Jacobian
165    for TensorTuple<TensorRank2<D, I, J>, TensorRank2<D, K, L>>
166{
167    fn fill_into(self, vector: &mut Vector) {
168        self.0
169            .into_iter()
170            .flatten()
171            .chain(self.1.into_iter().flatten())
172            .zip(vector.iter_mut())
173            .for_each(|(self_i, vector_i)| *vector_i = self_i)
174    }
175    fn fill_into_chained(self, other: Vector, vector: &mut Vector) {
176        self.0
177            .into_iter()
178            .flatten()
179            .chain(self.1.into_iter().flatten())
180            .chain(other)
181            .zip(vector.iter_mut())
182            .for_each(|(self_i, vector_i)| *vector_i = self_i)
183    }
184}
185
186impl<const D: usize, const I: usize, const J: usize, const K: usize, const L: usize> Solution
187    for TensorTuple<TensorRank2<D, I, J>, TensorRank2<D, K, L>>
188{
189    fn decrement_from(&mut self, other: &Vector) {
190        self.0
191            .iter_mut()
192            .flat_map(|x| x.iter_mut())
193            .chain(self.1.iter_mut().flat_map(|x| x.iter_mut()))
194            .zip(other.iter())
195            .for_each(|(self_i, vector_i)| *self_i -= vector_i)
196    }
197    fn decrement_from_chained(&mut self, other: &mut Vector, vector: Vector) {
198        self.0
199            .iter_mut()
200            .flat_map(|x| x.iter_mut())
201            .chain(self.1.iter_mut().flat_map(|x| x.iter_mut()))
202            .chain(other.iter_mut())
203            .zip(vector)
204            .for_each(|(entry_i, vector_i)| *entry_i -= vector_i)
205    }
206}
207
208impl<T1, T2> Sum for TensorTuple<T1, T2>
209where
210    T1: Tensor,
211    T2: Tensor,
212{
213    fn sum<Ii>(iter: Ii) -> Self
214    where
215        Ii: Iterator<Item = Self>,
216    {
217        iter.reduce(|mut acc, item| {
218            acc.0 += item.0;
219            acc.1 += item.1;
220            acc
221        })
222        .unwrap_or_else(Self::default)
223    }
224}
225
226impl<T1, T2> Div<TensorRank0> for TensorTuple<T1, T2>
227where
228    T1: Tensor,
229    T2: Tensor,
230{
231    type Output = Self;
232    fn div(mut self, tensor_rank_0: TensorRank0) -> Self::Output {
233        self /= tensor_rank_0;
234        self
235    }
236}
237
238impl<T1, T2> Div<&TensorRank0> for TensorTuple<T1, T2>
239where
240    T1: Tensor,
241    T2: Tensor,
242{
243    type Output = Self;
244    fn div(mut self, tensor_rank_0: &TensorRank0) -> Self::Output {
245        self /= tensor_rank_0;
246        self
247    }
248}
249
250impl<T1, T2> DivAssign<TensorRank0> for TensorTuple<T1, T2>
251where
252    T1: Tensor,
253    T2: Tensor,
254{
255    fn div_assign(&mut self, tensor_rank_0: TensorRank0) {
256        self.0 /= &tensor_rank_0;
257        self.1 /= tensor_rank_0;
258    }
259}
260
261impl<T1, T2> DivAssign<&TensorRank0> for TensorTuple<T1, T2>
262where
263    T1: Tensor,
264    T2: Tensor,
265{
266    fn div_assign(&mut self, tensor_rank_0: &TensorRank0) {
267        self.0 /= tensor_rank_0;
268        self.1 /= tensor_rank_0;
269    }
270}
271
272impl<T1, T2> Mul<TensorRank0> for TensorTuple<T1, T2>
273where
274    T1: Tensor,
275    T2: Tensor,
276{
277    type Output = Self;
278    fn mul(mut self, tensor_rank_0: TensorRank0) -> Self::Output {
279        self *= tensor_rank_0;
280        self
281    }
282}
283
284impl<T1, T2> Mul<&TensorRank0> for TensorTuple<T1, T2>
285where
286    T1: Tensor,
287    T2: Tensor,
288{
289    type Output = Self;
290    fn mul(mut self, tensor_rank_0: &TensorRank0) -> Self::Output {
291        self *= tensor_rank_0;
292        self
293    }
294}
295
296impl<T1, T2> Mul<TensorRank0> for &TensorTuple<T1, T2>
297where
298    T1: Tensor,
299    T2: Tensor,
300{
301    type Output = TensorTuple<T1, T2>;
302    fn mul(self, tensor_rank_0: TensorRank0) -> Self::Output {
303        //
304        // Cloning for now to avoid trait recursion nightmare.
305        //
306        TensorTuple(
307            self.0.clone() * tensor_rank_0,
308            self.1.clone() * tensor_rank_0,
309        )
310    }
311}
312
313impl<T1, T2> MulAssign<TensorRank0> for TensorTuple<T1, T2>
314where
315    T1: Tensor,
316    T2: Tensor,
317{
318    fn mul_assign(&mut self, tensor_rank_0: TensorRank0) {
319        self.0 *= &tensor_rank_0;
320        self.1 *= tensor_rank_0;
321    }
322}
323
324impl<T1, T2> MulAssign<&TensorRank0> for TensorTuple<T1, T2>
325where
326    T1: Tensor,
327    T2: Tensor,
328{
329    fn mul_assign(&mut self, tensor_rank_0: &TensorRank0) {
330        self.0 *= tensor_rank_0;
331        self.1 *= tensor_rank_0;
332    }
333}
334
335impl<T1, T2> Add for TensorTuple<T1, T2>
336where
337    T1: Tensor,
338    T2: Tensor,
339{
340    type Output = Self;
341    fn add(mut self, tensor_tuple: Self) -> Self::Output {
342        self += tensor_tuple;
343        self
344    }
345}
346
347impl<T1, T2> Add<&Self> for TensorTuple<T1, T2>
348where
349    T1: Tensor,
350    T2: Tensor,
351{
352    type Output = Self;
353    fn add(mut self, tensor_tuple: &Self) -> Self::Output {
354        self += tensor_tuple;
355        self
356    }
357}
358
359impl<T1, T2> AddAssign for TensorTuple<T1, T2>
360where
361    T1: Tensor,
362    T2: Tensor,
363{
364    fn add_assign(&mut self, tensor_tuple: Self) {
365        self.0 += tensor_tuple.0;
366        self.1 += tensor_tuple.1;
367    }
368}
369
370impl<T1, T2> AddAssign<&Self> for TensorTuple<T1, T2>
371where
372    T1: Tensor,
373    T2: Tensor,
374{
375    fn add_assign(&mut self, tensor_tuple: &Self) {
376        self.0 += &tensor_tuple.0;
377        self.1 += &tensor_tuple.1;
378    }
379}
380
381impl<T1, T2> Sub for TensorTuple<T1, T2>
382where
383    T1: Tensor,
384    T2: Tensor,
385{
386    type Output = Self;
387    fn sub(mut self, tensor_tuple: Self) -> Self::Output {
388        self -= tensor_tuple;
389        self
390    }
391}
392
393impl<T1, T2> Sub<&Self> for TensorTuple<T1, T2>
394where
395    T1: Tensor,
396    T2: Tensor,
397{
398    type Output = Self;
399    fn sub(mut self, tensor_tuple: &Self) -> Self::Output {
400        self -= tensor_tuple;
401        self
402    }
403}
404
405impl<T1, T2> Sub for &TensorTuple<T1, T2>
406where
407    T1: Tensor,
408    T2: Tensor,
409{
410    type Output = TensorTuple<T1, T2>;
411    fn sub(self, _tensor_tuple: Self) -> Self::Output {
412        unimplemented!("Avoiding trait recursion nightmare")
413    }
414}
415
416impl<T1, T2> SubAssign for TensorTuple<T1, T2>
417where
418    T1: Tensor,
419    T2: Tensor,
420{
421    fn sub_assign(&mut self, tensor_tuple: Self) {
422        self.0 -= tensor_tuple.0;
423        self.1 -= tensor_tuple.1;
424    }
425}
426
427impl<T1, T2> SubAssign<&Self> for TensorTuple<T1, T2>
428where
429    T1: Tensor,
430    T2: Tensor,
431{
432    fn sub_assign(&mut self, tensor_tuple: &Self) {
433        self.0 -= &tensor_tuple.0;
434        self.1 -= &tensor_tuple.1;
435    }
436}
437
438impl<const D: usize, const I: usize, const J: usize, const K: usize, const L: usize> Sub<Vector>
439    for TensorTuple<TensorRank2<D, I, J>, TensorRank2<D, K, L>>
440{
441    type Output = Self;
442    fn sub(mut self, vector: Vector) -> Self::Output {
443        self.0 = self.0 - vector.iter().take(D * D).copied().collect::<Vector>();
444        self.1 = self.1 - vector.iter().skip(D * D).copied().collect::<Vector>();
445        self
446    }
447}
448
449impl<const D: usize, const I: usize, const J: usize, const K: usize, const L: usize> Sub<&Vector>
450    for TensorTuple<TensorRank2<D, I, J>, TensorRank2<D, K, L>>
451{
452    type Output = Self;
453    fn sub(mut self, vector: &Vector) -> Self::Output {
454        self.0 = self.0 - vector.iter().take(D * D).copied().collect::<Vector>();
455        self.1 = self.1 - vector.iter().skip(D * D).copied().collect::<Vector>();
456        self
457    }
458}
459
460impl<T0, T1, T4, T5> Div<TensorTuple<T0, T1>> for &TensorTuple<T4, T5>
461where
462    T0: Tensor,
463    T1: Tensor,
464    T4: Tensor,
465    T5: Tensor,
466{
467    type Output = TensorTuple<T4, T5>;
468    fn div(self, _tensor_tuple: TensorTuple<T0, T1>) -> Self::Output {
469        unimplemented!()
470    }
471}