tn.draw: show abelian signature

Author: jcmgray

docs/changelog.md

@@ -12,6 +12,7 @@ Release notes for `quimb`.
 - [`schematic.Drawing`](quimb.schematic.drawing): add [`grid`](quimb.schematic.drawing.grid), [`grid3d`](quimb.schematic.drawing.grid3d), [`bezier`](quimb.schematic.drawing.bezier), [`star`](quimb.schematic.drawing.star), [`cross`](quimb.schematic.drawing.cross) and [`zigzag`](quimb.schematic.drawing.zigzag) methods.
 - [`schematic.Drawing`](quimb.schematic.drawing): add `relative` option to [`arrowhead`](quimb.schematic.drawing.arrowhead), `shorten` option to [`text_between`](quimb.schematic.drawing.text_between) and `text_left` and `text_right` options to [`line`](quimb.schematic.drawing.line).
 - refactor [`TEBDGen`](quimb.tensor.tensor_arbgeom_tebd.TEBDGen) and [`SimpleUpdateGen`](quimb.tensor.tensor_arbgeom_tebd.SimpleUpdateGen)
+- [`tn.draw()`](quimb.tensor.drawing.draw_tn): show abelian signature if using `symmray` arrays.
 
 (whats-new-1-11-2)=
 ## v1.11.2 (2025-07-30)

quimb/tensor/drawing.py

