drawing.py

"""Functionailty for drawing tensor networks.
"""
import collections
import importlib
import textwrap
import warnings

import numpy as np

from ..utils import autocorrect_kwargs, check_opt, valmap

# from ..schematic import average_color, darken_color, auto_colors, hash_to_color

HAS_FA2 = importlib.util.find_spec("fa2") is not None


@autocorrect_kwargs
def draw_tn(
    tn,
    color=None,
    *,
    show_inds=None,
    show_tags=None,
    output_inds=None,
    highlight_inds=(),
    highlight_tids=(),
    highlight_inds_color=(1.0, 0.2, 0.2),
    highlight_tids_color=(1.0, 0.2, 0.2),
    custom_colors=None,
    legend="auto",
    dim=2,
    fix=None,
    layout="auto",
    initial_layout="auto",
    refine_layout="auto",
    iterations="auto",
    k=None,
    pos=None,
    node_color=None,
    node_scale=1.0,
    node_size=None,
    node_alpha=1.0,
    node_shape="o",
    node_outline_size=None,
    node_outline_darkness=0.9,
    node_hatch="",
    edge_color=None,
    edge_scale=1.0,
    edge_alpha=1 / 2,
    multi_edge_spread=0.1,
    multi_tag_style="auto",
    show_left_inds=True,
    arrow_opts=None,
    label_color=None,
    font_size=10,
    font_size_inner=7,
    font_family="monospace",
    isdark=None,
    title=None,
    backend="matplotlib",
    figsize=(6, 6),
    margin=None,
    xlims=None,
    ylims=None,
    get=None,
    return_fig=False,
    ax=None,
):
    """Plot this tensor network as a networkx graph using matplotlib,
    with edge width corresponding to bond dimension.

    Parameters
    ----------
    color : sequence of tags, optional
        If given, uniquely color any tensors which have each of the tags.
        If some tensors have more than of the tags, only one color will show.
    output_inds : sequence of str, optional
        For hyper tensor networks explicitly specify which indices should be
        drawn as outer indices. If not set, the outer indices are assumed to be
        those that only appear on a single tensor.
    highlight_inds : iterable, optional
        Highlight these edges.
    highlight_tids : iterable, optional
        Highlight these nodes.
    highlight_inds_color
        What color to use for ``highlight_inds`` nodes.
    highlight_tids_color : tuple[float], optional
        What color to use for ``highlight_tids`` nodes.
    show_inds : {None, False, True, 'all', 'bond-size'}, optional
        Explicitly turn on labels for each tensors indices.
    show_tags : {None, False, True}, optional
        Explicitly turn on labels for each tensors tags.
    custom_colors : sequence of colors, optional
        Supply a custom sequence of colors to match the tags given
        in ``color``.
    title : str, optional
        Set a title for the axis.
    legend : "auto" or bool, optional
        Whether to draw a legend for the colored tags. If ``"auto"`` then
        only draw a legend if there are less than 20 tags.
    dim : {2, 2.5, 3}, optional
        What dimension to position the graph nodes in. 2.5 positions the nodes
        in 3D but then projects then down to 2D.
    fix : dict[tags_ind_or_tid], (float, float)], optional
        Used to specify actual relative positions for each tensor node.
        Each key should be a sequence of tags that uniquely identifies a
        tensor, a ``tid``, or a ``ind``, and each value should be a ``(x, y)``
        coordinate tuple.
    layout : str, optional
        How to layout the graph. Can be any of the following:

            - ``'auto'``: layout the graph using a networkx method then relax
              the layout using a force-directed algorithm.
            - a networkx layout method name, e.g. ``'kamada_kawai'``: just
              layout the graph using a networkx method, with no relaxation.
            - a graphviz method such as ``'dot'``, ``'neato'`` or ``'sfdp'``:
              layout the graph using ``pygraphviz``.

    initial_layout : {'auto', 'spectral', 'kamada_kawai', 'circular', \\
                      'planar', 'random', 'shell', 'bipartite', ...}, optional
        If ``layout == 'auto'`` The name of a networkx layout to use before
        iterating with the spring layout. Set `layout` directly or
        ``iterations=0`` if you don't want any spring relaxation.
    iterations : int, optional
        How many iterations to perform when when finding the best layout
        using node repulsion. Ramp this up if the graph is drawing messily.
    k : float, optional
        The optimal distance between nodes.
    pos : dict, optional
        Pre-computed positions for the nodes. If given, this will override
        ``layout``. The nodes shouuld be exactly the same as the nodes in the
        graph returned by ``draw(get='graph')``.
    node_color : tuple[float], optional
        Default color of nodes.
    node_scale : float, optional
        Scale the node sizes by this factor, in addition to the automatic
        scaling based on the number of tensors.
    node_size : None, float or dict, optional
        How big to draw the tensors. Can be a global single value, or a dict
        containing values for specific tags or tids. This is in absolute
        figure units. See ``node_scale`` simply scale the node sizes up or
        down.
    node_alpha : float, optional
        Transparency of the nodes.
    node_shape : None, str or dict, optional
        What shape to draw the tensors. Should correspond to a matplotlib
        scatter marker. Can be a global single value, or a dict containing
        values for specific tags or tids.
    node_outline_size : None, float or dict, optional
        The width of the border of each node. Can be a global single value, or
        a dict containing values for specific tags or tids.
    node_outline_darkness : float, optional
        Darkening of nodes outlines.
    edge_color : tuple[float], optional
        Default color of edges.
    edge_scale : float, optional
        How much to scale the width of the edges.
    edge_alpha : float, optional
        Set the alpha (opacity) of the drawn edges.
    multi_edge_spread : float, optional
        How much to spread the lines of multi-edges.
    show_left_inds : bool, optional
        Whether to show ``tensor.left_inds`` as incoming arrows.
    arrow_closeness : float, optional
        How close to draw the arrow to its target.
    arrow_length : float, optional
        The size of the arrow with respect to the edge.
    arrow_overhang : float, optional
        Varies the arrowhead between a triangle (0.0) and 'V' (1.0).
    arrow_linewidth : float, optional
        The width of the arrow line itself.
    label_color : tuple[float], optional
        Color to draw labels with.
    font_size : int, optional
        Font size for drawing tags and outer indices.
    font_size_inner : int, optional
        Font size for drawing inner indices.
    font_family : str, optional
        Font family to use for all labels.
    isdark : bool, optional
        Explicitly specify that the background is dark, and use slightly
        different default drawing colors. If not specified detects
        automatically from `matplotlib.rcParams`.
    figsize : tuple of int, optional
        The size of the drawing.
    margin : None or float, optional
        Specify an argument for ``ax.margin``, else the plot limits will try
        and be computed based on the node positions and node sizes.
    xlims : None or tuple, optional
        Explicitly set the x plot range.
    xlims : None or tuple, optional
        Explicitly set the y plot range.
    get : {None, 'pos', 'graph'}, optional
        If ``None`` then plot as normal, else if:

            - ``'pos'``, return the plotting positions of each ``tid`` and
              ``ind`` drawn as a node, this can supplied to subsequent calls as
              ``fix=pos`` to maintain positions, even as the graph structure
              changes.
            - ``'graph'``, return the ``networkx.Graph`` object. Note that this
              will potentially have extra nodes representing output and hyper
              indices.

    return_fig : bool, optional
        If True and ``ax is None`` then return the figure created rather than
        executing ``pyplot.show()``.
    ax : matplotlib.Axis, optional
        Draw the graph on this axis rather than creating a new figure.
    """
    import math

    import matplotlib as mpl
    import networkx as nx
    from matplotlib.colors import to_rgb, to_rgba

    from ..schematic import darken_color, hash_to_color

    check_opt(
        "multi_tag_style",
        multi_tag_style,
        ("auto", "pie", "nest", "average", "last"),
    )

    if output_inds is None:
        output_inds = set(tn.outer_inds())
    elif isinstance(output_inds, str):
        output_inds = {output_inds}
    else:
        output_inds = set(output_inds)

    # automatically decide whether to show tags and inds
    if show_inds is None:
        show_inds = len(tn.outer_inds()) <= 20
    show_inds = {False: "", True: "outer"}.get(show_inds, show_inds)

    if show_tags is None:
        show_tags = len(tn.tag_map) <= 20
    show_tags = {False: "", True: "tags"}.get(show_tags, show_tags)

    if isdark is None:
        isdark = sum(to_rgb(mpl.rcParams["figure.facecolor"])) / 3 < 0.5

    if isdark:
        default_draw_color = (0.55, 0.57, 0.60, 1.0)
        default_label_color = (0.85, 0.86, 0.87, 1.0)
    else:
        default_draw_color = (0.45, 0.47, 0.50, 1.0)
        default_label_color = (0.33, 0.34, 0.35, 1.0)

    if edge_color is None:
        edge_color = mpl.colors.to_rgba(default_draw_color, edge_alpha)
    elif edge_color is True:
        # hash edge to get color
        pass
    else:
        edge_color = mpl.colors.to_rgba(edge_color, edge_alpha)

    if node_color is None:
        node_color = mpl.colors.to_rgba(default_draw_color, node_alpha)
    else:
        node_color = mpl.colors.to_rgba(node_color, node_alpha)

    if label_color is None:
        label_color = default_label_color
    elif label_color == "inherit":
        label_color = mpl.rcParams["axes.labelcolor"]

    # get colors for tagged nodes
    colors = get_colors(color, custom_colors, node_alpha)

    if legend == "auto":
        legend = len(colors) <= 20

    highlight_tids_color = to_rgba(highlight_tids_color, node_alpha)
    highlight_inds_color = to_rgba(highlight_inds_color, edge_alpha)

    # set the size of the nodes and their border
    node_size = parse_dict_to_tids_or_inds(
        node_size,
        tn,
        default=1,
    )
    node_outline_size = parse_dict_to_tids_or_inds(
        node_outline_size,
        tn,
        default=1,
    )
    node_shape = parse_dict_to_tids_or_inds(node_shape, tn, default="o")
    node_hatch = parse_dict_to_tids_or_inds(node_hatch, tn, default="")

    # build the graph
    edges = collections.defaultdict(lambda: collections.defaultdict(list))
    nodes = collections.defaultdict(dict)

    # parse all indices / edges
    for ix, tids in tn.ind_map.items():
        tids = sorted(tids)

        isouter = ix in output_inds
        ishyper = isouter or (len(tids) != 2)
        ind_size = tn.ind_size(ix)
        edge_size = edge_scale * math.log2(ind_size)

        # compute a color for this index
        color = (
            highlight_inds_color
            if ix in highlight_inds
            else to_rgba(hash_to_color(ix))
            if edge_color is True
            else edge_color
        )

        # compute a label for this index
        if ishyper:
            # each tensor connects to the dummy node represeting the hyper edge
            pairs = [(tid, ix) for tid in tids]
            if isouter and len(tids) > 1:
                # 'hyper outer' index
                pairs.append((("outer", ix), ix))
            # hyper labels get put on dummy node
            label = ""

            nodes[ix]["ind"] = ix
            nodes[ix]["ind_size"] = ind_size
            # make actual node invisible
            nodes[ix]["color"] = (1.0, 1.0, 1.0, 1.0)
            nodes[ix]["size"] = 0.0
            nodes[ix]["outline_size"] = 0.0
            nodes[ix]["outline_color"] = (1.0, 1.0, 1.0, 1.0)
            nodes[ix]["marker"] = "."  # set this to avoid warning - size is 0
            nodes[ix]["hatch"] = ""

            # set these for plotly hover info
            nodes[ix]["tid"] = nodes[ix]["shape"] = nodes[ix]["tags"] = ""

            if ((show_inds == "outer") and isouter) or (show_inds == "all"):
                # show as outer index or inner index name
                nodes[ix]["label"] = ix
            elif show_inds == "bond-size":
                # show all bond sizes
                nodes[ix]["label"] = f"{tn.ind_size(ix)}"
            else:
                # labels hidden or inner edge
                nodes[ix]["label"] = ""

            nodes[ix]["label_fontsize"] = font_size_inner
            nodes[ix]["label_color"] = label_color
            nodes[ix]["label_fontfamily"] = font_family

        else:
            # standard edge
            pairs = [tuple(tids)]

            if show_inds == "all":
                # show inner index name
                label = ix
            elif show_inds == "bond-size":
                # show all bond sizes
                label = f"{ind_size}"
            else:
                # labels hidden or inner edge
                label = ""

        for pair in pairs:
            edges[pair]["color"].append(color)
            edges[pair]["ind"].append(ix)
            edges[pair]["ind_size"].append(ind_size)
            edges[pair]["edge_size"].append(edge_size)
            edges[pair]["label"].append(label)
            edges[pair]["label_fontsize"] = font_size_inner
            edges[pair]["label_color"] = label_color
            edges[pair]["label_fontfamily"] = font_family

            if isinstance(pair[0], tuple):
                # dummy hyper outer edge - no arrows
                edges[pair]["arrow_left"].append(False)
                edges[pair]["arrow_right"].append(False)
            else:
                # tensor side can always have an incoming arrow
                tl_left_inds = tn.tensor_map[pair[0]].left_inds
                edges[pair]["arrow_left"].append(
                    show_left_inds
                    and (tl_left_inds is not None)
                    and (ix in tl_left_inds)
                )
                if ishyper:
                    # hyper edge can't have an incoming arrow
                    edges[pair]["arrow_right"].append(False)
                else:
                    # standard edge can
                    tr_left_inds = tn.tensor_map[pair[1]].left_inds
                    edges[pair]["arrow_right"].append(
                        show_left_inds
                        and (tr_left_inds is not None)
                        and (ix in tr_left_inds)
                    )

    # parse all tensors / nodes
    for tid, t in tn.tensor_map.items():
        nodes[tid]["tid"] = tid
        nodes[tid]["tags"] = str(list(t.tags))
        nodes[tid]["shape"] = str(t.shape)
        nodes[tid]["size"] = node_size[tid]
        nodes[tid]["outline_size"] = node_outline_size[tid]
        nodes[tid]["marker"] = node_shape[tid]
        nodes[tid]["hatch"] = node_hatch[tid]

        if show_tags == "tags":
            node_label = ", ".join(map(str, t.tags))
            # make the tags appear with auto vertical extent
            nodes[tid]["label"] = "\n".join(
                textwrap.wrap(node_label, max(2 * len(node_label) ** 0.5, 16))
            )
        elif show_tags == "tids":
            nodes[tid]["label"] = str(tid)
        elif show_tags == "shape":
            nodes[tid]["label"] = nodes[tid]["shape"]
        else:
            nodes[tid]["label"] = ""

        nodes[tid]["label_fontsize"] = font_size
        nodes[tid]["label_color"] = label_color
        nodes[tid]["label_fontfamily"] = font_family

        if tid in highlight_tids:
            nodes[tid]["color"] = highlight_tids_color
            nodes[tid]["outline_color"] = darken_color(
                highlight_tids_color, node_outline_darkness
            )
        else:
            # collect all relevant tag colors
            multi_colors = []
            multi_outline_colors = []
            for tag in colors:
                if tag in t.tags:
                    multi_colors.append(colors[tag])
                    multi_outline_colors.append(
                        darken_color(colors[tag], node_outline_darkness)
                    )

            if len(multi_colors) >= 1:
                # set the basic color to the last tag
                nodes[tid]["color"] = multi_colors[-1]
                nodes[tid]["outline_color"] = multi_outline_colors[-1]
                if len(multi_colors) >= 2:
                    # have multiple relevant tags - store them, but some
                    # backends might support, so store alongside basic color
                    nodes[tid]["multi_colors"] = multi_colors
                    nodes[tid]["multi_outline_colors"] = multi_outline_colors
            else:
                # untagged node
                nodes[tid]["color"] = node_color
                nodes[tid]["outline_color"] = darken_color(
                    node_color, node_outline_darkness**2
                )

    G = nx.Graph()
    for edge, edge_data in edges.items():
        G.add_edge(*edge, **edge_data)
    for node, node_data in nodes.items():
        G.add_node(node, **node_data)

    if pos is None:
        pos = get_positions(
            tn=tn,
            G=G,
            fix=fix,
            layout=layout,
            initial_layout=initial_layout,
            refine_layout=refine_layout,
            k=k,
            dim=dim,
            iterations=iterations,
        )
    else:
        pos = _normalize_positions(pos)

    # compute a base size using the position and number of tensors
    # first get plot volume:
    node_packing_factor = tn.num_tensors**-0.45
    xs, ys, *zs = zip(*pos.values())
    xmin, xmax = min(xs), max(xs)
    ymin, ymax = min(ys), max(ys)
    # if there only a few tensors we don't want to limit the node size
    # because of flatness, also don't allow the plot volume to go to zero
    xrange = max(((xmax - xmin) / 2, node_packing_factor, 0.1))
    yrange = max(((ymax - ymin) / 2, node_packing_factor, 0.1))
    plot_volume = xrange * yrange
    if zs:
        zmin, zmax = min(zs[0]), max(zs[0])
        zrange = max(((zmax - zmin) / 2, node_packing_factor, 0.1))
        plot_volume *= zrange
    # in total we account for:
    #     - user specified scaling
    #     - number of tensors
    #     - how flat the plot area is (flatter requires smaller nodes)
    full_node_scale = 0.2 * node_scale * node_packing_factor * plot_volume**0.5

    default_outline_size = 6 * full_node_scale**0.5

    # update node size and position attributes
    for node, node_data in nodes.items():
        nodes[node]["size"] = G.nodes[node]["size"] = (
            full_node_scale * node_data["size"]
        )
        nodes[node]["outline_size"] = G.nodes[node]["outline_size"] = (
            default_outline_size * node_data["outline_size"]
        )
        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]

    if get == "pos":
        return pos
    if get == "graph,pos":
        return G, pos

    opts = {
        "colors": colors,
        "node_outline_darkness": node_outline_darkness,
        "title": title,
        "legend": legend,
        "multi_edge_spread": multi_edge_spread,
        "multi_tag_style": multi_tag_style,
        "arrow_opts": arrow_opts,
        "label_color": label_color,
        "font_family": font_family,
        "figsize": figsize,
        "margin": margin,
        "xlims": xlims,
        "ylims": ylims,
        "return_fig": return_fig,
        "ax": ax,
    }

    if get == "data":
        return edges, nodes, opts

    if backend == "matplotlib":
        return _draw_matplotlib(edges=edges, nodes=nodes, **opts)

    if backend == "matplotlib3d":
        return _draw_matplotlib3d(G, **opts)

    if backend == "plotly":
        return _draw_plotly(G, **opts)


