2. Drawing Tensor Networks

quimb has a lot of functionality for drawing tensor networks that can be useful for debugging, interactive development, and producing figures etc. This page is a general overview of various options, mostly centered around the method draw(). Underneath this calls networkx which itself uses matplotlib.

[1]:
%config InlineBackend.figure_formats = ['svg']
import quimb.tensor as qtn

We’ll use a 3D grid tensor network as our basic example.

[2]:
Lx = Ly = Lz = 4
D = 2
tn = qtn.TN3D_rand(Lx, Ly, Lz, D=D)

By default bonds are draw proportional to log2 of their dimension, whereas nodes are fixed in size.

[3]:
tn.draw()
_images/tensor-drawing_5_0.svg

By default index names are not shown and tensor tags are only shown for small tensors, these can both be controlled manually like so:

[4]:
qtn.PEPS.rand(3, 3, D).draw(show_tags=True, show_inds=True)
_images/tensor-drawing_7_0.svg

If you want to see inner index names (bonds) as well as the outer index names you need to use show_inds='all':

[5]:
qtn.PEPS.rand(3, 3, D).draw(show_tags=False, show_inds='all')
_images/tensor-drawing_9_0.svg

2.1. Coloring

The first argument to draw is color=, which can either be a single tag or a sequence of tags. All tensors with each tag will be colored the same, with later tags taking priority:

[6]:
# add the same tag to every tensor
tn.add_tag('CUBE')

# color that tag and each corner of our TN
color = ['CUBE'] + [
    f'I{i},{j},{k}'
    for i in (0, Lx - 1)
    for j in (0, Ly - 1)
    for k in (0, Lz - 1)
]

tn.draw(color=color)
_images/tensor-drawing_11_0.svg

If you have many tags or are simply only interested in the drawing the colors you can supply the legend=False option to turn off the legend.

Hint

quimb tries to produce a sequence of colors that are reasonably locally distigushable but also have some global ordering when using many colors. These are based on the palette designed with color blindness in mind by Bang Wong. You can supply custom colors with the custom_colors= kwarg.

2.2. Highlighting indices

You can visualize a subset of indices by supplying a sequence of them to the highlight_inds= kwarg like so:

[7]:
# get a central tensor and its indices
tag = f"I{Lx // 2},{Ly // 2},{Lz // 2}"
t = tn[tag]
inds = t.inds
tn.draw(color=tag, highlight_inds=inds)
_images/tensor-drawing_15_0.svg

The color can be controlled with highlight_inds_color.

2.3. Highlighting tids

While tensors can carry arbitrary tags and can usually be identified by these, it is sometimes useful to be able to highlight tensors based on their underlying tids - each of which is a unique integer representing a node in the hypergraph.

[8]:
# get the first plane of tensor tids
tids = list(tn.tensor_map.keys())[:Lx * Ly]
tids
[8]:
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

The color can be controlled with highlight_tids_color:

[9]:
tn.draw(highlight_tids=tids, highlight_tids_color=(1.0, 0.0, 0.5, 0.5))
_images/tensor-drawing_20_0.svg

2.4. Positioning tensors

2.4.1. Automatic layouts

The automatic layout strategy quimb adopts is to lay the tensors out using some relatively efficient scheme, before ‘relaxing’ the positions using a (slower) force repulsion algorithm into something usually more natural.

The iterations kwarg controls the number of force repulsion steps, set this to zero to use only the initial layout algorithm (the default of which is 'spectral'):

[10]:
tn.draw(iterations=0)
_images/tensor-drawing_22_0.svg

Another good choice for the initial layout that you might try if 'spectral' isn’t producing good results is 'kamada_kawai':

[11]:
tn.draw(iterations=0, initial_layout='kamada_kawai')
_images/tensor-drawing_24_0.svg

You should be able to specify most of the networkx layout algorithms:

[12]:
tn.draw(iterations=0, initial_layout='spiral')
_images/tensor-drawing_26_0.svg

2.4.2. Force Repulsion options

For the force repulsion layout, you can supply the spring constant k, which can have a significant effect on the layout:

[13]:
tn.draw(iterations=100, k=0.01)
_images/tensor-drawing_28_0.svg

You can also fix specific tensors (by either a tid or set of tags that uniquely identifies that tensor):

[14]:
fix = {
    'I0,0,0': (0, 0),
    'I0,0,1': (0, 1),
    'I1,0,0': (1, 0),
    'I1,0,1': (1, 1),
}

# when fixing tensors you often have to play with ``k``
tn.draw(k=0.001, fix=fix, color=fix.keys())
_images/tensor-drawing_30_0.svg

If you have `forceatlas2 (fa2) <https://github.com/bhargavchippada/forceatlas2>`__ installed then you can specify to use it rather than the slower networkx force repulsion algorithm at a certain threshold of nodes (by default 1000) with the option use_forceatlas2=1000.

2.4.3. Manually Specifying

