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#[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 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}