Refactor state latex drawer

qiskit/visualization/array.py

@@ -13,112 +13,57 @@
 Tools to create LaTeX arrays.
 """
 
-import math
-from fractions import Fraction
 import numpy as np
 
 from qiskit.exceptions import MissingOptionalLibraryError
 
 
-def _num_to_latex(num, precision=5):
-    """Takes a complex number as input and returns a latex representation
+def num_to_latex(raw_value, precision=15, coefficient=False, first_term=True):
+    """Convert a complex number to latex code suitable for a ket expression
 
     Args:
-        num (numerical): The number to be converted to latex.
-        precision (int): If the real or imaginary parts of num are not close
-                         to an integer, the number of decimal places to round to
-
+        raw_value (complex): Value to convert.
+        precision (int): Number of decimal places to round to (default 15).
+        coefficient (bool): Whether the number is to be used as a coefficient
+                            of a ket.
+        first_term (bool): If a coefficient, whether this number is the first
+                           coefficient in the expression.
     Returns:
-        str: Latex representation of num
+        str: latex code
     """
-    # Result is combination of maximum 4 strings in the form:
-    #     {common_facstring} ( {realstring} {operation} {imagstring}i )
-    # common_facstring: A common factor between the real and imaginary part
-    # realstring: The real part (inc. a negative sign if applicable)
-    # operation: The operation between the real and imaginary parts ('+' or '-')
-    # imagstring: Absolute value of the imaginary parts (i.e. not inc. any negative sign).
-    # This function computes each of these strings and combines appropriately.
-
-    r = np.real(num)
-    i = np.imag(num)
-    common_factor = None
-
-    # try to factor out common terms in imaginary numbers
-    if np.isclose(abs(r), abs(i)) and not np.isclose(r, 0) and not np.isclose(i, 0):
-        common_factor = abs(r)
-        r = r / common_factor
-        i = i / common_factor
-
-    common_terms = {
-        1 / math.sqrt(2): "\\tfrac{1}{\\sqrt{2}}",
-        1 / math.sqrt(3): "\\tfrac{1}{\\sqrt{3}}",
-        math.sqrt(2 / 3): "\\sqrt{\\tfrac{2}{3}}",
-        math.sqrt(3 / 4): "\\sqrt{\\tfrac{3}{4}}",
-        1 / math.sqrt(8): "\\tfrac{1}{\\sqrt{8}}",
-    }
-
-    def _proc_value(val):
-        # This function converts a real value to a latex string
-        # First, see if val is close to an integer:
-        val_mod = np.mod(val, 1)
-        if np.isclose(val_mod, 0) or np.isclose(val_mod, 1):
-            # If so, return that integer
-            return str(int(np.round(val)))
-        # Otherwise, see if it matches one of the common terms
-        for term, latex_str in common_terms.items():
-            if np.isclose(abs(val), term):
-                if val > 0:
-                    return latex_str
-                else:
-                    return "-" + latex_str
-        # try to factorise val nicely
-        frac = Fraction(val).limit_denominator()
-        num, denom = frac.numerator, frac.denominator
-        if abs(num) + abs(denom) < 20:
-            # If fraction is 'nice' return
-            if val > 0:
-                return f"\\tfrac{{{abs(num)}}}{{{abs(denom)}}}"
-            else:
-                return f"-\\tfrac{{{abs(num)}}}{{{abs(denom)}}}"
-        else:
-            # Failing everything else, return val as a decimal
-            return "{:.{}f}".format(val, precision).rstrip("0")
+    import sympy  # runtime import
 
-    # Get string (or None) for common factor between real and imag
-    if common_factor is None:
-        common_facstring = None
-    else:
-        common_facstring = _proc_value(common_factor)
+    raw_value = np.around(raw_value, precision)
+    value = sympy.nsimplify(raw_value, rational=False)
 
-    # Get string for real part
-    realstring = _proc_value(r)
+    if isinstance(value, sympy.core.numbers.Rational) and value.denominator > 50:
+        # Avoid showing ugly fractions (e.g. 50498971964399/62500000000000)
+        value = value.evalf()  # Display as float
 
