Repository: dgasmith/opt_einsum
Pull request #68 (2018-10-18 03:17:54)

Support more symbols for subscripts (#67)
* Add more symbol support based on hash

* Refactor the converting to subscripts routine a little bit

* Update tests

## Description
Make opt_einsum support anything hashable as subscripts, such as:
```python
contract(arr1, ['left edge', 'bond'], arr2, ['bond', 'right edge'])
```
See #67 

## Questions
- [x] I've tried to add the new feature to the document, but it seems the document does not have a chapter describing the inputs. Maybe I've left out something or maybe this should be solved in a different PR?

## Status
- [x] Ready to go

jcmgray: @dgasmith I think this feels like a natural extension/inclusion to me. 

I did think of an extra detail that might complicate things a bit - currently the ``int`` order specifies the output order, so:
```python
np.einsum(a, [1, 0])
```
performs the transpose and ``opt_einsum`` should match this behaviour. The current implementation will map this to ``'ab'`` which performs no transpose. Potential solutions:

1. Use ``get_symbol(s)`` if ``int`` supplied and allow no mixing with other hashables
2. Compute the output symbols explicitly if not specified, I think we might do this elsewhere already
3. Somehow sort the ``symbol_map`` - not clear how to do this for mixed hashable types


dgasmith: I agree that this turned out to be really quite elegant to integrate and should not cause much confusion. For the above issue, @liwt31 can you add a test like `opt_einsum.contract(np.zeros((2, 3)), [1, 0]).shape == (3, 2)` to make sure this always works? Might be good to add a test that uses `["b", "a"]` as our indices as well.

On the input issue would you mind starting a document that describes inputs? We can handle the canonical cases, if you just document this case it would be a great starting point.

@jcmgray Is allowing mixed hashable types really a flexibility that we need to handle?  

jcmgray: > @jcmgray Is allowing mixed hashable types really a flexibility that we need to handle?

No no, I think 99% use cases are just going to be strings. It would just have been nice and simple to essentially see integers with everything else as the same things: hashable.

So we just need to efficiently identify the integer case. Would just dispatching on the very first index be acceptable?:
```python
if isinstance(subscript_list[0][0], numbers.Integral):
    symbol_map = defaultdict(get_symbol)
else:
    symbols = (get_symbol(i) for i in itertools.count())
    symbol_map = defaultdict(lambda: next(symbols))
```
Or do we want the extra overhead of checking all the arguments are ``int`` or not?

For the docs, it will need to mention that the default order will be based on the input order rather than sorted order if non-integers are used.
dgasmith: I like it, and seems fairly readable with a short explanation. 
jcmgray: > I suppose we all prefer 1?

I think so. No use cases really come to mind where this would be critical, so best to be more stringent I guess. And it if does come up later, would be very easy to add.
dgasmith: This pull request **introduces 1 alert** when merging 33a09d6b2e23783a8ec1d0cb597f93dea7293c5f into 02d15b7fee68decec69950f9b04de993900d87bd - [view on LGTM.com](https://lgtm.com/projects/g/dgasmith/opt_einsum/rev/pr-3fc36f0aa0bd59bf0617ea83cc9f1d50521c8471)

**new alerts:**

* 1 for Unused import

---

*Comment posted by [LGTM.com](https://lgtm.com)*
dgasmith: I agree 1) sounds good for now.
liwt31: > On the input issue would you mind starting a document that describes inputs? We can handle the canonical cases, if you just document this case it would be a great starting point.

@dgasmith I started some document on the input format. The canonical cases are now left blank.
liwt31: I think that's all. Thank you for your comments and advice!
jcmgray: Looks good to me as well, thanks for this!

Patch:
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 7b8b92b..5d0c74c 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -147,5 +147,6 @@ Table of Contents
    :maxdepth: 1
    :caption: Help & Reference:
 
+   input_format
    api
    changelog
diff --git a/docs/source/input_format.rst b/docs/source/input_format.rst
new file mode 100644
index 0000000..fa7b1e7
--- /dev/null
+++ b/docs/source/input_format.rst
@@ -0,0 +1,41 @@
+============
+Input Format
+============
+
+The ``opt_einsum`` package is a drop-in replacement for the ``np.einsum`` function
+and supports all input formats that ``np.einsum`` supports.
+
+TBA: "Normal" inputs
+
+In ``opt_einsum``, you can also put anything hashable and comparable such as `str` in the subscript list.
+Note ``np.einsum`` currently does not support this syntax. A simple example of the syntax is:
+
+.. code:: python
+
+    >>> x, y, z = np.ones((1, 2)), np.ones((2, 2)), np.ones((2, 1))
+    >>> oe.contract(x, ('left', 'bond1'), y, ('bond1', 'bond2'), z, ('bond2', 'right'), ('left', 'right'))
+    array([[4.]])
+
+The subscripts need to be hashable so that ``opt_einsum`` can efficiently process them, and
+they should also be comparable because in this way we can keep a consistent behavior between integer 
+subscripts and other types of subscripts. For example:
+
+.. code:: python
+
+    >>> x = np.array([[0, 1], [2, 0]])
+    >>> oe.contract(x, (0, 1))  # original matrix
+    array([[0, 1],
+           [2, 0]])
+    >>> oe.contract(x, (1, 0)) # the transpose
+    array([[0, 2],
+           [1, 0]])
+    >>> oe.contract(x, ('a', 'b'))  # original matrix, consistent behavior
+    array([[0, 1],
+           [2, 0]])
+    >>> oe.contract(x, ('b', 'a')) # the transpose, consistent behavior
+    array([[0, 2],
+           [1, 0]])
+    >>> oe.contract(x, (0, 'a')) # relative sequence undefined, can't determine output
+    TypeError: For this input type lists must contain either Ellipsis or hashable and comparable object (e.g. int, str)
+
+
diff --git a/opt_einsum/parser.py b/opt_einsum/parser.py
index 910a977..8c93213 100644
--- a/opt_einsum/parser.py
+++ b/opt_einsum/parser.py
@@ -4,6 +4,7 @@
 A functionally equivalent parser of the numpy.einsum input parser
 """
 
+import itertools
 from collections import OrderedDict
 
 import numpy as np
@@ -187,6 +188,26 @@ def possibly_convert_to_numpy(x):
         return x
 
 
+def convert_subscripts(old_sub, symbol_map):
+    """Convert user custom subscripts list to subscript string according to `symbol_map`.
+
+    Examples
+    --------
+    >>>  oe.parser.convert_subscripts(['abc', 'def'], {'abc':'a', 'def':'b'})
+    'ab'
+    >>> oe.parser.convert_subscripts([Ellipsis, object], {object:'a'})
+    '...a'
+    """
+    new_sub = ""
+    for s in old_sub:
+        if s is Ellipsis:
+            new_sub += "..."
+        else:
+            # no need to try/except here because symbol_map has already been checked
+            new_sub += symbol_map[s]
+    return new_sub
+
+
 def parse_einsum_input(operands):
     """
     A reproduction of einsum c side einsum parsing in python.
@@ -230,28 +251,27 @@ def parse_einsum_input(operands):
 
         output_list = tmp_operands[-1] if len(tmp_operands) else None
         operands = [possibly_convert_to_numpy(x) for x in operand_list]
-        subscripts = ""
-        last = len(subscript_list) - 1
-        for num, sub in enumerate(subscript_list):
-            for s in sub:
-                if s is Ellipsis:
-                    subscripts += "..."
-                elif isinstance(s, int):
-                    subscripts += get_symbol(s)
-                else:
-                    raise TypeError("For this input type lists must contain either int or Ellipsis")
-            if num != last:
-                subscripts += ","
 
+        # build a map from user symbols to single-character symbols based on `get_symbol`
+        # The map retains the intrinsic order of user symbols
+        try:
+            # collect all user symbols
+            symbol_set = set(itertools.chain.from_iterable(subscript_list))
+        except TypeError:  # unhashable object
+            raise TypeError("For this input type lists must contain either Ellipsis or hashable and comparable object (e.g. int, str)")
+        # remove Ellipsis because it can not be compared with other objects
+        if Ellipsis in symbol_set:
+            symbol_set.remove(Ellipsis)
+        try:
+            # build the map based on sorted user symbols, retaining the order we lost in the `set`
+            symbol_map = {symbol: get_symbol(idx) for idx, symbol in enumerate(sorted(symbol_set))}
+        except TypeError:  # uncomparable object
+            raise TypeError("For this input type lists must contain either Ellipsis or hashable and comparable object (e.g. int, str)")
+
+        subscripts = ','.join(convert_subscripts(sub, symbol_map) for sub in subscript_list)
         if output_list is not None:
             subscripts += "->"
-            for s in output_list:
-                if s is Ellipsis:
-                    subscripts += "..."
-                elif isinstance(s, int):
-                    subscripts += get_symbol(s)
-                else:
-                    raise TypeError("For this input type lists must contain either int or Ellipsis")
+            subscripts += convert_subscripts(output_list, symbol_map)
 
     # Check for proper "->"
     if ("-" in subscripts) or (">" in subscripts):


Test patch:
diff --git a/opt_einsum/tests/test_input.py b/opt_einsum/tests/test_input.py
index 311356e..b82d59e 100644
--- a/opt_einsum/tests/test_input.py
+++ b/opt_einsum/tests/test_input.py
@@ -73,11 +73,12 @@ def test_type_errors():
     string = '...a->...a'
     views = build_views(string)
 
+    # Subscript list must contain Ellipsis or (hashable && comparable) object
     with pytest.raises(TypeError):
-        contract(views[0], [Ellipsis, 'a'], [Ellipsis, 0])
+        contract(views[0], [Ellipsis, 0], [Ellipsis, ['a']])
 
     with pytest.raises(TypeError):
-        contract(views[0], [Ellipsis, 0], [Ellipsis, 'a'])
+        contract(views[0], [Ellipsis, dict()], [Ellipsis, 'a'])
 
 
 def test_value_errors():
@@ -236,3 +237,19 @@ def test_large_int_input_format():
     string_output = contract(string, x, y, z)
     int_output = contract(x, (1000, 1001), y, (1001, 1002), z, (1002, 1003))
     assert np.allclose(string_output, int_output)
+    for i in range(10):
+        transpose_output = contract(x, (i+1, i))
+        assert np.allclose(transpose_output, x.T)
+
+
+def test_hashable_object_input_format():
+    string = 'ab,bc,cd'
+    x, y, z = build_views(string)
+    string_output = contract(string, x, y, z)
+    hash_output1 = contract(x, ('left', 'bond1'), y, ('bond1', 'bond2'), z, ('bond2', 'right'))
+    hash_output2 = contract(x, ('left', 'bond1'), y, ('bond1', 'bond2'), z, ('bond2', 'right'), ('left', 'right'))
+    assert np.allclose(string_output, hash_output1)
+    assert np.allclose(hash_output1, hash_output2)
+    for i in range(1, 10):
+        transpose_output = contract(x, ('b' * i, 'a' * i))
+        assert np.allclose(transpose_output, x.T)