cotengrust 0.1.0__tar.gz → 0.1.1__tar.gz

{cotengrust-0.1.0 → cotengrust-0.1.1}/.github/workflows/CI.yml

@@ -1,7 +1,7 @@
-# This file is autogenerated by maturin v0.15.2
+# This file is autogenerated by maturin v1.2.3
 # To update, run
 #
-#    maturin generate-ci github
+#    maturin generate-ci github --pytest
 #
 name: CI
 
@@ -41,6 +41,30 @@ jobs:
         with:
           name: wheels
           path: dist
+      - name: pytest
+        if: ${{ startsWith(matrix.target, 'x86_64') }}
+        shell: bash
+        run: |
+          set -e
+          ls dist/*
+          pip install cotengrust --find-links dist --force-reinstall
+          pip install pytest numpy cotengra
+          pytest --verbose
+      - name: pytest
+        if: ${{ !startsWith(matrix.target, 'x86') && matrix.target != 'ppc64' }}
+        uses: uraimo/run-on-arch-action@v2.5.0
+        with:
+          arch: ${{ matrix.target }}
+          distro: ubuntu22.04
+          githubToken: ${{ github.token }}
+          install: |
+            apt-get update
+            apt-get install -y --no-install-recommends python3 python3-pip
+            pip3 install -U pip pytest # numpy cotengra
+          run: |
+            set -e
+            pip3 install cotengrust --find-links dist --force-reinstall
+            pytest --verbose
 
   windows:
     runs-on: windows-latest
@@ -64,6 +88,15 @@ jobs:
         with:
           name: wheels
           path: dist
+      - name: pytest
+        if: ${{ !startsWith(matrix.target, 'aarch64') }}
+        shell: bash
+        run: |
+          set -e
+          ls dist/*
+          pip install cotengrust --find-links dist --force-reinstall
+          pip install pytest numpy cotengra
+          pytest --verbose
 
   macos:
     runs-on: macos-latest
@@ -86,6 +119,15 @@ jobs:
         with:
           name: wheels
           path: dist
+      - name: pytest
+        if: ${{ !startsWith(matrix.target, 'aarch64') }}
+        shell: bash
+        run: |
+          set -e
+          ls dist/*
+          pip install cotengrust --find-links dist --force-reinstall
+          pip install pytest numpy cotengra
+          pytest --verbose
 
   sdist:
     runs-on: ubuntu-latest
@@ -117,4 +159,4 @@ jobs:
           MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
         with:
           command: upload
-          args: --skip-existing *
+          args: --non-interactive --skip-existing *

{cotengrust-0.1.0 → cotengrust-0.1.1}/Cargo.lock

@@ -37,7 +37,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "cotengrust"
-version = "0.1.0"
+version = "0.1.1"
 dependencies = [
  "bit-set",
  "ordered-float",

{cotengrust-0.1.0 → cotengrust-0.1.1}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "cotengrust"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -19,4 +19,3 @@ rustc-hash = "1.1"
 codegen-units = 1
 lto = true
 opt-level = 3
-panic = "abort"

cotengrust-0.1.1/PKG-INFO

@@ -0,0 +1,249 @@
+Metadata-Version: 2.1
+Name: cotengrust
+Version: 0.1.1
+Classifier: Programming Language :: Rust
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+License-File: LICENSE
+Summary: Fast contraction ordering primitives for tensor networks.
+Author-email: Johnnie Gray <johnniemcgray@gmail.com>
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+
+# cotengrust
+
+`cotengrust` provides fast rust implementations of contraction ordering
+primitives for tensor networks or einsum expressions. The two main functions
+are:
+
+- `optimize_optimal(inputs, output, size_dict, **kwargs)`
+- `optimize_greedy(inputs, output, size_dict, **kwargs)`
+
+The optimal algorithm is an optimized version of the `opt_einsum` 'dp' 
+path - itself an implementation of https://arxiv.org/abs/1304.6112.
+
+
+## Installation
+
+`cotengrust` is available for most platforms from
+[PyPI](https://pypi.org/project/cotengrust/):
+
+```bash
+pip install cotengrust
+```
+
+or if you want to develop locally (which requires [pyo3](https://github.com/PyO3/pyo3) 
+and [maturin](https://github.com/PyO3/maturin)):
+
+```bash
+git clone https://github.com/jcmgray/cotengrust.git
+cd cotengrust
+maturin develop --release
+```
+(the release flag is very important for assessing performance!).
+
+
+## Usage
+
+If `cotengrust` is installed, then by default `cotengra` will use it for its
+greedy and optimal subroutines, notably subtree reconfiguration. You can also
+call the routines directly:
+
+```python
+import cotengra as ctg
+import cotengrust as ctgr
+
+# specify an 8x8 square lattice contraction
+inputs, output, shapes, size_dict = ctg.utils.lattice_equation([8, 8])
+
+# find the optimal 'combo' contraction path
+%%time
+path = ctgr.optimize_optimal(inputs, output, size_dict, minimize='combo')
+# CPU times: user 13.7 s, sys: 83.4 ms, total: 13.7 s
+# Wall time: 13.7 s
+
+# construct a contraction tree for further introspection
+tree = ctg.ContractionTree.from_path(
+    inputs, output, size_dict, path=path
+)
+tree.plot_rubberband()
+```
+![optimal-8x8-order](https://github.com/jcmgray/cotengrust/assets/8982598/f8e18ff2-5ace-4e46-81e1-06bffaef5e45)
+
+
+## API
+
+The optimize functions follow the api of the python implementations in `cotengra.pathfinders.path_basic.py`.
+
+```python
+def optimize_optimal(
+    inputs,
+    output,
+    size_dict,
+    minimize='flops',
+    cost_cap=2,
+    search_outer=False,
+    simplify=True,
+    use_ssa=False,
+):
+    """Find an optimal contraction ordering.
+
+    Parameters
+    ----------
+    inputs : Sequence[Sequence[str]]
+        The indices of each input tensor.
+    output : Sequence[str]
+        The indices of the output tensor.
+    size_dict : dict[str, int]
+        The size of each index.
+    minimize : str, optional
+        The cost function to minimize. The options are:
+
+        - "flops": minimize with respect to total operation count only
+          (also known as contraction cost)
+        - "size": minimize with respect to maximum intermediate size only
+          (also known as contraction width)
+        - 'write' : minimize the sum of all tensor sizes, i.e. memory written
+        - 'combo' or 'combo={factor}` : minimize the sum of
+          FLOPS + factor * WRITE, with a default factor of 64.
+        - 'limit' or 'limit={factor}` : minimize the sum of
+          MAX(FLOPS, alpha * WRITE) for each individual contraction, with a
+          default factor of 64.
+
+        'combo' is generally a good default in term of practical hardware
+        performance, where both memory bandwidth and compute are limited.
+    cost_cap : float, optional
+        The maximum cost of a contraction to initially consider. This acts like
+        a sieve and is doubled at each iteration until the optimal path can
+        be found, but supplying an accurate guess can speed up the algorithm.
+    search_outer : bool, optional
+        If True, consider outer product contractions. This is much slower but
+        theoretically might be required to find the true optimal 'flops'
+        ordering. In practical settings (i.e. with minimize='combo'), outer
+        products should not be required.
+    simplify : bool, optional
+        Whether to perform simplifications before optimizing. These are:
+
+            - ignore any indices that appear in all terms
+            - combine any repeated indices within a single term
+            - reduce any non-output indices that only appear on a single term
+            - combine any scalar terms
+            - combine any tensors with matching indices (hadamard products)
+
+        Such simpifications may be required in the general case for the proper
+        functioning of the core optimization, but may be skipped if the input
+        indices are already in a simplified form.
+    use_ssa : bool, optional
+        Whether to return the contraction path in 'single static assignment'
+        (SSA) format (i.e. as if each intermediate is appended to the list of
+        inputs, without removals). This can be quicker and easier to work with
+        than the 'linear recycled' format that `numpy` and `opt_einsum` use.
+
+    Returns
+    -------
+    path : list[list[int]]
+        The contraction path, given as a sequence of pairs of node indices. It
+        may also have single term contractions if `simplify=True`.
+    """
+    ...
+
+
+def optimize_greedy(
+    inputs,
+    output,
+    size_dict,
+    costmod=1.0,
+    temperature=0.0,
+    simplify=True,
+    use_ssa=False,
+):
+    """Find a contraction path using a (randomizable) greedy algorithm.
+
+    Parameters
+    ----------
+    inputs : Sequence[Sequence[str]]
+        The indices of each input tensor.
+    output : Sequence[str]
+        The indices of the output tensor.
+    size_dict : dict[str, int]
+        A dictionary mapping indices to their dimension.
+    costmod : float, optional
+        When assessing local greedy scores how much to weight the size of the
+        tensors removed compared to the size of the tensor added::
+
+            score = size_ab - costmod * (size_a + size_b)
+
+        This can be a useful hyper-parameter to tune.
+    temperature : float, optional
+        When asessing local greedy scores, how much to randomly perturb the
+        score. This is implemented as::
+
+            score -> sign(score) * log(|score|) - temperature * gumbel()
+
+        which implements boltzmann sampling.
+    simplify : bool, optional
+        Whether to perform simplifications before optimizing. These are:
+
+            - ignore any indices that appear in all terms
+            - combine any repeated indices within a single term
+            - reduce any non-output indices that only appear on a single term
+            - combine any scalar terms
+            - combine any tensors with matching indices (hadamard products)
+
+        Such simpifications may be required in the general case for the proper
+        functioning of the core optimization, but may be skipped if the input
+        indices are already in a simplified form.
+    use_ssa : bool, optional
+        Whether to return the contraction path in 'single static assignment'
+        (SSA) format (i.e. as if each intermediate is appended to the list of
+        inputs, without removals). This can be quicker and easier to work with
+        than the 'linear recycled' format that `numpy` and `opt_einsum` use.
+
+    Returns
+    -------
+    path : list[list[int]]
+        The contraction path, given as a sequence of pairs of node indices. It
+        may also have single term contractions if `simplify=True`.
+    """
+
+def optimize_simplify(
+    inputs,
+    output,
+    size_dict,
+    use_ssa=False,
+):
+    """Find the (partial) contracton path for simplifiactions only.
+
+    Parameters
+    ----------
+    inputs : Sequence[Sequence[str]]
+        The indices of each input tensor.
+    output : Sequence[str]
+        The indices of the output tensor.
+    size_dict : dict[str, int]
+        A dictionary mapping indices to their dimension.
+    use_ssa : bool, optional
+        Whether to return the contraction path in 'single static assignment'
+        (SSA) format (i.e. as if each intermediate is appended to the list of
+        inputs, without removals). This can be quicker and easier to work with
+        than the 'linear recycled' format that `numpy` and `opt_einsum` use.
+
+    Returns
+    -------
+    path : list[list[int]]
+        The contraction path, given as a sequence of pairs of node indices. It
+        may also have single term contractions.
+
+    """
+    ...
+
+def ssa_to_linear(ssa_path, n=None):
+    """Convert a SSA path to linear format."""
+    ...
+
+def find_subgraphs(inputs, output, size_dict,):
+    """Find all disconnected subgraphs of a specified contraction."""
+    ...
+```
+
+

cotengrust-0.1.1/README.md

@@ -0,0 +1,236 @@
+# cotengrust
+
+`cotengrust` provides fast rust implementations of contraction ordering
+primitives for tensor networks or einsum expressions. The two main functions
+are:
+
+- `optimize_optimal(inputs, output, size_dict, **kwargs)`
+- `optimize_greedy(inputs, output, size_dict, **kwargs)`
+
+The optimal algorithm is an optimized version of the `opt_einsum` 'dp' 
+path - itself an implementation of https://arxiv.org/abs/1304.6112.
+
+
+## Installation
+
+`cotengrust` is available for most platforms from
+[PyPI](https://pypi.org/project/cotengrust/):
+
+```bash
+pip install cotengrust
+```
+
+or if you want to develop locally (which requires [pyo3](https://github.com/PyO3/pyo3) 
+and [maturin](https://github.com/PyO3/maturin)):
+
+```bash
+git clone https://github.com/jcmgray/cotengrust.git
+cd cotengrust
+maturin develop --release
+```
+(the release flag is very important for assessing performance!).
+
+
+## Usage
+
+If `cotengrust` is installed, then by default `cotengra` will use it for its
+greedy and optimal subroutines, notably subtree reconfiguration. You can also
+call the routines directly:
+
+```python
+import cotengra as ctg
+import cotengrust as ctgr
+
+# specify an 8x8 square lattice contraction
+inputs, output, shapes, size_dict = ctg.utils.lattice_equation([8, 8])
+
+# find the optimal 'combo' contraction path
+%%time
+path = ctgr.optimize_optimal(inputs, output, size_dict, minimize='combo')
+# CPU times: user 13.7 s, sys: 83.4 ms, total: 13.7 s
+# Wall time: 13.7 s
+
+# construct a contraction tree for further introspection
+tree = ctg.ContractionTree.from_path(
+    inputs, output, size_dict, path=path
+)
+tree.plot_rubberband()
+```
+![optimal-8x8-order](https://github.com/jcmgray/cotengrust/assets/8982598/f8e18ff2-5ace-4e46-81e1-06bffaef5e45)
+
+
+## API
+
+The optimize functions follow the api of the python implementations in `cotengra.pathfinders.path_basic.py`.
+
+```python
+def optimize_optimal(
+    inputs,
+    output,
+    size_dict,
+    minimize='flops',
+    cost_cap=2,
+    search_outer=False,
+    simplify=True,
+    use_ssa=False,
+):
+    """Find an optimal contraction ordering.
+
+    Parameters
+    ----------
+    inputs : Sequence[Sequence[str]]
+        The indices of each input tensor.
+    output : Sequence[str]
+        The indices of the output tensor.
+    size_dict : dict[str, int]
+        The size of each index.
+    minimize : str, optional
+        The cost function to minimize. The options are:
+
+        - "flops": minimize with respect to total operation count only
+          (also known as contraction cost)
+        - "size": minimize with respect to maximum intermediate size only
+          (also known as contraction width)
+        - 'write' : minimize the sum of all tensor sizes, i.e. memory written
+        - 'combo' or 'combo={factor}` : minimize the sum of
+          FLOPS + factor * WRITE, with a default factor of 64.
+        - 'limit' or 'limit={factor}` : minimize the sum of
+          MAX(FLOPS, alpha * WRITE) for each individual contraction, with a
+          default factor of 64.
+
+        'combo' is generally a good default in term of practical hardware
+        performance, where both memory bandwidth and compute are limited.
+    cost_cap : float, optional
+        The maximum cost of a contraction to initially consider. This acts like
+        a sieve and is doubled at each iteration until the optimal path can
+        be found, but supplying an accurate guess can speed up the algorithm.
+    search_outer : bool, optional
+        If True, consider outer product contractions. This is much slower but
+        theoretically might be required to find the true optimal 'flops'
+        ordering. In practical settings (i.e. with minimize='combo'), outer
+        products should not be required.
+    simplify : bool, optional
+        Whether to perform simplifications before optimizing. These are:
+
+            - ignore any indices that appear in all terms
+            - combine any repeated indices within a single term
+            - reduce any non-output indices that only appear on a single term
+            - combine any scalar terms
+            - combine any tensors with matching indices (hadamard products)
+
+        Such simpifications may be required in the general case for the proper
+        functioning of the core optimization, but may be skipped if the input
+        indices are already in a simplified form.
+    use_ssa : bool, optional
+        Whether to return the contraction path in 'single static assignment'
+        (SSA) format (i.e. as if each intermediate is appended to the list of
+        inputs, without removals). This can be quicker and easier to work with
+        than the 'linear recycled' format that `numpy` and `opt_einsum` use.
+
+    Returns
+    -------
+    path : list[list[int]]
+        The contraction path, given as a sequence of pairs of node indices. It
+        may also have single term contractions if `simplify=True`.
+    """
+    ...
+
+
+def optimize_greedy(
+    inputs,
+    output,
+    size_dict,
+    costmod=1.0,
+    temperature=0.0,
+    simplify=True,
+    use_ssa=False,
+):
+    """Find a contraction path using a (randomizable) greedy algorithm.
+
+    Parameters
+    ----------
+    inputs : Sequence[Sequence[str]]
+        The indices of each input tensor.
+    output : Sequence[str]
+        The indices of the output tensor.
+    size_dict : dict[str, int]
+        A dictionary mapping indices to their dimension.
+    costmod : float, optional
+        When assessing local greedy scores how much to weight the size of the
+        tensors removed compared to the size of the tensor added::
+
+            score = size_ab - costmod * (size_a + size_b)
+
+        This can be a useful hyper-parameter to tune.
+    temperature : float, optional
+        When asessing local greedy scores, how much to randomly perturb the
+        score. This is implemented as::
+
+            score -> sign(score) * log(|score|) - temperature * gumbel()
+
+        which implements boltzmann sampling.
+    simplify : bool, optional
+        Whether to perform simplifications before optimizing. These are:
+
+            - ignore any indices that appear in all terms
+            - combine any repeated indices within a single term
+            - reduce any non-output indices that only appear on a single term
+            - combine any scalar terms
+            - combine any tensors with matching indices (hadamard products)
+
+        Such simpifications may be required in the general case for the proper
+        functioning of the core optimization, but may be skipped if the input
+        indices are already in a simplified form.
+    use_ssa : bool, optional
+        Whether to return the contraction path in 'single static assignment'
+        (SSA) format (i.e. as if each intermediate is appended to the list of
+        inputs, without removals). This can be quicker and easier to work with
+        than the 'linear recycled' format that `numpy` and `opt_einsum` use.
+
+    Returns
+    -------
+    path : list[list[int]]
+        The contraction path, given as a sequence of pairs of node indices. It
+        may also have single term contractions if `simplify=True`.
+    """
+
+def optimize_simplify(
+    inputs,
+    output,
+    size_dict,
+    use_ssa=False,
+):
+    """Find the (partial) contracton path for simplifiactions only.
+
+    Parameters
+    ----------
+    inputs : Sequence[Sequence[str]]
+        The indices of each input tensor.
+    output : Sequence[str]
+        The indices of the output tensor.
+    size_dict : dict[str, int]
+        A dictionary mapping indices to their dimension.
+    use_ssa : bool, optional
+        Whether to return the contraction path in 'single static assignment'
+        (SSA) format (i.e. as if each intermediate is appended to the list of
+        inputs, without removals). This can be quicker and easier to work with
+        than the 'linear recycled' format that `numpy` and `opt_einsum` use.
+
+    Returns
+    -------
+    path : list[list[int]]
+        The contraction path, given as a sequence of pairs of node indices. It
+        may also have single term contractions.
+
+    """
+    ...
+
+def ssa_to_linear(ssa_path, n=None):
+    """Convert a SSA path to linear format."""
+    ...
+
+def find_subgraphs(inputs, output, size_dict,):
+    """Find all disconnected subgraphs of a specified contraction."""
+    ...
+```
+

{cotengrust-0.1.0 → cotengrust-0.1.1}/pyproject.toml

@@ -1,16 +1,22 @@
-[build-system]
-requires = ["maturin>=0.15,<0.16"]
-build-backend = "maturin"
-
 [project]
 name = "cotengrust"
-requires-python = ">=3.7"
+version = "0.1.1"
+description = "Fast contraction ordering primitives for tensor networks."
+readme = "README.md"
+requires-python = ">=3.8"
 classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
+license = { file = "LICENSE" }
+authors = [
+    {name = "Johnnie Gray", email = "johnniemcgray@gmail.com"}
+]
 
+[build-system]
+requires = ["maturin>=0.15,<0.16"]
+build-backend = "maturin"
 
 [tool.maturin]
 features = ["pyo3/extension-module"]

{cotengrust-0.1.0 → cotengrust-0.1.1}/src/lib.rs

@@ -25,7 +25,7 @@ type SubContraction = (Legs, Score, BitPath);
 /// helper struct to build contractions from bottom up
 struct ContractionProcessor {
     nodes: Dict<Node, Legs>,
-    edges: Dict<Ix, Vec<Node>>,
+    edges: Dict<Ix, BTreeSet<Node>>,
     appearances: Vec<Count>,
     sizes: Vec<Score>,
     ssa: Node,
@@ -133,7 +133,7 @@ impl ContractionProcessor {
         size_dict: Dict<char, f32>,
     ) -> ContractionProcessor {
         let mut nodes: Dict<Node, Legs> = Dict::default();
-        let mut edges: Dict<Ix, Vec<Node>> = Dict::default();
+        let mut edges: Dict<Ix, BTreeSet<Node>> = Dict::default();
         let mut indmap: Dict<char, Ix> = Dict::default();
         let mut sizes: Vec<Score> = Vec::with_capacity(size_dict.len());
         let mut appearances: Vec<Count> = Vec::with_capacity(size_dict.len());
@@ -147,7 +147,7 @@ impl ContractionProcessor {
                     None => {
                         // index not parsed yet
                         indmap.insert(ind, c);
-                        edges.insert(c, vec![i as Node]);
+                        edges.insert(c, std::iter::once(i as Node).collect());
                         appearances.push(1);
                         sizes.push(f32::log(size_dict[&ind] as f32, 2.0));
                         legs.push((c, 1));
@@ -156,7 +156,7 @@ impl ContractionProcessor {
                     Some(&ix) => {
                         // index already present
                         appearances[ix as usize] += 1;
-                        edges.get_mut(&ix).unwrap().push(i as Node);
+                        edges.get_mut(&ix).unwrap().insert(i as Node);
                         legs.push((ix, 1));
                     }
                 };
@@ -204,11 +204,15 @@ impl ContractionProcessor {
     fn pop_node(&mut self, i: Node) -> Legs {
         let legs = self.nodes.remove(&i).unwrap();
         for (ix, _) in legs.iter() {
-            let nodes = self.edges.get_mut(&ix).unwrap();
-            if nodes.len() == 1 {
+            let enodes = match self.edges.get_mut(&ix) {
+                Some(enodes) => enodes,
+                // if repeated index, might have already been removed
+                None => continue,
+            };
+            enodes.remove(&i);
+            if enodes.len() == 0 {
+                // last node with this index -> remove from map
                 self.edges.remove(&ix);
-            } else {
-                nodes.retain(|&j| j != i);
             }
         }
         legs
@@ -221,8 +225,8 @@ impl ContractionProcessor {
         for (ix, _) in &legs {
             self.edges
                 .entry(*ix)
-                .and_modify(|nodes| nodes.push(i))
-                .or_insert(vec![i]);
+                .and_modify(|nodes| {nodes.insert(i);})
+                .or_insert(std::iter::once(i as Node).collect());
         }
         self.nodes.insert(i, legs);
         i
@@ -267,28 +271,27 @@ impl ContractionProcessor {
     /// combine and remove all scalars
     fn simplify_scalars(&mut self) {
         let mut scalars = Vec::new();
+        let mut j: Option<Node> = None;
+        let mut jndim: usize = 0;
         for (i, legs) in self.nodes.iter() {
-            if legs.len() == 0 {
+            let ndim = legs.len();
+            if ndim == 0 {
                 scalars.push(*i);
+            } else {
+                // also search for smallest other term to multiply into
+                if j.is_none() || ndim < jndim {
+                    j = Some(*i);
+                    jndim = ndim;
+                }
             }
         }
         if scalars.len() > 0 {
-            for &i in &scalars {
-                self.pop_node(i);
+            for p in 0..scalars.len() - 1 {
+                let i = scalars[p];
+                let j = scalars[p + 1];
+                let k = self.contract_nodes(i, j);
+                scalars[p + 1] = k;
             }
-            let (res, con) = match self.nodes.iter().min_by_key(|&(_, legs)| legs.len()) {
-                Some((&j, _)) => {
-                    let res = self.pop_node(j);
-                    let con: Vec<Node> = scalars.into_iter().chain(vec![j].into_iter()).collect();
-                    (res, con)
-                }
-                None => {
-                    let res = Vec::new();
-                    (res, scalars)
-                }
-            };
-            self.add_node(res);
-            self.ssa_path.push(con);
         }
     }
 
@@ -393,6 +396,8 @@ impl ContractionProcessor {
 
         // get the initial candidate contractions
         for ix_nodes in self.edges.values() {
+            // convert to vector for combinational indexing
+            let ix_nodes: Vec<Node> = ix_nodes.iter().cloned().collect();
             // for all combinations of nodes with a connected edge
             for ip in 0..ix_nodes.len() {
                 let i = ix_nodes[ip];
@@ -579,26 +584,70 @@ fn compute_con_cost_combo(
     (new_legs, new_score)
 }
 
+fn compute_con_cost_limit(
+    temp_legs: Legs,
+    appearances: &Vec<Count>,
+    sizes: &Vec<Score>,
+    iscore: Score,
+    jscore: Score,
+    factor: Score,
+) -> (Legs, Score) {
+    // remove indices that have reached final appearance
+    // and compute cost and size of local contraction
+    let mut new_legs: Legs = Legs::with_capacity(temp_legs.len());
+    let mut size: Score = 0.0;
+    let mut cost: Score = 0.0;
+    for (ix, ix_count) in temp_legs.into_iter() {
+        // all involved indices contribute to the cost
+        let d = sizes[ix as usize];
+        cost += d;
+        if ix_count != appearances[ix as usize] {
+            // not last appearance -> kept index contributes to new size
+            new_legs.push((ix, ix_count));
+            size += d;
+        }
+    }
+    // whichever is more expensive, the cost or the scaled write
+    let new_local_score = cost.max(factor + size);
+
+    // the total score including history
+    let new_score = logadd(logadd(iscore, jscore), new_local_score);
+
+    (new_legs, new_score)
+}
+
 impl ContractionProcessor {
     fn optimize_optimal_connected(
         &mut self,
         subgraph: Vec<Node>,
         minimize: Option<String>,
-        factor: Option<Score>,
         cost_cap: Option<Score>,
+        search_outer: Option<bool>,
     ) {
+        // parse the minimize argument
         let minimize = minimize.unwrap_or("flops".to_string());
-        let factor = f32::ln(factor.unwrap_or(64.0));
-        let compute_cost = match minimize.as_str() {
+        let mut minimize_split = minimize.split('-');
+        let minimize_type = minimize_split.next().unwrap();
+        let factor = minimize_split
+            .next()
+            .map_or(64.0, |s| s.parse::<f32>().unwrap())
+            .ln();
+        if minimize_split.next().is_some() {
+            // multiple hyphens -> raise error
+            panic!("invalid minimize: {:?}", minimize);
+        }
+        let compute_cost = match minimize_type {
             "flops" => compute_con_cost_flops,
             "size" => compute_con_cost_size,
             "write" => compute_con_cost_write,
             "combo" => compute_con_cost_combo,
+            "limit" => compute_con_cost_limit,
             _ => panic!(
-                "minimize must be one of 'flops', 'size', 'write', or 'combo', got {}",
+                "minimize must be one of 'flops', 'size', 'write', 'combo', or 'limit', got {}",
                 minimize
             ),
         };
+        let search_outer = search_outer.unwrap_or(false);
 
         // storage for each possible contraction to reach subgraph of size m
         let mut contractions: Vec<Dict<Subgraph, SubContraction>> =
@@ -624,7 +673,7 @@ impl ContractionProcessor {
 
         let mut ip: usize;
         let mut jp: usize;
-        let mut outer: bool;
+        let mut skip_because_outer: bool;
 
         let cost_cap_incr = f32::ln(2.0);
         let mut cost_cap = cost_cap.unwrap_or(cost_cap_incr);
@@ -647,7 +696,8 @@ impl ContractionProcessor {
                             let mut temp_legs: Legs = Vec::with_capacity(ilegs.len() + jlegs.len());
                             ip = 0;
                             jp = 0;
-                            outer = true;
+                            // if search_outer -> we will never skip
+                            skip_because_outer = !search_outer;
                             while ip < ilegs.len() && jp < jlegs.len() {
                                 if ilegs[ip].0 < jlegs[jp].0 {
                                     // index only appears in ilegs
@@ -662,10 +712,10 @@ impl ContractionProcessor {
                                     temp_legs.push((ilegs[ip].0, ilegs[ip].1 + jlegs[jp].1));
                                     ip += 1;
                                     jp += 1;
-                                    outer = false;
+                                    skip_because_outer = false;
                                 }
                             }
-                            if outer {
+                            if skip_because_outer {
                                 // no shared indices -> outer product
                                 continue;
                             }
@@ -683,7 +733,7 @@ impl ContractionProcessor {
                             );
 
                             if new_score > cost_cap {
-                                // contraction not allowed yet due to cost
+                                // contraction not allowed yet due to 'sieve'
                                 continue;
                             }
 
@@ -711,10 +761,10 @@ impl ContractionProcessor {
                             }
                         }
                     }
-                    // move new contractions from temp into the main storage, there
-                    // might be contractions for the same subgraph in this, but
-                    // because we check eagerly best_scores above, later entries
-                    // are guaranteed to be better
+                    // move new contractions from temp into the main storage,
+                    // there might be contractions for the same subgraph in
+                    // this, but because we check eagerly best_scores above,
+                    // later entries are guaranteed to be better
                     contractions_m_temp.drain(..).for_each(|(k, v)| {
                         contractions[m].insert(k, v);
                     });
@@ -722,7 +772,7 @@ impl ContractionProcessor {
             }
             cost_cap += cost_cap_incr;
         }
-        // can only ever be a single entry in contractions[nterms] -> the best one
+        // can only ever be a single entry in contractions[nterms] -> the best
         let (_, _, best_path) = contractions[nterms].values().next().unwrap();
 
         // convert from the bitpath to the actual (subgraph) node ids
@@ -738,17 +788,45 @@ impl ContractionProcessor {
     fn optimize_optimal(
         &mut self,
         minimize: Option<String>,
-        factor: Option<Score>,
         cost_cap: Option<Score>,
+        search_outer: Option<bool>,
     ) {
         for subgraph in self.subgraphs() {
-            self.optimize_optimal_connected(subgraph, minimize.clone(), factor, cost_cap);
+            self.optimize_optimal_connected(subgraph, minimize.clone(), cost_cap, search_outer);
         }
     }
 }
 
 // --------------------------- PYTHON FUNCTIONS ---------------------------- //
 
+#[pyfunction]
+#[pyo3()]
+fn ssa_to_linear(ssa_path: SSAPath, n: Option<usize>) -> SSAPath {
+    let n = match n {
+        Some(n) => n,
+        None => ssa_path.iter().map(|v| v.len()).sum::<usize>() + ssa_path.len() + 1,
+    };
+    let mut ids: Vec<Node> = (0..n).map(|i| i as Node).collect();
+    let mut path: SSAPath = Vec::with_capacity(2 * n - 1);
+    let mut ssa = n as Node;
+    for scon in ssa_path {
+        // find the locations of the ssa ids in the list of ids
+        let mut con: Vec<Node> = scon
+            .iter()
+            .map(|s| ids.binary_search(s).unwrap() as Node)
+            .collect();
+        // remove the ssa ids from the list
+        con.sort();
+        for j in con.iter().rev() {
+            ids.remove(*j as usize);
+        }
+        path.push(con);
+        ids.push(ssa);
+        ssa += 1;
+    }
+    path
+}
+
 #[pyfunction]
 #[pyo3()]
 fn find_subgraphs(
@@ -766,10 +844,16 @@ fn optimize_simplify(
     inputs: Vec<Vec<char>>,
     output: Vec<char>,
     size_dict: Dict<char, f32>,
+    use_ssa: Option<bool>,
 ) -> SSAPath {
+    let n = inputs.len();
     let mut cp = ContractionProcessor::new(inputs, output, size_dict);
     cp.simplify();
-    cp.ssa_path
+    if use_ssa.unwrap_or(false) {
+        cp.ssa_path
+    } else {
+        ssa_to_linear(cp.ssa_path, Some(n))
+    }
 }
 
 #[pyfunction]
@@ -781,15 +865,23 @@ fn optimize_greedy(
     costmod: Option<f32>,
     temperature: Option<f32>,
     simplify: Option<bool>,
+    use_ssa: Option<bool>,
 ) -> Vec<Vec<Node>> {
+    let n = inputs.len();
     let mut cp = ContractionProcessor::new(inputs, output, size_dict);
     if simplify.unwrap_or(true) {
+        // perform simplifications
         cp.simplify();
     }
+    // greddily contract each connected subgraph
     cp.optimize_greedy(costmod, temperature);
     // optimize any remaining disconnected terms
     cp.optimize_remaining_by_size();
-    cp.ssa_path
+    if use_ssa.unwrap_or(false) {
+        cp.ssa_path
+    } else {
+        ssa_to_linear(cp.ssa_path, Some(n))
+    }
 }
 
 #[pyfunction]
@@ -799,23 +891,32 @@ fn optimize_optimal(
     output: Vec<char>,
     size_dict: Dict<char, f32>,
     minimize: Option<String>,
-    factor: Option<Score>,
     cost_cap: Option<Score>,
+    search_outer: Option<bool>,
     simplify: Option<bool>,
+    use_ssa: Option<bool>,
 ) -> Vec<Vec<Node>> {
+    let n = inputs.len();
     let mut cp = ContractionProcessor::new(inputs, output, size_dict);
     if simplify.unwrap_or(true) {
+        // perform simplifications
         cp.simplify();
     }
-    cp.optimize_optimal(minimize, factor, cost_cap);
+    // optimally contract each connected subgraph
+    cp.optimize_optimal(minimize, cost_cap, search_outer);
     // optimize any remaining disconnected terms
     cp.optimize_remaining_by_size();
-    cp.ssa_path
+    if use_ssa.unwrap_or(false) {
+        cp.ssa_path
+    } else {
+        ssa_to_linear(cp.ssa_path, Some(n))
+    }
 }
 
 /// A Python module implemented in Rust.
 #[pymodule]
 fn cotengrust(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(ssa_to_linear, m)?)?;
     m.add_function(wrap_pyfunction!(find_subgraphs, m)?)?;
     m.add_function(wrap_pyfunction!(optimize_simplify, m)?)?;
     m.add_function(wrap_pyfunction!(optimize_greedy, m)?)?;

cotengrust-0.1.1/tests/test_cotengrust.py

@@ -0,0 +1,205 @@
+import pytest
+
+try:
+    import cotengra as ctg
+
+    ctg_missing = False
+except ImportError:
+    ctg_missing = True
+    ctg = None
+
+import cotengrust as ctgr
+
+
+requires_cotengra = pytest.mark.skipif(ctg_missing, reason="requires cotengra")
+
+
+@pytest.mark.parametrize("which", ["greedy", "optimal"])
+def test_basic_call(which):
+    inputs = [('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'a')]
+    output = ('b', 'd')
+    size_dict = {'a': 2, 'b': 3, 'c': 4, 'd': 5}
+    path = {
+        "greedy": ctgr.optimize_greedy,
+        "optimal": ctgr.optimize_optimal,
+    }[
+        which
+    ](inputs, output, size_dict)
+    assert all(len(con) <= 2 for con in path)
+
+
+def find_output_str(lhs):
+    tmp_lhs = lhs.replace(",", "")
+    return "".join(s for s in sorted(set(tmp_lhs)) if tmp_lhs.count(s) == 1)
+
+
+def eq_to_inputs_output(eq):
+    if "->" not in eq:
+        eq += "->" + find_output_str(eq)
+    inputs, output = eq.split("->")
+    inputs = inputs.split(",")
+    inputs = [list(s) for s in inputs]
+    output = list(output)
+    return inputs, output
+
+
+def get_rand_size_dict(inputs, d_min=2, d_max=3):
+    import random
+
+    size_dict = {}
+    for term in inputs:
+        for ix in term:
+            if ix not in size_dict:
+                size_dict[ix] = random.randint(d_min, d_max)
+    return size_dict
+
+
+# these are taken from opt_einsum
+test_case_eqs = [
+    # Test scalar-like operations
+    "a,->a",
+    "ab,->ab",
+    ",ab,->ab",
+    ",,->",
+    # Test hadamard-like products
+    "a,ab,abc->abc",
+    "a,b,ab->ab",
+    # Test index-transformations
+    "ea,fb,gc,hd,abcd->efgh",
+    "ea,fb,abcd,gc,hd->efgh",
+    "abcd,ea,fb,gc,hd->efgh",
+    # Test complex contractions
+    "acdf,jbje,gihb,hfac,gfac,gifabc,hfac",
+    "cd,bdhe,aidb,hgca,gc,hgibcd,hgac",
+    "abhe,hidj,jgba,hiab,gab",
+    "bde,cdh,agdb,hica,ibd,hgicd,hiac",
+    "chd,bde,agbc,hiad,hgc,hgi,hiad",
+    "chd,bde,agbc,hiad,bdi,cgh,agdb",
+    "bdhe,acad,hiab,agac,hibd",
+    # Test collapse
+    "ab,ab,c->",
+    "ab,ab,c->c",
+    "ab,ab,cd,cd->",
+    "ab,ab,cd,cd->ac",
+    "ab,ab,cd,cd->cd",
+    "ab,ab,cd,cd,ef,ef->",
+    # Test outer prodcuts
+    "ab,cd,ef->abcdef",
+    "ab,cd,ef->acdf",
+    "ab,cd,de->abcde",
+    "ab,cd,de->be",
+    "ab,bcd,cd->abcd",
+    "ab,bcd,cd->abd",
+    # Random test cases that have previously failed
+    "eb,cb,fb->cef",
+    "dd,fb,be,cdb->cef",
+    "bca,cdb,dbf,afc->",
+    "dcc,fce,ea,dbf->ab",
+    "fdf,cdd,ccd,afe->ae",
+    "abcd,ad",
+    "ed,fcd,ff,bcf->be",
+    "baa,dcf,af,cde->be",
+    "bd,db,eac->ace",
+    "fff,fae,bef,def->abd",
+    "efc,dbc,acf,fd->abe",
+    # Inner products
+    "ab,ab",
+    "ab,ba",
+    "abc,abc",
+    "abc,bac",
+    "abc,cba",
+    # GEMM test cases
+    "ab,bc",
+    "ab,cb",
+    "ba,bc",
+    "ba,cb",
+    "abcd,cd",
+    "abcd,ab",
+    "abcd,cdef",
+    "abcd,cdef->feba",
+    "abcd,efdc",
+    # Inner than dot
+    "aab,bc->ac",
+    "ab,bcc->ac",
+    "aab,bcc->ac",
+    "baa,bcc->ac",
+    "aab,ccb->ac",
+    # Randomly built test caes
+    "aab,fa,df,ecc->bde",
+    "ecb,fef,bad,ed->ac",
+    "bcf,bbb,fbf,fc->",
+    "bb,ff,be->e",
+    "bcb,bb,fc,fff->",
+    "fbb,dfd,fc,fc->",
+    "afd,ba,cc,dc->bf",
+    "adb,bc,fa,cfc->d",
+    "bbd,bda,fc,db->acf",
+    "dba,ead,cad->bce",
+    "aef,fbc,dca->bde",
+]
+
+
+@requires_cotengra
+@pytest.mark.parametrize("eq", test_case_eqs)
+@pytest.mark.parametrize("which", ["greedy", "optimal"])
+def test_manual_cases(eq, which):
+    inputs, output = eq_to_inputs_output(eq)
+    size_dict = get_rand_size_dict(inputs)
+    path = {
+        "greedy": ctgr.optimize_greedy,
+        "optimal": ctgr.optimize_optimal,
+    }[
+        which
+    ](inputs, output, size_dict)
+    assert all(len(con) <= 2 for con in path)
+    tree = ctg.ContractionTree.from_path(
+        inputs, output, size_dict, path=path, check=True
+    )
+    assert tree.is_complete()
+
+
+@requires_cotengra
+@pytest.mark.parametrize("seed", range(10))
+@pytest.mark.parametrize("which", ["greedy", "optimal"])
+def test_basic_rand(seed, which):
+    inputs, output, shapes, size_dict = ctg.utils.rand_equation(
+        n=10,
+        reg=4,
+        n_out=2,
+        n_hyper_in=1,
+        n_hyper_out=1,
+        d_min=2,
+        d_max=3,
+        seed=seed,
+    )
+    path = {
+        "greedy": ctgr.optimize_greedy,
+        "optimal": ctgr.optimize_optimal,
+    }[
+        which
+    ](inputs, output, size_dict)
+    assert all(len(con) <= 2 for con in path)
+    tree = ctg.ContractionTree.from_path(
+        inputs, output, size_dict, path=path, check=True
+    )
+    assert tree.is_complete()
+
+
+@requires_cotengra
+def test_optimal_lattice_eq():
+    inputs, output, _, size_dict = ctg.utils.lattice_equation(
+        [4, 5], d_max=3, seed=42
+    )
+
+    path = ctgr.optimize_optimal(inputs, output, size_dict, minimize='flops')
+    tree = ctg.ContractionTree.from_path(
+        inputs, output, size_dict, path=path
+    )
+    assert tree.contraction_cost() == 3628
+
+    path = ctgr.optimize_optimal(inputs, output, size_dict, minimize='size')
+    assert all(len(con) <= 2 for con in path)
+    tree = ctg.ContractionTree.from_path(
+        inputs, output, size_dict, path=path
+    )
+    assert tree.contraction_width() == pytest.approx(6.754887502163468)

cotengrust-0.1.0/PKG-INFO

@@ -1,8 +0,0 @@
-Metadata-Version: 2.1
-Name: cotengrust
-Version: 0.1.0
-Classifier: Programming Language :: Rust
-Classifier: Programming Language :: Python :: Implementation :: CPython
-Classifier: Programming Language :: Python :: Implementation :: PyPy
-License-File: LICENSE
-Requires-Python: >=3.7