@@ -379,26 +379,44 @@ def draw_tn(
                 # dummy hyper outer edge - no arrows
                 edges[pair]["arrow_left"].append(False)
                 edges[pair]["arrow_right"].append(False)
+                edges[pair]["label_left"].append(None)
+                edges[pair]["label_right"].append(None)
             else:
                 # tensor side can always have an incoming arrow
-                tl_left_inds = tn.tensor_map[pair[0]].left_inds
+                tl = tn.tensor_map[pair[0]]
+                tl_left_inds = tl.left_inds
                 edges[pair]["arrow_left"].append(
                     show_left_inds
                     and (tl_left_inds is not None)
                     and (ix in tl_left_inds)
                 )
+
+                if hasattr(tl.data, "signature"):
+                    sigl = tl.data.signature[tl.inds.index(ix)]
+                    edges[pair]["label_left"].append(sigl)
+                else:
+                    edges[pair]["label_left"].append(None)
+
                 if ishyper:
                     # hyper edge can't have an incoming arrow
                     edges[pair]["arrow_right"].append(False)
+                    edges[pair]["label_right"].append(None)
                 else:
                     # standard edge can
-                    tr_left_inds = tn.tensor_map[pair[1]].left_inds
+                    tr = tn.tensor_map[pair[1]]
+                    tr_left_inds = tr.left_inds
                     edges[pair]["arrow_right"].append(
                         show_left_inds
                         and (tr_left_inds is not None)
                         and (ix in tr_left_inds)
                     )
 
+                    if hasattr(tr.data, "signature"):
+                        sigr = tr.data.signature[tr.inds.index(ix)]
+                        edges[pair]["label_right"].append(sigr)
+                    else:
+                        edges[pair]["label_right"].append(None)
+
     # parse all tensors / nodes
     for tid, t in tn.tensor_map.items():
         nodes[tid]["tid"] = tid
@@ -519,7 +537,11 @@ def draw_tn(
         nodes[node]["coo"] = G.nodes[node]["coo"] = pos[node]
 
     for (i, j), edge_data in edges.items():
-        edges[i, j]["coos"] = G.edges[i, j]["coos"] = pos[i], pos[j]
+        edge_data["coos"] = G.edges[i, j]["coos"] = pos[i], pos[j]
+        edge_data["shorten"] = G.edges[i, j]["shorten"] = (
+            nodes[i]["size"],
+            nodes[j]["size"],
+        )
 
     if get == "pos":
         return pos
@@ -677,7 +699,7 @@ def _draw_matplotlib(
         fig = None
 
     arrow_opts = arrow_opts or {}
-    arrow_opts.setdefault("center", 3 / 4)
+    arrow_opts.setdefault("center", 0.8)
     arrow_opts.setdefault("linewidth", 1)
     arrow_opts.setdefault("width", 0.08)
     arrow_opts.setdefault("length", 0.12)
@@ -692,7 +714,10 @@ def _draw_matplotlib(
         labels = edge_data["label"]
         arrow_lefts = edge_data["arrow_left"]
         arrow_rights = edge_data["arrow_right"]
+        label_lefts = edge_data["label_left"]
+        label_rights = edge_data["label_right"]
         multiplicity = len(edge_colors)
+        shorten = edge_data["shorten"]
 
         if multiplicity > 1:
             offsets = np.linspace(
@@ -731,11 +756,25 @@ def _draw_matplotlib(
                     color=edge_data["label_color"],
                     fontfamily=edge_data["label_fontfamily"],
                 )
+            if label_lefts[m]:
+                line_opts["text_left"] = dict(
+                    text=label_lefts[m],
+                    fontsize=edge_data["label_fontsize"] + 3,
+                    color=edge_data["label_color"],
+                    fontfamily=edge_data["label_fontfamily"],
+                )
+            if label_rights[m]:
+                line_opts["text_right"] = dict(
+                    text=label_rights[m],
+                    fontsize=edge_data["label_fontsize"] + 3,
+                    color=edge_data["label_color"],
+                    fontfamily=edge_data["label_fontfamily"],
+                )
 
             if multiplicity > 1:
                 d.line_offset(offset=offsets[m], **line_opts)
             else:
-                d.line(**line_opts)
+                d.line(shorten=shorten, **line_opts)
 
     # draw the tensors
     for _, node_data in nodes.items():

quimb/tensor/tensor_1d.py

@@ -1971,7 +1971,7 @@ def add_MPS(self, other, inplace=False, **kwargs):
     add_MPS_ = functools.partialmethod(add_MPS, inplace=True)
 
     def permute_arrays(self, shape="lrp"):
-        """Permute the indices of each tensor in this MPS to match ``shape``.
+        """Ensure the arrays are stored internally in the specified order.
         This doesn't change how the overall object interacts with other tensor
         networks but may be useful for extracting the underlying arrays
         consistently. This is an inplace operation.

quimb/tensor/tensor_core.py

@@ -1725,6 +1725,9 @@ def check(self):
                 f"Tensor data contains non-finite values: {self.data}."
             )
 
+        if hasattr(self.data, "check"):
+            self.data.check()
+
     @property
     def owners(self):
         return self._owners
@@ -4539,6 +4542,10 @@ def check(self):
                     "Mismatched index dimension for index "
                     f"'{ix}' in tensors {ts}."
                 )
+            if len(ts) == 2 and hasattr(ts[0].data, "check_with"):
+                axa = ts[0].inds.index(ix)
+                axb = ts[1].inds.index(ix)
+                ts[0].data.check_with(ts[1].data, [axa], [axb])
 
     def add_tag(self, tag, where=None, which="all", record=None):
         """Add tag(s) to every tensor in this network, or if ``where`` is
@@ -5498,7 +5505,7 @@ def _select_local_tids(
                 )
 
                 # full-rank decompose the outer tensor
-                l, r = self.tensor_map[tid_out].split(
+                _, r = self.tensor_map[tid_out].split(
                     left_inds=None,
                     right_inds=[ix],
                     max_bond=None,
@@ -9593,7 +9600,7 @@ def split(self, left_inds, right_inds=None, **split_opts):
 
     def trace(self, left_inds, right_inds, **contract_opts):
         """Trace over ``left_inds`` joined with ``right_inds``"""
-        tn = self.reindex({u: l for u, l in zip(left_inds, right_inds)})
+        tn = self.reindex(dict(zip(left_inds, right_inds)))
         return tn.contract_tags(..., **contract_opts)
 
     def to_dense(self, *inds_seq, to_qarray=False, **contract_opts):