Fix #853: improve formatter for pretty-printing musical score diagrams

qualtran/bloqs/rotations/quantum_variable_rotation.py

@@ -191,6 +191,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
         if self.cost_dtype.signed:
             out[0] = bb.add(ZPowGate(exponent=1, eps=eps), q=out[0])
         offset = 1 + self.cost_dtype.num_frac - self.num_frac_rotations
+        offset = offset if isinstance(offset, int) else 1
         for i in range(num_rotations):
             power_of_two = i - self.num_frac_rotations
             exp = (2**power_of_two) * self.gamma * 2

qualtran/drawing/_show_funcs.py

@@ -1,3 +1,3 @@
 #  Copyright 2023 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,7 +15,9 @@
 """Convenience functions for showing rich displays in Jupyter notebook."""
 
 import os
-from typing import Dict, Optional, overload, Sequence, TYPE_CHECKING, Union
+import re
+import sympy
+from typing import Dict, Optional, Text, overload, Sequence, TYPE_CHECKING, Union, Tuple
 
 import IPython.display
 import ipywidgets
@@ -25,13 +27,12 @@
 from .bloq_counts_graph import format_counts_sigma, GraphvizCallGraph
 from .flame_graph import get_flame_graph_svg_data
 from .graphviz import PrettyGraphDrawer, TypedGraphDrawer
-from .musical_score import draw_musical_score, get_musical_score_data
+from .musical_score import MusicalScoreData, TextBox, draw_musical_score, get_musical_score_data
 from .qpic_diagram import qpic_diagram_for_bloq
 
 if TYPE_CHECKING:
     import networkx as nx
-    import sympy
-
+
 
 def show_bloq(bloq: 'Bloq', type: str = 'graph'):  # pylint: disable=redefined-builtin
     """Display a visual representation of the bloq in IPython.
@@ -49,7 +50,9 @@
     elif type.lower() == 'dtype':
         IPython.display.display(TypedGraphDrawer(bloq).get_svg())
     elif type.lower() == 'musical_score':
-        draw_musical_score(get_musical_score_data(bloq))
+        msd = get_musical_score_data(bloq)
+        pretty_msd = pretty_format_msd(msd)
+        draw_musical_score(pretty_msd)
     elif type.lower() == 'latex':
         show_bloq_via_qpic(bloq)
     else:
@@ -142,3 +145,81 @@
 
     IPython.display.display(Image(output_file_path, width=width, height=height))
     os.remove(output_file_path)
+
+
+def pretty_format_msd(msd: MusicalScoreData) -> MusicalScoreData:
+    """
+    Beautifies MSD to enable pretty diagrams
+
+    Args:
+        msd: A raw MSD
+
+    Returns:
+        new_msd: A pretty MSD
+
+    """
+
+    def symbols_to_identity(lbl: str) -> Tuple[str, str]:
+        """
+        Exchanges any symbols in the label for integer 1 or returns lbl if no symbols found.
+
+        Args:
+            lbl: The label to be processed.
+
+        Returns:
+            new_lbl: A label without symbols
+
+        """
+
+        pattern = r"Abs\((?P<symbol>[a-zA-Z])\)"
+        match = re.search(pattern, lbl)
+        if match:
+            symbol = match.group("symbol")
+            new_lbl = lbl.replace(symbol, "1")
+            return new_lbl, symbol
+        return lbl, ""
+
+    simpify_locals = {
+        "Min": sympy.Min,
+        "ceiling": sympy.ceiling,
+        "log2": lambda x: sympy.log(x, 2)
+    }
+
+    mult = 1
+    pretty_soqs = []
+    for soq_item in msd.soqs:
+        if isinstance(soq_item.symb, (TextBox, Text)):
+            try:
+                lbl_raw = soq_item.symb.text
+                lbl_no_symbols, symbol = symbols_to_identity(lbl_raw)
+                gate, base, exponent = sum([p.split("**", 1) for p in lbl_no_symbols.split("^")], [])
+
+                if len(base.split("*", 1)) > 1:
+                    mult_str, base = base.rsplit("*", 1)
+                    mult = sympy.sympify(mult_str, evaluate=True)
+
+                exponent = sympy.sympify(exponent, locals=simpify_locals, evaluate=True)
+                expression = str(base) + "**" + str(exponent)
+                expression = sympy.sympify(expression, locals=simpify_locals, evaluate=True)
+                new_lbl = str(gate) + "^" + symbol + "*" + str(expression * mult)
+
+                new_soq = soq_item.__class__(
+                    symb= TextBox(text=new_lbl) if isinstance(soq_item.symb, TextBox) else Text(text=new_lbl, fontsize=soq_item.symb.fontsize),
+                    rpos= soq_item.rpos,
+                    ident= soq_item.ident
+                )
+                pretty_soqs.append(new_soq)
+            except (ValueError, TypeError, NameError, sympy.SympifyError):
+                pretty_soqs.append(soq_item)
+        else:
+            pretty_soqs.append(soq_item)
+
+    pretty_msd = MusicalScoreData(
+        max_x=msd.max_x,
+        max_y=msd.max_y,
+        soqs=pretty_soqs,
+        hlines=msd.hlines,
+        vlines=msd.vlines
+    )
+
+    return pretty_msd