Repository: dgasmith/opt_einsum
Pull request #12 (2018-02-03 17:45:40)

Improvements for large networks of tensors
This PR contains a couple of improvements for efficiently contracting large networks of tensors, e.g. MPS, PEPS, MERA.

1. Allow ~1000 unique indices to be used. This is only for python3, which uses unicode strings by default. The only trick here really is that calls (e.g. pairwise contractions) to np.einsum must have their string mapped into [a-z,A-Z] with a bit of care taken about ordering.

2. Improve the performance significantly when finding the greedy path for many tensors. This is partly from refactoring ``find_contraction``, which was a hot spot, but mainly from ignoring any potential outer products (which are never advantageous to perform eagerly) while potential inner products are remaining.

For example, finding a path for 200 tensors with ~300 unique indices takes < ~1sec. Although one eventually might want to speed things along with a few manual hints, this is pretty useful functionality.
And from a few small tests, it doesn't seem like much overhead has been introduced for small and/or heavily connected graphs.

Comments obviously welcome!


Patch:
diff --git a/README.md b/README.md
index 9e4ed20..43ecb82 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,7 @@ Optimized Einsum: A tensor contraction order optimizer
 
  - [Optimizing numpy's einsum function](https://github.com/dgasmith/opt_einsum/blob/master/README.md#optimizing-numpys-einsum-function)
  - [Obtaining the path expression](https://github.com/dgasmith/opt_einsum/blob/master/README.md#obtaining-the-path-expression)
+ - [Reusing paths](https://github.com/dgasmith/opt_einsum/blob/master/README.md#reusing-paths-using-contract_expression)
  - [More details on paths](https://github.com/dgasmith/opt_einsum/blob/master/README.md#more-details-on-paths)
  - [Finding the optimal path](https://github.com/dgasmith/opt_einsum/blob/master/README.md#finding-the-optimal-path)
  - [Finding the opportunistic path](https://github.com/dgasmith/opt_einsum/blob/master/README.md#finding-the-opportunistic-path)
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index 4efb1a9..eda2a54 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -118,7 +118,6 @@ def contract_path(*operands, **kwargs):
 
     # Python side parsing
     input_subscripts, output_subscript, operands = parser.parse_einsum_input(operands)
-    subscripts = input_subscripts + '->' + output_subscript
 
     # Build a few useful list and sets
     input_list = input_subscripts.split(',')
@@ -155,7 +154,7 @@ def contract_path(*operands, **kwargs):
             if memory_limit == -1:
                 memory_arg = int(1e20)
             else:
-                raise ValidationError("Memory limit must be larger than 0, or -1")
+                raise ValueError("Memory limit must be larger than 0, or -1")
         else:
             memory_arg = int(memory_limit)
 
@@ -255,6 +254,26 @@ def contract_path(*operands, **kwargs):
     return path, path_print
 
 
+def _np_einsum(*operands, **kwargs):
+    """Numpy einsum, but with pre-parse for valid characters if string given.
+    """
+    if not isinstance(operands[0], str):
+        return np.einsum(*operands, **kwargs)
+
+    einsum_str, operands = operands[0], operands[1:]
+
+    # Do we need to temporarily map indices into [a-z,A-Z] range?
+    if not parser.has_valid_einsum_chars_only(einsum_str):
+
+        # Explicitly find output str first so as to maintain order
+        if '->' not in einsum_str:
+            einsum_str += '->' + parser.find_output_str(einsum_str)
+
+        einsum_str = parser.convert_to_valid_einsum_chars(einsum_str)
+
+    return np.einsum(einsum_str, *operands, **kwargs)
+
+
 # Rewrite einsum to handle different cases
 def contract(*operands, **kwargs):
     """
@@ -332,7 +351,7 @@ def contract(*operands, **kwargs):
 
     # If no optimization, run pure einsum
     if optimize_arg is False:
-        return np.einsum(*operands, **einsum_kwargs)
+        return _np_einsum(*operands, **einsum_kwargs)
 
     # Grab non-einsum kwargs
     use_blas = kwargs.pop('use_blas', True)
@@ -401,7 +420,7 @@ def _core_contract(operands, contraction_list, **einsum_kwargs):
             if (tensor_result != results_index) or handle_out:
                 if handle_out:
                     einsum_kwargs["out"] = out_array
-                new_view = np.einsum(tensor_result + '->' + results_index, new_view, **einsum_kwargs)
+                new_view = _np_einsum(tensor_result + '->' + results_index, new_view, **einsum_kwargs)
 
         # Call einsum
         else:
@@ -410,7 +429,7 @@ def _core_contract(operands, contraction_list, **einsum_kwargs):
                 einsum_kwargs["out"] = out_array
 
             # Do the contraction
-            new_view = np.einsum(einsum_str, *tmp_operands, **einsum_kwargs)
+            new_view = _np_einsum(einsum_str, *tmp_operands, **einsum_kwargs)
 
         # Append new items and derefernce what we can
         operands.append(new_view)
@@ -446,7 +465,7 @@ class ContractExpression:
         try:
             return _core_contract(list(arrays), self.contraction_list, out=out, **self.einsum_kwargs)
         except ValueError as err:
-            original_msg = "".join(err.args) if err.args else ""
+            original_msg = str(err.args) if err.args else ""
             msg = ("Internal error while evaluating `ContractExpression`. Note that few checks are performed"
                    " - the number and rank of the array arguments must match the original expression. "
                    "The internal error was: '%s'" % original_msg, )
diff --git a/opt_einsum/helpers.py b/opt_einsum/helpers.py
index 5072001..dfc88ca 100644
--- a/opt_einsum/helpers.py
+++ b/opt_einsum/helpers.py
@@ -4,6 +4,7 @@ Contains helper functions for opt_einsum testing scripts
 
 from __future__ import division, absolute_import, print_function
 
+import itertools
 import numpy as np
 
 chars = 'abcdefghijklmopq'
@@ -118,15 +119,13 @@ def find_contraction(positions, input_sets, output_set):
     ({'a', 'c'}, [{'a', 'c'}, {'a', 'c'}], {'b', 'd'}, {'a', 'b', 'c', 'd'})
     """
 
+    remaining = list(input_sets)
     idx_contract = set()
-    idx_remain = output_set.copy()
-    remaining = []
-    for ind, value in enumerate(input_sets):
-        if ind in positions:
-            idx_contract |= value
-        else:
-            remaining.append(value)
-            idx_remain |= value
+
+    for i in sorted(positions, reverse=True):
+        idx_contract |= remaining.pop(i)
+
+    idx_remain = set(itertools.chain(output_set, *remaining))
 
     new_result = idx_remain & idx_contract
     idx_removed = (idx_contract - new_result)
diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py
index d490703..e973826 100644
--- a/opt_einsum/parser.py
+++ b/opt_einsum/parser.py
@@ -3,11 +3,61 @@ A functionally equivalent parser of the numpy.einsum input parser
 """
 
 import numpy as np
+import sys
+
+einsum_symbols_base = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
+einsum_symbols = str(einsum_symbols_base)
+
+# boost the number of symbols using unicode if python3
+if sys.version_info[0] >= 3:
+    einsum_symbols += ''.join(map(chr, range(193, 688)))
+    einsum_symbols += ''.join(map(chr, range(913, 1367)))
 
-einsum_symbols = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
 einsum_symbols_set = set(einsum_symbols)
 
 
+def is_valid_einsum_char(x):
+    """Check if the character ``x`` is valid for numpy einsum.
+    """
+    return (x in einsum_symbols_base) or (x in ',->.')
+
+
+def has_valid_einsum_chars_only(einsum_str):
+    """Check if ``einsum_str`` contains only valid characters for numpy einsum.
+    """
+    return all(map(is_valid_einsum_char, einsum_str))
+
+
+def convert_to_valid_einsum_chars(einsum_str):
+    """Convert the str ``einsum_str`` to contain only the alphabetic characters
+    valid for numpy einsum.
+    """
+    # partition into valid and invalid sets
+    valid, invalid = set(), set()
+    for x in einsum_str:
+        (valid if is_valid_einsum_char(x) else invalid).add(x)
+
+    # get replacements for invalid chars that are not already used
+    available = (x for x in einsum_symbols if x not in valid)
+
+    # map invalid to available and replace in the inputs
+    replacer = dict(zip(invalid, available))
+    return "".join(replacer.get(x, x) for x in einsum_str)
+
+
+def find_output_str(subscripts):
+    """Find the output string for the inputs ``susbcripts``.
+    """
+    tmp_subscripts = subscripts.replace(",", "")
+    output_subscript = ""
+    for s in sorted(set(tmp_subscripts)):
+        if s not in einsum_symbols_set:
+            raise ValueError("Character %s is not a valid symbol." % s)
+        if tmp_subscripts.count(s) == 1:
+            output_subscript += s
+    return output_subscript
+
+
 def parse_einsum_input(operands):
     """
     A reproduction of einsum c side einsum parsing in python.
@@ -45,7 +95,7 @@ def parse_einsum_input(operands):
         for s in subscripts:
             if s in '.,->':
                 continue
-            if s not in einsum_symbols:
+            if s not in einsum_symbols_set:
                 raise ValueError("Character %s is not a valid symbol." % s)
 
     else:
@@ -135,13 +185,7 @@ def parse_einsum_input(operands):
             subscripts += "->" + output_sub.replace("...", out_ellipse)
         else:
             # Special care for outputless ellipses
-            output_subscript = ""
-            tmp_subscripts = subscripts.replace(",", "")
-            for s in sorted(set(tmp_subscripts)):
-                if s not in einsum_symbols:
-                    raise ValueError("Character %s is not a valid symbol." % s)
-                if tmp_subscripts.count(s) == 1:
-                    output_subscript += s
+            output_subscript = find_output_str(subscripts)
             normal_inds = ''.join(sorted(set(output_subscript) - set(out_ellipse)))
 
             subscripts += "->" + out_ellipse + normal_inds
@@ -150,15 +194,7 @@ def parse_einsum_input(operands):
     if "->" in subscripts:
         input_subscripts, output_subscript = subscripts.split("->")
     else:
-        input_subscripts = subscripts
-        # Build output subscripts
-        tmp_subscripts = subscripts.replace(",", "")
-        output_subscript = ""
-        for s in sorted(set(tmp_subscripts)):
-            if s not in einsum_symbols:
-                raise ValueError("Character %s is not a valid symbol." % s)
-            if tmp_subscripts.count(s) == 1:
-                output_subscript += s
+        input_subscripts, output_subscript = subscripts, find_output_str(subscripts)
 
     # Make sure output subscripts are in the input
     for char in output_subscript:
diff --git a/opt_einsum/paths.py b/opt_einsum/paths.py
index 0fda762..49de5a6 100644
--- a/opt_einsum/paths.py
+++ b/opt_einsum/paths.py
@@ -1,7 +1,7 @@
 """
 Contains the path technology behind opt_einsum in addition to several path helpers
 """
-
+from itertools import combinations
 from . import helpers
 
 
@@ -41,10 +41,7 @@ def optimal(input_sets, output_set, idx_dict, memory_limit):
         iter_results = []
 
         # Compute all unique pairs
-        comb_iter = []
-        for x in range(len(input_sets) - iteration):
-            for y in range(x + 1, len(input_sets) - iteration):
-                comb_iter.append((x, y))
+        comb_iter = tuple(combinations(range(len(input_sets) - iteration), 2))
 
         for curr in full_results:
             cost, positions, remaining = curr
@@ -84,6 +81,99 @@ def optimal(input_sets, output_set, idx_dict, memory_limit):
     return path
 
 
+def _parse_possible_contraction(positions, input_sets, output_set, idx_dict,
+                                memory_limit, path_cost, naive_cost):
+    """Compute the cost (removed size + flops) and resultant indices for
+    performing the contraction specified by ``positions``.
+
+    Parameters
+    ----------
+    positions : tuple of int
+        The locations of the proposed tensors to contract.
+    input_sets : list of sets
+        The indices found on each tensors.
+    output_set : set
+        The output indices of the expression.
+    idx_dict : dict
+        Mapping of each index to its size.
+    memory_limit : int
+        The total allowed size for an intermediary tensor.
+    path_cost : int
+        The contraction cost so far.
+    naive_cost : int
+        The cost of the unoptimized expression.
+
+    Returns
+    -------
+    cost : (int, int)
+        A tuple containing the size of any indices removed, and the flop cost.
+    positions : tuple of int
+        The locations of the proposed tensors to contract.
+    new_input_sets : list of sets
+        The resulting new list of indices if this proposed contraction is performed.
+
+    """
+    # Find the contraction
+    contract = helpers.find_contraction(positions, input_sets, output_set)
+    idx_result, new_input_sets, idx_removed, idx_contract = contract
+
+    # Sieve the results based on memory_limit
+    if helpers.compute_size_by_dict(idx_result, idx_dict) > memory_limit:
+        return None
+
+    # Build sort tuple
+    removed_size = helpers.compute_size_by_dict(idx_removed, idx_dict)
+    cost = helpers.flop_count(idx_contract, idx_removed,
+                              len(positions), idx_dict)
+    sort = (-removed_size, cost)
+
+    # Sieve based on total cost as well
+    if (path_cost + cost) > naive_cost:
+        return None
+
+    # Add contraction to possible choices
+    return [sort, positions, new_input_sets]
+
+
+def _update_other_results(results, best):
+    """Update the positions and provisional input_sets of ``results`` based on
+    performing the contraction result ``best``. Remove any involving the tensors
+    contracted.
+
+    Parameters
+    ----------
+    results :
+        List of contraction results produced by ``_parse_possible_contraction``.
+    best :
+        The best contraction of ``results`` i.e. the one that will be performed.
+
+    Returns
+    -------
+    mod_results :
+        The list of modifed results, updated with outcome of ``best`` contraction.
+    """
+    best_con = best[1]
+    bx, by = best_con
+    mod_results = []
+
+    for cost, (x, y), con_sets in results:
+
+        # ignore results involving tensors just contracted
+        if x in best_con or y in best_con:
+            continue
+
+        # update the input_sets
+        del con_sets[by - int(by > x) - int(by > y)]
+        del con_sets[bx - int(bx > x) - int(bx > y)]
+        con_sets.insert(-1, best[2][-1])
+
+        # update the position indices
+        mod_con = x - int(x > bx) - int(x > by), y - int(y > bx) - int(y > by)
+        mod_results.append((cost, mod_con, con_sets))
+
+    return mod_results
+
+
 def greedy(input_sets, output_set, idx_dict, memory_limit):
     """
     Finds the path by contracting the best pair until the input list is
@@ -127,49 +217,52 @@ def greedy(input_sets, output_set, idx_dict, memory_limit):
     idx_result, new_input_sets, idx_removed, idx_contract = contract
     naive_cost = helpers.flop_count(idx_contract, idx_removed, len(input_sets), idx_dict)
 
-    path_cost = 0
+    comb_iter = combinations(range(len(input_sets)), 2)
+    iteration_results = []
 
+    path_cost = 0
     path = []
-    for iteration in range(len(input_sets) - 1):
-        iteration_results = []
-        comb_iter = []
-
-        # Compute all unique pairs
-        for x in range(len(input_sets)):
-            for y in range(x + 1, len(input_sets)):
-                comb_iter.append((x, y))
 
+    for iteration in range(len(input_sets) - 1):
         for positions in comb_iter:
 
-            # Find the contraction
-            contract = helpers.find_contraction(positions, input_sets, output_set)
-            idx_result, new_input_sets, idx_removed, idx_contract = contract
-
-            # Sieve the results based on memory_limit
-            if helpers.compute_size_by_dict(idx_result, idx_dict) > memory_limit:
+            # always initially ignore outer products
+            if input_sets[positions[0]].isdisjoint(input_sets[positions[1]]):
                 continue
 
-            # Build sort tuple
-            removed_size = helpers.compute_size_by_dict(idx_removed, idx_dict)
-            cost = helpers.flop_count(idx_contract, idx_removed, len(positions), idx_dict)
-            sort = (-removed_size, cost)
+            result = _parse_possible_contraction(positions, input_sets, output_set, idx_dict,
+                                                 memory_limit, path_cost, naive_cost)
+            if result is not None:
+                iteration_results.append(result)
 
-            # Sieve based on total cost as well
-            if (path_cost + cost) > naive_cost:
-                continue
+        # If we did not find a new inner contraction remaining
+        if len(iteration_results) == 0:
 
-            # Add contraction to possible choices
-            iteration_results.append([sort, positions, new_input_sets])
+            # Then check the outer products
+            for positions in combinations(range(len(input_sets)), 2):
+                result = _parse_possible_contraction(positions, input_sets, output_set, idx_dict,
+                                                     memory_limit, path_cost, naive_cost)
+                if result is not None:
+                    iteration_results.append(result)
 
-        # If we did not find a new contraction contract remaining
-        if len(iteration_results) == 0:
-            path.append(tuple(range(len(input_sets))))
-            break
+            # If we still did not find any remaining contraction, contract everything
+            if len(iteration_results) == 0:
+                path.append(tuple(range(len(input_sets))))
+                break
 
         # Sort based on first index
         best = min(iteration_results, key=lambda x: x[0])
-        path.append(best[1])
+
+        # Now propagate as many results as possible to next iteration
+        iteration_results = _update_other_results(iteration_results, best)
+
+        # Next iteration only compute contractions with the new tensor
         input_sets = best[2]
+        new_tensor_pos = len(input_sets) - 1
+        comb_iter = ((i, new_tensor_pos) for i in range(new_tensor_pos))
+
+        # Update path and total cost
+        path.append(best[1])
         path_cost += best[0][1]
 
     return path


Test patch:
diff --git a/tests/test_contract.py b/tests/test_contract.py
index 6885522..c63f798 100644
--- a/tests/test_contract.py
+++ b/tests/test_contract.py
@@ -3,6 +3,7 @@ Tets a series of opt_einsum contraction paths to ensure the results are the same
 """
 
 from __future__ import division, absolute_import, print_function
+import sys
 
 import numpy as np
 from opt_einsum import contract, contract_path, helpers, contract_expression
@@ -109,6 +110,23 @@ def test_compare(string):
     assert np.allclose(ein, opt)
 
 
+@pytest.mark.skipif(sys.version_info[0] < 3, reason='requires python3')
+@pytest.mark.parametrize("string", tests)
+def test_compare_greek(string):
+    views = helpers.build_views(string)
+
+    ein = contract(string, *views, optimize=False, use_blas=False)
+
+    # convert to greek
+    string = ''.join(chr(ord(c) + 848) if c not in ',->.' else c for c in string)
+
+    opt = contract(string, *views, optimize='greedy', use_blas=False)
+    assert np.allclose(ein, opt)
+
+    opt = contract(string, *views, optimize='optimal', use_blas=False)
+    assert np.allclose(ein, opt)
+
+
 @pytest.mark.parametrize("string", tests)
 def test_compare_blas(string):
     views = helpers.build_views(string)
@@ -121,6 +139,33 @@ def test_compare_blas(string):
     assert np.allclose(ein, opt)
 
 
+@pytest.mark.skipif(sys.version_info[0] < 3, reason='requires python3')
+@pytest.mark.parametrize("string", tests)
+def test_compare_blas_greek(string):
+    views = helpers.build_views(string)
+
+    ein = contract(string, *views, optimize=False)
+
+    # convert to greek
+    string = ''.join(chr(ord(c) + 848) if c not in ',->.' else c for c in string)
+
+    opt = contract(string, *views, optimize='greedy')
+    assert np.allclose(ein, opt)
+
+    opt = contract(string, *views, optimize='optimal')
+    assert np.allclose(ein, opt)
+
+
+@pytest.mark.skipif(sys.version_info[0] < 3, reason='requires python3')
+def test_some_non_alphabet_maintains_order():
+    # 'c beta a' should automatically go to -> 'a c beta'
+    string = 'c' + chr(ord('b') + 848) + 'a'
+    # but beta will be temporarily replaced with 'b' for which 'cba->abc'
+    # so check manual output kicks in:
+    x = np.random.rand(2, 3, 4)
+    assert np.allclose(contract(string, x), contract('cxa', x))
+
+
 def test_printing():
     string = "bbd,bda,fc,db->acf"
     views = helpers.build_views(string)
diff --git a/tests/test_paths.py b/tests/test_paths.py
index 8270303..0d9310a 100644
--- a/tests/test_paths.py
+++ b/tests/test_paths.py
@@ -169,3 +169,9 @@ def test_greedy_edge_cases():
 
     path, path_str = oe.contract_path(expression, *tensors, path='greedy', memory_limit=-1)
     assert check_path(path, [(0, 1), (0, 2), (0, 1)])
+
+
+def test_can_optimize_outer_products():
+    a, b, c = [np.random.randn(10, 10) for _ in range(3)]
+    d = np.random.randn(10, 2)
+    assert oe.contract_path("ab,cd,ef,fg", a, b, c, d, path='greedy')[0] == [(2, 3), (0, 2), (0, 1)]