Circuit diagrams for unitary quantum kernels in LaTeX

python/cudaq/__init__.py

@@ -34,6 +34,7 @@
             print("Could not find a suitable cuQuantum Python package.")
         pass
 
+from .display import display_trace
 from .kernel.kernel_decorator import kernel, PyKernelDecorator
 from .kernel.kernel_builder import make_kernel, QuakeValue, PyKernel
 from .kernel.ast_bridge import globalAstRegistry, globalKernelRegistry
@@ -111,6 +112,8 @@
 AsyncStateResult = cudaq_runtime.AsyncStateResult
 vqe = cudaq_runtime.vqe
 draw = cudaq_runtime.draw
+displaySVG = display_trace.displaySVG
+getSVGstring = display_trace.getSVGstring
 
 ComplexMatrix = cudaq_runtime.ComplexMatrix
 to_qir = cudaq_runtime.get_qir

python/cudaq/display/display_trace.py

@@ -0,0 +1,30 @@
+# ============================================================================ #
+# Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                   #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #
+
+import cudaq
+
+
+def getSVGstring(kernel, *args):
+    from subprocess import check_output, STDOUT
+    from tempfile import TemporaryDirectory
+
+    latex_string = cudaq.draw("latex", kernel, *args)
+    with TemporaryDirectory() as tmpdirname:
+        with open(tmpdirname + "/cudaq-trace.tex", "w") as f:
+            f.write(latex_string)
+        # this needs `latex` and `quantikz` to be installed, e.g. through `apt install texlive-latex-extra`
+        check_output(["latex", "cudaq-trace"], cwd=tmpdirname, stderr=STDOUT)
+        check_output(["dvisvgm", "cudaq-trace"], cwd=tmpdirname, stderr=STDOUT)
+        with open(tmpdirname + "/cudaq-trace.svg", "rb") as f:
+            return str(f.read(), encoding="utf-8")
+
+
+def displaySVG(kernel, *args):
+    from IPython.display import SVG, display
+
+    display(SVG(data=getSVGstring(kernel, *args)))

python/cudaq/kernel/kernel_decorator.py

@@ -180,6 +180,29 @@ def __str__(self):
         self.compile()
         return str(self.module)
 
+    def _repr_svg_(self):
+        """
+        Return the SVG representation of `self` (:class:`PyKernelDecorator`).
+        This assumes no arguments are required to execute the kernel,
+        and `latex` (with `quantikz` package) and `dvisvgm` are installed,
+        and the temporary directory is writable.
+        If any of these assumptions fail, returns None.
+        """
+        self.compile()  # compile if not yet compiled
+        if self.argTypes is None or len(self.argTypes) != 0:
+            return None
+        from cudaq import getSVGstring
+
+        try:
+            from subprocess import CalledProcessError
+
+            try:
+                return getSVGstring(self)
+            except CalledProcessError:
+                return None
+        except ImportError:
+            return None
+
     def isCastable(self, fromTy, toTy):
         if F64Type.isinstance(toTy):
             return F32Type.isinstance(fromTy) or IntegerType.isinstance(fromTy)

python/runtime/cudaq/algorithms/py_draw.cpp

@@ -9,20 +9,24 @@
 #include "utils/OpaqueArguments.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
-#include <iostream>
 #include <pybind11/complex.h>
 #include <pybind11/stl.h>
