tnc/contractionpath/
paths.rs

1//! Contraction path finders.
2
3use crate::{
4    contractionpath::{ssa_replace_ordering, ContractionPath},
5    tensornetwork::tensor::CompositeTensor,
6};
7
8pub mod branchbound;
9pub mod cotengrust;
10#[cfg(feature = "cotengra")]
11pub mod hyperoptimization;
12#[cfg(feature = "cotengra")]
13pub mod tree_annealing;
14#[cfg(feature = "cotengra")]
15pub mod tree_reconfiguration;
16#[cfg(feature = "cotengra")]
17pub mod tree_tempering;
18pub mod weighted_branchbound;
19
20/// An optimizer for finding a contraction path.
21pub trait Pathfinder {
22    type Result: ContractionPathResult;
23
24    /// Finds a contraction path for the `tensor`.
25    ///
26    /// Uses `&mut self` to allow for internal state such as caching.
27    fn find_path(&mut self, tensor: &CompositeTensor) -> Self::Result;
28}
29
30/// The result of running a contraction [`Pathfinder`].
31pub trait ContractionPathResult {
32    /// Returns the best found contraction path in SSA format.
33    fn ssa_path(&self) -> &ContractionPath;
34
35    /// Returns the best found contraction path in ReplaceLeft format.
36    fn replace_path(&self) -> ContractionPath;
37
38    /// Returns the total op count of the best path found.
39    fn flops(&self) -> f64;
40
41    /// Returns the max memory (in number of elements) of the best path found.
42    fn size(&self) -> f64;
43}
44
45/// Basic result data from running a contraction [`Pathfinder`].
46#[derive(Debug, Clone, Default, PartialEq)]
47pub struct BasicContractionPathResult {
48    /// The found path in SSA format.
49    ssa_path: ContractionPath,
50    /// The computational cost of the found path.
51    flops: f64,
52    /// The peak memory of the found path.
53    size: f64,
54}
55
56impl ContractionPathResult for BasicContractionPathResult {
57    #[inline]
58    fn ssa_path(&self) -> &ContractionPath {
59        &self.ssa_path
60    }
61
62    #[inline]
63    fn replace_path(&self) -> ContractionPath {
64        ssa_replace_ordering(&self.ssa_path)
65    }
66
67    #[inline]
68    fn flops(&self) -> f64 {
69        self.flops
70    }
71
72    #[inline]
73    fn size(&self) -> f64 {
74        self.size
75    }
76}
77
78/// The cost metric to optimize for.
79#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
80pub enum CostType {
81    /// Number of flops or operations.
82    Flops,
83    /// Size of the biggest contraction.
84    Size,
85}
86
87pub(crate) fn validate_path(path: &ContractionPath) {
88    let mut contracted = Vec::<usize>::new();
89    for nested in path.nested.values() {
90        validate_path(nested);
91    }
92
93    for (u, v) in &path.toplevel {
94        assert!(
95            !contracted.contains(u),
96            "Contracting already contracted tensors: {u:?}, path: {path:?}"
97        );
98        contracted.push(*v);
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    use crate::path;
107
108    #[test]
109    #[should_panic(
110        expected = "Contracting already contracted tensors: 1, path: ContractionPath { nested: {}, toplevel: [(0, 1), (1, 2)] }"
111    )]
112    fn test_validate_paths() {
113        let invalid_path = path![(0, 1), (1, 2)];
114        validate_path(&invalid_path);
115    }
116}
tnc/contractionpath/paths.rs

tnc/contractionpath/
paths.rs
