tnc/tensornetwork/
tensor.rs

1use std::iter::zip;
2use std::num::TryFromIntError;
3use std::ops::{BitAnd, BitOr, BitXor, BitXorAssign, Sub};
4
5use approx::AbsDiffEq;
6use bytemuck::{TransparentWrapper, TransparentWrapperAlloc};
7use rustc_hash::FxHashMap;
8use serde::{Deserialize, Serialize};
9
10use crate::tensornetwork::tensordata::TensorData;
11use crate::utils::datastructures::UnionFind;
12
13/// Unique index of a leg.
14pub type EdgeIndex = usize;
15
16/// Index of a tensor in a tensor network.
17pub type TensorIndex = usize;
18
19/// An abstract tensor. This can either be a [`CompositeTensor`] or a [`LeafTensor`].
20#[derive(Debug, Clone, PartialEq, TransparentWrapper, Serialize, Deserialize)]
21#[repr(transparent)]
22pub struct Tensor(TensorRepr);
23
24/// A composite tensor that has other tensors as children, similar to a tensor
25/// network.
26#[derive(Debug, Clone, PartialEq, TransparentWrapper, Serialize, Deserialize)]
27#[repr(transparent)]
28pub struct CompositeTensor(TensorRepr);
29
30/// A single leaf tensor.
31#[derive(Debug, Clone, PartialEq, TransparentWrapper, Serialize, Deserialize)]
32#[repr(transparent)]
33pub struct LeafTensor(TensorRepr);
34
35/// The type of a tensor.
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
37pub enum TensorType {
38    Composite,
39    Leaf,
40}
41
42/// Abstract representation of a tensor.
43#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
44pub struct TensorRepr {
45    /// The type of this tensor.
46    kind: TensorType,
47
48    /// The inner tensors that make up this tensor. If non-empty, this tensor is
49    /// called a *composite* tensor.
50    tensors: Vec<Tensor>,
51
52    /// The legs of the tensor. Each leg should have a unique id. Connected tensors
53    /// are recognized by having at least one leg id in common.
54    legs: Vec<EdgeIndex>,
55
56    /// The bond dimensions of the legs, same length and order as `legs`. It is
57    /// assumed (but not checked!) that the bond dimensions of different tensors that
58    /// connect to the same leg match.
59    bond_dims: Vec<u64>,
60
61    /// The data of the tensor.
62    tensordata: TensorData,
63}
64
65impl Tensor {
66    /// Returns the kind of this tensor.
67    #[inline]
68    pub fn kind(&self) -> TensorType {
69        self.0.kind
70    }
71
72    /// Returns true if the tensor is a leaf tensor, without any nested tensors.
73    ///
74    /// # Examples
75    /// ```
76    /// # use tnc::tensornetwork::tensor::{CompositeTensor, LeafTensor, Tensor};
77    /// # use rustc_hash::FxHashMap;
78    /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6)]);
79    /// let leaf: Tensor = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims).into();
80    /// assert_eq!(leaf.is_leaf(), true);
81    /// let comp: Tensor = CompositeTensor::new(vec![leaf]).into();
82    /// assert_eq!(comp.is_leaf(), false);
83    /// ```
84    #[inline]
85    pub fn is_leaf(&self) -> bool {
86        self.kind() == TensorType::Leaf
87    }
88
89    /// Returns a reference to this tensor as a [`LeafTensor`] if it is one.
90    #[inline]
91    pub fn as_leaf(&self) -> Option<&LeafTensor> {
92        self.is_leaf()
93            .then(|| LeafTensor::wrap_ref(Self::peel_ref(self)))
94    }
95
96    /// Returns this tensor as a [`LeafTensor`] if it is one.
97    #[inline]
98    pub fn into_leaf(self) -> Option<LeafTensor> {
99        self.is_leaf().then(|| LeafTensor::wrap(Self::peel(self)))
100    }
101
102    /// Returns true if the tensor is composite.
103    ///
104    /// # Examples
105    /// ```
106    /// # use tnc::tensornetwork::tensor::{CompositeTensor, LeafTensor, Tensor};
107    /// # use rustc_hash::FxHashMap;
108    /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6)]);
109    /// let leaf: Tensor = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims).into();
110    /// assert_eq!(leaf.is_composite(), false);
111    /// let comp: Tensor = CompositeTensor::new(vec![leaf]).into();
112    /// assert_eq!(comp.is_composite(), true);
113    /// ```
114    #[inline]
115    pub fn is_composite(&self) -> bool {
116        self.kind() == TensorType::Composite
117    }
118
119    /// Returns a reference to this tensor as a [`CompositeTensor`] if it is one.
120    #[inline]
121    pub fn as_composite(&self) -> Option<&CompositeTensor> {
122        self.is_composite()
123            .then(|| CompositeTensor::wrap_ref(Self::peel_ref(self)))
124    }
125
126    /// Returns this tensor as a [`CompositeTensor`] if it is one.
127    #[inline]
128    pub fn into_composite(self) -> Option<CompositeTensor> {
129        self.is_composite()
130            .then(|| CompositeTensor::wrap(Self::peel(self)))
131    }
132}
133
134pub trait TensorList {
135    fn into_tensors(self) -> Vec<Tensor>;
136    fn as_tensors(&self) -> &[Tensor];
137}
138
139impl TensorList for Vec<Tensor> {
140    #[inline]
141    fn into_tensors(self) -> Vec<Tensor> {
142        self
143    }
144
145    #[inline]
146    fn as_tensors(&self) -> &[Tensor] {
147        self
148    }
149}
150
151impl From<LeafTensor> for Tensor {
152    #[inline]
153    fn from(value: LeafTensor) -> Self {
154        Self::wrap(LeafTensor::peel(value))
155    }
156}
157
158impl TensorList for Vec<LeafTensor> {
159    #[inline]
160    fn into_tensors(self) -> Vec<Tensor> {
161        Tensor::wrap_vec(LeafTensor::peel_vec(self))
162    }
163
164    #[inline]
165    fn as_tensors(&self) -> &[Tensor] {
166        Tensor::wrap_slice(LeafTensor::peel_slice(self))
167    }
168}
169
170impl From<CompositeTensor> for Tensor {
171    #[inline]
172    fn from(value: CompositeTensor) -> Self {
173        Self::wrap(CompositeTensor::peel(value))
174    }
175}
176
177impl TensorList for Vec<CompositeTensor> {
178    #[inline]
179    fn into_tensors(self) -> Vec<Tensor> {
180        Tensor::wrap_vec(CompositeTensor::peel_vec(self))
181    }
182
183    #[inline]
184    fn as_tensors(&self) -> &[Tensor] {
185        Tensor::wrap_slice(CompositeTensor::peel_slice(self))
186    }
187}
188
189impl CompositeTensor {
190    /// Creates a new composite tensor with the given nested tensors.
191    #[inline]
192    pub fn new<T>(tensors: T) -> Self
193    where
194        T: TensorList,
195    {
196        Self(TensorRepr {
197            kind: TensorType::Composite,
198            tensors: tensors.into_tensors(),
199            legs: Vec::new(),
200            bond_dims: Vec::new(),
201            tensordata: TensorData::None,
202        })
203    }
204
205    /// Returns a reference to this composite tensor as a [`Tensor`].
206    #[inline]
207    pub fn as_tensor(&self) -> &Tensor {
208        Tensor::wrap_ref(Self::peel_ref(self))
209    }
210
211    /// Returns the children tensors.
212    ///
213    /// # Examples
214    /// ```
215    /// # use tnc::tensornetwork::tensor::{CompositeTensor, LeafTensor, Tensor};
216    /// # use rustc_hash::FxHashMap;
217    /// # use approx::assert_abs_diff_eq;
218    /// let bond_dims = FxHashMap::from_iter([(0, 17), (1, 19), (2, 8)]);
219    /// let v1 = LeafTensor::new_from_map(vec![0, 1], &bond_dims);
220    /// let v2 = LeafTensor::new_from_map(vec![1, 2], &bond_dims);
221    /// let tn = CompositeTensor::new(vec![v1.clone(), v2.clone()]);
222    /// for (tensor, ref_tensor) in std::iter::zip(tn.tensors(), vec![v1, v2]){
223    ///    assert_abs_diff_eq!(tensor, ref_tensor.as_tensor());
224    /// }
225    /// ```
226    #[inline]
227    pub fn tensors(&self) -> &Vec<Tensor> {
228        &self.0.tensors
229    }
230
231    /// Get the ith tensor.
232    ///
233    /// # Examples
234    /// ```
235    /// # use tnc::tensornetwork::tensor::{CompositeTensor, LeafTensor, Tensor};
236    /// # use rustc_hash::FxHashMap;
237    /// # use approx::assert_abs_diff_eq;
238    /// let bond_dims = FxHashMap::from_iter([(0, 17), (1, 19), (2, 8)]);
239    /// let v1 = LeafTensor::new_from_map(vec![0, 1], &bond_dims);
240    /// let v2 = LeafTensor::new_from_map(vec![1, 2], &bond_dims);
241    /// let tn = CompositeTensor::new(vec![v1.clone(), v2]);
242    /// assert_abs_diff_eq!(tn.tensor(0), v1.as_tensor());
243    /// ```
244    #[inline]
245    pub fn tensor(&self, i: TensorIndex) -> &Tensor {
246        &self.0.tensors[i]
247    }
248
249    /// Converts this tensor into the vec of its children.
250    #[inline]
251    pub fn into_tensors(self) -> Vec<Tensor> {
252        self.0.tensors
253    }
254
255    /// Returns true if the tensor is empty. This means, it doesn't have any
256    /// children.
257    ///
258    /// # Examples
259    /// ```
260    /// # use tnc::tensornetwork::tensor::CompositeTensor;
261    /// let tensor = CompositeTensor::default();
262    /// assert_eq!(tensor.is_empty(), true);
263    /// ```
264    #[inline]
265    pub fn is_empty(&self) -> bool {
266        self.0.tensors.is_empty()
267    }
268
269    /// Returns the number of direct children of this tensor.
270    ///
271    /// # Examples
272    /// ```
273    /// # use tnc::tensornetwork::tensor::CompositeTensor;
274    /// let tensor = CompositeTensor::default();
275    /// assert_eq!(tensor.len(), 0);
276    /// ```
277    #[inline]
278    pub fn len(&self) -> usize {
279        self.0.tensors.len()
280    }
281
282    /// Gets a nested [`Tensor`] based on the `nested_indices` which specify the
283    /// index of the tensor at each level of the hierarchy.
284    ///
285    /// # Examples
286    /// ```
287    /// # use tnc::tensornetwork::tensor::{CompositeTensor, LeafTensor};
288    /// # use rustc_hash::FxHashMap;
289    /// let bond_dims = FxHashMap::from_iter([(0, 17), (1, 19), (2, 8), (3, 2), (4, 1)]);
290    /// let mut v1 = LeafTensor::new_from_map(vec![0, 1], &bond_dims);
291    /// let mut v2 = LeafTensor::new_from_map(vec![1, 2], &bond_dims);
292    /// let mut v3 = LeafTensor::new_from_map(vec![2, 3], &bond_dims);
293    /// let mut v4 = LeafTensor::new_from_map(vec![3, 4], &bond_dims);
294    /// let tn1 = CompositeTensor::new(vec![v1, v2]);
295    /// let tn2 = CompositeTensor::new(vec![v3.clone(), v4]);
296    /// let nested_tn = CompositeTensor::new(vec![tn1, tn2]);
297    ///
298    /// let found = nested_tn.nested_tensor(&[1, 0]);
299    /// assert!(found.is_leaf());
300    /// let found = found.as_leaf().unwrap();
301    /// assert_eq!(found.legs(), v3.legs());
302    /// ```
303    pub fn nested_tensor(&self, nested_indices: &[usize]) -> &Tensor {
304        let mut tensor = self.as_tensor();
305        for index in nested_indices {
306            tensor = tensor.as_composite().unwrap().tensor(*index);
307        }
308        tensor
309    }
310
311    /// Returns the total number of leaf tensors in the hierarchy.
312    pub fn total_num_tensors(&self) -> usize {
313        self.0
314            .tensors
315            .iter()
316            .map(|t| match t.kind() {
317                TensorType::Composite => t.as_composite().unwrap().total_num_tensors(),
318                TensorType::Leaf => 1,
319            })
320            .sum()
321    }
322
323    /// Pushes additional `tensor` into this composite tensor.
324    #[inline]
325    pub fn push_tensor<T>(&mut self, tensor: T)
326    where
327        T: Into<Tensor>,
328    {
329        self.0.tensors.push(tensor.into());
330    }
331
332    /// Pushes additional `tensors` into this composite tensor.
333    #[inline]
334    pub fn push_tensors<T>(&mut self, tensors: T)
335    where
336        T: TensorList,
337    {
338        let mut tensors = tensors.into_tensors();
339        self.0.tensors.append(&mut tensors);
340    }
341
342    /// Reserves space for at least `additional` more tensors to be pushed to this
343    /// composite tensor without reallocation.
344    #[inline]
345    pub fn reserve(&mut self, additional: usize) {
346        self.0.tensors.reserve(additional);
347    }
348
349    /// Returns whether all tensors inside this tensor are connected. This currently
350    /// requires all children to be leaf tensors.
351    ///
352    /// # Examples
353    /// ```
354    /// # use tnc::tensornetwork::tensor::{CompositeTensor, LeafTensor};
355    /// # use rustc_hash::FxHashMap;
356    /// // Create a tensor network with two connected tensors
357    /// let bond_dims = FxHashMap::from_iter([(0, 17), (1, 19), (2, 8), (3, 5)]);
358    /// let v1 = LeafTensor::new_from_map(vec![0, 1], &bond_dims);
359    /// let v2 = LeafTensor::new_from_map(vec![1, 2], &bond_dims);
360    /// let mut tn = CompositeTensor::new(vec![v1, v2]);
361    /// assert!(tn.is_connected());
362    ///
363    /// // Introduce a new tensor that is not connected
364    /// let v3 = LeafTensor::new_from_map(vec![3], &bond_dims);
365    /// tn.push_tensor(v3);
366    /// assert!(!tn.is_connected());
367    /// ```
368    pub fn is_connected(&self) -> bool {
369        let num_tensors = self.len();
370        let mut uf = UnionFind::new(num_tensors);
371
372        for t1_id in 0..num_tensors {
373            for t2_id in (t1_id + 1)..num_tensors {
374                let t1 = self
375                    .tensor(t1_id)
376                    .as_leaf()
377                    .expect("Expected all children to be leaves");
378                let t2 = self
379                    .tensor(t2_id)
380                    .as_leaf()
381                    .expect("Expected all children to be leaves");
382                if !(t1 & t2).legs().is_empty() {
383                    uf.union(t1_id, t2_id);
384                }
385            }
386        }
387
388        uf.count_sets() == 1
389    }
390
391    /// Get output legs after tensor contraction.
392    pub fn external_tensor(&self) -> LeafTensor {
393        self.tensors()
394            .iter()
395            .fold(LeafTensor::default(), |acc, tensor| {
396                let tensor = match tensor.kind() {
397                    TensorType::Composite => &tensor.as_composite().unwrap().external_tensor(),
398                    TensorType::Leaf => tensor.as_leaf().unwrap(),
399                };
400                &acc ^ tensor
401            })
402    }
403}
404
405impl Default for CompositeTensor {
406    fn default() -> Self {
407        Self(TensorRepr {
408            kind: TensorType::Composite,
409            tensors: Vec::new(),
410            legs: Vec::new(),
411            bond_dims: Vec::new(),
412            tensordata: TensorData::None,
413        })
414    }
415}
416
417impl AbsDiffEq for CompositeTensor {
418    type Epsilon = f64;
419
420    fn default_epsilon() -> Self::Epsilon {
421        f64::EPSILON
422    }
423
424    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
425        if self.len() != other.len() {
426            return false;
427        }
428        for (tensor, other_tensor) in zip(self.tensors(), other.tensors()) {
429            if !Tensor::abs_diff_eq(tensor, other_tensor, epsilon) {
430                return false;
431            }
432        }
433        true
434    }
435}
436
437impl LeafTensor {
438    /// Constructs a leaf tensor object with the given `legs` (edge ids) and
439    /// corresponding `bond_dims`. The tensor doesn't have underlying data.
440    #[inline]
441    pub(crate) fn new(legs: Vec<EdgeIndex>, bond_dims: Vec<u64>) -> Self {
442        Self::new_with_data(legs, bond_dims, TensorData::None)
443    }
444
445    /// Constructs a leaf tensor object with the given `legs` (edge ids),
446    /// corresponding `bond_dims` and `data`.
447    #[inline]
448    pub(crate) fn new_with_data(
449        legs: Vec<EdgeIndex>,
450        bond_dims: Vec<u64>,
451        data: TensorData,
452    ) -> Self {
453        assert_eq!(legs.len(), bond_dims.len());
454        Self(TensorRepr {
455            kind: TensorType::Leaf,
456            legs,
457            tensors: Vec::new(),
458            bond_dims,
459            tensordata: data,
460        })
461    }
462
463    /// Constructs a leaf tensor with the given edge ids and a mapping of edge ids
464    /// to corresponding bond dimensions.
465    ///
466    /// # Examples
467    /// ```
468    /// # use tnc::tensornetwork::tensor::LeafTensor;
469    /// # use rustc_hash::FxHashMap;
470    /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6)]);
471    /// let tensor = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims);
472    /// assert_eq!(tensor.legs(), &[1, 2, 3]);
473    /// assert_eq!(tensor.bond_dims(), &[2, 4, 6]);
474    /// ```
475    #[inline]
476    pub fn new_from_map(legs: Vec<EdgeIndex>, bond_dims_map: &FxHashMap<EdgeIndex, u64>) -> Self {
477        let bond_dims = legs.iter().map(|l| bond_dims_map[l]).collect();
478        Self::new(legs, bond_dims)
479    }
480
481    /// Constructs a leaf tensor with the given edge ids and the same bond dimension
482    /// for all edges.
483    ///
484    /// # Examples
485    /// ```
486    /// # use tnc::tensornetwork::tensor::LeafTensor;
487    /// let tensor = LeafTensor::new_from_const(vec![1, 2, 3], 2);
488    /// assert_eq!(tensor.legs(), &[1, 2, 3]);
489    /// assert_eq!(tensor.bond_dims(), &[2, 2, 2]);
490    /// ```
491    #[inline]
492    pub fn new_from_const(legs: Vec<EdgeIndex>, bond_dim: u64) -> Self {
493        let bond_dims = vec![bond_dim; legs.len()];
494        Self::new(legs, bond_dims)
495    }
496
497    /// Returns a reference to this leaf tensor as a [`Tensor`].
498    #[inline]
499    pub fn as_tensor(&self) -> &Tensor {
500        Tensor::wrap_ref(Self::peel_ref(self))
501    }
502
503    /// Returns a new leaf tensor with the same legs and bond dimensions as this
504    /// tensor, but without any data.
505    #[inline]
506    pub fn shallow_clone(&self) -> Self {
507        Self::new(self.legs().clone(), self.bond_dims().clone())
508    }
509
510    /// Returns edge ids of the tensor.
511    ///
512    /// # Examples
513    /// ```
514    /// # use tnc::tensornetwork::tensor::LeafTensor;
515    /// let tensor = LeafTensor::new_from_const(vec![1, 2, 3], 3);
516    /// assert_eq!(tensor.legs(), &[1, 2, 3]);
517    /// ```
518    #[inline]
519    pub fn legs(&self) -> &Vec<EdgeIndex> {
520        &self.0.legs
521    }
522
523    /// Returns an iterator of tuples of leg ids and their corresponding bond size.
524    #[inline]
525    pub fn edges(&self) -> impl Iterator<Item = (EdgeIndex, u64)> + '_ {
526        std::iter::zip(
527            self.0.legs.iter().copied(),
528            self.0.bond_dims.iter().copied(),
529        )
530    }
531
532    /// Getter for bond dimensions.
533    #[inline]
534    pub fn bond_dims(&self) -> &Vec<u64> {
535        &self.0.bond_dims
536    }
537
538    /// Returns the shape of tensor. This is the same as the bond dimensions, but as
539    /// `usize`. The conversion can fail, hence a [`Result`] is returned.
540    pub fn shape(&self) -> Result<Vec<usize>, TryFromIntError> {
541        self.0.bond_dims.iter().map(|&dim| dim.try_into()).collect()
542    }
543
544    /// Returns the number of dimensions.
545    ///
546    /// # Examples
547    /// ```
548    /// # use tnc::tensornetwork::tensor::LeafTensor;
549    /// # use rustc_hash::FxHashMap;
550    /// let bond_dims = FxHashMap::from_iter([(1, 4), (2, 6), (3, 2)]);
551    /// let tensor = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims);
552    /// assert_eq!(tensor.dims(), 3);
553    /// ```
554    #[inline]
555    pub fn dims(&self) -> usize {
556        self.0.legs.len()
557    }
558
559    /// Returns the number of elements. This is a f64 to avoid overflow in large
560    /// tensors.
561    ///
562    /// # Examples
563    /// ```
564    /// # use tnc::tensornetwork::tensor::LeafTensor;
565    /// # use rustc_hash::FxHashMap;
566    /// let bond_dims = FxHashMap::from_iter([(1, 5), (2, 15), (3, 8)]);
567    /// let tensor = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims);
568    /// assert_eq!(tensor.size(), 600.0);
569    /// ```
570    #[inline]
571    pub fn size(&self) -> f64 {
572        self.0.bond_dims.iter().map(|v| *v as f64).product()
573    }
574
575    /// Converts this tensor into the leg ids and bond dimensions it is made of.
576    #[inline]
577    pub fn into_legs(self) -> Vec<EdgeIndex> {
578        self.0.legs
579    }
580
581    /// Converts this tensor into the data it contains.
582    #[inline]
583    pub fn into_data(self) -> TensorData {
584        self.0.tensordata
585    }
586
587    /// Converts this tensor into the leg ids, bond dimensions and data it is made
588    /// of.
589    pub fn into_inner(self) -> (Vec<EdgeIndex>, Vec<u64>, TensorData) {
590        (self.0.legs, self.0.bond_dims, self.0.tensordata)
591    }
592
593    /// Getter for tensor data.
594    #[inline]
595    pub fn tensor_data(&self) -> &TensorData {
596        &self.0.tensordata
597    }
598
599    /// Setter for tensor data.
600    ///
601    /// # Examples
602    ///
603    /// ```
604    /// # use tnc::tensornetwork::tensor::LeafTensor;
605    /// # use tnc::tensornetwork::tensordata::TensorData;
606    /// let mut tensor = LeafTensor::new_from_const(vec![0, 1], 2);
607    /// let tensordata = TensorData::Gate((String::from("x"), vec![], false));
608    /// tensor.set_tensor_data(tensordata);
609    /// ```
610    #[inline]
611    pub fn set_tensor_data(&mut self, tensordata: TensorData) {
612        self.0.tensordata = tensordata;
613    }
614
615    /// Returns the tensor with legs in `self` that are not in `other`.
616    ///
617    /// # Examples
618    /// ```
619    /// # use tnc::tensornetwork::tensor::LeafTensor;
620    /// # use rustc_hash::FxHashMap;
621    /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6), (4, 3), (5, 9)]);
622    /// let tensor1 = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims);
623    /// let tensor2 = LeafTensor::new_from_map(vec![4, 2, 5], &bond_dims);
624    /// let diff_tensor = &tensor1 - &tensor2;
625    /// assert_eq!(diff_tensor.legs(), &[1, 3]);
626    /// assert_eq!(diff_tensor.bond_dims(), &[2, 6]);
627    /// ```
628    #[must_use]
629    pub fn difference(&self, other: &Self) -> Self {
630        let mut new_legs = Vec::with_capacity(self.legs().len());
631        let mut new_bond_dims = Vec::with_capacity(new_legs.capacity());
632        for (leg, dim) in self.edges() {
633            if !other.legs().contains(&leg) {
634                new_legs.push(leg);
635                new_bond_dims.push(dim);
636            }
637        }
638        Self::new(new_legs, new_bond_dims)
639    }
640
641    /// Returns the tensor with union of legs in both `self` and `other`.
642    ///
643    /// # Examples
644    /// ```
645    /// # use tnc::tensornetwork::tensor::LeafTensor;
646    /// # use rustc_hash::FxHashMap;
647    /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6), (4, 3), (5, 9)]);
648    /// let tensor1 = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims);
649    /// let tensor2 = LeafTensor::new_from_map(vec![4, 2, 5], &bond_dims);
650    /// let union_tensor = &tensor1 | &tensor2;
651    /// assert_eq!(union_tensor.legs(), &[1, 2, 3, 4, 5]);
652    /// assert_eq!(union_tensor.bond_dims(), &[2, 4, 6, 3, 9]);
653    /// ```
654    #[must_use]
655    pub fn union(&self, other: &Self) -> Self {
656        let mut new_legs = Vec::with_capacity(self.legs().len() + other.legs().len());
657        let mut new_bond_dims = Vec::with_capacity(new_legs.capacity());
658        new_legs.extend_from_slice(self.legs());
659        new_bond_dims.extend_from_slice(self.bond_dims());
660        for (leg, dim) in other.edges() {
661            if !self.legs().contains(&leg) {
662                new_legs.push(leg);
663                new_bond_dims.push(dim);
664            }
665        }
666        Self::new(new_legs, new_bond_dims)
667    }
668
669    /// Returns the tensor with intersection of legs in `self` and `other`.
670    ///
671    /// # Examples
672    /// ```
673    /// # use tnc::tensornetwork::tensor::LeafTensor;
674    /// # use rustc_hash::FxHashMap;
675    /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6), (4, 3), (5, 9)]);
676    /// let tensor1 = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims);
677    /// let tensor2 = LeafTensor::new_from_map(vec![4, 2, 5], &bond_dims);
678    /// let intersection_tensor = &tensor1 & &tensor2;
679    /// assert_eq!(intersection_tensor.legs(), &[2]);
680    /// assert_eq!(intersection_tensor.bond_dims(), &[4]);
681    /// ```
682    #[must_use]
683    pub fn intersection(&self, other: &Self) -> Self {
684        let mut new_legs = Vec::with_capacity(self.legs().len().min(other.legs().len()));
685        let mut new_bond_dims = Vec::with_capacity(new_legs.capacity());
686        for (leg, dim) in self.edges() {
687            if other.legs().contains(&leg) {
688                new_legs.push(leg);
689                new_bond_dims.push(dim);
690            }
691        }
692        Self::new(new_legs, new_bond_dims)
693    }
694
695    /// Returns the tensor with symmetrical difference of legs in `self` and `other`.
696    ///
697    /// # Examples
698    /// ```
699    /// # use tnc::tensornetwork::tensor::LeafTensor;
700    /// # use rustc_hash::FxHashMap;
701    /// let bond_dims = FxHashMap::from_iter([(1, 2), (2, 4), (3, 6), (4, 3), (5, 9)]);
702    /// let tensor1 = LeafTensor::new_from_map(vec![1, 2, 3], &bond_dims);
703    /// let tensor2 = LeafTensor::new_from_map(vec![4, 2, 5], &bond_dims);
704    /// let sym_dif_tensor = &tensor1 ^ &tensor2;
705    /// assert_eq!(sym_dif_tensor.legs(), &[1, 3, 4, 5]);
706    /// assert_eq!(sym_dif_tensor.bond_dims(), &[2, 6, 3, 9]);
707    /// ```
708    #[must_use]
709    pub fn symmetric_difference(&self, other: &Self) -> Self {
710        let mut new_legs = Vec::with_capacity(self.legs().len() + other.legs().len());
711        let mut new_bond_dims = Vec::with_capacity(new_legs.capacity());
712        for (leg, dim) in self.edges() {
713            if !other.legs().contains(&leg) {
714                new_legs.push(leg);
715                new_bond_dims.push(dim);
716            }
717        }
718        for (leg, dim) in other.edges() {
719            if !self.legs().contains(&leg) {
720                new_legs.push(leg);
721                new_bond_dims.push(dim);
722            }
723        }
724        Self::new(new_legs, new_bond_dims)
725    }
726}
727
728impl Default for LeafTensor {
729    fn default() -> Self {
730        Self(TensorRepr {
731            kind: TensorType::Leaf,
732            tensors: Vec::new(),
733            legs: Vec::new(),
734            bond_dims: Vec::new(),
735            tensordata: TensorData::None,
736        })
737    }
738}
739
740impl BitOr for &LeafTensor {
741    type Output = LeafTensor;
742    #[inline]
743    fn bitor(self, rhs: Self) -> Self::Output {
744        self.union(rhs)
745    }
746}
747
748impl BitAnd for &LeafTensor {
749    type Output = LeafTensor;
750    #[inline]
751    fn bitand(self, rhs: Self) -> Self::Output {
752        self.intersection(rhs)
753    }
754}
755
756impl BitXor for &LeafTensor {
757    type Output = LeafTensor;
758    #[inline]
759    fn bitxor(self, rhs: Self) -> Self::Output {
760        self.symmetric_difference(rhs)
761    }
762}
763
764impl Sub for &LeafTensor {
765    type Output = LeafTensor;
766    #[inline]
767    fn sub(self, rhs: Self) -> Self::Output {
768        self.difference(rhs)
769    }
770}
771
772impl BitXorAssign<&LeafTensor> for LeafTensor {
773    #[inline]
774    fn bitxor_assign(&mut self, rhs: &Self) {
775        *self = self.symmetric_difference(rhs);
776    }
777}
778
779impl AbsDiffEq for LeafTensor {
780    type Epsilon = f64;
781
782    fn default_epsilon() -> Self::Epsilon {
783        f64::EPSILON
784    }
785
786    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
787        if self.legs() != other.legs() {
788            return false;
789        }
790        if self.bond_dims() != other.bond_dims() {
791            return false;
792        }
793
794        TensorData::abs_diff_eq(self.tensor_data(), other.tensor_data(), epsilon)
795    }
796}
797
798impl AbsDiffEq for Tensor {
799    type Epsilon = f64;
800
801    fn default_epsilon() -> Self::Epsilon {
802        f64::default_epsilon()
803    }
804
805    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
806        match (self.kind(), other.kind()) {
807            (TensorType::Leaf, TensorType::Leaf) => {
808                let self_leaf = self.as_leaf().unwrap();
809                let other_leaf = other.as_leaf().unwrap();
810                LeafTensor::abs_diff_eq(self_leaf, other_leaf, epsilon)
811            }
812            (TensorType::Composite, TensorType::Composite) => {
813                let self_comp = self.as_composite().unwrap();
814                let other_comp = other.as_composite().unwrap();
815                CompositeTensor::abs_diff_eq(self_comp, other_comp, epsilon)
816            }
817            _ => false,
818        }
819    }
820}
821
822#[cfg(test)]
823mod tests {
824    use super::*;
825
826    use std::iter::zip;
827
828    use rustc_hash::FxHashMap;
829
830    use crate::tensornetwork::tensordata::TensorData;
831
832    macro_rules! assert_matches {
833        ($left:expr, $pattern:pat) => {
834            match $left {
835                $pattern => (),
836                _ => panic!(
837                    "Expected pattern {} but got {:?}",
838                    stringify!($pattern),
839                    $left
840                ),
841            }
842        };
843    }
844
845    mod leaf {
846        use super::*;
847
848        #[test]
849        fn default() {
850            let tensor = LeafTensor::default();
851            assert!(tensor.legs().is_empty());
852            assert!(tensor.bond_dims().is_empty());
853            assert_matches!(tensor.tensor_data(), TensorData::None);
854        }
855
856        #[test]
857        fn new() {
858            let tensor = LeafTensor::new(vec![2, 4, 5], vec![4, 2, 6]);
859            assert_eq!(tensor.legs(), &[2, 4, 5]);
860            assert_eq!(tensor.bond_dims(), &[4, 2, 6]);
861            assert_matches!(tensor.tensor_data(), TensorData::None);
862        }
863
864        #[test]
865        fn new_from_map() {
866            let bond_dims = FxHashMap::from_iter([(1, 1), (2, 4), (3, 7), (4, 2), (5, 6)]);
867            let tensor = LeafTensor::new_from_map(vec![2, 4, 5], &bond_dims);
868            assert_eq!(tensor.legs(), &[2, 4, 5]);
869            assert_eq!(tensor.bond_dims(), &[4, 2, 6]);
870            assert_matches!(tensor.tensor_data(), TensorData::None);
871        }
872
873        #[test]
874        fn new_from_const() {
875            let tensor = LeafTensor::new_from_const(vec![9, 2, 5, 1], 3);
876            assert_eq!(tensor.legs(), &[9, 2, 5, 1]);
877            assert_eq!(tensor.bond_dims(), &[3, 3, 3, 3]);
878            assert_matches!(tensor.tensor_data(), TensorData::None);
879        }
880    }
881
882    mod composite {
883        use super::*;
884
885        #[test]
886        fn default() {
887            let tensor = CompositeTensor::default();
888            assert!(tensor.is_empty());
889            assert_eq!(tensor.len(), 0);
890        }
891
892        #[test]
893        fn external_tensor() {
894            let bond_dims = FxHashMap::from_iter([
895                (2, 2),
896                (3, 4),
897                (4, 6),
898                (5, 8),
899                (6, 10),
900                (7, 12),
901                (8, 14),
902                (9, 16),
903            ]);
904            let tensor_1 = LeafTensor::new_from_map(vec![2, 3, 4], &bond_dims);
905            let tensor_2 = LeafTensor::new_from_map(vec![2, 3, 5], &bond_dims);
906            let tensor_12 = CompositeTensor::new(vec![tensor_1, tensor_2]);
907
908            let tensor_3 = LeafTensor::new_from_map(vec![6, 7, 8], &bond_dims);
909            let tensor_4 = LeafTensor::new_from_map(vec![6, 8, 9], &bond_dims);
910            let tensor_34 = CompositeTensor::new(vec![tensor_3, tensor_4]);
911
912            let tensor_1234 = CompositeTensor::new(vec![tensor_12, tensor_34]);
913
914            let external = tensor_1234.external_tensor();
915            assert_eq!(external.legs(), &[4, 5, 7, 9]);
916            assert_eq!(external.bond_dims(), &[6, 8, 12, 16]);
917        }
918
919        #[test]
920        fn test_push_tensor() {
921            let bond_dims =
922                FxHashMap::from_iter([(2, 17), (3, 1), (4, 11), (8, 3), (9, 20), (7, 7), (10, 14)]);
923            let ref_tensor_1 = LeafTensor::new_from_map(vec![8, 4, 9], &bond_dims);
924            let ref_tensor_2 = LeafTensor::new_from_map(vec![7, 10, 2], &bond_dims);
925
926            let mut tensor = CompositeTensor::default();
927
928            // Push tensor 1
929            let tensor_1 = LeafTensor::new_from_map(vec![8, 4, 9], &bond_dims);
930            tensor.push_tensor(tensor_1);
931
932            for (sub_tensor, ref_tensor) in zip(tensor.tensors(), [&ref_tensor_1]) {
933                let sub_tensor = sub_tensor.as_leaf().unwrap();
934                assert_eq!(sub_tensor.legs(), ref_tensor.legs());
935                assert_eq!(sub_tensor.bond_dims(), ref_tensor.bond_dims());
936            }
937
938            // Push tensor 2
939            let tensor_2 = LeafTensor::new_from_map(vec![7, 10, 2], &bond_dims);
940            tensor.push_tensor(tensor_2);
941
942            for (sub_tensor, ref_tensor) in zip(tensor.tensors(), [&ref_tensor_1, &ref_tensor_2]) {
943                let sub_tensor = sub_tensor.as_leaf().unwrap();
944                assert_eq!(sub_tensor.legs(), ref_tensor.legs());
945                assert_eq!(sub_tensor.bond_dims(), ref_tensor.bond_dims());
946            }
947        }
948
949        #[test]
950        fn test_push_tensors() {
951            let bond_dims =
952                FxHashMap::from_iter([(2, 17), (3, 1), (4, 11), (8, 3), (9, 20), (7, 7), (10, 14)]);
953            let ref_tensor_1 = LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims);
954            let ref_tensor_2 = LeafTensor::new_from_map(vec![8, 4, 9], &bond_dims);
955            let ref_tensor_3 = LeafTensor::new_from_map(vec![7, 10, 2], &bond_dims);
956
957            let tensor_1 = LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims);
958            let tensor_2 = LeafTensor::new_from_map(vec![8, 4, 9], &bond_dims);
959            let tensor_3 = LeafTensor::new_from_map(vec![7, 10, 2], &bond_dims);
960            let mut tensor = CompositeTensor::default();
961            tensor.push_tensors(vec![tensor_1, tensor_2, tensor_3]);
962
963            for (sub_tensor, ref_tensor) in zip(
964                tensor.tensors(),
965                &vec![ref_tensor_1, ref_tensor_2, ref_tensor_3],
966            ) {
967                let sub_tensor = sub_tensor.as_leaf().unwrap();
968                assert_eq!(sub_tensor.legs(), ref_tensor.legs());
969                assert_eq!(sub_tensor.bond_dims(), ref_tensor.bond_dims());
970            }
971        }
972    }
973}
tnc/tensornetwork/tensor.rs

tnc/tensornetwork/
tensor.rs
