{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nRandomized measurements can estimate Rényi-2 entropy using random local Clifford rotations and repeated shots.\\n\\nExercises:\\n- Replace the analytic probabilities with sampled bitstrings (multinomial) to mimic finite shots.\\n- Compare a product state vs. Bell state.\\n'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Randomized measurements can estimate Rényi-2 entropy using random local Clifford rotations and repeated shots.\n",
    "\n",
    "Exercises:\n",
    "- Replace the analytic probabilities with sampled bitstrings (multinomial) to mimic finite shots.\n",
    "- Compare a product state vs. Bell state.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True ⟨S2⟩ (toy): 1.0000000000000007\n",
      "Estimated ⟨S2⟩ (toy): 1.0\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from qiskit import QuantumCircuit\n",
    "from qiskit.quantum_info import Statevector, DensityMatrix, partial_trace\n",
    "rng = np.random.default_rng(7)\n",
    "\n",
    "def random_unitary_for_1q():\n",
    "    # A random unitary on 1 qubit via Euler angles\n",
    "    a,b,c = rng.random(3)*2*np.pi\n",
    "    qc = QuantumCircuit(1)\n",
    "    qc.rx(a,0); qc.ry(b,0); qc.rz(c,0)\n",
    "    return qc\n",
    "\n",
    "def renyi2_from_probs(p):\n",
    "    return -np.log2(np.sum(p**2))\n",
    "\n",
    "prep = QuantumCircuit(2)\n",
    "prep.h(0); prep.cx(0,1)\n",
    "\n",
    "M = 200\n",
    "true_entropy = []\n",
    "vals = []\n",
    "for _ in range(M):\n",
    "    #Apply random unitaries\n",
    "    Ua, Ub = random_unitary_for_1q(), random_unitary_for_1q()\n",
    "    circ = prep.compose(Ua, qubits=[0]).compose(Ub, qubits=[1])\n",
    "\n",
    "    #Compute true entropy\n",
    "    psi = Statevector.from_instruction(circ)\n",
    "    rhoA = partial_trace(DensityMatrix(psi), [1]).data # Trace out qubit 1\n",
    "    evals = np.linalg.eigvalsh(rhoA)\n",
    "    true_entropy.append(renyi2_from_probs(evals))\n",
    "\n",
    "    #Compute toy model probabilities\n",
    "    proba = np.abs(psi.data)**2\n",
    "    proba_A = np.array([np.sum(proba[0::2]), np.sum(proba[1::2])])  # Marginalize over qubit 1\n",
    "    proba_A /= np.sum(proba_A)  # Normalize in case of bad numerical precision\n",
    "    X = 2*(proba_A[0]**2  - proba_A[0]*proba_A[1] + proba_A[1]**2) # Formula from DOI: 10.1126/science.aau4963 (Brydges–Elben–Roos–Zoller)\n",
    "    vals.append(X)\n",
    "\n",
    "print(\"True ⟨S2⟩ (toy):\", np.mean(true_entropy))\n",
    "print(\"Estimated ⟨S2⟩ (toy):\", -np.log2(np.mean(vals)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
