tnc/tensornetwork/tensor.rs
1use std::iter::zip;
2use std::num::TryFromIntError;
3use std::ops::{BitAnd, BitOr, BitXor, BitXorAssign, Sub};
4
5use float_cmp::{ApproxEq, F64Margin};
6use rustc_hash::FxHashMap;
7use serde::{Deserialize, Serialize};
8
9use crate::tensornetwork::tensordata::TensorData;
10use crate::utils::datastructures::UnionFind;
11
12/// Unique index of a leg.
13pub type EdgeIndex = usize;
14
15/// Index of a tensor in a tensor network.
16pub type TensorIndex = usize;
17
18/// Abstract representation of a tensor. Can be a *composite* tensor (that is, a
19/// tensor network) or a *leaf* tensor.
20#[derive(Default, Debug, Clone, Serialize, Deserialize)]
21pub struct Tensor {
22 /// The inner tensors that make up this tensor. If non-empty, this tensor is
23 /// called a *composite* tensor.
24 pub(crate) tensors: Vec<Tensor>,
25
26 /// The legs of the tensor. Each leg should have a unique id. Connected tensors
27 /// are recognized by having at least one leg id in common.
28 pub(crate) legs: Vec<EdgeIndex>,
29
30 /// The bond dimensions of the legs, same length and order as `legs`. It is
31 /// assumed (but not checked!) that the bond dimensions of different tensors that
32 /// connect to the same leg match.
33 pub(crate) bond_dims: Vec<u64>,
34
35 /// The data of the tensor.
36 pub(crate) tensordata: TensorData,
37}
38
39impl Tensor {
40 /// Constructs a Tensor object with the given `legs` (edge ids) and corresponding
41 /// `bond_dims`. The tensor doesn't have underlying data.
42 #[inline]
43 pub(crate) fn new(legs: Vec<EdgeIndex>, bond_dims: Vec<u64>) -> Self {
44 assert_eq!(legs.len(), bond_dims.len());
45 Self {
46 legs,
47 bond_dims,
48 ..Default::default()
49 }
50 }
51
52 /// Constructs a Tensor using with the given edge ids and a mapping of edge ids
53 /// to corresponding bond dimension.
54 ///
55 /// # Examples
56 /// ```
57 /// # use tnc::tensornetwork::tensor::Tensor;
58 /// # use rustc_hash::FxHashMap;
59 /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6)]);
60 /// let tensor = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
61 /// assert_eq!(tensor.legs(), &[1, 2, 3]);
62 /// assert_eq!(tensor.bond_dims(), &[2, 4, 6]);
63 /// ```
64 #[inline]
65 pub fn new_from_map(legs: Vec<EdgeIndex>, bond_dims_map: &FxHashMap<EdgeIndex, u64>) -> Self {
66 let bond_dims = legs.iter().map(|l| bond_dims_map[l]).collect();
67 Self::new(legs, bond_dims)
68 }
69
70 /// Constructs a Tensor with the given edge ids and the same bond dimension for
71 /// all edges.
72 ///
73 /// # Examples
74 /// ```
75 /// # use tnc::tensornetwork::tensor::Tensor;
76 /// let tensor = Tensor::new_from_const(vec![1, 2, 3], 2);
77 /// assert_eq!(tensor.legs(), &[1, 2, 3]);
78 /// assert_eq!(tensor.bond_dims(), &[2, 2, 2]);
79 /// ```
80 #[inline]
81 pub fn new_from_const(legs: Vec<EdgeIndex>, bond_dim: u64) -> Self {
82 let bond_dims = vec![bond_dim; legs.len()];
83 Self::new(legs, bond_dims)
84 }
85
86 /// Creates a new composite tensor with the given nested tensors.
87 #[inline]
88 pub fn new_composite(tensors: Vec<Self>) -> Self {
89 Self {
90 tensors,
91 ..Default::default()
92 }
93 }
94
95 /// Returns edge ids of Tensor object.
96 ///
97 /// # Examples
98 /// ```
99 /// # use tnc::tensornetwork::tensor::Tensor;
100 /// let tensor = Tensor::new_from_const(vec![1, 2, 3], 3);
101 /// assert_eq!(tensor.legs(), &[1, 2, 3]);
102 /// ```
103 #[inline]
104 pub fn legs(&self) -> &Vec<EdgeIndex> {
105 &self.legs
106 }
107
108 /// Internal method to set legs. Needs pub(crate) for contraction order finding for hierarchies.
109 #[inline]
110 pub(crate) fn set_legs(&mut self, legs: Vec<EdgeIndex>) {
111 self.legs = legs;
112 }
113
114 /// Returns an iterator of tuples of leg ids and their corresponding bond size.
115 #[inline]
116 pub fn edges(&self) -> impl Iterator<Item = (&EdgeIndex, &u64)> + '_ {
117 std::iter::zip(&self.legs, &self.bond_dims)
118 }
119
120 /// Returns the nested tensors of a composite tensor.
121 ///
122 /// # Examples
123 ///
124 /// ```
125 /// # use tnc::tensornetwork::tensor::Tensor;
126 /// # use rustc_hash::FxHashMap;
127 /// let bond_dims = FxHashMap::from_iter([(0, 17), (1, 19), (2, 8)]);
128 /// let v1 = Tensor::new_from_map(vec![0, 1], &bond_dims);
129 /// let v2 = Tensor::new_from_map(vec![1, 2], &bond_dims);
130 /// let tn = Tensor::new_composite(vec![v1.clone(), v2.clone()]);
131 /// for (tensor, ref_tensor) in std::iter::zip(tn.tensors(), vec![v1, v2]){
132 /// assert_eq!(tensor.legs(), ref_tensor.legs());
133 /// }
134 /// ```
135 #[inline]
136 pub fn tensors(&self) -> &Vec<Self> {
137 &self.tensors
138 }
139
140 /// Gets a nested `Tensor` based on the `nested_indices` which specify the index
141 /// of the tensor at each level of the hierarchy.
142 ///
143 /// # Examples
144 /// ```
145 /// # use tnc::tensornetwork::tensor::Tensor;
146 /// # use rustc_hash::FxHashMap;
147 /// let bond_dims = FxHashMap::from_iter([(0, 17), (1, 19), (2, 8), (3, 2), (4, 1)]);
148 /// let mut v1 = Tensor::new_from_map(vec![0, 1], &bond_dims);
149 /// let mut v2 = Tensor::new_from_map(vec![1, 2], &bond_dims);
150 /// let mut v3 = Tensor::new_from_map(vec![2, 3], &bond_dims);
151 /// let mut v4 = Tensor::new_from_map(vec![3, 4], &bond_dims);
152 /// let tn1 = Tensor::new_composite(vec![v1, v2]);
153 /// let tn2 = Tensor::new_composite(vec![v3.clone(), v4]);
154 /// let nested_tn = Tensor::new_composite(vec![tn1, tn2]);
155 ///
156 /// assert_eq!(nested_tn.nested_tensor(&[1, 0]).legs(), v3.legs());
157 /// ```
158 pub fn nested_tensor(&self, nested_indices: &[usize]) -> &Tensor {
159 let mut tensor = self;
160 for index in nested_indices {
161 tensor = tensor.tensor(*index);
162 }
163 tensor
164 }
165
166 /// Returns the total number of leaf tensors in the hierarchy.
167 pub fn total_num_tensors(&self) -> usize {
168 if self.is_composite() {
169 self.tensors.iter().map(Self::total_num_tensors).sum()
170 } else {
171 1
172 }
173 }
174
175 /// Get ith Tensor.
176 ///
177 /// # Examples
178 ///
179 /// ```
180 /// # use tnc::tensornetwork::tensor::Tensor;
181 /// # use rustc_hash::FxHashMap;
182 /// let bond_dims = FxHashMap::from_iter([(0, 17), (1, 19), (2, 8)]);
183 /// let v1 = Tensor::new_from_map(vec![0, 1], &bond_dims);
184 /// let v2 = Tensor::new_from_map(vec![1, 2], &bond_dims);
185 /// let tn = Tensor::new_composite(vec![v1.clone(), v2]);
186 /// assert_eq!(tn.tensor(0).legs(), v1.legs());
187 /// ```
188 #[inline]
189 pub fn tensor(&self, i: TensorIndex) -> &Self {
190 &self.tensors[i]
191 }
192
193 /// Getter for bond dimensions.
194 #[inline]
195 pub fn bond_dims(&self) -> &Vec<u64> {
196 assert!(self.is_leaf());
197 &self.bond_dims
198 }
199
200 /// Returns the shape of tensor. This is the same as the bond dimensions, but as
201 /// `usize`. The conversion can fail, hence a [`Result`] is returned.
202 pub fn shape(&self) -> Result<Vec<usize>, TryFromIntError> {
203 self.bond_dims.iter().map(|&dim| dim.try_into()).collect()
204 }
205
206 /// Returns the number of dimensions.
207 ///
208 /// # Examples
209 /// ```
210 /// # use tnc::tensornetwork::tensor::Tensor;
211 /// # use rustc_hash::FxHashMap;
212 /// let bond_dims = FxHashMap::from_iter([(1, 4), (2, 6), (3, 2)]);
213 /// let tensor = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
214 /// assert_eq!(tensor.dims(), 3);
215 /// ```
216 #[inline]
217 pub fn dims(&self) -> usize {
218 self.legs.len()
219 }
220
221 /// Returns the number of elements. This is a f64 to avoid overflow in large
222 /// tensors.
223 ///
224 /// # Examples
225 /// ```
226 /// # use tnc::tensornetwork::tensor::Tensor;
227 /// # use rustc_hash::FxHashMap;
228 /// let bond_dims = FxHashMap::from_iter([(1, 5), (2, 15), (3, 8)]);
229 /// let tensor = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
230 /// assert_eq!(tensor.size(), 600.0);
231 /// ```
232 #[inline]
233 pub fn size(&self) -> f64 {
234 self.bond_dims.iter().map(|v| *v as f64).product()
235 }
236
237 /// Returns true if Tensor is a leaf tensor, without any nested tensors.
238 ///
239 /// # Examples
240 /// ```
241 /// # use tnc::tensornetwork::tensor::Tensor;
242 /// # use rustc_hash::FxHashMap;
243 /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6)]);
244 /// let tensor = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
245 /// assert_eq!(tensor.is_leaf(), true);
246 /// let comp = Tensor::new_composite(vec![tensor]);
247 /// assert_eq!(comp.is_leaf(), false);
248 /// ```
249 #[inline]
250 pub fn is_leaf(&self) -> bool {
251 self.tensors.is_empty()
252 }
253
254 /// Returns true if Tensor is composite.
255 ///
256 /// # Examples
257 /// ```
258 /// # use tnc::tensornetwork::tensor::Tensor;
259 /// # use rustc_hash::FxHashMap;
260 /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6)]);
261 /// let tensor = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
262 /// assert_eq!(tensor.is_composite(), false);
263 /// let comp = Tensor::new_composite(vec![tensor]);
264 /// assert_eq!(comp.is_composite(), true);
265 /// ```
266 #[inline]
267 pub fn is_composite(&self) -> bool {
268 !self.tensors.is_empty()
269 }
270
271 /// Returns true if Tensor is empty. This means, it doesn't have any subtensors,
272 /// has no legs and is doesn't have any data (e.g., is not a scalar).
273 ///
274 /// # Examples
275 /// ```
276 /// # use tnc::tensornetwork::tensor::Tensor;
277 /// let tensor = Tensor::default();
278 /// assert_eq!(tensor.is_empty(), true);
279 /// ```
280 #[inline]
281 pub fn is_empty(&self) -> bool {
282 self.tensors.is_empty()
283 && self.legs.is_empty()
284 && matches!(*self.tensor_data(), TensorData::Uncontracted)
285 }
286
287 /// Pushes additional `tensor` into this tensor, which must be a composite tensor.
288 #[inline]
289 pub fn push_tensor(&mut self, tensor: Self) {
290 assert!(
291 self.legs.is_empty() && matches!(self.tensordata, TensorData::Uncontracted),
292 "Cannot push tensors into a leaf tensor"
293 );
294 self.tensors.push(tensor);
295 }
296
297 /// Pushes additional `tensors` into this tensor, which must be a composite tensor.
298 #[inline]
299 pub fn push_tensors(&mut self, mut tensors: Vec<Self>) {
300 assert!(
301 self.legs.is_empty() && matches!(self.tensordata, TensorData::Uncontracted),
302 "Cannot push tensors into a leaf tensor"
303 );
304 self.tensors.append(&mut tensors);
305 }
306
307 /// Getter for tensor data.
308 #[inline]
309 pub fn tensor_data(&self) -> &TensorData {
310 &self.tensordata
311 }
312
313 /// Setter for tensor data.
314 ///
315 /// # Examples
316 ///
317 /// ```
318 /// # use tnc::tensornetwork::tensor::Tensor;
319 /// # use tnc::tensornetwork::tensordata::TensorData;
320 /// let mut tensor = Tensor::new_from_const(vec![0, 1], 2);
321 /// let tensordata = TensorData::Gate((String::from("x"), vec![], false));
322 /// tensor.set_tensor_data(tensordata);
323 /// ```
324 #[inline]
325 pub fn set_tensor_data(&mut self, tensordata: TensorData) {
326 assert!(
327 self.is_leaf() || matches!(tensordata, TensorData::Uncontracted),
328 "Cannot add data to composite tensor"
329 );
330 self.tensordata = tensordata;
331 }
332
333 /// Returns whether all tensors inside this tensor are connected.
334 /// This only checks the top-level, not recursing into composite tensors.
335 ///
336 /// # Examples
337 /// ```
338 /// # use tnc::tensornetwork::tensor::Tensor;
339 /// # use rustc_hash::FxHashMap;
340 /// // Create a tensor network with two connected tensors
341 /// let bond_dims = FxHashMap::from_iter([(0, 17), (1, 19), (2, 8), (3, 5)]);
342 /// let v1 = Tensor::new_from_map(vec![0, 1], &bond_dims);
343 /// let v2 = Tensor::new_from_map(vec![1, 2], &bond_dims);
344 /// let mut tn = Tensor::new_composite(vec![v1, v2]);
345 /// assert!(tn.is_connected());
346 ///
347 /// // Introduce a new tensor that is not connected
348 /// let v3 = Tensor::new_from_map(vec![3], &bond_dims);
349 /// tn.push_tensor(v3);
350 /// assert!(!tn.is_connected());
351 /// ```
352 pub fn is_connected(&self) -> bool {
353 let num_tensors = self.tensors.len();
354 let mut uf = UnionFind::new(num_tensors);
355
356 for t1_id in 0..num_tensors {
357 for t2_id in (t1_id + 1)..num_tensors {
358 let t1 = &self.tensors[t1_id];
359 let t2 = &self.tensors[t2_id];
360 if !(t1 & t2).legs.is_empty() {
361 uf.union(t1_id, t2_id);
362 }
363 }
364 }
365
366 uf.count_sets() == 1
367 }
368
369 /// Returns `Tensor` with legs in `self` that are not in `other`.
370 ///
371 /// # Examples
372 /// ```
373 /// # use tnc::tensornetwork::tensor::Tensor;
374 /// # use rustc_hash::FxHashMap;
375 /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6), (4, 3), (5, 9)]);
376 /// let tensor1 = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
377 /// let tensor2 = Tensor::new_from_map(vec![4, 2, 5], &bond_dims);
378 /// let diff_tensor = &tensor1 - &tensor2;
379 /// assert_eq!(diff_tensor.legs(), &[1, 3]);
380 /// assert_eq!(diff_tensor.bond_dims(), &[2, 6]);
381 /// ```
382 #[must_use]
383 pub fn difference(&self, other: &Self) -> Self {
384 let mut new_legs = Vec::with_capacity(self.legs.len());
385 let mut new_bond_dims = Vec::with_capacity(new_legs.capacity());
386 for (leg, dim) in self.edges() {
387 if !other.legs.contains(leg) {
388 new_legs.push(*leg);
389 new_bond_dims.push(*dim);
390 }
391 }
392 Self::new(new_legs, new_bond_dims)
393 }
394
395 /// Returns `Tensor` with union of legs in both `self` and `other`.
396 ///
397 /// # Examples
398 /// ```
399 /// # use tnc::tensornetwork::tensor::Tensor;
400 /// # use rustc_hash::FxHashMap;
401 /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6), (4, 3), (5, 9)]);
402 /// let tensor1 = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
403 /// let tensor2 = Tensor::new_from_map(vec![4, 2, 5], &bond_dims);
404 /// let union_tensor = &tensor1 | &tensor2;
405 /// assert_eq!(union_tensor.legs(), &[1, 2, 3, 4, 5]);
406 /// assert_eq!(union_tensor.bond_dims(), &[2, 4, 6, 3, 9]);
407 /// ```
408 #[must_use]
409 pub fn union(&self, other: &Self) -> Self {
410 let mut new_legs = Vec::with_capacity(self.legs.len() + other.legs.len());
411 let mut new_bond_dims = Vec::with_capacity(new_legs.capacity());
412 new_legs.extend_from_slice(&self.legs);
413 new_bond_dims.extend_from_slice(&self.bond_dims);
414 for (leg, dim) in other.edges() {
415 if !self.legs.contains(leg) {
416 new_legs.push(*leg);
417 new_bond_dims.push(*dim);
418 }
419 }
420 Self::new(new_legs, new_bond_dims)
421 }
422
423 /// Returns `Tensor` with intersection of legs in `self` and `other`.
424 ///
425 /// # Examples
426 /// ```
427 /// # use tnc::tensornetwork::tensor::Tensor;
428 /// # use rustc_hash::FxHashMap;
429 /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6), (4, 3), (5, 9)]);
430 /// let tensor1 = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
431 /// let tensor2 = Tensor::new_from_map(vec![4, 2, 5], &bond_dims);
432 /// let intersection_tensor = &tensor1 & &tensor2;
433 /// assert_eq!(intersection_tensor.legs(), &[2]);
434 /// assert_eq!(intersection_tensor.bond_dims(), &[4]);
435 /// ```
436 #[must_use]
437 pub fn intersection(&self, other: &Self) -> Self {
438 let mut new_legs = Vec::with_capacity(self.legs.len().min(other.legs.len()));
439 let mut new_bond_dims = Vec::with_capacity(new_legs.capacity());
440 for (leg, dim) in self.edges() {
441 if other.legs.contains(leg) {
442 new_legs.push(*leg);
443 new_bond_dims.push(*dim);
444 }
445 }
446 Self::new(new_legs, new_bond_dims)
447 }
448
449 /// Returns `Tensor` with symmetrical difference of legs in `self` and `other`.
450 ///
451 /// # Examples
452 /// ```
453 /// # use tnc::tensornetwork::tensor::Tensor;
454 /// # use rustc_hash::FxHashMap;
455 /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6), (4, 3), (5, 9)]);
456 /// let tensor1 = Tensor::new_from_map(vec![1, 2, 3], &bond_dims);
457 /// let tensor2 = Tensor::new_from_map(vec![4, 2, 5], &bond_dims);
458 /// let sym_dif_tensor = &tensor1 ^ &tensor2;
459 /// assert_eq!(sym_dif_tensor.legs(), &[1, 3, 4, 5]);
460 /// assert_eq!(sym_dif_tensor.bond_dims(), &[2, 6, 3, 9]);
461 /// ```
462 #[must_use]
463 pub fn symmetric_difference(&self, other: &Self) -> Self {
464 let mut new_legs = Vec::with_capacity(self.legs.len() + other.legs.len());
465 let mut new_bond_dims = Vec::with_capacity(new_legs.capacity());
466 for (leg, dim) in self.edges() {
467 if !other.legs.contains(leg) {
468 new_legs.push(*leg);
469 new_bond_dims.push(*dim);
470 }
471 }
472 for (leg, dim) in other.edges() {
473 if !self.legs.contains(leg) {
474 new_legs.push(*leg);
475 new_bond_dims.push(*dim);
476 }
477 }
478 Self::new(new_legs, new_bond_dims)
479 }
480
481 /// Get output legs after tensor contraction
482 pub fn external_tensor(&self) -> Tensor {
483 if self.is_leaf() {
484 return self.clone();
485 }
486
487 let mut ext_tensor = Self::default();
488 for tensor in &self.tensors {
489 let new_tensor = if tensor.is_composite() {
490 &tensor.external_tensor()
491 } else {
492 tensor
493 };
494 ext_tensor = &ext_tensor ^ new_tensor;
495 }
496
497 ext_tensor
498 }
499}
500
501impl ApproxEq for &Tensor {
502 type Margin = F64Margin;
503
504 fn approx_eq<M: Into<Self::Margin>>(self, other: Self, margin: M) -> bool {
505 let margin = margin.into();
506 if self.legs != other.legs {
507 return false;
508 }
509 if self.bond_dims != other.bond_dims {
510 return false;
511 }
512 if self.tensors.len() != other.tensors.len() {
513 return false;
514 }
515 for (tensor, other_tensor) in zip(&self.tensors, &other.tensors) {
516 if !tensor.approx_eq(other_tensor, margin) {
517 return false;
518 }
519 }
520
521 self.tensordata.approx_eq(&other.tensordata, margin)
522 }
523}
524
525impl BitOr for &Tensor {
526 type Output = Tensor;
527 #[inline]
528 fn bitor(self, rhs: &Tensor) -> Tensor {
529 self.union(rhs)
530 }
531}
532
533impl BitAnd for &Tensor {
534 type Output = Tensor;
535 #[inline]
536 fn bitand(self, rhs: &Tensor) -> Tensor {
537 self.intersection(rhs)
538 }
539}
540
541impl BitXor for &Tensor {
542 type Output = Tensor;
543 #[inline]
544 fn bitxor(self, rhs: &Tensor) -> Tensor {
545 self.symmetric_difference(rhs)
546 }
547}
548
549impl Sub for &Tensor {
550 type Output = Tensor;
551 #[inline]
552 fn sub(self, rhs: &Tensor) -> Tensor {
553 self.difference(rhs)
554 }
555}
556
557impl BitXorAssign<&Tensor> for Tensor {
558 #[inline]
559 fn bitxor_assign(&mut self, rhs: &Tensor) {
560 *self = self.symmetric_difference(rhs);
561 }
562}
563
564#[cfg(test)]
565mod tests {
566 use super::*;
567
568 use std::iter::zip;
569
570 use rustc_hash::FxHashMap;
571
572 use crate::tensornetwork::tensordata::TensorData;
573
574 macro_rules! assert_matches {
575 ($left:expr, $pattern:pat) => {
576 match $left {
577 $pattern => (),
578 _ => panic!(
579 "Expected pattern {} but got {:?}",
580 stringify!($pattern),
581 $left
582 ),
583 }
584 };
585 }
586
587 #[test]
588 fn test_empty_tensor() {
589 let tensor = Tensor::default();
590 assert!(tensor.tensors.is_empty());
591 assert!(tensor.legs.is_empty());
592 assert!(tensor.bond_dims.is_empty());
593 assert!(tensor.is_empty());
594 }
595
596 #[test]
597 fn test_new() {
598 let tensor = Tensor::new(vec![2, 4, 5], vec![4, 2, 6]);
599 assert_eq!(tensor.legs(), &[2, 4, 5]);
600 assert_eq!(tensor.bond_dims(), &[4, 2, 6]);
601 assert_matches!(tensor.tensor_data(), TensorData::Uncontracted);
602 }
603
604 #[test]
605 fn test_new_from_map() {
606 let bond_dims = FxHashMap::from_iter([(1, 1), (2, 4), (3, 7), (4, 2), (5, 6)]);
607 let tensor = Tensor::new_from_map(vec![2, 4, 5], &bond_dims);
608 assert_eq!(tensor.legs(), &[2, 4, 5]);
609 assert_eq!(tensor.bond_dims(), &[4, 2, 6]);
610 assert_matches!(tensor.tensor_data(), TensorData::Uncontracted);
611 }
612
613 #[test]
614 fn test_new_from_const() {
615 let tensor = Tensor::new_from_const(vec![9, 2, 5, 1], 3);
616 assert_eq!(tensor.legs(), &[9, 2, 5, 1]);
617 assert_eq!(tensor.bond_dims(), &[3, 3, 3, 3]);
618 assert_matches!(tensor.tensor_data(), TensorData::Uncontracted);
619 }
620
621 #[test]
622 fn test_external_tensor() {
623 let bond_dims = FxHashMap::from_iter([
624 (2, 2),
625 (3, 4),
626 (4, 6),
627 (5, 8),
628 (6, 10),
629 (7, 12),
630 (8, 14),
631 (9, 16),
632 ]);
633 let tensor_1 = Tensor::new_from_map(vec![2, 3, 4], &bond_dims);
634 let tensor_2 = Tensor::new_from_map(vec![2, 3, 5], &bond_dims);
635 let tensor_12 = Tensor::new_composite(vec![tensor_1, tensor_2]);
636
637 let tensor_3 = Tensor::new_from_map(vec![6, 7, 8], &bond_dims);
638 let tensor_4 = Tensor::new_from_map(vec![6, 8, 9], &bond_dims);
639 let tensor_34 = Tensor::new_composite(vec![tensor_3, tensor_4]);
640
641 let tensor_1234 = Tensor::new_composite(vec![tensor_12, tensor_34]);
642
643 let external = tensor_1234.external_tensor();
644 assert_eq!(external.legs(), &[4, 5, 7, 9]);
645 assert_eq!(external.bond_dims(), &[6, 8, 12, 16]);
646 }
647
648 #[test]
649 fn test_push_tensor() {
650 let bond_dims =
651 FxHashMap::from_iter([(2, 17), (3, 1), (4, 11), (8, 3), (9, 20), (7, 7), (10, 14)]);
652 let ref_tensor_1 = Tensor::new_from_map(vec![8, 4, 9], &bond_dims);
653 let ref_tensor_2 = Tensor::new_from_map(vec![7, 10, 2], &bond_dims);
654
655 let mut tensor = Tensor::default();
656
657 // Push tensor 1
658 let tensor_1 = Tensor::new_from_map(vec![8, 4, 9], &bond_dims);
659 tensor.push_tensor(tensor_1);
660
661 for (sub_tensor, ref_tensor) in zip(tensor.tensors(), [&ref_tensor_1]) {
662 assert_eq!(sub_tensor.legs(), ref_tensor.legs());
663 assert_eq!(sub_tensor.bond_dims(), ref_tensor.bond_dims());
664 }
665
666 // Push tensor 2
667 let tensor_2 = Tensor::new_from_map(vec![7, 10, 2], &bond_dims);
668 tensor.push_tensor(tensor_2);
669
670 for (sub_tensor, ref_tensor) in zip(tensor.tensors(), [&ref_tensor_1, &ref_tensor_2]) {
671 assert_eq!(sub_tensor.legs(), ref_tensor.legs());
672 assert_eq!(sub_tensor.bond_dims(), ref_tensor.bond_dims());
673 }
674
675 // Test that other fields are unchanged
676 assert_matches!(tensor.tensor_data(), TensorData::Uncontracted);
677 assert!(tensor.legs().is_empty());
678 }
679
680 #[test]
681 #[should_panic(expected = "Cannot push tensors into a leaf tensor")]
682 fn test_push_tensor_to_leaf() {
683 let bond_dims =
684 FxHashMap::from_iter([(2, 17), (3, 1), (4, 11), (8, 3), (9, 20), (7, 7), (10, 14)]);
685 let mut leaf_tensor = Tensor::new_from_map(vec![4, 3, 2], &bond_dims);
686 let pushed_tensor = Tensor::new_from_map(vec![8, 4, 9], &bond_dims);
687 leaf_tensor.push_tensor(pushed_tensor);
688 }
689
690 #[test]
691 fn test_push_tensors() {
692 let bond_dims =
693 FxHashMap::from_iter([(2, 17), (3, 1), (4, 11), (8, 3), (9, 20), (7, 7), (10, 14)]);
694 let ref_tensor_1 = Tensor::new_from_map(vec![4, 3, 2], &bond_dims);
695 let ref_tensor_2 = Tensor::new_from_map(vec![8, 4, 9], &bond_dims);
696 let ref_tensor_3 = Tensor::new_from_map(vec![7, 10, 2], &bond_dims);
697
698 let tensor_1 = Tensor::new_from_map(vec![4, 3, 2], &bond_dims);
699 let tensor_2 = Tensor::new_from_map(vec![8, 4, 9], &bond_dims);
700 let tensor_3 = Tensor::new_from_map(vec![7, 10, 2], &bond_dims);
701 let mut tensor = Tensor::default();
702 tensor.push_tensors(vec![tensor_1, tensor_2, tensor_3]);
703
704 assert_matches!(tensor.tensor_data(), TensorData::Uncontracted);
705
706 for (sub_tensor, ref_tensor) in zip(
707 tensor.tensors(),
708 &vec![ref_tensor_1, ref_tensor_2, ref_tensor_3],
709 ) {
710 assert_eq!(sub_tensor.legs(), ref_tensor.legs());
711 assert_eq!(sub_tensor.bond_dims(), ref_tensor.bond_dims());
712 }
713 }
714
715 #[test]
716 #[should_panic(expected = "Cannot push tensors into a leaf tensor")]
717 fn test_push_tensors_to_leaf() {
718 let bond_dims =
719 FxHashMap::from_iter([(2, 17), (3, 1), (4, 11), (8, 3), (9, 20), (7, 7), (10, 14)]);
720 let mut leaf_tensor = Tensor::new_from_map(vec![4, 3, 2], &bond_dims);
721 let pushed_tensor_1 = Tensor::new_from_map(vec![8, 4, 9], &bond_dims);
722 let pushed_tensor_2 = Tensor::new_from_map(vec![7, 10, 2], &bond_dims);
723
724 leaf_tensor.push_tensors(vec![pushed_tensor_1, pushed_tensor_2]);
725 }
726}