Add helper to plot counts from results

docs/tutorials/Solving_a_MWIS.ipynb

@@ -436,7 +436,7 @@
    "id": "24",
    "metadata": {},
    "source": [
-    "Finally we plot the full distribution of measured bitstrings, highlighting the exact MWIS solution `0110` in green. An adiabatic run that stayed in the ground state should return `0110` with overwhelming probability."
+    "Finally we plot the full distribution of measured bitstrings, highlighting the exact MWIS solution `0110`. An adiabatic run that stayed in the ground state should return `0110` with overwhelming probability."
    ]
   },
   {
@@ -446,33 +446,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from qoolqit.visualization import plot_bitstrings\n",
     "\n",
     "SOLUTION = \"0110\"\n",
     "\n",
-    "def plot_distribution(counts, solution, top=None):\n",
-    "    \"\"\"Bar plot of a bitstring-count distribution, highlighting the solution.\n",
-    "\n",
-    "    Args:\n",
-    "        counts (dict[str, int]): Mapping from measured bitstring to its count.\n",
-    "        solution (str): The bitstring to highlight (the exact MWIS answer).\n",
-    "        top (int | None): If given, only show the `top` most frequent bitstrings.\n",
-    "    \"\"\"\n",
-    "    counts = dict(sorted(counts.items(), key=lambda kv: kv[1], reverse=True))\n",
-    "    if top is not None:\n",
-    "        counts = dict(list(counts.items())[:top])\n",
-    "\n",
-    "    colors = [\"tab:green\" if b == solution else \"tab:blue\" for b in counts]\n",
-    "    plt.figure(figsize=(12, 5))\n",
-    "    plt.bar(counts.keys(), counts.values(), width=0.6, color=colors)\n",
-    "    plt.xlabel(\"bitstring\")\n",
-    "    plt.ylabel(\"counts\")\n",
-    "    plt.title(f\"Measurement distribution (solution {solution} in green)\")\n",
-    "    plt.xticks(rotation=\"vertical\")\n",
-    "    plt.tight_layout()\n",
-    "    plt.show()\n",
-    "\n",
-    "plot_distribution(counts, SOLUTION, top=20)"
+    "plot_bitstrings(counts, highlight={SOLUTION: \"#00C887\"})\n"
    ]
   },
   {
@@ -518,7 +497,7 @@
     "print(\"Most frequent bitstring:\", max(counts_analog, key=counts_analog.get))\n",
     "print(f\"P({SOLUTION}) = {p_solution:.2%}\")\n",
     "\n",
-    "plot_distribution(counts_analog, SOLUTION, top=20)"
+    "plot_bitstrings(counts, highlight={SOLUTION: \"#00C887\"})"
    ]
   },
   {
@@ -550,7 +529,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "qoolqit",
+   "display_name": "devqoolqit (3.14.6.final.0)",
    "language": "python",
    "name": "python3"
   },
@@ -564,7 +543,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.11"
+   "version": "3.14.6"
   }
  },
  "nbformat": 4,

docs/tutorials/solving_a_qubo.ipynb