You can also simply specify all positions manually using the fix kwarg. Here’s that illustrated with a axonometric projection:

[15]:
import math
import itertools

def get_3d_pos(i, j, k, a=22, b=45, p=0.2):
    return (
        + i * math.cos(math.pi * a / 180) + j * math.cos(math.pi * b / 180) / 2**p,
        - i * math.sin(math.pi * a / 180) + j * math.sin(math.pi * b / 180) / 2**p + k
    )

pos = {
    f'I{i},{j},{k}': get_3d_pos(i, j, k)
    for i in range(Lx)
    for j in range(Ly)
    for k in range(Lz)
}


tn.draw(fix=pos, color=pos.keys(), legend=False)
_images/tensor-drawing_33_0.svg

If you want to retrieve an automatic positioning, e.g. for repeated use in an animation, you can pass the get='pos' option, which simply returns the positions as a dict mapping each tid to a 2D coordinate:

[16]:
pos = tn.draw(get='pos')
pos[0], pos[1], pos[2]
[16]:
(array([-1.00139021,  0.06992426]),
 array([-0.81934865,  0.31153009]),
 array([-0.58950712,  0.58652939]))

2.5. Hyper-edges

Hyper edges (indices which appear on 3 or more tensors) are represented as separate ‘nodes’ of zero size - since they are equivalent to placing a multi-dimensional COPY-tensor is such locations.

[17]:
htn = qtn.HTN3D_classical_ising_partition_function(3, 3, 3, beta=0.22)
htn.draw()
_images/tensor-drawing_37_0.svg

Another way to visualize such hyperedges, using ‘rubber bands’, is provided by `hypernetx <https://github.com/pnnl/HyperNetX>`__ - both the ind_map of a tensor network and the pos generate by draw are directly compatible:

[18]:
import hypernetx

H = hypernetx.Hypergraph(htn.ind_map)
hypernetx.draw(H, pos=htn.draw(get='pos'))
_images/tensor-drawing_39_0.svg

2.6. Spanning trees

Various algorithms in quimb make use of a tree generated by spanning out from a particular region.

[19]:
span_opts = {
    'max_distance': 3,
    'distance_sort': 'min',
    'ndim_sort': 'max',
}

qtn.TN2D_rand(7, 7, 3).draw_tree_span(
    tags=['I2,3', 'I2,2'],  which='any', **span_opts
)
_images/tensor-drawing_41_0.svg

2.7. Interaction with matplotlib

You can either add other stuff to the figure that quimb creates, or you can supply a matplotlib axis to add the tensor network drawing to directly.

The return_fig=True option allows you to modify the figure or save it to file:

[20]:
fig = tn.draw(return_fig=True)
fig.set_facecolor('yellow')
_images/tensor-drawing_43_0.svg

This could be saved with e.g.:

fig.savefig('my-tn-drawing.png', bbox_inches='tight', dpi=300)

The ax=ax option allows you to add to an existing plot:

[21]:
import matplotlib.pyplot as plt

fig, axs = plt.subplots(10, 10)

for ax in axs.flat:
    tn = qtn.TN_rand_reg(n=12, reg=3, D=2)
    tn.draw(tn.tags, ax=ax, legend=False, show_tags=False, node_size=10)
    ax.axis('off')
_images/tensor-drawing_45_0.svg

2.8. ‘Publication style’ figures

There are some likely settings to tweak to generate neat ‘publication style’ figures, but probably the most noticeable setting is explicitly laying out the nodes with fix.

Here we demonstrate drawing a PEPS with various options, in particular, embedding it in a 3D space so no edges are overlapping.

[22]:
psi = qtn.PEPS.rand(6, 6, 4)

# fix the site tensors in one plane
fix = {
    psi.site_tag(i, j): get_3d_pos(i, j, 1)
    for i, j in psi.gen_site_coos()
}
# fix the site inidices in plane below
fix.update({
    psi.site_ind(i, j): get_3d_pos(i, j, 0.5)
    for i, j in psi.gen_site_coos()
})

# specific a tensor and its neighbors
tag0 = 'I2,2'
tags = ['I2,1', 'I1,2', 'I2,3', 'I3,2']

# create some arrows
psi.canonize_around_(tag0, max_distance=1)

# draw, with some manual style settings
psi.draw(
    color=(tag0, *tags),
    custom_colors=[(0.8, 0.3, 0.7)] + [(0.3, 0.8, 0.2)] * 4,
    fix=fix,
    edge_color='black',
    edge_alpha=1.0,
    edge_scale=1.0,
    arrow_overhang=1.0,
    arrow_linewidth=4,
    node_size=400,
    node_outline_darkness=0.0,
    node_outline_size=2.5,
    node_hatch={tag: '////' for tag in tags},
    node_shape={tag: 'h' for tag in tags} | {tag0: 's'},
    legend=False,
)
_images/tensor-drawing_47_0.svg