{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Compute the CHSH S-parameter for a Bell state under various noise models.\n",
    "\n",
    "Exercises:\n",
    "- Replace the analytic scaling with simulated Kraus noise channels.\n",
    "- Try amplitude damping and compare.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "S = 2.8284271247461903  (Tsirelson bound!)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from qiskit import QuantumCircuit\n",
    "from qiskit.quantum_info import SparsePauliOp\n",
    "from qiskit_aer import AerSimulator\n",
    "from qiskit_aer.primitives import EstimatorV2 as AerEstimator  \n",
    "from qiskit.transpiler import generate_preset_pass_manager   \n",
    "\n",
    "backend = AerSimulator(method=\"statevector\")\n",
    "pm = generate_preset_pass_manager(backend=backend, optimization_level=1)\n",
    "est = AerEstimator.from_backend(backend) \n",
    "\n",
    "\n",
    "def E(theta_a, theta_b):\n",
    "    #correlator <psi| ZZ |psi> (= p_00 + p_11 - p_10 - p_01) with rotations by where |psi> = U(theta_a) \\otimes U(theta_b)|Φ+>, that would be \"the device preparation state\"\n",
    "    qc = QuantumCircuit(2)\n",
    "    qc.h(0); qc.cx(0,1) #Bell state\n",
    "    qc.ry(theta_a,0); qc.ry(theta_b,1) # modelize crudely the internal state of the device\n",
    "    isa_qc = pm.run(qc)\n",
    "    ob = SparsePauliOp(\"ZZ\")\n",
    "    isa_ZZ = ob.apply_layout(isa_qc.layout)\n",
    "    pub = [(isa_qc,isa_ZZ)]\n",
    "    return AerEstimator().run(pub).result()[0].data.evs\n",
    "\n",
    "a0, a1 = 0, np.pi/2\n",
    "b0, b1 = np.pi/4, -np.pi/4\n",
    "S = E(a0,b0)+E(a0,b1)+E(a1,b0)-E(a1,b1)\n",
    "if S >= 2:\n",
    "    if abs(S - 2*np.sqrt(2)) < 1e-3:\n",
    "        print(\"S =\", S, \" (Tsirelson bound!)\")\n",
    "    elif S < 2*np.sqrt(2):\n",
    "        print(\"S =\", S, \" (CHSH violation!)\")\n",
    "    else:\n",
    "        print(\"S =\", S, \" (error!)\")\n",
    "elif S < 2:\n",
    "    print(\"S =\", S, \" (no CHSH violation)\")\n",
    "else:\n",
    "    print(\"S =\", S, \" (error!)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
