{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Pre-Trajectory Sampling with Batch Execution (PTSBE)\n",
        "\n",
        "PTSBE (Pre-Trajectory Sampling with Batch Execution) is a method for sampling from noisy quantum circuits efficiently. Instead of simulating the full density matrix and sampling once per shot, PTSBE:\n",
        "\n",
        "1. Traces the kernel to get the gate sequence and qubit layout.\n",
        "2. Extracts noise sites by matching the noise model to the trace (each noisy gate becomes a *noise site* with a set of Kraus outcomes, e.g. $I$, $X$, $Y$, $Z$ for depolarization).\n",
        "3. Generates trajectories — each trajectory is one possible *realization* of noise (one Kraus outcome per site). A *sampling strategy* decides which trajectories to use (e.g. by sampling trajectories proporitional to their likelihood).\n",
        "4. Allocates shots across trajectories (e.g. by error likelihood).\n",
        "5. Runs batches — for each trajectory, the circuit is run as a noiseless circuit with that trajectory's outcomes applied; results are collected.\n",
        "6. Aggregates all per-trajectory counts into a single `SampleResult`.\n",
        "\n",
        "Given sufficient trajectory and shot samples PTSBE will be equivalent to standard trajectory based sampling up to sampling noise. However, it has the additional advantage of being able to batch many shots per trajectory and to control simulation cost via the number of trajectories. This notebook runs the full workflow with a single API call: `cudaq.ptsbe.sample()`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Set up the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import cudaq\n",
        "\n",
        "cudaq.set_target(\"nvidia\")\n",
        "cudaq.set_random_seed(42)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Define the circuit and noise model\n",
        "\n",
        "Define a kernel and attach a noise model. Each gate you add to the noise model becomes a noise site when that gate appears in the circuit. In this example, for single-qubit gates (e.g. $H$) we use a `DepolarizationChannel` with one qubit and for $CNOT$ we add the **qubit pair** `[control, target]` with `Depolarization2` (two-qubit depolarization channel). Here we use a small Bell-style circuit to demonstrate the"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "depol_1q = 0.05\n",
        "depol_2q = 0.03\n",
        "\n",
        "@cudaq.kernel\n",
        "def bell_with_noise():\n",
        "    q = cudaq.qvector(2)\n",
        "    h(q[0])\n",
        "    x.ctrl(q[0], q[1])\n",
        "    mz(q)\n",
        "\n",
        "noise = cudaq.NoiseModel()\n",
        "noise.add_all_qubit_channel(\"h\", cudaq.DepolarizationChannel(depol_1q))\n",
        "noise.add_all_qubit_channel(\"x\", cudaq.Depolarization2(depol_2q), num_controls=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Inline noise with `apply_noise`\n",
        "\n",
        "Instead of (or in addition to) attaching noise via a `NoiseModel`, you can place noise at specific points in the kernel with `cudaq.apply_noise`. PTSBE traces these calls and includes them as noise sites alongside any model-attached noise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Inline apply_noise result:\n",
            "{ 00:5007 01:79 10:69 11:4845 }\n",
            "\n"
          ]
        }
      ],
      "source": [
        "@cudaq.kernel\n",
        "def bell_inline_noise():\n",
        "    q = cudaq.qvector(2)\n",
        "    h(q[0])\n",
        "    cudaq.apply_noise(cudaq.DepolarizationChannel, depol_1q, q[0])\n",
        "    x.ctrl(q[0], q[1])\n",
        "    cudaq.apply_noise(cudaq.Depolarization2, depol_2q, q[0], q[1])\n",
        "    mz(q)\n",
        "\n",
        "result_inline = cudaq.ptsbe.sample(\n",
        "    bell_inline_noise,\n",
        "    shots_count=10000,\n",
        ")\n",
        "print(\"Inline apply_noise result:\")\n",
        "print(result_inline)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Run PTSBE sampling\n",
        "\n",
        "Call `cudaq.ptsbe.sample()` with the kernel, noise model, and shot count. This example uses `OrderedSamplingStrategy` plus `LOW_WEIGHT_BIAS` shot allocation to emphasize likely low-error trajectories with a smaller shot budget.\n",
        "\n",
        "Optional arguments are:\n",
        "\n",
        "- `sampling_strategy` — how trajectories are chosen:\n",
        "  - `ProbabilisticSamplingStrategy(seed=...)` (default): Performs Monte Carlo sampling over trajectories.\n",
        "  - `ExhaustiveSamplingStrategy()`: use all possible trajectories (every combination of Kraus outcomes per noise site).\n",
        "  - `OrderedSamplingStrategy()`: Select the top-$k$ trajectories by probability (highest first), up to `max_trajectories`.\n",
        "- `shot_allocation` — how shots are split across the chosen trajectories:\n",
        "  - `PROPORTIONAL` (default): allocate shots in proportion to each sampled trajectory's weighting.\n",
        "  - `UNIFORM`: give each trajectory the same number of shots.\n",
        "  - `LOW_WEIGHT_BIAS`: bias more shots toward low-weight (fewer errors) trajectories, optional `bias_strength` (default 2.0).\n",
        "  - `HIGH_WEIGHT_BIAS`: bias more shots toward high-weight trajectories, optional `bias_strength` (default 2.0).\n",
        "- `max_trajectories`: cap the number of trajectories (useful for large shot counts).\n",
        "- `return_execution_data`: If ``True``, the result includes trace instructions and per-trajectory data (``result.ptsbe_execution_data``). This is currently an experimental API and subject to change in future releases. See section 5 below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Comparison with same shots=10000, max_trajectories=16\n",
            "PTSBE (ordered + low-weight bias)    shots= 10000  counts={ 00:5044 01:5 10:5 11:4946 }\n",
            "\n",
            "PTSBE (ordered + proportional)       shots= 10000  counts={ 00:4919 01:45 10:59 11:4977 }\n",
            "\n",
            "Standard noisy sample                shots= 10000  counts={ 00:5029 01:93 10:74 11:4804 }\n",
            "\n"
          ]
        }
      ],
      "source": [
        "shots = 10000\n",
        "max_traj = 16\n",
        "strategy = cudaq.ptsbe.OrderedSamplingStrategy()\n",
        "\n",
        "def print_result(label, sample_result):\n",
        "    print(f\"{label:<36} shots={sample_result.get_total_shots():>6}  counts={sample_result}\")\n",
        "\n",
        "biased_alloc = cudaq.ptsbe.ShotAllocationStrategy(\n",
        "    cudaq.ptsbe.ShotAllocationType.LOW_WEIGHT_BIAS,\n",
        "    bias_strength=3.0,\n",
        "    seed=42,\n",
        ")\n",
        "result_biased = cudaq.ptsbe.sample(\n",
        "    bell_with_noise,\n",
        "    noise_model=noise,\n",
        "    shots_count=shots,\n",
        "    sampling_strategy=strategy,\n",
        "    shot_allocation=biased_alloc,\n",
        "    max_trajectories=max_traj,\n",
        ")\n",
        "\n",
        "unbiased_alloc = cudaq.ptsbe.ShotAllocationStrategy(\n",
        "    cudaq.ptsbe.ShotAllocationType.PROPORTIONAL,\n",
        "    seed=42,\n",
        ")\n",
        "result_unbiased = cudaq.ptsbe.sample(\n",
        "    bell_with_noise,\n",
        "    noise_model=noise,\n",
        "    shots_count=shots,\n",
        "    sampling_strategy=strategy,\n",
        "    shot_allocation=unbiased_alloc,\n",
        "    max_trajectories=max_traj,\n",
        ")\n",
        "\n",
        "result_standard = cudaq.sample(\n",
        "    bell_with_noise,\n",
        "    noise_model=noise,\n",
        "    shots_count=shots,\n",
        ")\n",
        "\n",
        "print(f\"Comparison with same shots={shots}, max_trajectories={max_traj}\")\n",
        "print_result(\"PTSBE (ordered + low-weight bias)\", result_biased)\n",
        "print_result(\"PTSBE (ordered + proportional)\", result_unbiased)\n",
        "print_result(\"Standard noisy sample\", result_standard)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Larger circuit for execution data\n",
        "\n",
        "To demonstrate PTSBE's execution data on a larger circuit, we define a 12-qubit GHZ state with depolarization on every gate. This creates 12 noise sites with a large trajectory space."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "n_qubits = 12\n",
        "\n",
        "@cudaq.kernel\n",
        "def ghz(n: int):\n",
        "    q = cudaq.qvector(n)\n",
        "    h(q[0])\n",
        "    for i in range(1, n):\n",
        "        x.ctrl(q[i - 1], q[i])\n",
        "    mz(q)\n",
        "\n",
        "ghz_depol_1q = 0.01\n",
        "ghz_depol_2q = 0.01\n",
        "\n",
        "ghz_noise = cudaq.NoiseModel()\n",
        "ghz_noise.add_all_qubit_channel(\"h\", cudaq.DepolarizationChannel(ghz_depol_1q))\n",
        "ghz_noise.add_all_qubit_channel(\"x\", cudaq.Depolarization2(ghz_depol_2q), num_controls=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Inspecting trajectories with execution data\n",
        "\n",
        "Pass `return_execution_data=True` to get the PTSBE execution data via `result.ptsbe_execution_data`. This reveals *why* PTSBE is efficient: most probability mass concentrates on a few low-error trajectories, so a small number of circuit simulations captures the bulk of the physics. The highest-probability trajectory (no errors at any noise site) typically receives the vast majority of shots, while multi-error trajectories are exponentially suppressed.\n",
        "\n",
        "Note: this is an experimental API and may change in future releases."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Trajectories: 147, Total shots: 1,000,000\n",
            "[2026-03-16 15:10:06.241] [warning] [PTSBESampleResult.cpp:20] PTSBE execution data API is experimental and may change in a future release.\n",
            "\n",
            "Top 5 trajectories (highest probability):\n",
            "  #1: p_theory=0.886385, p_empirical=0.888297, shots=888,297, errors=0, cumulative shots=88.8%\n",
            "  #2: p_theory=0.002984, p_empirical=0.002378, shots=2,378, errors=1, cumulative shots=89.1%\n",
            "  #3: p_theory=0.002984, p_empirical=0.002333, shots=2,333, errors=1, cumulative shots=89.3%\n",
            "  #4: p_theory=0.002984, p_empirical=0.004269, shots=4,269, errors=1, cumulative shots=89.7%\n",
            "  #5: p_theory=0.000597, p_empirical=0.000434, shots=434, errors=1, cumulative shots=89.8%\n",
            "\n",
            "Lowest-probability trajectory:\n",
            "  prob=2.71e-10, shots=359, errors=3\n",
            "\n",
            "Highest-probability trajectory (id=0):\n",
            "  prob=0.886385, shots=888,297\n",
            "  Kraus selections:\n",
            "    site 1 [depolarization_channel on q[0]]: K0 (no-error)\n",
            "    site 3 [depolarization2 on q[1, 0]]: K0 (no-error)\n",
            "    site 5 [depolarization2 on q[2, 1]]: K0 (no-error)\n",
            "    site 7 [depolarization2 on q[3, 2]]: K0 (no-error)\n",
            "    site 9 [depolarization2 on q[4, 3]]: K0 (no-error)\n",
            "    site 11 [depolarization2 on q[5, 4]]: K0 (no-error)\n",
            "    site 13 [depolarization2 on q[6, 5]]: K0 (no-error)\n",
            "    site 15 [depolarization2 on q[7, 6]]: K0 (no-error)\n",
            "    site 17 [depolarization2 on q[8, 7]]: K0 (no-error)\n",
            "    site 19 [depolarization2 on q[9, 8]]: K0 (no-error)\n",
            "    site 21 [depolarization2 on q[10, 9]]: K0 (no-error)\n",
            "    site 23 [depolarization2 on q[11, 10]]: K0 (no-error)\n",
            "\n",
            "Lowest-probability trajectory (id=108):\n",
            "  prob=2.71e-10, shots=359\n",
            "  Kraus selections:\n",
            "    site 1 [depolarization_channel on q[0]]: K0 (no-error)\n",
            "    site 3 [depolarization2 on q[1, 0]]: K13 (error)\n",
            "    site 5 [depolarization2 on q[2, 1]]: K12 (error)\n",
            "    site 7 [depolarization2 on q[3, 2]]: K0 (no-error)\n",
            "    site 9 [depolarization2 on q[4, 3]]: K0 (no-error)\n",
            "    site 11 [depolarization2 on q[5, 4]]: K7 (error)\n",
            "    site 13 [depolarization2 on q[6, 5]]: K0 (no-error)\n",
            "    site 15 [depolarization2 on q[7, 6]]: K0 (no-error)\n",
            "    site 17 [depolarization2 on q[8, 7]]: K0 (no-error)\n",
            "    site 19 [depolarization2 on q[9, 8]]: K0 (no-error)\n",
            "    site 21 [depolarization2 on q[10, 9]]: K0 (no-error)\n",
            "    site 23 [depolarization2 on q[11, 10]]: K0 (no-error)\n",
            "\n",
            "Trajectories grouped by error count:\n",
            "  0 errors: 1 trajectories, 888,297 shots (88.8%)\n",
            "  1 errors: 126 trajectories, 103,948 shots (10.4%)\n",
            "  2 errors: 19 trajectories, 7,396 shots (0.7%)\n",
            "  3 errors: 1 trajectories, 359 shots (0.0%)\n"
          ]
        }
      ],
      "source": [
        "from collections import Counter\n",
        "\n",
        "exec_shots = 1_000_000\n",
        "max_traj = 256\n",
        "result_with_data = cudaq.ptsbe.sample(\n",
        "    ghz,\n",
        "    n_qubits,\n",
        "    noise_model=ghz_noise,\n",
        "    shots_count=exec_shots,\n",
        "    max_trajectories=max_traj,\n",
        "    return_execution_data=True,\n",
        ")\n",
        "assert result_with_data.has_execution_data()\n",
        "data = result_with_data.ptsbe_execution_data\n",
        "\n",
        "trajs = sorted(data.trajectories, key=lambda t: t.probability, reverse=True)\n",
        "print(f\"Trajectories: {len(trajs)}, Total shots: {exec_shots:,}\\n\")\n",
        "\n",
        "# `t.probability` is the per-shot probability mass in the full trajectory distribution.\n",
        "print(\"Top 5 trajectories (highest probability):\")\n",
        "cumulative = 0\n",
        "for rank, t in enumerate(trajs[:5], 1):\n",
        "    cumulative += t.num_shots\n",
        "    n_errors = sum(1 for s in t.kraus_selections if s.is_error)\n",
        "    p_theory = t.probability\n",
        "    p_empirical = t.num_shots / exec_shots\n",
        "    print(f\"  #{rank}: p_theory={p_theory:.6f}, p_empirical={p_empirical:.6f}, \"\n",
        "          f\"shots={t.num_shots:,}, errors={n_errors}, \"\n",
        "          f\"cumulative shots={100 * cumulative / exec_shots:.1f}%\")\n",
        "\n",
        "# Lowest-probability trajectory\n",
        "lowest = trajs[-1]\n",
        "n_errors_low = sum(1 for s in lowest.kraus_selections if s.is_error)\n",
        "print(f\"\\nLowest-probability trajectory:\")\n",
        "print(f\"  prob={lowest.probability:.2e}, shots={lowest.num_shots}, errors={n_errors_low}\")\n",
        "\n",
        "noise_instructions = {i: inst for i, inst in enumerate(data.instructions)\n",
        "                      if inst.type == cudaq.ptsbe.TraceInstructionType.Noise}\n",
        "\n",
        "def fmt_selection(sel):\n",
        "    # `channel.params` carries the noise rates passed to the channel, and\n",
        "    # `channel.channel` is the underlying `cudaq.KrausChannel` with `.noise_type`,\n",
        "    # `.parameters`, and `.get_ops()` for the Kraus matrices.\n",
        "    channel = noise_instructions[sel.circuit_location]\n",
        "    label = \"error\" if sel.is_error else \"no-error\"\n",
        "    params = \", \".join(f\"{p:g}\" for p in channel.params)\n",
        "    n_ops = len(channel.channel.get_ops())\n",
        "    return (f\"    site {sel.circuit_location} [{channel.name}(p={params}) on q{channel.targets}]: \"\n",
        "            f\"K{sel.kraus_operator_index} / {n_ops} ops ({label})\")\n",
        "\n",
        "highest = trajs[0]\n",
        "print(f\"\\nHighest-probability trajectory (id={highest.trajectory_id}):\")\n",
        "print(f\"  prob={highest.probability:.6f}, shots={highest.num_shots:,}\")\n",
        "print(f\"  Kraus selections:\")\n",
        "for sel in highest.kraus_selections:\n",
        "    print(fmt_selection(sel))\n",
        "\n",
        "print(f\"\\nLowest-probability trajectory (id={lowest.trajectory_id}):\")\n",
        "print(f\"  prob={lowest.probability:.2e}, shots={lowest.num_shots}\")\n",
        "print(f\"  Kraus selections:\")\n",
        "for sel in lowest.kraus_selections:\n",
        "    print(fmt_selection(sel))\n",
        "\n",
        "# Error count histogram\n",
        "error_counts = Counter(\n",
        "    sum(1 for s in t.kraus_selections if s.is_error) for t in trajs\n",
        ")\n",
        "print(\"\\nTrajectories grouped by error count:\")\n",
        "for n_err in sorted(error_counts):\n",
        "    n_traj = error_counts[n_err]\n",
        "    total = sum(t.num_shots for t in trajs\n",
        "                if sum(1 for s in t.kraus_selections if s.is_error) == n_err)\n",
        "    print(f\"  {n_err} errors: {n_traj} trajectories, \"\n",
        "          f\"{total:,} shots ({100 * total / exec_shots:.1f}%)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Performance of PTSBE vs standard noisy sampling\n",
        "\n",
        "The table below times compares the performance of standard trajectory sampling with PTSBE on the 12-qubit GHZ circuit at increasing shot counts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "     shots    standard (s)     PTSBE (s)   speedup\n",
            "----------------------------------------------------\n",
            "     1,000           0.043         0.024      1.8x\n",
            "    10,000           0.119         0.024      4.9x\n",
            "   100,000           0.739         0.029     25.8x\n",
            " 1,000,000           5.800         0.051    113.1x\n"
          ]
        }
      ],
      "source": [
        "import time\n",
        "\n",
        "shot_counts = [1_000, 10_000, 100_000, 1_000_000]\n",
        "max_traj = 64\n",
        "\n",
        "print(f\"{'shots':>10}  {'standard (s)':>14}  {'PTSBE (s)':>12}  {'speedup':>8}\")\n",
        "print(\"-\" * 52)\n",
        "\n",
        "for shots in shot_counts:\n",
        "    # Standard noisy sampling (one trajectory per shot)\n",
        "    t0 = time.perf_counter()\n",
        "    cudaq.sample(ghz, n_qubits, noise_model=ghz_noise, shots_count=shots)\n",
        "    t_standard = time.perf_counter() - t0\n",
        "\n",
        "    # PTSBE (batched trajectories)\n",
        "    t0 = time.perf_counter()\n",
        "    cudaq.ptsbe.sample(\n",
        "        ghz, n_qubits,\n",
        "        noise_model=ghz_noise,\n",
        "        shots_count=shots,\n",
        "        max_trajectories=max_traj,\n",
        "    )\n",
        "    t_ptsbe = time.perf_counter() - t0\n",
        "\n",
        "    speedup = t_standard / t_ptsbe if t_ptsbe > 0 else float('inf')\n",
        "    print(f\"{shots:>10,}  {t_standard:>14.3f}  {t_ptsbe:>12.3f}  {speedup:>7.1f}x\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv (3.12.3)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}
