Organized visualizers to hopefully allow easier implementations of ne…

src/visualize/base/__init__.py

@@ -1,2 +1,7 @@
-from .graph import *
-from .sphere import *
+from .graph import graph
+from .sphere import sphere
+from .color_bar import color_bar
+from .light_mode import light_mode
+from .hamming_distance import hamming_distance
+from .get_qsphere_coordinates import get_qsphere_coordinates
+from .qsphere_latitude_finder import qsphere_latitude_finder

src/visualize/base/color_bar.py

@@ -0,0 +1,13 @@
+import numpy as np
+
+
+def color_bar(plt, _text, _accent, colors, norm) -> None:
+    cbar = plt.colorbar(
+        plt.cm.ScalarMappable(cmap=colors, norm=norm), ax=plt.gca(), shrink=0.55
+    )
+    cbar.set_label("Phase Angle", rotation=270, labelpad=15, color=_accent)
+    cbar.set_ticks([2 * np.pi, (3 * np.pi) / 2, np.pi, np.pi / 2, 0])
+    cbar.ax.yaxis.set_tick_params(color=_text)
+    cbar.outline.set_edgecolor(_text)
+    cbar.set_ticklabels(["2π", "3π / 2", "π", "π / 2", "0"], color=_text)
+    return

src/visualize/base/get_qsphere_coordinates.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+
+def get_qsphere_coordinates(num_qubits: int, lat_vals):
+    coords = []
+    phi = []
+    theta = []
+
+    for i in range(len(lat_vals)):
+        temp_arr = np.linspace(
+            2 * (np.pi) / len(lat_vals[i]), 2 * (np.pi), len(lat_vals[i])
+        )
+        theta.append(temp_arr)
+
+    phi = np.linspace(0, np.pi, num_qubits + 1)
+
+    for i in range(len(phi)):
+        for j in range(len(theta[i])):
+            x1 = 1 * np.sin(phi[i]) * np.cos(theta[i][j])
+            y1 = 1 * np.sin(phi[i]) * np.sin(theta[i][j])
+            z1 = 1 * np.cos(phi[i])
+            x, y, z = [0, x1], [0, y1], [0, z1]
+            coords.append([x, y, z])
+
+    return coords

src/visualize/base/hamming_distance.py

@@ -0,0 +1,2 @@
+def hamming_distance(l1: str, l2: str) -> int:
+    return l1.count("1") == l2.count("1")

src/visualize/base/light_mode.py

@@ -0,0 +1,8 @@
+from . import theme
+
+
+def light_mode(light_mode: bool) -> None:
+    if light_mode:
+        theme.TEXT_COLOR = "black"
+        theme.ACCENT_COLOR = "black"
+        theme.BACKGROUND_COLOR = "white"

src/visualize/base/qsphere_latitude_finder.py

@@ -0,0 +1,30 @@
+from collections import deque
+from .hamming_distance import hamming_distance
+
+
+def qsphere_latitude_finder(num_qubits: int, state_list):
+    latitude_values = [[]]
+    for _ in range(num_qubits - 1):
+        latitude_values.append([])
+    latitude_values.append([])
+
+    queue_of_state = deque(state_list)
+    latitude_values[0].append(queue_of_state.popleft())
+    latitude_values[-1].append(queue_of_state.pop())
+
+    bit_representation = "0" * (num_qubits - 1) + "1"
+
+    for i in range(1, len(bit_representation)):
+        latitude_values[i].append(bit_representation)
+        queue_of_state.remove(bit_representation)
+        list_temp = list(bit_representation)
+        list_temp[i - 1] = "1"
+        bit_representation = "".join(list_temp)
+
+    while queue_of_state:
+        bit_representation = queue_of_state.popleft()
+        for i in range(1, len(latitude_values) - 1):
+            if hamming_distance(bit_representation, latitude_values[i][0]):
+                latitude_values[i].append(bit_representation)
+
+    return latitude_values

src/visualize/base/sphere.py

@@ -2,7 +2,7 @@
 import numpy as np
 
 
-def sphere(_background):
+def sphere(_background: str):
     plt.clf()
     plt.close()
     plt.clf()

src/visualize/base/theme.py

@@ -0,0 +1,5 @@
+LIGHT_GREY = "lightgrey"
+GREY = "grey"
+TEXT_COLOR = "white"
+ACCENT_COLOR = "#39c0ba"
+BACKGROUND_COLOR = "#2e3037"

src/visualize/bloch.py

@@ -1,31 +1,20 @@
 import matplotlib.pyplot as plt
 import numpy as np