@@ -365,46 +365,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from collections import Counter\n",
+    "from qoolqit.visualization import plot_bitstrings\n",
     "\n",
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "\n",
-    "def plot_distribution(counter, solutions, bins=10):\n",
-    "    counter = Counter(counter)\n",
-    "    counter = dict(counter.most_common(bins))\n",
-    "    color = [\n",
-    "        \"tab:green\" if key in solutions.tolist() else \"tab:blue\" for key in counter\n",
-    "    ]\n",
-    "    fig, ax = plt.subplots()\n",
-    "    ax.set_xlabel(\"Bitstrings\")\n",
-    "    ax.set_ylabel(\"Counts\")\n",
-    "    ax.bar(\n",
-    "        range(len(counter)), counter.values(), color=color, tick_label=counter.keys()\n",
-    "    )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "26",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "plot_distribution(counter, marked_bitstrings)"
+    "highlight = {b: \"#00C887\" for b in marked_bitstrings}\n",
+    "plot_bitstrings(counter, highlight=highlight)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "27",
+   "id": "26",
    "metadata": {},
    "source": [
     "As we can see, the bitstrings we had marked as the optimal solutions of this QUBO problem were the ones sampled with the highest probability, meaning the the QUBO problem was successfully solved with the quantum program we defined."
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "28",
+   "id": "27",
    "metadata": {},
    "source": [
     "## Advanced Compilation \n",
@@ -418,7 +395,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "29",
+   "id": "28",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -428,7 +405,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "30",
+   "id": "29",
    "metadata": {},
    "source": [
     "Concretely, we can see the beneficial effect of rescaling the drive duration on the simulation results:"
@@ -437,19 +414,19 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "31",
+   "id": "30",
    "metadata": {},
    "outputs": [],
    "source": [
     "job = emulator.run(program)\n",
     "results = job.results()\n",
     "counter = results.final_bitstrings\n",
-    "plot_distribution(counter, marked_bitstrings)"
+    "plot_bitstrings(counter, highlight=highlight)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "32",
+   "id": "31",
    "metadata": {},
    "source": [
     "Here the execution was relatively fast and easy, but for larger QUBO instances, or for QPU execution (which might have some queue), see the [Execution](https://docs.pasqal.com/qoolqit/qoolqitDoc/fundamentals/execution/execution/) section of the QoolQit documentation."
@@ -458,7 +435,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "qoolqit",
+   "display_name": "devqoolqit (3.14.6.final.0)",
    "language": "python",
    "name": "python3"
   },
@@ -472,7 +449,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.11"
+   "version": "3.14.6"
   }
  },
  "nbformat": 4,

qoolqit/visualization.py

@@ -0,0 +1,79 @@
+"""Visualization helpers for QoolQit results."""
+
+from __future__ import annotations
+
+from collections import Counter
+
+import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
+
+DEFAULT_BAR_COLOR = "#397378"
+
+
+def plot_bitstrings(
+    counts: dict[str, int],
+    top: int | None = None,
+    normalize: bool = False,
+    color: str = DEFAULT_BAR_COLOR,
+    highlight: dict[str, str] | None = None,
+    label: str | None = None,
+    ax: Axes | None = None,
+) -> None:
+    """Plot bitstring counts, optionally highlighting selected ones.
+
+    Arguments:
+        counts: Mapping of bitstrings to counts.
+        top: Plot only the top N counts.
+        normalize: Normalize counts to probabilities. Defaults to False.
+        color: Bar color.
+        highlight: Mapping of bitstrings to highlight colors. Highlighted
+            outcomes get their bar and tick label colored accordingly.
+        label: Legend label for the bars. Call ax.legend() to show it.
+        ax: Axes to draw on. Creates new axes if omitted.
+    """
+
+    if not counts:
+        raise ValueError("counts cannot be empty")
+
+    total = sum(counts.values())
+    if normalize and total == 0:
+        raise ValueError("cannot plot normalized counts with zero total counts")
+
+    if top is not None and top <= 0:
+        raise ValueError("top must be a positive integer")
+
+    # most_common(None) returns all entries, sorted by decreasing count
+    selected_counts = Counter(counts).most_common(top)
+    bitstrings = [bitstring for bitstring, _ in selected_counts]
+    values = [count / total if normalize else count for _, count in selected_counts]
+
+    highlight = highlight or {}
+
+    # Create the plot if no axes are provided
+    if ax is None:
+        _, ax = plt.subplots(figsize=(12, 5))
+
+    positions = range(len(bitstrings))
+    bar_colors = [highlight.get(bitstring, color) for bitstring in bitstrings]
+    ax.bar(positions, values, width=0.65, color=bar_colors)
+
+    if label is not None:
+        # A zero-height bar draws nothing but gives the legend a swatch in
+        # `color`, regardless of which bitstrings are highlighted.
+        ax.bar(0, 0, color=color, label=label)
+
+    ax.set_xticks(list(positions))
+    ax.set_xticklabels(bitstrings)
+
+    # Highlighted outcomes are also marked on their tick label
+    for tick_label in ax.get_xticklabels():
+        if tick_label.get_text() in highlight:
+            tick_label.set_color(highlight[tick_label.get_text()])
+            tick_label.set_fontweight("bold")
+
+    ax.tick_params(axis="x", labelrotation=90)
+    ax.grid(axis="y", linestyle="--", alpha=0.4)
+
+    # Default labels; the caller can override via ax.set_xlabel/ax.set_ylabel
+    ax.set_ylabel("Probability" if normalize else "Counts")
+    ax.set_xlabel("Bitstrings")

tests/test_visualization.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import pytest
+from matplotlib.colors import to_rgba
+
+from qoolqit.visualization import DEFAULT_BAR_COLOR, plot_bitstrings
+
+
+def test_plot_bitstrings_errors() -> None:
+    with pytest.raises(ValueError, match="counts cannot be empty"):
+        plot_bitstrings(counts={})
+    with pytest.raises(ValueError, match="cannot plot normalized counts with zero total counts"):
+        plot_bitstrings(counts={"000": 0}, normalize=True)
+    with pytest.raises(ValueError, match="top must be a positive integer"):
+        plot_bitstrings(counts={"000": 1, "001": 2}, top=0)
+
+
+def test_plot_bitstrings_default_bar_color() -> None:
+    _, ax = plt.subplots()
+    plot_bitstrings(counts={"000": 1, "001": 2}, ax=ax)
+
+    bars = ax.containers[0]
+    assert all(bar.get_facecolor() == to_rgba(DEFAULT_BAR_COLOR) for bar in bars)
+
+
+def test_plot_bitstrings_label_shows_up_in_legend() -> None:
+    _, ax = plt.subplots()
+    plot_bitstrings(counts={"000": 1, "001": 2}, label="run 1", ax=ax)
+
+    assert ax.get_legend() is None
+    legend = ax.legend()
+    assert legend.legend_handles[0].get_facecolor() == to_rgba(DEFAULT_BAR_COLOR)
+    assert legend.get_texts()[0].get_text() == "run 1"
+
+
+def test_plot_bitstrings_highlight_colors_the_bar() -> None:
+    _, ax = plt.subplots()
+    plot_bitstrings(counts={"000": 1, "001": 2}, highlight={"001": "tab:red"}, ax=ax)
+
+    bars = dict(zip(["001", "000"], ax.containers[0]))
+    assert bars["001"].get_facecolor() == to_rgba("tab:red")
+    assert bars["000"].get_facecolor() == to_rgba(DEFAULT_BAR_COLOR)
+
+
+def test_plot_bitstrings_legend_uses_base_color_even_if_first_bar_highlighted() -> None:
+    _, ax = plt.subplots()
+    # "001" has the higher count so it plots first, and is also highlighted.
+    plot_bitstrings(counts={"000": 1, "001": 2}, highlight={"001": "tab:red"}, label="run 1", ax=ax)
+
+    legend = ax.legend()
+    assert legend.legend_handles[0].get_facecolor() == to_rgba(DEFAULT_BAR_COLOR)
+
+
+def test_plot_bitstrings_two_calls_on_same_axes_keep_their_own_highlights() -> None:
+    _, ax = plt.subplots()
+    plot_bitstrings(counts={"000": 1, "001": 2}, highlight={"001": "tab:red"}, ax=ax)
+    plot_bitstrings(counts={"000": 1, "001": 2}, color="C2", ax=ax)
+
+    first_call_bars = dict(zip(["001", "000"], ax.containers[0]))
+    assert first_call_bars["001"].get_facecolor() == to_rgba("tab:red")