{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n### Overview\\nSimulate a simple entanglement purification step (BBPSSW-like) on isotropic states and track fidelity gains.\\n\\nExercises:\\n- Chain multiple rounds and observe convergence threshold behavior.\\n'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "### Overview\n",
    "Simulate a simple entanglement purification step (BBPSSW-like) on isotropic states and track fidelity gains.\n",
    "\n",
    "Exercises:\n",
    "- Chain multiple rounds and observe convergence threshold behavior.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F_in=0.50  F_out≈0.500  p_succ≈0.556\n",
      "F_in=0.60  F_out≈0.620  p_succ≈0.609\n",
      "F_in=0.70  F_out≈0.735  p_succ≈0.680\n",
      "F_in=0.80  F_out≈0.838  p_succ≈0.769\n",
      "F_in=0.90  F_out≈0.926  p_succ≈0.876\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "def bbpssw_round(F):\n",
    "    # Toy mapping for isotropic states under BBPSSW (not exact; illustrative)\n",
    "    # Successful output fidelity (heuristic):\n",
    "    F_out = (F**2 + ((1-F)/3)**2) / (F**2 + 2*F*(1-F)/3 + 5*((1-F)/3)**2)\n",
    "    p_succ = F**2 + 2*F*(1-F)/3 + 5*((1-F)/3)**2\n",
    "    return F_out, p_succ\n",
    "for F in [0.5,0.6,0.7,0.8,0.9]:\n",
    "    Fo, ps = bbpssw_round(F)\n",
    "    print(f\"F_in={F:.2f}  F_out≈{Fo:.3f}  p_succ≈{ps:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bell States Density Matrices:\n",
      "Phi+:\n",
      " DensityMatrix([[0.5+0.j, 0. +0.j, 0. +0.j, 0.5+0.j],\n",
      "               [0. +0.j, 0. +0.j, 0. +0.j, 0. +0.j],\n",
      "               [0. +0.j, 0. +0.j, 0. +0.j, 0. +0.j],\n",
      "               [0.5+0.j, 0. +0.j, 0. +0.j, 0.5+0.j]],\n",
      "              dims=(2, 2))\n",
      "Phi-:\n",
      " DensityMatrix([[ 0.5+0.j,  0. +0.j,  0. +0.j, -0.5+0.j],\n",
      "               [ 0. +0.j,  0. +0.j,  0. +0.j,  0. +0.j],\n",
      "               [ 0. +0.j,  0. +0.j,  0. +0.j,  0. +0.j],\n",
      "               [-0.5+0.j,  0. +0.j,  0. +0.j,  0.5+0.j]],\n",
      "              dims=(2, 2))\n",
      "Psi+:\n",
      " DensityMatrix([[0. +0.j, 0. +0.j, 0. +0.j, 0. +0.j],\n",
      "               [0. +0.j, 0.5+0.j, 0.5+0.j, 0. +0.j],\n",
      "               [0. +0.j, 0.5+0.j, 0.5+0.j, 0. +0.j],\n",
      "               [0. +0.j, 0. +0.j, 0. +0.j, 0. +0.j]],\n",
      "              dims=(2, 2))\n",
      "Psi-:\n",
      " DensityMatrix([[ 0. +0.j,  0. +0.j,  0. +0.j,  0. +0.j],\n",
      "               [ 0. +0.j,  0.5+0.j, -0.5+0.j,  0. +0.j],\n",
      "               [ 0. +0.j, -0.5+0.j,  0.5+0.j,  0. +0.j],\n",
      "               [ 0. +0.j,  0. +0.j,  0. +0.j,  0. +0.j]],\n",
      "              dims=(2, 2))\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from qiskit import QuantumCircuit\n",
    "from qiskit.quantum_info import Statevector, DensityMatrix, partial_trace\n",
    "\n",
    "qc_phip = QuantumCircuit(2)\n",
    "qc_phip.h(0)\n",
    "qc_phip.cx(0, 1)\n",
    "rho_phip = DensityMatrix.from_instruction(qc_phip)\n",
    "\n",
    "qc_phim = QuantumCircuit(2)\n",
    "qc_phim.x(0)\n",
    "qc_phim.h(0)\n",
    "qc_phim.cx(0, 1)\n",
    "rho_phim = DensityMatrix.from_instruction(qc_phim)\n",
    "\n",
    "\n",
    "qc_psip = QuantumCircuit(2)\n",
    "qc_psip.x(1)\n",
    "qc_psip.h(0)\n",
    "qc_psip.cx(0, 1)\n",
    "rho_psip = DensityMatrix.from_instruction(qc_psip)\n",
    "\n",
    "\n",
    "qc_psim = QuantumCircuit(2)\n",
    "qc_psim.x(0)\n",
    "qc_psim.x(1)\n",
    "qc_psim.h(0)\n",
    "qc_psim.cx(0, 1)\n",
    "rho_psim = DensityMatrix.from_instruction(qc_psim)\n",
    "\n",
    "print(\"Bell States Density Matrices:\")\n",
    "print(\"Phi+:\\n\", rho_phip)\n",
    "print(\"Phi-:\\n\", rho_phim)\n",
    "print(\"Psi+:\\n\", rho_psip)\n",
    "print(\"Psi-:\\n\", rho_psim)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F_in=0.50  F_out=0.900\n",
      "F_in=0.60  F_out=0.953\n",
      "F_in=0.70  F_out=0.980\n",
      "F_in=0.80  F_out=0.993\n",
      "F_in=0.90  F_out=0.999\n"
     ]
    }
   ],
   "source": [
    "rho_id   = DensityMatrix(np.eye(4) / 4)\n",
    "\n",
    "def isotropic_state(F):\n",
    "\t\"\"\"Create an isotropic state with fidelity F to |Φ+>.\"\"\"\n",
    "\treturn F * rho_phip + (1 - F) / 3 * (rho_phim + rho_psip + rho_psim)\n",
    "\n",
    "def bbpssw_circuit():\n",
    "\t\"\"\"Construct the BBPSSW purification circuit.\"\"\"\n",
    "\tcirc = QuantumCircuit(4, 2)\n",
    "\t# CNOTs\n",
    "\tcirc.cx(0, 2)\n",
    "\tcirc.cx(1, 3)\n",
    "\treturn circ\n",
    "\n",
    "def simulate_bbpssw(F):\n",
    "\t\"\"\"Simulate one round of BBPSSW purification on isotropic states.\"\"\"\n",
    "\t# Prepare initial state\n",
    "\trho1 = isotropic_state(F)\n",
    "\trho2 = isotropic_state(F)\n",
    "\trho_in = rho1.tensor(rho2)\n",
    "\tproj = rho_id.tensor(rho_phip+rho_phim)\n",
    "\n",
    "\t# Create BBPSSW circuit\n",
    "\tcirc = bbpssw_circuit()\n",
    "\n",
    "\t# Apply circuit to the state\n",
    "\trho_out = DensityMatrix(circ).evolve(rho_in)\n",
    "\n",
    "\t# Trace out measured qubits (2 and 3) over only successful outcomes and renormalize\n",
    "\trho_reduced = partial_trace(DensityMatrix(proj.data @ rho_out.data), [2, 3])\n",
    "\tnorm = np.trace(rho_reduced.data)\n",
    "\trho_reduced = DensityMatrix(rho_reduced.data / norm)\n",
    "\n",
    "\t# Calculate fidelity with |Φ+>\n",
    "\tfid = np.real(np.trace(rho_reduced.data @ rho_phip.data))\n",
    "\n",
    "\treturn fid, rho_reduced\n",
    "\n",
    "for F in [0.5, 0.6, 0.7, 0.8, 0.9]:\n",
    "\tFo, rho_final = simulate_bbpssw(F)\n",
    "\tprint(f\"F_in={F:.2f}  F_out={Fo:.3f}\")"
   ]
  },
  {
   "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
}