-from .base import sphere
+from .base import sphere, theme, light_mode
 from ..tools import probability, amplitude
 
 
 def bloch(
-    quantumstate=np.array,
+    quantumstate: any,
     path: str = "BlochSphere.png",
     save: bool = False,
     show: bool = True,
-    darkmode: bool = True,
+    light: bool = False,
 ):
-    if np.log2(len(quantumstate)) > 1:
-        exit(
-            f"Error: BlochSphere() --", f"BlochSphere only calculates 1 qubit circuits."
-        )
     amplitutes = amplitude(quantumstate)
     phase_angles = probability(quantumstate, False)
-    if darkmode:
-        _text = "white"
-        _accent = "#39c0ba"
-        _background = "#2e3037"
-    else:
-        _text = "black"
-        _accent = "black"
-        _background = "white"
-    ax = sphere(_background)
+    light_mode(light)
+    ax = sphere(theme.BACKGROUND_COLOR)
     ax.quiver(1, 0, 0, 0.75, 0, 0, color="lightgray")
     ax.text(2, 0, 0, "+x", color="gray")
     ax.quiver(0, 1, 0, 0, 0.75, 0, color="lightgray")
@@ -42,9 +31,9 @@ def bloch(
     y = 1 * np.sin(theta) * np.sin(phi)
     z = 1 * np.cos(theta)
     xs, ys, zs = [0, x], [0, y], [0, z]
-    ax.plot3D(xs, ys, zs, color=_accent, markevery=100)
-    ax.scatter(xs[1], ys[1], zs[1], s=5, color=_accent)
-    ax.text(xs[1] * 1.15, ys[1] * 1.15, zs[1] * 1.15, "|ψ⟩", color=_text)
+    ax.plot3D(xs, ys, zs, color=theme.ACCENT_COLOR, markevery=100)
+    ax.scatter(xs[1], ys[1], zs[1], s=5, color=theme.ACCENT_COLOR)
+    ax.text(xs[1] * 1.15, ys[1] * 1.15, zs[1] * 1.15, "|ψ⟩", color=theme.ACCENT_COLOR)
     plt.tight_layout()
     plt.axis("off")
     if save:

src/visualize/probability.py

@@ -1,7 +1,6 @@
 import matplotlib.pyplot as plt
-from .base.graph import graph
+from .base import graph, light_mode, theme
 from ..tools import probability as prob
-from ..tools.base import convert_state
 import numpy as np
 
 
@@ -10,30 +9,28 @@ def probability(
     path: str = "probabilities.png",
     save: bool = False,
     show: bool = True,
-    darkmode: bool = True,
+    light: bool = False,
 ):
-
-    num_qubits = int(np.log2(len(convert_state(state))))
+    probabilities = prob(state)
+    num_qubits = int(np.log2(probabilities.size))
     state_list = [format(i, "b").zfill(num_qubits) for i in range(2**num_qubits)]
-    percents = [i * 100 for i in prob(state)]
-    if darkmode:
-        _text = "white"
-        _accent = "#39c0ba"
-        _background = "#2e3037"
-    else:
-        _text = "black"
-        _accent = "black"
-        _background = "white"
+    percents = [i * 100 for i in probabilities]
+
     plt.clf()
     plt.close()
-    ax = graph(_text, _background, num_qubits)
+
+    light_mode(light)
+    ax = graph(theme.TEXT_COLOR, theme.BACKGROUND_COLOR, num_qubits)
     ax.bar(state_list, percents, color="#39c0ba")
-    plt.xlabel("Computational basis states", color=_accent)
-    plt.ylabel("Probability (%)", labelpad=5, color=_accent)
-    plt.title("Probabilities", pad=10, color=_accent)
+
+    plt.xlabel("Computational basis states", color=theme.ACCENT_COLOR)
+    plt.ylabel("Probability (%)", labelpad=5, color=theme.ACCENT_COLOR)
+    plt.title("Probabilities", pad=10, color=theme.ACCENT_COLOR)
     plt.tight_layout()
+
     if save:
         plt.savefig(path)
     if show:
         plt.show()
+
     return

src/visualize/q_sphere.py

@@ -1,105 +1,57 @@
-from collections import deque
 import matplotlib.pyplot as plt
 import numpy as np
-from .base.sphere import sphere
+from .base import (
+    sphere,
+    color_bar,
+    theme,
+    light_mode,
+    qsphere_latitude_finder,
+    get_qsphere_coordinates,
+)
 from ..tools import probability, phaseangle
 
 
-def hamming_distance(l1: str, l2: str):
-    return l1.count("1") == l2.count("1")
-
-
-def latitude_finder(num_qubits: int, state_list):
-    latitude_values = [[]]
-    for _ in range(num_qubits - 1):
-        latitude_values.append([])
-    latitude_values.append([])
-    queue_of_state = deque(state_list)
-    latitude_values[0].append(queue_of_state.popleft())
-    latitude_values[-1].append(queue_of_state.pop())
-    bit_representation = "0" * (num_qubits - 1) + "1"
-    for i in range(1, len(bit_representation)):
-        latitude_values[i].append(bit_representation)
-        queue_of_state.remove(bit_representation)
-        list_temp = list(bit_representation)
-        list_temp[i - 1] = "1"
-        bit_representation = "".join(list_temp)
-    while queue_of_state:
-        bit_representation = queue_of_state.popleft()
-        for i in range(1, len(latitude_values) - 1):
-            if hamming_distance(bit_representation, latitude_values[i][0]):
-                latitude_values[i].append(bit_representation)
-    return latitude_values
-
-
-def get_coords(num_qubits, lat_vals):
-    coords = []
-    phi = []
-    theta = []
-    for i in range(len(lat_vals)):
-        temp_arr = np.linspace(
-            2 * (np.pi) / len(lat_vals[i]), 2 * (np.pi), len(lat_vals[i])
-        )
-        theta.append(temp_arr)
-    phi = np.linspace(0, np.pi, num_qubits + 1)
-    for i in range(len(phi)):
-        for j in range(len(theta[i])):
-            x1 = 1 * np.sin(phi[i]) * np.cos(theta[i][j])
-            y1 = 1 * np.sin(phi[i]) * np.sin(theta[i][j])
-            z1 = 1 * np.cos(phi[i])
-            x, y, z = [0, x1], [0, y1], [0, z1]
-            coords.append([x, y, z])
-
-    return coords
-
-
 def q_sphere(
-    quantumstate,
+    circuit: any,
     path: str = "qsphere.png",
     save: bool = False,
     show: bool = True,
-    darkmode: bool = True,
+    light: bool = False,
 ):
-    num_qubits = int(np.log2((len(quantumstate.state))))
-    probs = probability(quantumstate)
-    angle = phaseangle(quantumstate)
+    num_qubits = int(np.log2((len(circuit.state))))
+    probs = probability(circuit)
+    angle = phaseangle(circuit)
+
     state_list = [format(i, "b").zfill(num_qubits) for i in range(2**num_qubits)]
+
     prob_dict = {state_list[i]: probs[i] for i in range(len(state_list))}
     phase_dict = {state_list[i]: angle[i] for i in range(len(state_list))}
-    lat_vals = latitude_finder(num_qubits, state_list)
-    if darkmode:
-        _text = "white"
-        _accent = "#39c0ba"
-        _background = "#2e3037"
-    else:
-        _text = "black"
-        _accent = "black"
-        _background = "white"
-    ax = sphere(_background)
-    coords = get_coords(num_qubits, lat_vals)
+    lat_vals = qsphere_latitude_finder(num_qubits, state_list)
+
+    light_mode(light)
+    ax = sphere(theme.BACKGROUND_COLOR)
+    coords = get_qsphere_coordinates(num_qubits, lat_vals)
     ham_states = [item for sublist in lat_vals for item in sublist]
     colors = plt.get_cmap("hsv")
     norm = plt.Normalize(0, np.pi * 2)
+    color_bar(plt, theme.TEXT_COLOR, theme.ACCENT_COLOR, colors, norm)
+
     for i, j in zip(coords, ham_states):
         cur_prob = prob_dict[j]
         cur_phase = phase_dict[j]
         if cur_prob > 0:
             x, y, z = i[0], i[1], i[2]
             ax.plot3D(x, y, z, color=colors(norm(cur_phase)))
             ax.scatter(x[1], y[1], z[1], s=5, color=colors(norm(cur_phase)))
-            ax.text(x[1] * 1.15, y[1] * 1.15, z[1] * 1.15, f"|{j}>", color=_text)
-    cbar = plt.colorbar(
-        plt.cm.ScalarMappable(cmap=colors, norm=norm), ax=plt.gca(), shrink=0.55
-    )
-    cbar.set_label("Phase Angle", rotation=270, labelpad=15, color=_accent)
-    cbar.set_ticks([2 * np.pi, (3 * np.pi) / 2, np.pi, np.pi / 2, 0])
-    cbar.ax.yaxis.set_tick_params(color=_text)
-    cbar.outline.set_edgecolor(_text)
-    cbar.set_ticklabels(["2π", "3π / 2", "π", "π / 2", "0"], color=_text)
+            ax.text(
+                x[1] * 1.15, y[1] * 1.15, z[1] * 1.15, f"|{j}>", color=theme.TEXT_COLOR
+            )
+
     plt.tight_layout()
     plt.axis("off")
     if save:
         plt.savefig(path)
     if show:
         plt.show()
+
     return