def parse_dict_to_tids_or_inds(spec, tn, default="__NONE__"):
    """Parse a dictionary possibly containing a mix of tags, tids and inds, to
    a dictionary with only sinlge tids and inds as keys. If a tag or set of
    tags are given as a key, all matching tensor tids will receive the value.
    """
    #
    if (spec is not None) and (not isinstance(spec, dict)):
        # assume new default value for everything
        return collections.defaultdict(lambda: spec)

    # allow not specifying a default value
    if default != "__NONE__":
        new = collections.defaultdict(lambda: default)
    else:
        new = {}

    if spec is None:
        return new

    # parse the special values
    for k, v in spec.items():
        if (
            # given as tid
            (isinstance(k, int) and k in tn.tensor_map)
            or
            # given as ind
            (isinstance(k, str) and k in tn.ind_map)
        ):
            # already a tid
            new[k] = v
            continue

        try:
            for tid in tn._get_tids_from_tags(k):
                new[tid] = v
        except KeyError:
            # just ignore keys that don't match any tensor
            pass

    return new


def _add_legend_matplotlib(
    ax, colors, legend, node_outline_darkness, label_color, font_family
):
    import matplotlib.pyplot as plt

    # create legend
    if colors and legend:
        handles = []
        for color in colors.values():
            ecolor = tuple(
                (1.0 if i == 3 else node_outline_darkness) * c
                for i, c in enumerate(color)
            )
            handles += [
                plt.Line2D(
                    [0],
                    [0],
                    marker="o",
                    color=color,
                    markeredgecolor=ecolor,
                    markeredgewidth=1,
                    linestyle="",
                    markersize=10,
                )
            ]

        # needed in case '_' is the first character
        lbls = [f" {lbl}" for lbl in colors]

        legend = ax.legend(
            handles,
            lbls,
            ncol=max(round(len(handles) / 20), 1),
            loc="center left",
            bbox_to_anchor=(1, 0.5),
            labelcolor=label_color,
            prop={"family": font_family},
        )
        # do this manually as otherwise can't make only face transparent
        legend.get_frame().set_alpha(None)
        legend.get_frame().set_facecolor((0.0, 0.0, 0.0, 0.0))
        legend.get_frame().set_edgecolor((0.6, 0.6, 0.6, 0.2))


