jcmgray / cotengra

Hyper optimized contraction trees for large tensor networks and einsums
https://cotengra.readthedocs.io
Apache License 2.0
174 stars 32 forks source link

Hypergraph plot fails after contraction #36

Closed nathanieltornow closed 1 month ago

nathanieltornow commented 1 month ago

Hi! :)

While playing around with the Hypergraph class, I found an issue (probably) with the Hypergraph.contract() method. In particular, after contracting two nodes, the Hypergraph.plot() function fails.

Here is an example:

import cotengra as ctg

inputs = [
    ('a', 'b', 'x'),
    ('b', 'c', 'd'),
    ('c', 'e', 'y'),
    ('e', 'a', 'd'),
]
output = ('x', 'y')
size_dict = {'x': 2, 'y': 3, 'a': 4, 'b': 5, 'c': 6, 'd': 7, 'e': 8}

hg = ctg.HyperGraph(inputs, output, size_dict)
hg.contract(0, 1) # comment this line to make the code work
hg.plot()

I would expect this code to plot the updated hypergraph with contracted nodes {0, 1} into the new node 4. However, code fails with:

Traceback (most recent call last):
  File "/home/nate/test/test.py", line 14, in <module>
    hg.plot()
  File "/home/nate/qvm/.venv/lib/python3.11/site-packages/cotengra/plot.py", line 15, in wrapped
    fig, ax = fn(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^
  File "/home/nate/qvm/.venv/lib/python3.11/site-packages/cotengra/plot.py", line 55, in new_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/nate/qvm/.venv/lib/python3.11/site-packages/cotengra/plot.py", line 1290, in plot_hypergraph
    hypergraph_compute_plot_info_G(
  File "/home/nate/qvm/.venv/lib/python3.11/site-packages/cotengra/plot.py", line 455, in hypergraph_compute_plot_info_G
    color = _node_colorer(nd)
            ^^^^^^^^^^^^^^^^^
  File "/home/nate/qvm/.venv/lib/python3.11/site-packages/cotengra/plot.py", line 432, in _node_colorer
    return node_colors[nd]
           ~~~~~~~~~~~^^^^
IndexError: list index out of range

Is this expected behavior? If not, there is probably some missing "bookkeeping" in the contract() method to ensure indices are correct.

Thank you! Nate

(I'm on cotengra version 0.6.2)

jcmgray commented 1 month ago

Thanks for the issue @nathanieltornow! should be fixed by https://github.com/jcmgray/cotengra/commit/e19e7db2d4796b83df5f296251a038981b7dbe3e.