Create heatmap.py

quread/heatmap.py

@@ -0,0 +1,110 @@
+# quread/heatmap.py
+from __future__ import annotations
+
+import csv
+import io
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+@dataclass
+class HeatmapConfig:
+    rows: int = 8          # chip rows
+    cols: int = 8          # chip cols
+    missing_value: float = 0.0
+
+
+def _default_qubit_coords(n_qubits: int, rows: int, cols: int) -> Dict[int, Tuple[int, int]]:
+    """
+    Simple default mapping: q0.. mapped row-major across the chip grid.
+    Override later with user-uploaded mapping if needed.
+    """
+    m: Dict[int, Tuple[int, int]] = {}
+    for q in range(n_qubits):
+        r = q // cols
+        c = q % cols
+        if r >= rows:
+            break
+        m[q] = (r, c)
+    return m
+
+
+def _parse_history_rows(csv_text: str) -> List[dict]:
+    """
+    Expects the CSV you generate in Task 2A:
+      step,gate,target,control,theta
+    """
+    f = io.StringIO(csv_text or "")
+    reader = csv.DictReader(f)
+    rows = []
+    for r in reader:
+        rows.append({
+            "step": int(r.get("step", "0") or 0),
+            "gate": (r.get("gate") or "").strip(),
+            "target": r.get("target", "").strip(),
+            "control": r.get("control", "").strip(),
+            "theta": (r.get("theta") or "").strip(),
+        })
+    return rows
+
+
+def make_activity_heatmap(
+    csv_text: str,
+    n_qubits: int,
+    cfg: Optional[HeatmapConfig] = None,
+    qubit_coords: Optional[Dict[int, Tuple[int, int]]] = None,
+) -> plt.Figure:
+    """
+    Builds a heatmap counting how many times each qubit participates in an operation.
+    - Single-qubit op increments its target qubit.
+    - CNOT increments both control and target.
+    """
+    cfg = cfg or HeatmapConfig()
+    coords = qubit_coords or _default_qubit_coords(n_qubits, cfg.rows, cfg.cols)
+
+    grid = np.full((cfg.rows, cfg.cols), cfg.missing_value, dtype=float)
+
+    # count per qubit
+    counts = np.zeros((n_qubits,), dtype=float)
+
+    rows = _parse_history_rows(csv_text)
+    for r in rows:
+        gate = r["gate"].upper()
+
+        # target
+        if r["target"] != "":
+            t = int(r["target"])
+            if 0 <= t < n_qubits:
+                counts[t] += 1.0
+
+        # control for CNOT
+        if gate == "CNOT" and r["control"] != "":
+            c = int(r["control"])
+            if 0 <= c < n_qubits:
+                counts[c] += 1.0
+
+    # place into chip grid
+    for q, (rr, cc) in coords.items():
+        if 0 <= q < n_qubits and 0 <= rr < cfg.rows and 0 <= cc < cfg.cols:
+            grid[rr, cc] = counts[q]
+
+    # plot
+    fig, ax = plt.subplots(figsize=(6, 5))
+    im = ax.imshow(grid, interpolation="nearest")
+    ax.set_title("Qubit activity heatmap (from circuit CSV)")
+    ax.set_xlabel("Chip column")
+    ax.set_ylabel("Chip row")
+    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+
+    # annotate qubit ids on mapped cells
+    for q, (rr, cc) in coords.items():
+        if 0 <= rr < cfg.rows and 0 <= cc < cfg.cols:
+            ax.text(cc, rr, f"q{q}", ha="center", va="center", fontsize=9)
+
+    ax.set_xticks(range(cfg.cols))
+    ax.set_yticks(range(cfg.rows))
+    fig.tight_layout()
+    return fig