def _draw_matplotlib(
    edges,
    nodes,
    *,
    colors=None,
    node_outline_darkness=0.9,
    title=None,
    legend=True,
    multi_edge_spread=0.1,
    multi_tag_style="auto",
    arrow_opts=None,
    label_color=None,
    font_family="monospace",
    figsize=(6, 6),
    margin=None,
    xlims=None,
    ylims=None,
    return_fig=False,
    ax=None,
):
    import matplotlib.pyplot as plt

    from quimb.schematic import Drawing, average_color

    d = Drawing(figsize=figsize, ax=ax)
    if ax is None:
        fig = d.fig
        ax = d.ax
        fig.patch.set_alpha(0.0)
        ax.patch.set_alpha(0.0)
    else:
        fig = None

    arrow_opts = arrow_opts or {}
    arrow_opts.setdefault("center", 3 / 4)
    arrow_opts.setdefault("linewidth", 1)
    arrow_opts.setdefault("width", 0.08)
    arrow_opts.setdefault("length", 0.12)

    if title is not None:
        ax.set_title(str(title))

    for _, edge_data in edges.items():
        cooa, coob = edge_data["coos"]
        edge_colors = edge_data["color"]
        edge_sizes = edge_data["edge_size"]
        labels = edge_data["label"]
        arrow_lefts = edge_data["arrow_left"]
        arrow_rights = edge_data["arrow_right"]
        multiplicity = len(edge_colors)

        if multiplicity > 1:
            offsets = np.linspace(
                +multiplicity * multi_edge_spread / 2,
                -multiplicity * multi_edge_spread / 2,
                multiplicity,
            )
        else:
            offsets = None

        for m in range(multiplicity):
            line_opts = dict(
                cooa=cooa,
                coob=coob,
                linewidth=edge_sizes[m],
                color=edge_colors[m],
            )

            arrowhead, reverse = {
                (False, False): (None, False),  # no arrow
                (False, True): (True, False),  # arrowhead to right
                (True, False): (True, True),  # arrowhead to left
                (True, True): (True, "both"),  # arrowheads both sides
            }[arrow_lefts[m], arrow_rights[m]]

            if arrowhead:
                line_opts["arrowhead"] = dict(
                    reverse=reverse,
                    **arrow_opts,
                )

            if labels[m]:
                line_opts["text"] = dict(
                    text=labels[m],
                    fontsize=edge_data["label_fontsize"],
                    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)

    # draw the tensors
    for _, node_data in nodes.items():
        patch_opts = dict(
            coo=node_data["coo"],
            radius=node_data["size"],
            facecolor=node_data["color"],
            edgecolor=node_data["outline_color"],
            linewidth=node_data["outline_size"],
            hatch=node_data["hatch"],
        )
        marker = node_data["marker"]

        if "multi_colors" in node_data:
            # tensor has multiple tags which are colored

            if multi_tag_style in ("pie", "auto"):
                # draw a mini pie chart
                if marker not in ("o", "."):
                    warnings.warn(
                        "Can only draw multi-colored nodes as circles."
                    )

                angles = np.linspace(
                    0, 360, len(node_data["multi_colors"]) + 1
                )
                for i, (color, outline_color) in enumerate(
                    zip(
                        node_data["multi_colors"],
                        node_data["multi_outline_colors"],
                    )
                ):
                    patch_opts["facecolor"] = color
                    patch_opts["edgecolor"] = outline_color
                    d.wedge(
                        theta1=angles[i] - 67.5,
                        theta2=angles[i + 1] - 67.5,
                        **patch_opts,
                    )
            elif multi_tag_style == "nest":
                # draw nested markers of decreasing size
                radii = np.linspace(
                    node_data["size"], 0, len(node_data["multi_colors"]) + 1
                )
                for i, (color, outline_color) in enumerate(
                    zip(
                        node_data["multi_colors"],
                        node_data["multi_outline_colors"],
                    )
                ):
                    patch_opts["facecolor"] = color
                    patch_opts["edgecolor"] = outline_color
                    d.marker(
                        marker=marker,
                        **{**patch_opts, "radius": radii[i], "linewidth": 0},
                    )
            elif multi_tag_style == "last":
                # draw a single marker with last tag
                patch_opts["facecolor"] = node_data["multi_colors"][-1]
                patch_opts["edgecolor"] = node_data["multi_outline_colors"][-1]
                d.marker(marker=marker, **patch_opts)
            else:  # multi_tag_style == "average":
                # draw a single marker with average color
                patch_opts["facecolor"] = average_color(
                    node_data["multi_colors"]
                )
                patch_opts["edgecolor"] = average_color(
                    node_data["multi_outline_colors"]
                )
                d.marker(marker=marker, **patch_opts)

        else:
            d.marker(marker=marker, **patch_opts)

        if node_data["label"]:
            d.text(
                node_data["coo"],
                node_data["label"],
                fontsize=node_data["label_fontsize"],
                color=node_data["label_color"],
                fontfamily=node_data["label_fontfamily"],
            )

    _add_legend_matplotlib(
        ax, colors, legend, node_outline_darkness, label_color, font_family
    )

    if fig is None:
        # ax was supplied, don't modify and simply return
        return
    else:
        # axes and figure were created
        if xlims is not None:
            ax.set_xlim(xlims)
        if ylims is not None:
            ax.set_ylim(ylims)
        if margin is not None:
            ax.margins(margin)

    if return_fig:
        return fig
    else:
        plt.show()
        plt.close(fig)


