tnc/contractionpath/paths/
cotengrust.rs

1use std::iter::zip;
2
3use cotengrust::{optimize_greedy_rust, optimize_optimal_rust, optimize_random_greedy_rust};
4use itertools::Itertools;
5use rustc_hash::FxHashMap;
6
7use crate::contractionpath::contraction_cost::contract_path_cost;
8use crate::contractionpath::paths::{
9    BasicContractionPathResult, ContractionPathResult, Pathfinder,
10};
11use crate::contractionpath::{ssa_replace_ordering, ContractionPath, SimplePath};
12use crate::tensornetwork::tensor::{CompositeTensor, LeafTensor, TensorType};
13
14/// The optimization method to use.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum OptMethod {
17    /// Searches for the optimal path, quite slow.
18    Optimal,
19    /// Uses a greedy algorithm to find a path.
20    Greedy,
21    /// Tries multiple greedy paths and selects the best one.
22    RandomGreedy(usize),
23}
24
25/// A contraction path finder using the `cotengrust` library.
26#[derive(Debug, Clone)]
27pub struct Cotengrust {
28    opt_method: OptMethod,
29}
30
31impl Cotengrust {
32    /// Creates a new Cotengrust optimizer using the specified optimization method.
33    #[inline]
34    pub fn new(opt_method: OptMethod) -> Self {
35        Self { opt_method }
36    }
37
38    /// Finds a contraction path for a "classical" tensor network, i.e. the inputs
39    /// are all leaf tensors.
40    fn optimize_single(&self, inputs: &[LeafTensor], output: &LeafTensor) -> SimplePath {
41        // Check if the inputs are empty (cotengrust does not handle this gracefully)
42        if inputs.is_empty() {
43            return SimplePath::default();
44        }
45
46        // Convert the inputs to the cotengra format
47        let (inputs, output, size_dict) = tensor_legs_to_digit(inputs, output);
48
49        // Find the contraction path
50        let path = match &self.opt_method {
51            OptMethod::Greedy => optimize_greedy_rust(
52                inputs,
53                output,
54                size_dict,
55                None,
56                None,
57                None,
58                Some(42),
59                false,
60                true,
61            ),
62            &OptMethod::RandomGreedy(ntrials) => {
63                optimize_random_greedy_rust(
64                    inputs,
65                    output,
66                    size_dict,
67                    ntrials,
68                    None,
69                    None,
70                    None,
71                    Some(42),
72                    false,
73                    true,
74                )
75                .0
76            }
77            OptMethod::Optimal => {
78                optimize_optimal_rust(inputs, output, size_dict, None, None, None, false, true)
79            }
80        };
81
82        // Convert the path back to our format
83        path.into_iter()
84            .map(|pair| {
85                let [a, b] = pair[..] else {
86                    panic!("Expected two indices in contraction path pair")
87                };
88                (a as _, b as _)
89            })
90            .collect_vec()
91    }
92}
93
94/// Converts tensor leg inputs to chars. Creates new inputs, outputs and size_dict that can be fed to Cotengra.
95fn tensor_legs_to_digit(
96    inputs: &[LeafTensor],
97    output: &LeafTensor,
98) -> (Vec<Vec<char>>, Vec<char>, FxHashMap<char, f32>) {
99    fn leg_to_char(leg: usize) -> char {
100        char::from_u32(leg.try_into().unwrap()).unwrap()
101    }
102    let mut new_inputs = vec![Vec::new(); inputs.len()];
103    let new_output = output.legs().iter().copied().map(leg_to_char).collect();
104    let mut new_size_dict = FxHashMap::default();
105
106    for (tensor, labels) in zip(inputs, new_inputs.iter_mut()) {
107        labels.reserve_exact(tensor.legs().len());
108        for (leg, dim) in tensor.edges() {
109            let character = leg_to_char(leg);
110            labels.push(character);
111            new_size_dict.insert(character, dim as f32);
112        }
113    }
114    (new_inputs, new_output, new_size_dict)
115}
116
117impl Pathfinder for Cotengrust {
118    type Result = BasicContractionPathResult;
119
120    fn find_path(&mut self, tensor: &CompositeTensor) -> BasicContractionPathResult {
121        // Handle nested tensors first
122        let mut nested_paths = FxHashMap::default();
123        let leaves = tensor
124            .tensors()
125            .iter()
126            .enumerate()
127            .map(|(index, t)| match t.kind() {
128                TensorType::Composite => {
129                    let composite = t.as_composite().unwrap();
130                    let mut ct = Cotengrust::new(self.opt_method);
131                    let result = ct.find_path(composite);
132                    nested_paths.insert(index, result.ssa_path().clone());
133                    composite.external_tensor()
134                }
135                TensorType::Leaf => t.clone().into_leaf().unwrap(),
136            })
137            .collect_vec();
138
139        // Now handle the outer tensor
140        let external_tensor = tensor.external_tensor();
141        let outer_path = self.optimize_single(&leaves, &external_tensor);
142        let best_path = ContractionPath {
143            nested: nested_paths,
144            toplevel: outer_path,
145        };
146        let replace_path = ssa_replace_ordering(&best_path);
147
148        // Compute the cost
149        let (op_cost, mem_cost) = contract_path_cost(tensor.tensors(), &replace_path, true);
150        BasicContractionPathResult {
151            ssa_path: best_path,
152            flops: op_cost,
153            size: mem_cost,
154        }
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    use crate::path;
163
164    fn setup_simple() -> CompositeTensor {
165        let bond_dims =
166            FxHashMap::from_iter([(0, 5), (1, 2), (2, 6), (3, 8), (4, 1), (5, 3), (6, 4)]);
167        CompositeTensor::new(vec![
168            LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
169            LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
170            LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
171        ])
172    }
173
174    fn setup_complex() -> CompositeTensor {
175        let bond_dims = FxHashMap::from_iter([
176            (0, 27),
177            (1, 18),
178            (2, 12),
179            (3, 15),
180            (4, 5),
181            (5, 3),
182            (6, 18),
183            (7, 22),
184            (8, 45),
185            (9, 65),
186            (10, 5),
187            (11, 17),
188        ]);
189        CompositeTensor::new(vec![
190            LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
191            LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
192            LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
193            LeafTensor::new_from_map(vec![6, 8, 9], &bond_dims),
194            LeafTensor::new_from_map(vec![10, 8, 9], &bond_dims),
195            LeafTensor::new_from_map(vec![5, 1, 0], &bond_dims),
196        ])
197    }
198
199    fn setup_simple_inner_product() -> CompositeTensor {
200        let bond_dims =
201            FxHashMap::from_iter([(0, 5), (1, 2), (2, 6), (3, 8), (4, 1), (5, 3), (6, 4)]);
202        CompositeTensor::new(vec![
203            LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
204            LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
205            LeafTensor::new_from_map(vec![0, 1, 5], &bond_dims),
206            LeafTensor::new_from_map(vec![1, 6], &bond_dims),
207        ])
208    }
209
210    fn setup_simple_outer_product() -> CompositeTensor {
211        let bond_dims = FxHashMap::from_iter([(0, 3), (1, 2), (2, 2)]);
212        CompositeTensor::new(vec![
213            LeafTensor::new_from_map(vec![0], &bond_dims),
214            LeafTensor::new_from_map(vec![1], &bond_dims),
215            LeafTensor::new_from_map(vec![2], &bond_dims),
216        ])
217    }
218
219    fn setup_complex_outer_product() -> CompositeTensor {
220        let bond_dims = FxHashMap::from_iter([(0, 5), (1, 4)]);
221        CompositeTensor::new(vec![
222            LeafTensor::new_from_map(vec![0], &bond_dims),
223            LeafTensor::new_from_map(vec![0], &bond_dims),
224            LeafTensor::new_from_map(vec![1], &bond_dims),
225            LeafTensor::new_from_map(vec![1], &bond_dims),
226        ])
227    }
228
229    #[test]
230    fn test_contract_order_greedy_simple() {
231        let tn = setup_simple();
232        let mut opt = Cotengrust::new(OptMethod::Greedy);
233        let result = opt.find_path(&tn);
234
235        assert_eq!(
236            result,
237            BasicContractionPathResult {
238                ssa_path: path![(0, 1), (3, 2)],
239                flops: 600.,
240                size: 538.
241            }
242        );
243    }
244
245    #[test]
246    fn test_contract_order_greedy_simple_inner() {
247        let tn = setup_simple_inner_product();
248        let mut opt = Cotengrust::new(OptMethod::Greedy);
249        let result = opt.find_path(&tn);
250
251        assert_eq!(
252            result,
253            BasicContractionPathResult {
254                ssa_path: path![(0, 1), (2, 3), (4, 5)],
255                flops: 228.,
256                size: 121.
257            }
258        );
259    }
260
261    #[test]
262    fn test_contract_order_greedy_simple_outer() {
263        let tn = setup_simple_outer_product();
264        let mut opt = Cotengrust::new(OptMethod::Greedy);
265        let result = opt.find_path(&tn);
266
267        assert_eq!(
268            result,
269            BasicContractionPathResult {
270                ssa_path: path![(2, 1), (0, 3)],
271                flops: 16.,
272                size: 19.
273            }
274        );
275    }
276
277    #[test]
278    fn test_contract_order_greedy_complex_outer() {
279        let tn = setup_complex_outer_product();
280        let mut opt = Cotengrust::new(OptMethod::Greedy);
281        let result = opt.find_path(&tn);
282
283        assert_eq!(
284            result,
285            BasicContractionPathResult {
286                ssa_path: path![(0, 1), (2, 3), (5, 4)],
287                flops: 10.,
288                size: 11.
289            }
290        );
291    }
292
293    #[test]
294    fn test_contract_order_greedy_complex() {
295        let tn = setup_complex();
296        let mut opt = Cotengrust::new(OptMethod::Greedy);
297        let result = opt.find_path(&tn);
298
299        assert_eq!(
300            result,
301            BasicContractionPathResult {
302                ssa_path: path![(1, 5), (3, 4), (6, 0), (7, 2), (9, 8)],
303                flops: 529815.,
304                size: 89478.
305            }
306        );
307    }
308}
tnc/contractionpath/paths/cotengrust.rs

tnc/contractionpath/paths/
cotengrust.rs
