tnc/tensornetwork/
contraction.rs

1//! Functionality to contract tensor networks.
2use itertools::Itertools;
3use log::debug;
4use ndarray::Axis;
5use tblis::{tensor_mult, TensorView};
6
7use crate::{
8    contractionpath::ContractionPath,
9    tensornetwork::{
10        tensor::{CompositeTensor, LeafTensor, Tensor},
11        tensordata::{DataTensor, TensorData},
12    },
13};
14
15/// Fully contracts `tn` based on the given `contract_path` using ReplaceLeft format.
16/// Returns the resulting tensor.
17///
18/// # Examples
19/// ```
20/// # use tnc::{
21/// #   contractionpath::paths::{cotengrust::{Cotengrust, OptMethod}, Pathfinder, ContractionPathResult},
22/// #   builders::sycamore_circuit::sycamore_circuit,
23/// #   tensornetwork::tensor::Tensor,
24/// #   tensornetwork::contraction::contract_tensor_network,
25/// # };
26/// # use rand::rngs::StdRng;
27/// # use rand::SeedableRng;
28/// let mut r = StdRng::seed_from_u64(42);
29/// let mut r_tn = sycamore_circuit(2, 1, &mut r).into_expectation_value_network();
30/// let mut opt = Cotengrust::new(OptMethod::Greedy);
31/// let result = opt.find_path(&r_tn);
32/// let opt_path = result.replace_path();
33/// let result = contract_tensor_network(r_tn, &opt_path);
34/// ```
35pub fn contract_tensor_network(tn: CompositeTensor, contract_path: &ContractionPath) -> LeafTensor {
36    debug!(len = tn.len(); "Start contracting tensor network");
37
38    // Wrap the tensors into options, so we can take out used tensors
39    let mut tensors = tn.into_tensors().into_iter().map(Some).collect_vec();
40
41    // Contract child composite tensors first
42    for (index, inner_path) in &contract_path.nested {
43        let tc = tensors[*index]
44            .take()
45            .and_then(Tensor::into_composite)
46            .unwrap();
47        let contracted = contract_tensor_network(tc, inner_path);
48        tensors[*index] = Some(contracted.into());
49    }
50
51    // Contract all leaf tensors
52    for (i, j) in &contract_path.toplevel {
53        let ti = tensors[*i].take().and_then(Tensor::into_leaf).unwrap();
54        let tj = tensors[*j].take().and_then(Tensor::into_leaf).unwrap();
55        let contracted = contract_tensors(ti, tj);
56        tensors[*i] = Some(contracted.into());
57    }
58    debug!("Completed tensor network contraction");
59
60    // Remove all the None values and return the final tensor
61    tensors
62        .into_iter()
63        .flatten()
64        .exactly_one()
65        .unwrap()
66        .into_leaf()
67        .unwrap()
68}
69
70fn contract_tensors(tensor_a: LeafTensor, tensor_b: LeafTensor) -> LeafTensor {
71    let mut tensor_symmetric_difference = &tensor_a ^ &tensor_b;
72
73    let (a_legs, _, a_data) = tensor_a.into_inner();
74    let (b_legs, _, b_data) = tensor_b.into_inner();
75
76    let result = contract_ndarrays(
77        tensor_symmetric_difference.legs(),
78        &a_legs,
79        a_data.into_data(),
80        &b_legs,
81        b_data.into_data(),
82    );
83
84    tensor_symmetric_difference.set_tensor_data(TensorData::Matrix(result));
85    tensor_symmetric_difference
86}
87
88fn contract_ndarrays(
89    out_labels: &[usize],
90    a_labels: &[usize],
91    a_data: DataTensor,
92    b_labels: &[usize],
93    b_data: DataTensor,
94) -> DataTensor {
95    assert_eq!(a_labels.len(), a_data.ndim());
96    assert_eq!(b_labels.len(), b_data.ndim());
97
98    // Find output shape
99    let mut out_shape = Vec::with_capacity(out_labels.len());
100    for label in out_labels {
101        if let Some(a_index) = a_labels.iter().position(|l| l == label) {
102            out_shape.push(a_data.len_of(Axis(a_index)));
103        } else if let Some(b_index) = b_labels.iter().position(|l| l == label) {
104            out_shape.push(b_data.len_of(Axis(b_index)));
105        } else {
106            panic!("Out label {label} not found in input a ({a_labels:?}) or b ({b_labels:?})");
107        }
108    }
109
110    // Contract with TBLIS
111    let a_view = TensorView::new(a_labels, a_data.shape(), a_data.strides(), a_data.as_ptr());
112    let b_view = TensorView::new(b_labels, b_data.shape(), b_data.strides(), b_data.as_ptr());
113    let out_data = tensor_mult(out_labels, &out_shape, a_view, b_view);
114
115    DataTensor::from_shape_vec(out_shape, out_data).unwrap()
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    use approx::assert_abs_diff_eq;
123    use num_complex::Complex64;
124    use rustc_hash::FxHashMap;
125    use serde::Deserialize;
126
127    use crate::{path, tensornetwork::tensordata::TensorData};
128
129    #[derive(Debug, Deserialize)]
130    struct TestTensor {
131        legs: Vec<usize>,
132        shape: Vec<u64>,
133        data: Vec<Complex64>,
134    }
135
136    type TestData = FxHashMap<String, TestTensor>;
137
138    static TEST_DATA: &str = include_str!("contraction_test_data.json");
139
140    fn load_test_data() -> TestData {
141        serde_json::from_str(TEST_DATA).unwrap()
142    }
143
144    fn pop_test_tensor(name: &str, data: &mut TestData) -> LeafTensor {
145        let test_tensor = data.remove(name).unwrap();
146        let mut tensor = LeafTensor::new(test_tensor.legs, test_tensor.shape);
147        tensor.set_tensor_data(TensorData::new_from_data(
148            &tensor.shape().unwrap(),
149            test_tensor.data,
150        ));
151        tensor
152    }
153
154    #[test]
155    fn test_tensor_contraction() {
156        let mut data = load_test_data();
157        // t1 is of shape [3, 2, 7]
158        let t1 = pop_test_tensor("A", &mut data);
159        // t2 is of shape [7, 8, 6]
160        let t2 = pop_test_tensor("B", &mut data);
161        // t3 is of shape [3, 5, 8]
162        let t3 = pop_test_tensor("C", &mut data);
163        // t12 is of shape [8, 6, 3, 2]
164        let t12 = pop_test_tensor("AxB", &mut data);
165        // t23 is of shape [3, 5, 7, 6]
166        let t23 = pop_test_tensor("BxC", &mut data);
167
168        let out = contract_tensors(t2.clone(), t1);
169        assert_abs_diff_eq!(&out, &t12, epsilon = 1e-14);
170
171        let out = contract_tensors(t3, t2);
172        assert_abs_diff_eq!(&out, &t23, epsilon = 1e-14);
173    }
174
175    #[test]
176    fn test_tn_contraction() {
177        let mut data = load_test_data();
178        // t1 is of shape [3, 2, 7]
179        let t1 = pop_test_tensor("A", &mut data);
180        // t2 is of shape [7, 8, 6]
181        let t2 = pop_test_tensor("B", &mut data);
182        // t3 is of shape [3, 5, 8]
183        let t3 = pop_test_tensor("C", &mut data);
184        // tout is of shape [5, 6, 2]
185        let tout = pop_test_tensor("ABxC", &mut data);
186
187        let tn = CompositeTensor::new(vec![t1, t2, t3]);
188        let contract_path = path![(1, 0), (2, 1)];
189
190        let result = contract_tensor_network(tn, &contract_path);
191        assert_abs_diff_eq!(&result, &tout, epsilon = 1e-14);
192    }
193
194    #[test]
195    fn test_outer_product_contraction() {
196        let bond_dims = FxHashMap::from_iter([(0, 3), (1, 2)]);
197        let mut t1 = LeafTensor::new_from_map(vec![0], &bond_dims);
198        let mut t2 = LeafTensor::new_from_map(vec![1], &bond_dims);
199        t1.set_tensor_data(TensorData::new_from_data(
200            &[3],
201            vec![
202                Complex64::new(1.0, 0.0),
203                Complex64::new(2.0, 5.0),
204                Complex64::new(3.0, -1.0),
205            ],
206        ));
207        t2.set_tensor_data(TensorData::new_from_data(
208            &[2],
209            vec![Complex64::new(-4.0, 2.0), Complex64::new(0.0, -1.0)],
210        ));
211        let t3 = CompositeTensor::new(vec![t1, t2]);
212        let contract_path = path![(0, 1)];
213
214        let mut tn_ref = LeafTensor::new_from_map(vec![0, 1], &bond_dims);
215        tn_ref.set_tensor_data(TensorData::new_from_data(
216            &[3, 2],
217            vec![
218                Complex64::new(-4.0, 2.0),
219                Complex64::new(0.0, -1.0),
220                Complex64::new(-18.0, -16.0),
221                Complex64::new(5.0, -2.0),
222                Complex64::new(-10.0, 10.0),
223                Complex64::new(-1.0, -3.0),
224            ],
225        ));
226
227        let result = contract_tensor_network(t3, &contract_path);
228        assert_abs_diff_eq!(&result, &tn_ref);
229    }
230
231    #[test]
232    fn dimension_order() {
233        let mut ket0 = LeafTensor::new_from_const(vec![0], 2);
234        ket0.set_tensor_data(TensorData::new_from_data(
235            &[2],
236            vec![Complex64::ONE, Complex64::ZERO],
237        ));
238
239        let mut mat = LeafTensor::new_from_const(vec![1, 0], 2);
240        mat.set_tensor_data(TensorData::new_from_data(
241            &[2, 2],
242            vec![
243                Complex64::new(1.0, 0.0),
244                Complex64::new(2.0, 0.0),
245                Complex64::new(3.0, 0.0),
246                Complex64::new(4.0, 0.0),
247            ],
248        ));
249
250        let tn = CompositeTensor::new(vec![ket0, mat]);
251        let contract_path = path![(0, 1)];
252
253        let mut tn_ref = LeafTensor::new_from_const(vec![1], 2);
254        tn_ref.set_tensor_data(TensorData::new_from_data(
255            &[2],
256            vec![Complex64::new(1.0, 0.0), Complex64::new(3.0, 0.0)],
257        ));
258
259        let result = contract_tensor_network(tn, &contract_path);
260        assert_abs_diff_eq!(result, &tn_ref);
261    }
262}
tnc/tensornetwork/contraction.rs

tnc/tensornetwork/
contraction.rs