def _linearize_graph_data(G, multi_tag_style="auto"):
    from ..schematic import average_color

    edge_source = collections.defaultdict(list)
    for _, _, edge_data in G.edges(data=True):
        cooa, coob = edge_data["coos"]
        x0, y0, *maybe_z0 = cooa
        x1, y1, *maybe_z1 = coob
        edge_source["x0"].append(x0)
        edge_source["y0"].append(y0)
        edge_source["x1"].append(x1)
        edge_source["y1"].append(y1)
        if maybe_z0:
            edge_source["z0"].extend(maybe_z0)
            edge_source["z1"].extend(maybe_z1)

        # we just aggregate all multi-edges into one
        edge_source["color"].append(average_color(edge_data["color"]))
        edge_source["edge_size"].append(sum(edge_data["edge_size"]))
        edge_source["ind"].append(" ".join(edge_data["ind"]))
        edge_source["ind_size"].append(np.prod(edge_data["ind_size"]))
        edge_source["label"].append(" ".join(edge_data["label"]))

    node_source = collections.defaultdict(list)
    for _, node_data in G.nodes(data=True):
        if "ind" in node_data:
            continue

        x, y, *maybe_z = node_data["coo"]

        if "multi_colors" not in node_data:
            # single marker
            mcs = [node_data["color"]]
            mocs = [node_data["outline_color"]]
            szs = [node_data["size"]]
            os = node_data["outline_size"]
        elif multi_tag_style == "average":
            # plot a single marker with average color
            mcs = [average_color(node_data["multi_colors"])]
            mocs = [average_color(node_data["multi_outline_colors"])]
            szs = [node_data["size"]]
            os = node_data["outline_size"]
        elif multi_tag_style == "last":
            # plot a single marker with last tag
            mcs = [node_data["multi_colors"][-1]]
            mocs = [node_data["multi_outline_colors"][-1]]
            szs = [node_data["size"]]
            os = node_data["outline_size"]
        else:  # multi_tag_style in ("auto", "nest"):
            # plot multiple nested markers
            mcs = node_data["multi_colors"]
            mocs = node_data["multi_outline_colors"]
            szs = np.linspace(node_data["size"], 0, len(mcs) + 1)
            os = 0.0

        for mc, moc, sz in zip(mcs, mocs, szs):
            node_source["x"].append(x)
            node_source["y"].append(y)
            if maybe_z:
                node_source["z"].extend(maybe_z)

            node_source["color"].append(mc)
            node_source["outline_color"].append(moc)
            node_source["size"].append(sz)
            node_source["outline_size"].append(os)

            for k in ("hatch", "tags", "shape", "tid", "label"):
                node_source[k].append(node_data.get(k, None))

    return dict(edge_source), dict(node_source)