+#include <string>
+#include <tuple>
+#include <vector>
 
 namespace cudaq {
 
 void pyAltLaunchKernel(const std::string &, MlirModule, OpaqueArguments &,
                        const std::vector<std::string> &);
 
-/// @brief Run `cudaq::draw` on the provided kernel.
-std::string pyDraw(py::object &kernel, py::args args) {
-
+namespace {
+std::tuple<std::string, MlirModule, OpaqueArguments *>
+getKernelLaunchParameters(py::object &kernel, py::args args) {
   if (py::len(kernel.attr("arguments")) != args.size())
-    throw std::runtime_error("Invalid number of arguments passed to draw.");
+    throw std::runtime_error("Invalid number of arguments passed to draw:" +
+                             std::to_string(args.size()) + " expected " +
+                             std::to_string(py::len(kernel.attr("arguments"))));
 
   if (py::hasattr(kernel, "compile"))
     kernel.attr("compile")();
@@ -32,17 +36,55 @@ std::string pyDraw(py::object &kernel, py::args args) {
   args = simplifiedValidateInputArguments(args);
   auto *argData = toOpaqueArgs(args, kernelMod, kernelName);
 
+  return {kernelName, kernelMod, argData};
+}
+} // namespace
+
+/// @brief Run `cudaq::draw` on the provided kernel.
+std::string pyDraw(py::object &kernel, py::args args) {
+  auto [kernelName, kernelMod, argData] =
+      getKernelLaunchParameters(kernel, args);
+
   return details::extractTrace([&]() mutable {
     pyAltLaunchKernel(kernelName, kernelMod, *argData, {});
     delete argData;
   });
 }
 
+/// @brief Run `cudaq::draw`'s string overload on the provided kernel.
+std::string pyDraw(std::string format, py::object &kernel, py::args args) {
+  if (format == "ascii") {
+    return pyDraw(kernel, args);
+  } else if (format == "latex") {
+    auto [kernelName, kernelMod, argData] =
+        getKernelLaunchParameters(kernel, args);
+
+    return details::extractTraceLatex([&]() mutable {
+      pyAltLaunchKernel(kernelName, kernelMod, *argData, {});
+      delete argData;
+    });
+  } else {
+    throw std::runtime_error("Invalid format passed to draw.");
+  }
+}
+
 /// @brief Bind the draw cudaq function
 void bindPyDraw(py::module &mod) {
-  mod.def(
-      "draw", &pyDraw,
-      R"#(Return a UTF-8 encoded string representing drawing of the execution 
+  mod.def("draw",
+          py::overload_cast<std::string, py::object &, py::args>(&pyDraw),
+          R"#(Return a string representing the drawing of the execution path, 
+in the format specified as the first argument. If the format is 
+'ascii', the output will be a UTF-8 encoded string. If the format 
+is 'latex', the output will be a LaTeX string.
+
+Args:
+  format (str): The format of the output. Can be 'ascii' or 'latex'.
+  kernel (:class:`Kernel`): The :class:`Kernel` to draw.
+  *arguments (Optional[Any]): The concrete values to evaluate the kernel 
+    function at. Leave empty if the kernel doesn't accept any arguments.)#")
+      .def(
+          "draw", py::overload_cast<py::object &, py::args>(&pyDraw),
+          R"#(Return a UTF-8 encoded string representing drawing of the execution 
 path, i.e., the trace, of the provided `kernel`.
 
 Args:

python/tests/display/CMakeLists.txt

@@ -0,0 +1,14 @@
+# ============================================================================ #
+# Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                   #
+# All rights reserved.                                                         #
+# # This source code and the accompanying materials are made available under #
+# the terms of the Apache License 2.0 which accompanies this distribution. #
+# ============================================================================ #
+
+add_test(
+  NAME cudaq-py-display-draw
+  COMMAND ${PYTHON_EXECUTABLE} test_draw.py
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+set_tests_properties(
+  cudaq-py-display-draw PROPERTIES ENVIRONMENT
+                                    "PYTHONPATH=${CMAKE_BINARY_DIR}/python")

python/tests/display/test_draw.py

@@ -0,0 +1,113 @@
+# ============================================================================ #
+# Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                   #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #
+
+import cudaq
+import numpy as np
+import os
+import pytest
+
+
+@cudaq.kernel
+def bar(qvec: cudaq.qview):
+    # FIXME https://github.com/NVIDIA/cuda-quantum/issues/1734
+    # rx(np.e, qvec[0])
+    rx(2.71828182845904523536028, qvec[0])
+    ry(np.pi, qvec[1])
+    # FIXME https://github.com/NVIDIA/cuda-quantum/issues/1734
+    # cudaq.adjoint(rz, np.pi, qvec[2])
+    rz(-np.pi, qvec[2])
+
+
+@cudaq.kernel
+def zaz(qub: cudaq.qubit):
+    sdg(qub)
+
+
+@cudaq.kernel
+def kernel():
+    q = cudaq.qvector(4)
+    h(q)
+    x.ctrl(q[0], q[1])
+    y.ctrl(q[0], q[1], q[2])
+    y.ctrl(q[2], q[0], q[1])
+    y.ctrl(q[1], q[2], q[0])
+    z(q[2])
+    r1(3.14159, q[0])
+    tdg(q[1])
+    s(q[2])
+    swap.ctrl(q[0], q[2])
+    swap.ctrl(q[1], q[2])
+    swap.ctrl(q[0], q[1])
+    swap.ctrl(q[0], q[2])
+    swap.ctrl(q[1], q[2])
+    swap.ctrl(q[3], q[0], q[1])
+    swap.ctrl(q[0], q[3], q[1], q[2])
+    swap.ctrl(q[1], q[0], q[3])
+    swap.ctrl(q[1], q[2], q[0], q[3])
+    bar(q)
+    cudaq.control(zaz, q[1], q[0])
+    cudaq.adjoint(bar, q)
+
+
+def test_draw():
+    """Test draw function, mainly copied from draw_tester.cpp"""
+    # fmt: off
+    expected_str = R"""
+     ╭───╮               ╭───╮╭───────────╮                          ╭───────╮»
+q0 : ┤ h ├──●────●────●──┤ y ├┤ r1(3.142) ├──────╳─────╳──╳─────╳──●─┤>      ├»
+     ├───┤╭─┴─╮  │  ╭─┴─╮╰─┬─╯╰──┬─────┬──╯      │     │  │     │  │ │       │»
+q1 : ┤ h ├┤ x ├──●──┤ y ├──●─────┤ tdg ├─────────┼──╳──╳──┼──╳──╳──╳─┤●      ├»
+     ├───┤╰───╯╭─┴─╮╰─┬─╯  │     ╰┬───┬╯   ╭───╮ │  │     │  │  │  │ │  swap │»
+q2 : ┤ h ├─────┤ y ├──●────●──────┤ z ├────┤ s ├─╳──╳─────╳──╳──┼──╳─│       │»
+     ├───┤     ╰───╯              ╰───╯    ╰───╯                │  │ │       │»
+q3 : ┤ h ├──────────────────────────────────────────────────────●──●─┤>      ├»
+     ╰───╯                                                           ╰───────╯»
+
+################################################################################
+
+╭───────╮╭───────────╮    ╭─────╮   ╭────────────╮
+┤>      ├┤ rx(2.718) ├────┤ sdg ├───┤ rx(-2.718) ├
+│       │├───────────┤    ╰──┬──╯   ├────────────┤
+┤●      ├┤ ry(3.142) ├───────●──────┤ ry(-3.142) ├
+│  swap │├───────────┴╮╭───────────╮╰────────────╯
+┤●      ├┤ rz(-3.142) ├┤ rz(3.142) ├──────────────
+│       │╰────────────╯╰───────────╯              
+┤>      ├─────────────────────────────────────────
+╰───────╯                                         
+"""
+    # fmt: on
+    expected_str = expected_str[1:]
+    produced_string = cudaq.draw(kernel)
+    assert expected_str == produced_string
+
+
+def test_draw_latex():
+    """Test draw function, mainly copied from draw_tester.cpp"""
+    # fmt: off
+    expected_str = R"""
+\documentclass{minimal}
+\usepackage{quantikz}
+\begin{document}
+\begin{quantikz}
+  \lstick{$q_0$} & \gate{H} & \ctrl{1} & \ctrl{2} & \ctrl{1} & \gate{Y} & \gate{R_1(3.142)} & \qw & \swap{2} & \qw & \swap{1} & \swap{2} & \qw & \swap{1} & \ctrl{2} & \swap{3} & \swap{3} & \gate{R_x(2.718)} & \gate{S^\dag} & \gate{R_x(-2.718)} & \qw \\
+  \lstick{$q_1$} & \gate{H} & \gate{X} & \ctrl{1} & \gate{Y} & \ctrl{-1} & \gate{T^\dag} & \qw & \qw & \swap{1} & \targX{} & \qw & \swap{1} & \targX{} & \swap{1} & \ctrl{2} & \ctrl{2} & \gate{R_y(3.142)} & \ctrl{-1} & \gate{R_y(-3.142)} & \qw \\
+  \lstick{$q_2$} & \gate{H} & \qw & \gate{Y} & \ctrl{-1} & \ctrl{-2} & \gate{Z} & \gate{S} & \targX{} & \targX{} & \qw & \targX{} & \targX{} & \qw & \targX{} & \qw & \ctrl{-2} & \gate{R_z(-3.142)} & \gate{R_z(3.142)} & \qw & \qw \\
+  \lstick{$q_3$} & \gate{H} & \qw & \qw & \qw & \qw & \qw & \qw & \qw & \qw & \qw & \qw & \qw & \ctrl{-3} & \ctrl{-2} & \targX{} & \targX{} & \qw & \qw & \qw & \qw \\
+\end{quantikz}
+\end{document}
+"""
+    # fmt: on
+    expected_str = expected_str[1:]
+    produced_string = cudaq.draw("latex", kernel)
+    assert expected_str == produced_string
+
+
+# leave for gdb debugging
+if __name__ == "__main__":
+    loc = os.path.abspath(__file__)
+    pytest.main([loc, "-s"])