# Copyright 2021-2024 Cambridge Quantum Computing Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Lambeq drawing
==============
Functionality for drawing lambeq diagrams. This work is based on DisCoPy
(https://discopy.org/) which is released under the BSD 3-Clause "New"
or "Revised" License.
"""
from __future__ import annotations
from math import sqrt
import os
from tempfile import NamedTemporaryFile, TemporaryDirectory
from typing import TYPE_CHECKING
from PIL import Image
from lambeq.backend import grammar, quantum
from lambeq.backend.drawing.drawable import (BOX_HEIGHT, BoxNode,
DrawableDiagram,
DrawablePregroup,
LEDGE,
WireEndpointType)
from lambeq.backend.drawing.drawing_backend import (DEFAULT_ASPECT,
DEFAULT_MARGINS,
DrawingBackend)
from lambeq.backend.drawing.helpers import drawn_as_spider, needs_asymmetry
from lambeq.backend.drawing.mat_backend import MatBackend
from lambeq.backend.drawing.text_printer import PregroupTextPrinter
from lambeq.backend.drawing.tikz_backend import TikzBackend
from lambeq.backend.grammar import Diagram
if TYPE_CHECKING:
from IPython.core.display import HTML as HTML_ty
[docs]def draw(diagram: Diagram, **params) -> None:
"""Draw a grammar diagram.
Parameters
----------
diagram: Diagram
Diagram to draw.
draw_as_nodes : bool, optional
Whether to draw boxes as nodes, default is `False`.
color : string, optional
Color of the box or node, default is white (`'#ffffff'`) for
boxes and red (`'#ff0000'`) for nodes.
textpad : pair of floats, optional
Padding between text and wires, default is `(0.1, 0.1)`.
draw_type_labels : bool, optional
Whether to draw type labels, default is `True`.
draw_box_labels : bool, optional
Whether to draw box labels, default is `True`.
aspect : string, optional
Aspect ratio, one of `['auto', 'equal']`.
margins : tuple, optional
Margins, default is `(0.05, 0.05)`.
nodesize : float, optional
BoxNode size for spiders and controlled gates.
fontsize : int, optional
Font size for the boxes, default is `12`.
fontsize_types : int, optional
Font size for the types, default is `12`.
figsize : tuple, optional
Figure size.
path : str, optional
Where to save the image, if `None` we call `plt.show()`.
to_tikz : bool, optional
Whether to output tikz code instead of matplotlib.
asymmetry : float, optional
Make a box and its dagger mirror images, default is
`.25 * any(box.is_dagger for box in diagram.boxes)`.
foliated : bool, default: False
If true, each box of the diagram is drawn in a separate
layer. By default boxes are compressed upwards into
available space.
"""
params['asymmetry'] = params.get(
'asymmetry', .25 * needs_asymmetry(diagram))
drawable = DrawableDiagram.from_diagram(diagram,
params.get('foliated', False))
drawable.scale_and_pad(params.get('scale', (1, 1)),
params.get('pad', (0, 0)))
if 'backend' in params:
backend: DrawingBackend = params.pop('backend')
elif params.get('to_tikz', False):
backend = TikzBackend(
use_tikzstyles=params.get('use_tikzstyles', None))
else:
backend = MatBackend(figsize=params.get('figsize', None))
min_size = 0.01
max_v = max([v for point in ([point.coordinates for point in
drawable.wire_endpoints + drawable.boxes]) for v in point]
+ [min_size])
params['nodesize'] = round(params.get('nodesize', 1.) / sqrt(max_v), 3)
backend = _draw_wires(backend, drawable, **params)
backend.draw_spiders(drawable, **params)
for node in drawable.boxes:
if isinstance(node.obj, (quantum.Ket, quantum.Bra, quantum.Bit)):
backend = _draw_brakets(backend, drawable, node, **params)
elif isinstance(node.obj, quantum.Discard):
backend = _draw_discard(backend, drawable, node, **params)
elif isinstance(node.obj, quantum.Measure):
backend = _draw_measure(backend, drawable, node, **params)
elif isinstance(node.obj, quantum.Controlled):
backend = _draw_controlled_gate(backend, drawable, node, **params)
elif not drawn_as_spider(node.obj):
backend = _draw_box(backend, drawable, node, **params)
backend.output(
path=params.get('path', None),
baseline=0,
tikz_options=params.get('tikz_options', None),
show=params.get('show', True),
margins=params.get('margins', DEFAULT_MARGINS))
[docs]def draw_pregroup(diagram: Diagram, **params) -> None:
""" Draw a pregroup grammar diagram.
A pregroup diagram is structured as:
(State @ State ... State) >> (Cups and Swaps)
Parameters
----------
diagram: Diagram
Diagram to draw.
draw_as_nodes : bool, optional
Whether to draw boxes as nodes, default is `False`.
color : string, optional
Color of the box or node, default is white (`'#ffffff'`) for
boxes and red (`'#ff0000'`) for nodes.
textpad : pair of floats, optional
Padding between text and wires, default is `(0.1, 0.1)`.
aspect : string, optional
Aspect ratio, one of `['auto', 'equal']`.
margins : tuple, optional
Margins, default is `(0.05, 0.05)`.
fontsize : int, optional
Font size for the boxes, default is `12`.
fontsize_types : int, optional
Font size for the types, default is `12`.
figsize : tuple, optional
Figure size.
path : str, optional
Where to save the image, if `None` we call `plt.show()`.
to_tikz : bool, optional
Whether to output tikz code instead of matplotlib.
"""
if not diagram.is_pregroup:
raise ValueError('Diagram is not a valid pregroup diagram.')
drawable = DrawablePregroup.from_diagram(diagram)
drawable.scale_and_pad(params.get('scale', (1, 1)),
params.get('pad', (0, 0)))
if 'backend' in params:
backend: DrawingBackend = params.pop('backend')
elif params.get('to_tikz', False):
backend = TikzBackend(
use_tikzstyles=params.get('use_tikzstyles', None))
else:
backend = MatBackend(figsize=params.get('figsize', None))
backend = _draw_wires(backend, drawable, **params)
backend.draw_spiders(drawable, **params)
for node in drawable.boxes:
if not drawn_as_spider(node.obj):
backend = _draw_pregroup_state(backend, node, **params)
backend.output(
path=params.get('path', None),
baseline=len(drawable.boxes) / 2 or .5,
tikz_options=params.get('tikz_options', None),
show=params.get('show', True),
margins=params.get('margins', DEFAULT_MARGINS),
aspect=params.get('aspect', DEFAULT_ASPECT))
[docs]def render_as_str(diagram: Diagram,
word_spacing: int = 2,
use_at_separator: bool = False,
compress_layers: bool = True,
use_ascii: bool = False) -> str:
"""Render a grammar diagram as text.
Presently only implemented for pregroup diagrams.
Parameters
----------
diagram: Diagram
Diagram to draw.
word_spacing : int, default: 2
The number of spaces between the words of the diagrams.
use_at_separator : bool, default: False
Whether to represent types using @ as the monoidal product.
Otherwise, use the unicode dot character.
compress_layers : bool, default: True
Whether to draw boxes in the same layer when they can occur
simultaneously, otherwise, draw one box per layer.
use_ascii: bool, default: False
Whether to draw using ASCII characters only, for
compatibility reasons.
Returns
-------
str
Drawing of diagram in string format.
"""
if diagram.is_pregroup:
text_printer = PregroupTextPrinter(word_spacing,
use_at_separator,
compress_layers,
use_ascii)
else:
# TODO: Add text/CLI drawing for non-pregroup diagrams.
raise NotImplementedError('Text drawing is only supported for'
' pregroups. Provided diagram is not a'
' pregroup diagram.')
return text_printer.diagram2str(diagram)
[docs]def to_gif(diagrams: list[Diagram],
path: str | None = None,
timestep: int = 500,
loop: bool = False,
**params) -> str | HTML_ty:
"""Build a GIF stepping through the given diagrams.
Parameters
----------
diagrams: list of Diagrams
Sequence of diagrams to draw.
path : str
Where to save the image, if `None` a gif gets created.
timestep : int, optional
Time step in milliseconds, default is `500`.
loop : bool, optional
Whether to loop, default is `False`
params : any, optional
Passed to `Diagram.draw`.
Returns
-------
IPython.display.HTML or str
HTML to display the generated GIF
"""
steps, frames = diagrams, []
path = path or os.path.basename(NamedTemporaryFile(
suffix='.gif', prefix='tmp_', dir='.').name)
with TemporaryDirectory() as directory:
for i, _diagram in enumerate(steps):
tmp_path = os.path.join(directory, f'{i}.png')
_diagram.draw(path=tmp_path, **params)
frames.append(Image.open(tmp_path))
if loop:
frames = frames + frames[::-1]
frames[0].save(path, format='GIF', append_images=frames[1:],
save_all=True, duration=timestep,
**{'loop': 0} if loop else {}) # type: ignore[arg-type]
try:
from IPython.display import HTML
return HTML(f'<img src="{path}">')
except ImportError:
return f'<img src="{path}">'
[docs]def draw_equation(*terms: grammar.Diagram,
symbol: str = '=',
space: float = 1,
path: str | None = None,
**params) -> None:
"""Draw an equation with multiple diagrams.
Parameters
----------
terms: list of Diagrams
Diagrams in equation.
symbol: str
Symbol separating equations. '=' by default.
space: float
Amount of space between adjacent diagrams.
path : str, optional
Where to save the image, if `None` we call `plt.show()`.
**params:
Additional drawing parameters, passed to :meth:`draw`.
"""
params['asymmetry'] = params.get(
'asymmetry', .25 * any(needs_asymmetry(d) for d in terms))
pad = params.get('pad', (0, 0))
scale_x, scale_y = params.get('scale', (1, 1))
if 'backend' in params:
backend: DrawingBackend = params.pop('backend')
elif params.get('to_tikz', False):
backend = TikzBackend(
use_tikzstyles=params.get('use_tikzstyles', None))
else:
backend = MatBackend(figsize=params.get('figsize', None))
for i, term in enumerate(terms):
term.draw(**dict(
params, show=False, path=None,
backend=backend, scale=(scale_x, scale_y), pad=pad))
pad = (backend.max_width + space, 0)
if i < len(terms) - 1:
backend.draw_text(symbol, pad[0], 0)
pad = (pad[0] + space, pad[1])
return backend.output(
path=path,
baseline=0,
tikz_options=params.get('tikz_options', None),
show=params.get('show', True),
margins=params.get('margins', DEFAULT_MARGINS),
aspect=params.get('aspect', DEFAULT_ASPECT))
def _draw_box(backend: DrawingBackend,
drawable_diagram: DrawableDiagram,
drawable_box: BoxNode,
asymmetry: float,
**params) -> DrawingBackend:
"""Draw a box on a given backend.
Parameters
----------
backend: DrawingBackend
A lambeq drawing backend.
drawable_diagram: DrawableDiagram
A drawable diagram.
drawable_box: BoxNode
A BoxNode to be drawn. Must be in `drawable_diagram`.
asymmetry: float
Amount of asymmetry, used to represent transposes,
conjugates and daggers,
**params:
Additional drawing parameters. See `drawing.draw`.
Returns
-------
backend: DrawingBackend
Drawing backend updated with the box's graphic.
"""
box = drawable_box.obj
if not box.dom and not box.cod:
left, right = drawable_box.x, drawable_box.x
left, right = drawable_box.get_x_lims(drawable_diagram)
height = drawable_box.y - BOX_HEIGHT / 2
points = [[left, height], [right, height],
[right, height + BOX_HEIGHT], [left, height + BOX_HEIGHT]]
conjd = bool(box.z)
daggd = isinstance(box, grammar.Daggered)
trand = conjd and daggd
if trand:
points[0][0] -= asymmetry
elif conjd:
points[3][0] -= asymmetry
elif daggd:
points[1][0] += asymmetry
else:
points[2][0] += asymmetry
backend.draw_polygon(*points)
if params.get('draw_box_labels', True):
backend.draw_text(box.name, drawable_box.x, drawable_box.y,
ha='center', va='center',
fontsize=params.get('fontsize', None))
return backend
def _draw_pregroup_state(backend: DrawingBackend,
drawable_box: BoxNode,
**params) -> DrawingBackend:
"""Draw a pregroup word state on a given backend.
Parameters
----------
backend: DrawingBackend
A lambeq drawing backend.
drawable_box: BoxNode
A BoxNode to be drawn.
**params:
Additional drawing parameters. See `drawing.draw`.
Returns
-------
backend: DrawingBackend
Drawing backend updated with the box's graphic.
"""
box = drawable_box.obj
left = drawable_box.x
right = left + 2
height = drawable_box.y - BOX_HEIGHT / 2
points = [[left, height], [right, height],
[right, height + BOX_HEIGHT], [(left + right) / 2, height + 0.6],
[left, height + BOX_HEIGHT]]
backend.draw_polygon(*points)
backend.draw_text(box.name, drawable_box.x + 1, drawable_box.y,
ha='center', va='center',
fontsize=params.get('fontsize', None))
return backend
def _draw_wires(backend: DrawingBackend,
drawable_diagram: DrawableDiagram,
**params) -> DrawingBackend:
"""Draw all wires of a diagram on a given backend.
Parameters
----------
backend: DrawingBackend
A lambeq drawing backend.
drawable_diagram: DrawableDiagram
A drawable diagram.
**params:
Additional drawing parameters. See :meth:`draw`.
Returns
-------
backend: DrawingBackend
Drawing backend updated with the wires' graphic.
"""
for src_idx, tgt_idx in drawable_diagram.wires:
source = drawable_diagram.wire_endpoints[src_idx]
target = drawable_diagram.wire_endpoints[tgt_idx]
backend.draw_wire(
source.coordinates, target.coordinates)
if (params.get('draw_type_labels', True) and source.kind in
{WireEndpointType.INPUT, WireEndpointType.COD}):
i, j = source.coordinates
pad_i, pad_j = params.get('textpad', (.1, .1))
pad_j = 0 if source.kind == WireEndpointType.INPUT else pad_j
backend.draw_text(
str(source.obj), i + pad_i, j - pad_j,
fontsize=params.get('fontsize_types',
params.get('fontsize', None)),
verticalalignment='top')
return backend
def _draw_brakets(backend: DrawingBackend,
drawable_diagram: DrawableDiagram,
drawable_box: BoxNode,
**params) -> DrawingBackend:
"""Draw Bras and Kets on a given backend.
Parameters
----------
backend: DrawingBackend
A lambeq drawing backend.
drawable_diagram: DrawableDiagram
A drawable diagram.
drawable_box: BoxNode
A BoxNode to be drawn. Must be in `drawable_diagram`.
**params:
Additional drawing parameters. See `drawing.draw`.
Returns
-------
backend: DrawingBackend
Drawing backend updated with the box's graphic.
"""
box = drawable_box.obj
assert isinstance(box, (quantum.Ket, quantum.Bra, quantum.Bit))
is_bra = isinstance(box, quantum.Bra)
factor = -1 if is_bra else 1
left, right = drawable_box.get_x_lims(drawable_diagram)
height = drawable_box.y - factor * .25
points = [[left, height], [right, height],
[(left+right) / 2, height + factor * .5]]
backend.draw_polygon(*points)
backend.draw_text(box.name,
drawable_box.x, drawable_box.y,
ha='center', va='center',
fontsize=params.get('fontsize', None))
return backend
def _draw_discard(backend: DrawingBackend,
drawable_diagram: DrawableDiagram,
drawable_box: BoxNode,
**params) -> DrawingBackend:
"""Draw a Discards on a given backend.
Parameters
----------
backend: DrawingBackend
A lambeq drawing backend.
drawable_diagram: DrawableDiagram
A drawable diagram.
drawable_box: BoxNode
A BoxNode to be drawn. Must be in `drawable_diagram`.
**params:
Additional drawing parameters. See `drawing.draw`.
Returns
-------
backend: DrawingBackend
Drawing backend updated with the box's graphic.
"""
left, right = drawable_box.get_x_lims(drawable_diagram)
height = drawable_box.y + 0.25
for j in range(3):
source = (left + .1 * j, height - .1 * j)
target = (right - .1 * j, height - .1 * j)
backend.draw_wire(source, target)
return backend
def _draw_measure(backend: DrawingBackend,
drawable_diagram: DrawableDiagram,
drawable_box: BoxNode,
**params) -> DrawingBackend:
"""Draw a Measure box.
Parameters
----------
backend: DrawingBackend
A lambeq drawing backend.
drawable_diagram: DrawableDiagram
A drawable diagram.
drawable_box: BoxNode
A BoxNode to be drawn. Must be in `drawable_diagram`.
**params:
Additional drawing parameters. See `drawing.draw`.
Returns
-------
backend: DrawingBackend
Drawing backend updated with the box's graphic.
"""
backend = _draw_box(backend,
drawable_diagram,
drawable_box,
draw_box_labels=False,
**params)
i, j = drawable_box.x, drawable_box.y
backend.draw_wire((i - .15, j - .1), (i, j + .1), bend_in=True)
backend.draw_wire((i, j + .1), (i + .15, j - .1), bend_out=True)
backend.draw_wire((i, j - .1), (i + .05, j + .15), style='->')
return backend
def _draw_controlled_gate(backend: DrawingBackend,
drawable_diagram: DrawableDiagram,
drawable_box: BoxNode,
**params) -> DrawingBackend:
""" Draw a Controlled gate.
Parameters
----------
backend: DrawingBackend
A lambeq drawing backend.
drawable_diagram: DrawableDiagram
A drawable diagram.
drawable_box: BoxNode
A BoxNode to be drawn. Must be in `drawable_diagram`.
**params:
Additional drawing parameters. See `drawing.draw`.
Returns
-------
backend: DrawingBackend
Drawing backend updated with the box's graphic.
"""
assert isinstance(drawable_box.obj, quantum.Controlled)
box = drawable_box.obj
distance = box.distance
c_size = len(box.controlled.dom)
all_wires_pos_x = sorted(set([
drawable_diagram.wire_endpoints[wire].x
for wire in drawable_box.cod_wires + drawable_box.dom_wires]))
all_wires_pos_y = sorted(set([
drawable_diagram.wire_endpoints[wire].y
for wire in drawable_box.cod_wires + drawable_box.dom_wires]))
middle_wire_pos_y = (min(all_wires_pos_y) + max(all_wires_pos_y)) / 2
# This is the index of the control location
index = 0 if distance > 0 else -1
sign = 1 if distance > 0 else -1
# Extract the location of the control and draw black dot
control_dot_coordinates = (all_wires_pos_x[index], middle_wire_pos_y)
backend.draw_node(*control_dot_coordinates,
color='black',
shape='circle',
nodesize=params.get('nodesize', 1))
control_wire_endpoint_coordinates = (all_wires_pos_x[index + distance],
middle_wire_pos_y)
controlled_middle_coordinates = (
all_wires_pos_x[index + distance] + sign * (c_size - 1) / 2,
middle_wire_pos_y)
# The target boundary is the point where the wire hits the box
target_boundary = control_wire_endpoint_coordinates
if box.controlled == quantum.X:
# CX gets drawn as a circled plus sign.
backend.draw_node(
*controlled_middle_coordinates,
shape='circle', color='white', edgecolor='black',
nodesize=2 * params.get('nodesize', 1))
backend.draw_node(
*controlled_middle_coordinates, shape='plus',
nodesize=2 * params.get('nodesize', 1))
# Draw the vertical line through the controlled box
backend.draw_wire(
(controlled_middle_coordinates[0], min(all_wires_pos_y)),
(controlled_middle_coordinates[0], max(all_wires_pos_y)))
else: # controlled box is not a CX gate
# If the controlled box is a regular box, we need to shift the
# endpoint controll wire to the left or right depending on the
# sign of the distance. This is indicated by shift_boundary
shift_boundary = True
# Get the connected wires of the controlled box
if sign > 0:
b_start = index + distance
new_dom_wires = drawable_box.dom_wires[b_start:b_start+c_size]
new_cod_wires = drawable_box.cod_wires[b_start:b_start+c_size]
elif sign < 0:
new_dom_wires = drawable_box.dom_wires[:c_size]
new_cod_wires = drawable_box.cod_wires[:c_size]
# Create a new box node for the controlled box
controlled_box_node = BoxNode(box.controlled,
*controlled_middle_coordinates,
new_dom_wires,
new_cod_wires)
if isinstance(box.controlled, quantum.Controlled): # nested control
backend = _draw_controlled_gate(
backend, drawable_diagram, controlled_box_node, **params)
next_box: quantum.Controlled | quantum.Box = box.controlled
while isinstance(next_box, quantum.Controlled):
if box.distance * next_box.distance < 0:
shift_boundary = False
break
next_box = next_box.controlled
if next_box == quantum.X:
shift_boundary = False
else:
backend = _draw_box(backend,
drawable_diagram,
controlled_box_node,
**params)
if shift_boundary:
target_boundary = (
control_wire_endpoint_coordinates[0] - sign * LEDGE,
control_wire_endpoint_coordinates[1])
# draw vertical line through control dot
backend.draw_wire((all_wires_pos_x[index], all_wires_pos_y[0]),
(all_wires_pos_x[index], all_wires_pos_y[-1]))
# draw all the other vertical wires
extra_offset = 1 if distance > 0 else len(box.controlled.dom)
for i in range(extra_offset, extra_offset + abs(distance) - 1):
backend.draw_wire((all_wires_pos_x[i], all_wires_pos_y[0]),
(all_wires_pos_x[i], all_wires_pos_y[-1]))
# TODO change bend_in and bend_out for tikz backend
backend.draw_wire(control_dot_coordinates,
target_boundary,
bend_in=True,
bend_out=True)
return backend