def _draw_matplotlib3d(G, **kwargs):
    import matplotlib.pyplot as plt

    edge_source, node_source = _linearize_graph_data(
        G, multi_tag_style=kwargs["multi_tag_style"]
    )

    ax = kwargs.pop("ax")
    if ax is None:
        fig = plt.figure(figsize=kwargs["figsize"])
        fig.patch.set_alpha(0.0)
        ax = plt.axes([0, 0, 1, 1], projection="3d")
        ax.patch.set_alpha(0.0)

        xmin = min(node_source["x"])
        xmax = max(node_source["x"])
        ymin = min(node_source["y"])
        ymax = max(node_source["y"])
        zmin = min(node_source["z"])
        zmax = max(node_source["z"])
        xyzmin = min((xmin, ymin, zmin))
        xyzmax = max((xmax, ymax, zmax))

        ax.set_xlim(xyzmin, xyzmax)
        ax.set_ylim(xyzmin, xyzmax)
        ax.set_zlim(xyzmin, xyzmax)
        ax.set_aspect("equal")
        ax.axis("off")

    # draw the edges
    # TODO: multiedges and left_inds
    for i in range(len(edge_source["x0"])):
        x0, x1 = edge_source["x0"][i], edge_source["x1"][i]
        xm = (x0 + x1) / 2
        y0, y1 = edge_source["y0"][i], edge_source["y1"][i]
        ym = (y0 + y1) / 2
        z0, z1 = edge_source["z0"][i], edge_source["z1"][i]
        zm = (z0 + z1) / 2
        ax.plot3D(
            [x0, x1],
            [y0, y1],
            [z0, z1],
            c=edge_source["color"][i],
            linewidth=edge_source["edge_size"][i],
        )
        label = edge_source["label"][i]
        if label:
            ax.text(
                xm,
                ym,
                zm,
                s=label,
                ha="center",
                va="center",
                color=edge_source["color"][i],
                fontsize=6,
            )

    node_source["color"] = [rgba[:3] for rgba in node_source["color"]]
    node_source["size"] = [100000 * s**2 for s in node_source["size"]]
    node_source["linewdith"] = [lw / 50 for lw in node_source["outline_size"]]

    # draw the nodes
    ax.scatter3D(
        xs="x",
        ys="y",
        zs="z",
        c="color",
        s="size",
        alpha=1.0,
        marker="o",
        data=node_source,
        depthshade=False,
        edgecolor=node_source["outline_color"],
        linewidth=node_source["outline_size"],
    )

    for _, node_data in G.nodes(data=True):
        label = node_data["label"]
        if label:
            ax.text(
                *node_data["coo"],
                s=label,
                ha="center",
                va="center",
                color=node_data["label_color"],
                fontsize=node_data["label_fontsize"],
                fontfamily=node_data["label_fontfamily"],
            )

    _add_legend_matplotlib(
        ax,
        kwargs["colors"],
        kwargs["legend"],
        kwargs["node_outline_darkness"],
        kwargs["label_color"],
        kwargs["font_family"],
    )

    if kwargs["return_fig"]:
        return fig
    else:
        plt.show()
        plt.close(fig)