-    # Get string for both imaginary part and operation between real and imaginary parts
-    if i > 0:
-        operation = "+"
-        imagstring = _proc_value(i)
-    else:
-        operation = "-"
-        imagstring = _proc_value(-i)
-    if imagstring == "1":
-        imagstring = ""  # Don't want to return '1i', just 'i'
-
-    # Now combine the strings appropriately:
-    if imagstring == "0":
-        return realstring  # realstring already contains the negative sign (if needed)
-    if realstring == "0":
-        # imagstring needs the negative sign adding
-        if operation == "-":
-            return f"-{imagstring}i"
-        else:
-            return f"{imagstring}i"
-    if common_facstring is not None:
-        return f"{common_facstring}({realstring} {operation} {imagstring}i)"
-    else:
-        return f"{realstring} {operation} {imagstring}i"
+    if isinstance(value, sympy.core.numbers.Float):
+        value = round(value, precision)
+
+    element = sympy.latex(value, full_prec=False)
+
+    if not coefficient:
+        return element
+
+    if isinstance(value, sympy.core.Add):
+        # element has two terms
+        element = f"({element})"
+
+    if element == "1":
+        element = ""
+    elif element == "-1":
+        element = "-"
 
+    if not first_term and not element.startswith("-"):
+        element = f"+{element}"
 
-def _matrix_to_latex(matrix, precision=5, prefix="", max_size=(8, 8)):
+    return element
+
+
+def _matrix_to_latex(matrix, precision=10, prefix="", max_size=(8, 8)):
     """Latex representation of a complex numpy array (with maximum dimension 2)
 
     Args:
@@ -134,13 +79,7 @@ def _matrix_to_latex(matrix, precision=5, prefix="", max_size=(8, 8)):
 
     Returns:
         str: Latex representation of the matrix
-
-    Raises:
-        ValueError: If minimum value in max_size < 3
     """
-    if min(max_size) < 3:
-        raise ValueError("""Smallest value in max_size must be greater than or equal to 3""")
-
     out_string = f"\n{prefix}\n"
     out_string += "\\begin{bmatrix}\n"
 
@@ -149,7 +88,7 @@ def _elements_to_latex(elements):
         # string from it; Each element separated by `&`
         el_string = ""
         for el in elements:
-            num_string = _num_to_latex(el, precision=precision)
+            num_string = num_to_latex(el, precision=precision)
             el_string += num_string + " & "
         el_string = el_string[:-2]  # remove trailing ampersands
         return el_string
@@ -197,7 +136,7 @@ def _rows_to_latex(rows, max_width):
     return out_string
 
 
-def array_to_latex(array, precision=5, prefix="", source=False, max_size=8):
+def array_to_latex(array, precision=10, prefix="", source=False, max_size=8):
     """Latex representation of a complex numpy array (with dimension 1 or 2)
 
     Args:
@@ -224,6 +163,7 @@ def array_to_latex(array, precision=5, prefix="", source=False, max_size=8):
     Raises:
         TypeError: If array can not be interpreted as a numerical numpy array.
         ValueError: If the dimension of array is not 1 or 2.
+        ValueError: If the smallest value in ``max_size`` is less than 3.
         MissingOptionalLibraryError: If ``source`` is ``False`` and ``IPython.display.Latex`` cannot be
                      imported.
     """
@@ -236,13 +176,17 @@ def array_to_latex(array, precision=5, prefix="", source=False, max_size=8):
         or types that can be converted to such arrays"""
         ) from err
 
-    if array.ndim <= 2:
-        if isinstance(max_size, int):
-            max_size = (max_size, max_size)
-        outstr = _matrix_to_latex(array, precision=precision, prefix=prefix, max_size=max_size)
-    else:
+    if isinstance(max_size, int):
+        max_size = (max_size, max_size)
+
+    if min(max_size) < 3:
+        raise ValueError("Smallest value in max_size must be greater than or equal to 3")
+
+    if array.ndim > 2:
         raise ValueError("array_to_latex can only convert numpy ndarrays of dimension 1 or 2")
 
+    outstr = _matrix_to_latex(array, precision=precision, prefix=prefix, max_size=max_size)
+
     if source is False:
         try:
             from IPython.display import Latex