def _draw_plotly(G, **kwargs):
    import plotly.graph_objects as go

    edge_source, node_source = _linearize_graph_data(
        G, multi_tag_style=kwargs["multi_tag_style"]
    )

    fig = go.Figure()
    fig.update_xaxes(visible=False)
    fig.update_yaxes(visible=False)
    fig.update_layout(
        width=100 * kwargs["figsize"][0],
        height=100 * kwargs["figsize"][1],
        margin=dict(l=10, r=10, b=10, t=10),
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        showlegend=False,
        scene=dict(
            xaxis=dict(visible=False),
            yaxis=dict(visible=False),
            zaxis=dict(visible=False),
        ),
    )

    for i in range(len(edge_source["x0"])):
        x0, x1 = edge_source["x0"][i], edge_source["x1"][i]
        y0, y1 = edge_source["y0"][i], edge_source["y1"][i]
        xm, ym = (x0 + x1) / 2, (y0 + y1) / 2
        *rgb, alpha = edge_source["color"][i]
        edge_kwargs = dict(
            x=[x0, xm, x1],
            y=[y0, ym, y1],
            opacity=alpha,
            line=dict(
                color=to_rgba_str(rgb, 1.0),
                width=edge_source["edge_size"][i],
            ),
            customdata=[[edge_source["ind"][i], edge_source["ind_size"][i]]]
            * 2,
            # show ind and ind_size on hover:
            hovertemplate="%{customdata[0]}<br>size: %{customdata[1]}",
            mode="lines",
            name="",
        )
        if "z0" in edge_source:
            z0, z1 = edge_source["z0"][i], edge_source["z1"][i]
            zm = (z0 + z1) / 2
            edge_kwargs["z"] = [z0, zm, z1]
            # edges appear much thinner in 3D
            edge_kwargs["line"]["width"] *= 2
            fig.add_trace(go.Scatter3d(**edge_kwargs))
        else:
            fig.add_trace(go.Scatter(**edge_kwargs))

    node_kwargs = dict(
        x=node_source["x"],
        y=node_source["y"],
        marker=dict(
            opacity=1.0,
            color=list(map(to_rgba_str, node_source["color"])),
            size=[300 * s for s in node_source["size"]],
            line=dict(
                color=list(map(to_rgba_str, node_source["outline_color"])),
                width=2,
            ),
        ),
        customdata=list(
            zip(node_source["tid"], node_source["shape"], node_source["tags"])
        ),
        hovertemplate=(
            "tid: %{customdata[0]}<br>"
            "shape: %{customdata[1]}<br>"
            "tags: %{customdata[2]}"
        ),
        mode="markers",
        name="",
    )
    if "z" in node_source:
        node_kwargs["z"] = node_source["z"]
        fig.add_trace(go.Scatter3d(**node_kwargs))
    else:
        fig.add_trace(go.Scatter(**node_kwargs))
    fig.show()


# ---------------------------- layout functions ----------------------------- #


def _normalize_positions(pos):
    # normalize to unit square
    xmin = ymin = zmin = float("inf")
    xmax = ymax = zmax = float("-inf")
    for x, y, *maybe_z in pos.values():
        xmin = min(xmin, x)
        xmax = max(xmax, x)
        ymin = min(ymin, y)
        ymax = max(ymax, y)
        for z in maybe_z:
            zmin = min(zmin, z)
            zmax = max(zmax, z)

    # maintain aspect ratio:
    # center each dimension separately
    xmid, ymid, zmid = (xmin + xmax) / 2, (ymin + ymax) / 2, (zmin + zmax) / 2
    # but scale all dimensions by the largest range
    xdiameter, ydiameter, zdiameter = xmax - xmin, ymax - ymin, zmax - zmin
    radius = max((xdiameter, ydiameter, zdiameter)) / 2

    for node, (x, y, *maybe_z) in pos.items():
        pos[node] = (
            (x - xmid) / radius,
            (y - ymid) / radius,
            *((z - zmid) / radius for z in maybe_z),
        )

    return pos


def _rotate(xy, theta):
    """Return a rotated set of points."""
    s = np.sin(theta)
    c = np.cos(theta)

    xyr = np.empty_like(xy)
    xyr[:, 0] = c * xy[:, 0] - s * xy[:, 1]
    xyr[:, 1] = s * xy[:, 0] + c * xy[:, 1]

    return xyr


def _span(xy):
    """Return the vertical span of the points."""
    return xy[:, 1].max() - xy[:, 1].min()


def _massage_pos(pos, nangles=360, flatten=False):
    """Rotate a position dict's points to cover a small vertical span"""
    xy = np.empty((len(pos), 2))
    for i, (x, y) in enumerate(pos.values()):
        xy[i, 0] = x
        xy[i, 1] = y

    thetas = np.linspace(0, 2 * np.pi, nangles, endpoint=False)
    rxys = (_rotate(xy, theta) for theta in thetas)
    rxy0 = min(rxys, key=lambda rxy: _span(rxy))

    if flatten:
        rxy0[:, 1] /= 2

    return dict(zip(pos, rxy0))


def phyllotaxis_points(n):
    """J. Kogan, "A New Computationally Efficient Method for Spacing Points on
    a Sphere," Rose-Hulman Undergraduate Mathematics Journal, 18(2), 2017
    Article 5. scholar.rose-hulman.edu/rhumj/vol18/iss2/5.
    """

    def spherical_coordinate(x, y):
        return [np.cos(x) * np.cos(y), np.sin(x) * np.cos(y), np.sin(y)]

    x = 0.1 + 1.2 * n
    pts = []
    start = -1.0 + 1.0 / (n - 1.0)
    increment = (2.0 - 2.0 / (n - 1.0)) / (n - 1.0)
    for j in range(n):
        s = start + j * increment
        pts.append(
            spherical_coordinate(
                s * x,
                np.pi
                / 2.0
                * np.copysign(1, s)
                * (1.0 - np.sqrt(1.0 - abs(s))),
            )
        )
    return pts


def layout_single_tensor(tn, dim=2):
    """Manually layout indices around a tensor either in a circle or sphere."""
    ((tid, t),) = tn.tensor_map.items()

    if dim == 2.5:
        dim = 3
        project_back_to_2d = True
    else:
        project_back_to_2d = False

    pos = {tid: (0.0,) * dim}
    if dim == 2:
        # fix around a circle
        angles = np.linspace(0, 2 * np.pi, t.ndim, endpoint=False)
        for ind, angle in zip(t.inds, angles):
            pos[ind] = (-np.cos(angle), np.sin(angle))
    else:
        # fix around a sphere
        for ind, coo in zip(t.inds, phyllotaxis_points(t.ndim)):
            pos[ind] = coo

    if project_back_to_2d:
        pos = {k: v[:2] for k, v in pos.items()}

    return pos


def layout_networkx(
    G,
    layout="kamada_kawai",
    pos0=None,
    fixed=None,
    dim=2,
    **kwargs,
):
    import networkx as nx

    layout_fn = getattr(nx, layout + "_layout")

    if pos0 is not None:
        if layout not in ("spring", "kamada_kawai"):
            warnings.warn(
                "Initial positions supplied but layout is not spring-based, "
                "so `pos0` is being ignored."
            )
        else:
            kwargs["pos"] = pos0

    if fixed is not None:
        if layout != "spring":
            warnings.warn(
                "Fixed positions supplied but layout is not spring-based, "
                "so `fixed` is being ignored."
            )
        else:
            kwargs["fixed"] = fixed

    return layout_fn(G, dim=dim, **kwargs)


def layout_pygraphviz(
    G,
    layout="neato",
    pos0=None,
    fixed=None,
    dim=2,
    iterations=None,
    k=None,
    **kwargs,
):
    # TODO: max iters
    # TODO: spring parameter
    # TODO: work out why pos0 and fix don't work
    import pygraphviz as pgv

    if k is not None:
        warnings.warn(
            "`k` is being ignored as layout is being done by pygraphviz."
        )

    aG = pgv.AGraph()

    # create nodes
    if pos0 is not None:
        fixed = fixed or set()
        for node, coo in pos0.items():
            pos = ",".join((f"{w:f}" for w in coo))
            pin = "true" if node in fixed else "false"
            aG.add_node(str(node), pos=pos, pin=pin)

        warnings.warn(
            "Initial and fixed positions don't seem "
            "to work currently with pygraphviz."
        )

    # create edges
    mapping = {}
    for nodea, nodeb in G.edges():
        s_nodea = str(nodea)
        s_nodeb = str(nodeb)
        mapping[s_nodea] = nodea
        mapping[s_nodeb] = nodeb
        aG.add_edge(s_nodea, s_nodeb)

    # layout options
    if iterations is not None:
        kwargs["maxiter"] = iterations
    if dim == 2.5:
        kwargs["dim"] = 3
        kwargs["dimen"] = 2
    else:
        kwargs["dim"] = kwargs["dimen"] = dim
    args = " ".join(f"-G{k}={v}" for k, v in kwargs.items())

    # run layout algorithm
    aG.layout(prog=layout, args=args)

    # extract layout
    pos = {}
    for snode, node in mapping.items():
        spos = aG.get_node(snode).attr["pos"]
        pos[node] = tuple(map(float, spos.split(",")))

    pos = _normalize_positions(pos)
    if dim < 3:
        pos = _massage_pos(pos)

    return pos


def get_positions(
    tn,
    G,
    *,
    dim=2,
    fix=None,
    layout="auto",
    initial_layout="auto",
    refine_layout="auto",
    iterations="auto",
    k=None,
):
    if (tn.num_tensors == 1) and (fix is None):
        # single tensor, layout manually
        return layout_single_tensor(tn, dim=dim)

    if layout != "auto":
        # don't use two step layout with relaxation
        initial_layout = layout
        iterations = 0

    if fix is None:
        fix = dict()
    else:
        fix = parse_dict_to_tids_or_inds(fix, tn)
        # find range with which to scale spectral points with
        xmin, xmax, ymin, ymax = (
            f(fix.values(), key=lambda xy: xy[i])[i]
            for f, i in [(min, 0), (max, 0), (min, 1), (max, 1)]
        )
        if xmin == xmax:
            xmin, xmax = xmin - 1, xmax + 1
        if ymin == ymax:
            ymin, ymax = ymin - 1, ymax + 1
        xymin, xymax = min(xmin, ymin), max(xmax, ymax)

    if all(node in fix for node in G.nodes):
        # everything is already fixed -> simply normalize
        return _normalize_positions(fix)

    if initial_layout == "auto":
        # automatically select
        if len(G) <= 100:
            # usually nicest
            initial_layout = "kamada_kawai"
        else:
            # faster, but not as nice
            initial_layout = "spectral"

    if refine_layout == "auto":
        # automatically select
        refine_layout = "spring"
        # if len(G) <= 100:
        #     # usually nicest
        #     refine_layout = "fdp"
        # else:
        #     # faster, but not as nice
        #     refine_layout = "sfdp"

    if iterations == "auto":
        # the smaller the graph, the more iterations we can afford
        iterations = max(200, 1000 - len(G))

    if dim == 2.5:
        dim = 3
        project_back_to_2d = True
    else:
        project_back_to_2d = False

    # use spectral or other layout as starting point
    if initial_layout in ("neato", "fdp", "sfdp", "dot"):
        pos0 = layout_pygraphviz(G, initial_layout, dim=dim)
    else:
        pos0 = layout_networkx(G, initial_layout, dim=dim)

    # scale points to fit with specified positions
    if fix:
        # but update with fixed positions
        pos0.update(
            valmap(
                lambda xy: np.array(
                    (
                        2 * (xy[0] - xymin) / (xymax - xymin) - 1,
                        2 * (xy[1] - xymin) / (xymax - xymin) - 1,
                    )
                ),
                fix,
            )
        )
        fixed = fix.keys()
    else:
        fixed = None

    # and then relax remaining using spring layout
    if iterations:
        if refine_layout == "spring":
            pos = layout_networkx(
                G,
                "spring",
                pos0=pos0,
                fixed=fixed,
                k=k,
                dim=dim,
                iterations=iterations,
            )
        elif refine_layout in ("fdp", "sfdp", "neato"):
            # XXX: currently doesn't seem to work with pos0 and fixed
            pos = layout_pygraphviz(
                G,
                refine_layout,
                pos0=pos0,
                fixed=fixed,
                k=k,
                dim=dim,
                iterations=iterations,
            )
        else:
            raise ValueError(f"Unknown refining layout {refine_layout}.")
    else:
        # no relaxation
        pos = pos0

    if project_back_to_2d:
        # ignore z-coordinate
        pos = {k: v[:2] for k, v in pos.items()}
        dim = 2

    # map all to range [-1, +1], but preserving aspect ratio
    pos = _normalize_positions(pos)

    if (not fix) and (dim == 2):
        # finally rotate them to cover a small vertical span
        pos = _massage_pos(pos)

    return pos


# ----------------------------- color functions ----------------------------- #


def get_colors(color, custom_colors=None, alpha=None):
    """Generate a sequence of rgbs for tag(s) ``color``."""
    from matplotlib.colors import to_rgba

    from ..schematic import auto_colors

    if color is None:
        return dict()

    if isinstance(color, str):
        color = (color,)

    if custom_colors is not None:
        rgbs = [to_rgba(c, alpha=alpha) for c in custom_colors]
        return dict(zip(color, rgbs))

    nc = len(color)
    if nc <= 7:
        rgbs = auto_colors(nc, alpha=alpha, default_sequence=True)
        return dict(zip(color, rgbs))

    rgbs = auto_colors(nc, alpha)
    return dict(zip(color, rgbs))


def to_rgba_str(color, alpha=None):
    from matplotlib.colors import to_rgba

    rgba = to_rgba(color, alpha)
    r = int(rgba[0] * 255) if isinstance(rgba[0], float) else rgba[0]
    g = int(rgba[1] * 255) if isinstance(rgba[1], float) else rgba[1]
    b = int(rgba[2] * 255) if isinstance(rgba[2], float) else rgba[2]
    return f"rgba({r}, {g}, {b}, {rgba[3]})"


def auto_color_html(s):
    """Automatically hash and color a string for HTML display."""
    from ..schematic import hash_to_color

    if not isinstance(s, str):
        s = str(s)
    return f'<b style="color: {hash_to_color(s)};">{s}</b>'


# ---------------------------- tensor functions ----------------------------- #


def visualize_tensor(tensor, **kwargs):
    """Visualize all entries of a tensor, with indices mapped into the plane
    and values mapped into a color wheel.

    Parameters
    ----------
    tensor : Tensor
        The tensor to visualize.
    skew_factor : float, optional
        When there are more than two dimensions, a factor to scale the
        rotations by to avoid overlapping data points.
    size_map : bool, optional
        Whether to map the tensor value magnitudes to marker size.
    size_scale : float, optional
        An overall factor to scale the marker size by.
    alpha_map : bool, optional
        Whether to map the tensor value magnitudes to marker alpha.
    alpha_pow : float, optional
        The power to raise the magnitude to when mapping to alpha.
    alpha : float, optional
        The overall alpha to use for all markers if ``not alpha_map``.
    show_lattice : bool, optional
        Show a small grey dot for every 'lattice' point regardless of value.
    lattice_opts : dict, optional
        Options to pass to ``maplotlib.Axis.scatter`` for the lattice points.
    linewidths : float, optional
        The linewidth to use for the markers.
    marker : str, optional
        The marker to use for the markers.
    figsize : tuple, optional
        The size of the figure to create, if ``ax`` is not provided.
    ax : matplotlib.Axis, optional
        The axis to draw to. If not provided, a new figure will be created.

    Returns
    -------
    fig : matplotlib.Figure
        The figure containing the plot, or ``None`` if ``ax`` was provided.
    ax : matplotlib.Axis
        The axis containing the plot.
    """
    import xyzpy as xyz

    kwargs.setdefault("legend", True)
    kwargs.setdefault("compass", True)
    kwargs.setdefault("compass_labels", tensor.inds)
    return xyz.visualize_tensor(tensor.data, **kwargs)


def choose_squarest_grid(x):
    p = x**0.5
    if p.is_integer():
        m = n = int(p)
    else:
        m = int(round(p))
        p = int(p)
        n = p if m * p >= x else p + 1
    return m, n


def visualize_tensors(
    tn,
    mode="network",
    r=None,
    r_scale=1.0,
    figsize=None,
    **visualize_opts,
):
    """Visualize all the entries of every tensor in this network.

    Parameters
    ----------
    tn : TensorNetwork
        The tensor network to visualize.
    mode : {'network', 'grid', 'row', 'col'}, optional
        How to arrange each tensor's visualization.

        - ``'network'``: arrange each tensor's visualization according to the
            automatic layout given by ``draw``.
        - ``'grid'``: arrange each tensor's visualization in a grid.
        - ``'row'``: arrange each tensor's visualization horizontally.
        - ``'col'``: arrange each tensor's visualization vertically.

    r : float, optional
        The absolute radius of each tensor's visualization, when
        ``mode='network'``.
    r_scale : float, optional
        A relative scaling factor for the radius of each tensor's
        visualization, when ``mode='network'``.
    figsize : tuple, optional
        The size of the figure to create, if ``ax`` is not provided.
    visualize_opts
        Supplied to ``visualize_tensor``.
    """
    from matplotlib import pyplot as plt

    if figsize is None:
        figsize = (2 * tn.num_tensors**0.4, 2 * tn.num_tensors**0.4)
    if r is None:
        r = 1.0 / tn.num_tensors**0.5
    r *= r_scale

    max_mag = None
    visualize_opts.setdefault("max_mag", max_mag)
    visualize_opts.setdefault("size_scale", r)

    if mode == "network":
        fig = plt.figure(figsize=figsize)
        pos = tn.draw(get="pos")
        for tid, (x, y) in pos.items():
            if tid not in tn.tensor_map:
                # hyper indez
                continue
            x = (x + 1) / 2 - r / 2
            y = (y + 1) / 2 - r / 2
            ax = fig.add_axes((x, y, r / 2, r / 2))
            tn.tensor_map[tid].visualize(ax=ax, **visualize_opts)
    else:
        if mode == "grid":
            px, py = choose_squarest_grid(tn.num_tensors)
        elif mode == "row":
            px, py = tn.num_tensors, 1
            figsize = (2 * figsize[0], figsize[1] / 2)
        elif mode == "col":
            px, py = 1, tn.num_tensors
            figsize = (figsize[0] / 2, 2 * figsize[1])

        fig, axs = plt.subplots(py, px, figsize=figsize)
        for i, t in enumerate(tn):
            t.visualize(ax=axs.flat[i], **visualize_opts)
        for ax in axs.flat[i:]:
            ax.set_axis_off()

    # transparent background
    fig.patch.set_alpha(0.0)

    plt.show()
    plt.close()