polar_cq_decoder_random_codeword.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91b17d07-1deb-4f90-8c02-c8f621b11241",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: jax in /opt/conda/lib/python3.9/site-packages (0.4.9)\n",
      "Requirement already satisfied: numpy>=1.21 in /opt/conda/lib/python3.9/site-packages (from jax) (1.22.4)\n",
      "Requirement already satisfied: scipy>=1.7 in /opt/conda/lib/python3.9/site-packages (from jax) (1.9.0)\n",
      "Requirement already satisfied: ml-dtypes>=0.1.0 in /opt/conda/lib/python3.9/site-packages (from jax) (0.1.0)\n",
      "Requirement already satisfied: opt-einsum in /opt/conda/lib/python3.9/site-packages (from jax) (3.3.0)\n",
      "Requirement already satisfied: jaxlib in /opt/conda/lib/python3.9/site-packages (0.4.9)\n",
      "Requirement already satisfied: ml-dtypes>=0.1.0 in /opt/conda/lib/python3.9/site-packages (from jaxlib) (0.1.0)\n",
      "Requirement already satisfied: scipy>=1.7 in /opt/conda/lib/python3.9/site-packages (from jaxlib) (1.9.0)\n",
      "Requirement already satisfied: numpy>=1.21 in /opt/conda/lib/python3.9/site-packages (from jaxlib) (1.22.4)\n"
     ]
    }
   ],
   "source": [
    "!pip install jax\n",
    "!pip install jaxlib\n",
    "import numpy as np\n",
    "from numpy import linalg as LA\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import grad, jit, vmap\n",
    "from scipy.linalg import sqrtm\n",
    "from numpy.linalg import matrix_rank\n",
    "from jax.numpy import linalg as jla\n",
    "from scipy.linalg import logm, expm\n",
    "\n",
    "from scipy.linalg import block_diag\n",
    "import copy\n",
    "import numba\n",
    "from numba import int64, float64, jit, njit, vectorize\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d2d2b686-fb30-4086-8a85-f2906c786f19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def qubit(x):\n",
    "  if x==0:\n",
    "    return np.array([1,0])\n",
    "  elif x==1:\n",
    "    return np.array([0,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "30504cfc-eb70-48d4-8e02-7fb70fde63e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pauli(x):\n",
    "    if x==0:\n",
    "        return np.array([[1,0],[0,1]])\n",
    "    if x==1:\n",
    "        return np.array([[0,1],[1,0]])\n",
    "    if x==2:\n",
    "        return np.array([[0,-1j],[+1j,0]])\n",
    "    if x==3:\n",
    "        return np.array([[1,0],[0,-1]])\n",
    "    else:\n",
    "        print('invalid')\n",
    "\n",
    "# modified rho based on delta, gamma, and input x rather than theta and p\n",
    "def rhom(delta,g,x):\n",
    "  return pauli(x)@np.array([[delta, g],[g, 1-delta]])@pauli(x)\n",
    "\n",
    "# \"bitnode\" take input states parameterized by (d1, g1) and (d2, g2) and returns the post-measurement state parameters and probabilities at a bitnode\n",
    "def bitnode(d1,d2,g1,g2):\n",
    "  # the state given root value z=0 is rho0 = np.kron(rhom(d1,g1,0),rhom(d2,g2,0))\n",
    "  # the state given root value z=1 is rho1 = np.kron(rhom(d1,g1,1),rhom(d2,g2,1))\n",
    "  \n",
    "  # we find the paired measurements by taking the eigenvectors of the difference matrix rhofull\n",
    "  rhofull = np.kron(rhom(d1,g1,0),rhom(d2,g2,0))-np.kron(rhom(d1,g1,1),rhom(d2,g2,1))\n",
    "  \n",
    "  evals, evecs = np.linalg.eigh(rhofull)\n",
    "  # fix eigenvector v0 \n",
    "  v0 = evecs[:,0]\n",
    "  \n",
    "  # symmetry operator Un = np.kron(pauli(1),pauli(1))\n",
    "  Un = np.kron(pauli(1),pauli(1))\n",
    "  # check if the second eigenvector evecs[:,1] is orthogonal to Un@evecs[:,1]\n",
    "  x = evecs[:,1]@(Un@evecs[:,1])\n",
    " \n",
    "  #print(\"evals\")\n",
    "  #print(evals)\n",
    "  if np.abs(x)<10e-10:\n",
    "    v1 = evecs[:,1]\n",
    "  # if not orthogonal, combine evecs[:,1], evecs[:,2] to create v1 s.t. v1@(Un@v1)= 0\n",
    "  if np.abs(x)>=10e-10:\n",
    "    vec1, vec2 = evecs[:,1], evecs[:,2]\n",
    "    b11, b12, b21, b22 = np.inner(vec1, (Un@vec1).conj()), np.inner(vec2, (Un@vec1).conj()), np.inner(vec1, (Un@vec2).conj()), np.inner(vec2, (Un@vec2).conj())\n",
    "    \n",
    "    alpha = (-b12-b21-np.sqrt((b12+b21)**2-4*b11*b22))/(2*b22)\n",
    "    v1 = vec1+alpha*vec2\n",
    "    v1 = v1/np.sqrt(v1@v1)\n",
    "    ##\n",
    "  #print(\"v0\")\n",
    "  #print(v0)\n",
    "  #print(\"v1\")\n",
    "  #print(v1)\n",
    "  ##\n",
    "  # the paired measurement is then given by {|v0><v0| + Un|v0><v0|Un, |v1><v1| + Un|v1><v1|Un}\n",
    "  ## find new state parameters (d1a, g1a) for measurement outcome 0\n",
    "  # find probability p0 of observing measurement  outcome 0\n",
    "  p0 = v0@np.kron(rhom(d1,g1,0),rhom(d2,g2,0))@v0+v0@Un@np.kron(rhom(d1,g1,0),rhom(d2,g2,0))@Un@v0\n",
    "  d1a, g1a = v0@np.kron(rhom(d1,g1,0),rhom(d2,g2,0))@v0/(p0+10e-21), v0@np.kron(rhom(d1,g1,0),rhom(d2,g2,0))@(Un@v0)/(p0+10e-21)\n",
    "  ## find new state parameters (d1b, g1b) for measurement outcome 1\n",
    "  # find probability p1 of observing measurement  outcome 1  \n",
    "  p1 = v1@np.kron(rhom(d1,g1,0),rhom(d2,g2,0))@v1+v1@Un@np.kron(rhom(d1,g1,0),rhom(d2,g2,0))@Un@v1\n",
    "  d2a, g2a = v1@np.kron(rhom(d1,g1,0),rhom(d2,g2,0))@v1/(p1+10e-21), v1@np.kron(rhom(d1,g1,0),rhom(d2,g2,0))@(Un@v1)/(p1+10e-21)\n",
    "  #print(\"pvec\")\n",
    "  #print([p0, p1])\n",
    "  return np.array([[1-d1a, g1a], [1-d2a, g2a]]), np.array([p0, p1])\n",
    "\n",
    "# \"checknode\" take input states parameterized by (d1, g1) and (d2, g2) and returns the post-measurement state parameters and probabilities at a checknode\n",
    "def checknode(d1,d2,g1,g2):\n",
    "  # rho0, rho1 correspond to the states at a check node when z=0 (z=1) respectively\n",
    "  rho0, rho1 = 1/2*(np.kron(rhom(d1,g1,0),rhom(d2,g2,0)) + np.kron(rhom(d1,g1,1),rhom(d2,g2,1))), 1/2*(np.kron(rhom(d1,g1,0),rhom(d2,g2,1)) + np.kron(rhom(d1,g1,1),rhom(d2,g2,0)))\n",
    "  # for check node combining, the optimal choice of eigenvectors appears to always be generated by v0 and v1\n",
    "  v0 = 1/np.sqrt(2)*np.array([1,0,0,1])\n",
    "  v1 = 1/np.sqrt(2)*np.array([-1, 0, 0, 1])\n",
    "  # symmetry operator for a check node\n",
    "  Un = np.kron(pauli(1),pauli(0))\n",
    "  ## find new state parameters (d1a, g1a) for measurement outcome 0\n",
    "  # find probability p0 of observing measurement  outcome 0\n",
    "  p0 = v0@rho0@v0+v0@Un@rho0@Un@v0\n",
    "  d1a, g1a = v0@rho0@v0/(p0+10e-21), v0@rho0@(Un@v0)/(p0+10e-21)\n",
    "  ## find new gamma, delta for second outcome\n",
    "  p1 = v1@rho0@v1+v1@Un@rho0@Un@v1\n",
    "  d2a, g2a = v1@rho0@v1/(p1+10e-21), v1@rho0@(Un@v1)/(p1+10e-21)\n",
    "  # return new gamma, delta pairs as well as respective probabilities tra and tr2a\n",
    "  # d1a=min([d1a,1-d1a])\n",
    "  # d2a=min([d2a,1-d2a])\n",
    "  return np.array([[d1a, g1a], [d2a, g2a]]), np.array([p0, p1])\n",
    "\n",
    "H = 1/np.sqrt(2)*np.array([[1,-1],[1,1]])\n",
    "sgz = np.array([[1,0],[0,-1]])\n",
    "sgx = np.array([[0,1],[1,0]])\n",
    " \n",
    "# pure state channel outputs\n",
    "def rho(th):\n",
    "  return H@np.array([[np.cos(th/2)**2, np.cos(th/2)*np.sin(th/2)],\n",
    "                   [np.cos(th/2)*np.sin(th/2), np.sin(th/2)**2]])@np.transpose(H)\n",
    "\n",
    "# output of the channel given input z, parameter th, and flip parameter p\n",
    "def W(z,th,p):\n",
    "  if z == 0:\n",
    "    return (1-p)*rho(th)+p*np.array([[1/2, 0], [0, 1/2]])\n",
    "  if z == 1:\n",
    "    return (1-p)*sgx@rho(th)@sgx+p*np.array([[1/2, 0], [0, 1/2]])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4af22e42-ea50-487a-a838-084336fe4aac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def helstrom(density_mat1,unitary):\n",
    "  r1=density_mat1\n",
    "  u=unitary\n",
    "  r2=u@r1@np.conjugate(np.transpose(u))\n",
    "  l,v=LA.eig(r1-r2)\n",
    "  return l,v\n",
    "\n",
    "\n",
    "def helstrom_success(rho1,unitary):\n",
    "  l,vec=helstrom(rho1,unitary)\n",
    "  u=unitary\n",
    "  rho2=u@rho1@np.conjugate(np.transpose(u))\n",
    "  v_pos_eig = np.array(vec[:,l>0])\n",
    "  v_pos_eigh= np.conjugate(np.transpose(v_pos_eig))\n",
    "  p1=np.trace(v_pos_eigh @ rho1 @ v_pos_eig)\n",
    "  p2=1-np.trace(v_pos_eigh@ rho2 @ v_pos_eig)\n",
    "  return 0.5*(p1+p2)\n",
    "\n",
    "def helstrom_success_vec(X):\n",
    "  o=[]\n",
    "  l=len(X[0])\n",
    "  for i in range(l):\n",
    "    d,g=X[0][i],X[1][i]\n",
    "    r=rhom(d,g,0)\n",
    "    h=helstrom_success(r,pauli(1))\n",
    "    o.append(h)\n",
    "  return o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d89c53eb-12bb-4e6a-922b-5a56f84a75a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def rhom_jit(delta:float64,g:float64) -> (float64[:,:]):\n",
    "    return np.array([[delta, g],[g, 1-delta]],dtype=float64)\n",
    "\n",
    "# \"bitnode\" take input states parameterized by (d1, g1) and (d2, g2) and returns the post-measurement state parameters and probabilities at a bitnode\n",
    "@njit\n",
    "def bitnode_jit(d1:float64,d2:float64,g1:float64,g2:float64) ->(float64[:,:],float64[:]):\n",
    "  # the state given root value z=0 is rho0 = np.kron(rhom(d1,g1,0),rhom(d2,g2,0))\n",
    "  # the state given root value z=1 is rho1 = np.kron(rhom(d1,g1,1),rhom(d2,g2,1))\n",
    "  \n",
    "  # we find the paired measurements by taking the eigenvectors of the difference matrix rhofull\n",
    "  x=np.array([[0,1],[1,0]],dtype=float64)\n",
    "  \n",
    "  rhofull = np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))-np.kron(x@rhom_jit(d1,g1)@x,x@rhom_jit(d2,g2)@x)\n",
    " \n",
    "  evals, evecs = np.linalg.eigh(rhofull)\n",
    "  # fix eigenvector v0 \n",
    "  v0 = evecs[:,0]\n",
    "  \n",
    "  # symmetry operator Un = np.kron(pauli(1),pauli(1))\n",
    "  Un = np.kron(x,x)\n",
    "  # check if the second eigenvector evecs[:,1] is orthogonal to Un@evecs[:,1]\n",
    "  x1 = evecs[:,1]@(Un@evecs[:,1])\n",
    " \n",
    "  #print(\"evals\")\n",
    "  #print(evals)\n",
    "  if np.abs(x1)<10e-10:\n",
    "    v1 = evecs[:,1]\n",
    "  # if not orthogonal, combine evecs[:,1], evecs[:,2] to create v1 s.t. v1@(Un@v1)= 0\n",
    "  if np.abs(x1)>=10e-10:\n",
    "    vec1, vec2 = evecs[:,1], evecs[:,2]\n",
    "    b11, b12, b21, b22 = np.dot(vec1, (Un@vec1).conj()), np.dot(vec2, (Un@vec1).conj()), np.dot(vec1, (Un@vec2).conj()), np.dot(vec2, (Un@vec2).conj())\n",
    "    \n",
    "    alpha = (-b12-b21-np.sqrt((b12+b21)**2-4*b11*b22))/(2*b22)\n",
    "    v1 = vec1+alpha*vec2\n",
    "    v1 = v1/np.sqrt(v1@v1)\n",
    "    ##\n",
    "  #print(\"v0\")\n",
    "  #print(v0)\n",
    "  #print(\"v1\")\n",
    "  #print(v1)\n",
    "  ##\n",
    "  # the paired measurement is then given by {|v0><v0| + Un|v0><v0|Un, |v1><v1| + Un|v1><v1|Un}\n",
    "  ## find new state parameters (d1a, g1a) for measurement outcome 0\n",
    "  # find probability p0 of observing measurement  outcome 0\n",
    "  p0 = v0@np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))@v0+v0@Un@np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))@Un@v0\n",
    "  d1a, g1a = v0@np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))@v0/(p0+10e-21), v0@np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))@(Un@v0)/(p0+10e-21)\n",
    "  ## find new state parameters (d1b, g1b) for measurement outcome 1\n",
    "  # find probability p1 of observing measurement  outcome 1  \n",
    "  p1 = v1@np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))@v1+v1@Un@np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))@Un@v1\n",
    "  d2a, g2a = v1@np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))@v1/(p1+10e-21), v1@np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2))@(Un@v1)/(p1+10e-21)\n",
    "  #print(\"pvec\")\n",
    "  #print([p0, p1])\n",
    "  d1a=min([d1a,1-d1a])\n",
    "  d2a=min([d2a,1-d2a])\n",
    "  return np.array([[d1a, g1a], [d2a, g2a]]), np.array([p0, p1])\n",
    "\n",
    "# \"checknode\" take input states parameterized by (d1, g1) and (d2, g2) and returns the post-measurement state parameters and probabilities at a checknode\n",
    "@njit\n",
    "def checknode_jit(d1:float64,d2:float64,g1:float64,g2:float64) ->(float64[:,:],float64[:]):\n",
    "  x=np.array([[0,1],[1,0]],dtype=float64)\n",
    "  I=np.array([[1,0],[0,1]],dtype=float64)\n",
    "  # rho0, rho1 correspond to the states at a check node when z=0 (z=1) respectively\n",
    "  rho0, rho1 = 1/2*(np.kron(rhom_jit(d1,g1),rhom_jit(d2,g2)) + np.kron(x@rhom_jit(d1,g1)@x,x@rhom_jit(d2,g2)@x)), 1/2*(np.kron(rhom_jit(d1,g1),x@rhom_jit(d2,g2)@x) + np.kron(x@rhom_jit(d1,g1)@x,rhom_jit(d2,g2)))\n",
    "  # for check node combining, the optimal choice of eigenvectors appears to always be generated by v0 and v1\n",
    "  v0 = 1/np.sqrt(2)*np.array([1,0,0,1])\n",
    "  v1 = 1/np.sqrt(2)*np.array([-1, 0, 0, 1])\n",
    "  # symmetry operator for a check node\n",
    "  Un = np.kron(x,I)\n",
    "  ## find new state parameters (d1a, g1a) for measurement outcome 0\n",
    "  # find probability p0 of observing measurement  outcome 0\n",
    "  p0 = v0@rho0@v0+v0@Un@rho0@Un@v0\n",
    "  d1a, g1a = v0@rho0@v0/(p0+10e-21), v0@rho0@(Un@v0)/(p0+10e-21)\n",
    "  ## find new gamma, delta for second outcome\n",
    "  p1 = v1@rho0@v1+v1@Un@rho0@Un@v1\n",
    "  d2a, g2a = v1@rho0@v1/(p1+10e-21), v1@rho0@(Un@v1)/(p1+10e-21)\n",
    "  # return new gamma, delta pairs as well as respective probabilities tra and tr2a\n",
    "  d1a=min([d1a,1-d1a])\n",
    "  d2a=min([d2a,1-d2a])\n",
    "  return np.array([[d1a, g1a], [d2a, g2a]]), np.array([p0, p1])\n",
    "  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3b7e60e9-7fbd-4440-9b49-cd1b1cb7e493",
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def bitnode_jit2(d1:float64,d2:float64,g1:float64,g2:float64,pr=None)->(float64[:,:]):\n",
    "  rho,pb=bitnode_jit(d1,d2,g1,g2)\n",
    "  if pr==None:\n",
    "    s=pb[0]\n",
    "  else:\n",
    "    s=pr\n",
    "  o=int(np.random.random()>s)\n",
    "  #choice([0,1],p=[s,1-s])\n",
    "  if o==0:\n",
    "    return rho[0]\n",
    "  else:\n",
    "    return rho[1]\n",
    "\n",
    "@njit\n",
    "def checknode_jit1(d1:float64,d2:float64,g1:float64,g2:float64,pr=None)->(float64[:,:]):\n",
    "  rho,pc=checknode_jit(d1,d2,g1,g2)\n",
    "  if pr==None:\n",
    "    s=pc[0]\n",
    "  else:\n",
    "    s=pr\n",
    "  o=int(np.random.random()>s)\n",
    "  if o==0:\n",
    "    return rho[0]\n",
    "  else:\n",
    "    return rho[1]\n",
    "\n",
    "@njit\n",
    "def bitnode_vec_jit(d1:float64[:],d2:float64[:],g1:float64[:],g2:float64[:],pr_vec=None,perm=None)->(float64[:],float64[:]):\n",
    "    l=np.shape(d1)[0]\n",
    "    if perm==None:\n",
    "      p=np.random.permutation(l)\n",
    "    else:\n",
    "      p=perm\n",
    "    d=np.zeros(l)\n",
    "    g=np.zeros(l)\n",
    "    if pr_vec==None:\n",
    "      for i in range(l):\n",
    "        d[i],g[i]=bitnode_jit2(d1[i],d2[p[i]],g1[i],g2[p[i]])\n",
    "        \n",
    "    else:\n",
    "      for i in range(l):\n",
    "        d[i],g[i]=bitnode_jit2(d1[i],d2[p[i]],g1[i],g2[p[i]],pr_vec[i])\n",
    "        \n",
    "    return d,g\n",
    "@njit\n",
    "def checknode_vec_jit(d1:float64[:],d2:float64[:],g1:float64[:],g2:float64[:],pr_vec=None,perm=None)->(float64[:],float64[:]):\n",
    "    l=np.shape(d1)[0]\n",
    "    if perm==None:\n",
    "      p=np.random.permutation(l)\n",
    "    else:\n",
    "      p=perm\n",
    "    d=np.zeros(l)\n",
    "    g=np.zeros(l)\n",
    "    if pr_vec==None:\n",
    "      for i in range(l):\n",
    "        d[i],g[i]=checknode_jit1(d1[i],d2[p[i]],g1[i],g2[p[i]])\n",
    "        \n",
    "    else:\n",
    "      for i in range(l):\n",
    "        d[i],g[i]=checknode_jit1(d1[i],d2[p[i]],g1[i],g2[p[i]],pr_vec[i])\n",
    "        \n",
    "    return d,g\n",
    "    \n",
    "    \n",
    "@njit   \n",
    "def bitnode_power_jit(d:float64[:],g:float64[:],k:int64,pr_vec=None,perm=None)->(float64[:],float64[:]):\n",
    "  if k==1:\n",
    "    return d,g\n",
    "  else:\n",
    "    d1,g1=bitnode_vec_jit(d,d,g,g,pr_vec,perm)\n",
    "    if k>2:\n",
    "      for i in range(k-2):\n",
    "        d1,g1=bitnode_vec_jit(d,d1,g,g1,pr_vec,perm)\n",
    "    return d1,g1\n",
    "\n",
    "  \n",
    "@njit\n",
    "def checknode_power_jit(d:float64[:],g:float64[:],k:int64,pr_vec=None,perm=None)->(float64[:],float64[:]):\n",
    "  if k==1:\n",
    "    return d,g\n",
    "  else:\n",
    "    d1,g1=checknode_vec_jit(d,d,g,g,pr_vec,perm)\n",
    "    if k>2:\n",
    "      for i in range(k-2):\n",
    "        d1,g1=checknode_vec_jit(d,d1,g,g1,pr_vec,perm)\n",
    "    return d1,g1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b4b1362e-c9ba-46b6-b3be-7a50245456ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def polar(n:int64,d:float64[:],g=float64[:])->(float64[:,:,:]):\n",
    " # s=0\n",
    "  #for i in range(n+1):\n",
    "  #  s=s+2**i\n",
    "  \n",
    "  old_stage=np.zeros(shape=(1,2,len(d)),dtype=float64)\n",
    "  old_stage[0][0],old_stage[0][1]=d,g\n",
    " # o=np.zeros(shape=(s,2,len(x)))\n",
    "  for i in range(n):\n",
    "    new_stage=np.zeros(shape=(2**(i+1),2,len(d)))\n",
    "    for j in range(2**(i)):\n",
    "      new_stage[2*j][0],new_stage[2*j][1]=checknode_vec_jit(old_stage[j][0],old_stage[j][0],old_stage[j][1],old_stage[j][1])\n",
    "      new_stage[2*j+1][0],new_stage[2*j+1][1]=bitnode_vec_jit(old_stage[j][0],old_stage[j][0],old_stage[j][1],old_stage[j][1])\n",
    "    old_stage=new_stage\n",
    "  return old_stage\n",
    "\n",
    "\n",
    "def polar_design(biterrd,d):\n",
    "    '''\n",
    "    Design polar code from additive channel noise scores\n",
    "\n",
    "          Arguments:\n",
    "                  biterrd (float[:]): Numpy array of channel noise scores (e.g., error rates)\n",
    "                  d (float): Sum constraint of total noise for unfrozen channels\n",
    "\n",
    "          Returns:\n",
    "                  f (float[:]): Numpy array indicating frozen bits by 0 and info bits by 0.5\n",
    "    '''\n",
    "    # Sort into increasing order and compute cumulative sum\n",
    "    order = np.argsort(biterrd)\n",
    "    SE = biterrd[order]\n",
    "    CSE = np.cumsum(SE)\n",
    "\n",
    "    # Find best frozen bits\n",
    "    k = np.sum(CSE<d)\n",
    "    f = np.zeros(len(biterrd))\n",
    "    f[order[:k]] = 0.5\n",
    "    return f\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "976e9a4a-9c69-4da4-990f-abca81664750",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bitnode_gen_unitary(d1,d2,g1,g2):\n",
    "  rhofull = np.kron(rhom(d1,g1,0),rhom(d2,g2,0))-np.kron(rhom(d1,g1,1),rhom(d2,g2,1))\n",
    "  \n",
    "  evals, evecs = LA.eigh(rhofull)\n",
    "  # fix eigenvector v0 \n",
    "  v0 = evecs[:,0]\n",
    "  \n",
    "  # symmetry operator Un = np.kron(pauli(1),pauli(1))\n",
    "  Un = np.kron(pauli(1),pauli(1))\n",
    "  # check if the second eigenvector evecs[:,1] is orthogonal to Un@evecs[:,1]\n",
    "  x = evecs[:,1]@(Un@evecs[:,1])\n",
    " \n",
    "  #print(\"evals\")\n",
    "  #print(evals)\n",
    "  if np.abs(x)<10e-10:\n",
    "    v1 = evecs[:,1]\n",
    "  # if not orthogonal, combine evecs[:,1], evecs[:,2] to create v1 s.t. v1@(Un@v1)= 0\n",
    "  if np.abs(x)>=10e-10:\n",
    "    vec1, vec2 = evecs[:,1], evecs[:,2]\n",
    "    b11, b12, b21, b22 = np.inner(vec1, (Un@vec1).conj()), np.inner(vec2, (Un@vec1).conj()), np.inner(vec1, (Un@vec2).conj()), np.inner(vec2, (Un@vec2).conj())\n",
    "    \n",
    "    alpha = (-b12-b21-np.sqrt((b12+b21)**2-4*b11*b22))/(2*b22)\n",
    "    v1 = vec1+alpha*vec2\n",
    "    v1 = v1/np.sqrt(v1@v1)\n",
    "  vb=np.zeros((4,4))\n",
    "  vb[:,0]=Un@v0\n",
    "  vb[:,1]=Un@v1\n",
    "  vb[:,2]=v0\n",
    "  vb[:,3]=v1\n",
    "  return np.transpose(vb)\n",
    "\n",
    "def checknode_gen_unitary(d1,d2,g1,g2):\n",
    "  v0 = 1/np.sqrt(2)*np.array([1,0,0,1])\n",
    "  v1 = 1/np.sqrt(2)*np.array([-1, 0, 0, 1])\n",
    "  ux=np.kron(pauli(1),np.eye(2))\n",
    "  vcf=np.zeros((4,4))\n",
    "  vcf[:,0]=v0\n",
    "  vcf[:,1]=v1\n",
    "  vcf[:,2]=ux@v0\n",
    "  vcf[:,3]=ux@v1\n",
    "  return np.transpose(vcf)\n",
    "\n",
    "CU=checknode_gen_unitary(0.1,0.1,0,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5a498795-0b2e-451c-84e4-a5db13460140",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_permutation(qubit_list,number_qubits,conditional_list1=[],conditional_list2=[]):\n",
    "  n=number_qubits\n",
    "  l=len(qubit_list)\n",
    "  l1=len(conditional_list1)\n",
    "  l2=len(conditional_list2)\n",
    "  p=np.zeros(n,dtype=int)\n",
    "  for j in range(l):\n",
    "    p[qubit_list[j]]=l1+l2++j\n",
    "  for j in range(l1):\n",
    "    p[conditional_list1[j]]=j\n",
    "  for j in range(l2):\n",
    "    p[conditional_list2[j]]=l1+j\n",
    "  c=0\n",
    "  for j in range(n):\n",
    "    if j not in qubit_list and j not in conditional_list1 and j not in conditional_list2:\n",
    "      p[j]=l+l1+l2+c\n",
    "      c=c+1\n",
    "  p1=np.zeros((len(p),len(p)))\n",
    "  for i in range(len(p)):\n",
    "    p1[:,i][p[i]]=1\n",
    "  p2=p1.T\n",
    "  pi=np.zeros(n,dtype=int)\n",
    "  for i in range(n):\n",
    "    pi[i]=int(np.where(p1.T[:,i]==1)[0][0])\n",
    "  return p,pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f56e042a-ab28-48da-a534-8226e016401d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_unitary(qubit_list,unitary,q_state,number_qubits):\n",
    "  n=number_qubits\n",
    "  l=len(qubit_list)\n",
    "  u=unitary\n",
    "  q=np.reshape(q_state,2*np.ones(n,dtype=int))\n",
    "  p,pi=apply_permutation(qubit_list,n)\n",
    "  q1=np.moveaxis(q,np.array(range(n)),p)\n",
    "  q2=np.reshape(q1,(2**l,2**(n-l)))\n",
    "  q_out=u@q2\n",
    "  q_out=np.reshape(q_out,2*np.ones(n,dtype=int))\n",
    "  q_out=np.moveaxis(q_out,np.array(range(n)),pi)\n",
    "  q_out=np.reshape(q_out,np.shape(q_state))\n",
    "  return q_out\n",
    "\n",
    "def apply_unitary_inverse(qubit_list,unitary,q_state,number_qubits):\n",
    "  n=number_qubits\n",
    "  l=len(qubit_list)\n",
    "  u=unitary\n",
    "  q=np.reshape(q_state,2*np.ones(n,dtype=int))\n",
    "  p,pi=apply_permutation(qubit_list,n)\n",
    "  q1=np.moveaxis(q,np.array(range(n)),p)\n",
    "  q2=np.reshape(q1,(2**l,2**(n-l)))\n",
    "  q_out=np.transpose(u)@q2\n",
    "  q_out=np.reshape(q_out,2*np.ones(n,dtype=int))\n",
    "  q_out=np.moveaxis(q_out,np.array(range(n)),pi)\n",
    "  q_out=np.reshape(q_out,np.shape(q_state))\n",
    "  return q_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f0cb8bfa-832d-4a7c-aad1-0a3314a09de1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_conditional_unitary(qubit_list,conditional_list1,conditional_list2,dg_list1,dg_list2,q_state,number_qubits):\n",
    "  n=number_qubits\n",
    "  l=len(qubit_list)\n",
    "  l1=len(conditional_list1)\n",
    "  l2=len(conditional_list2)\n",
    "  q=np.reshape(q_state,2*np.ones(n,dtype=int))\n",
    "  p=np.zeros(n,dtype=int)\n",
    "  p,pi=apply_permutation(qubit_list,n,conditional_list1,conditional_list2)\n",
    "  q1=np.moveaxis(q,np.array(range(n)),p)\n",
    "  q2=np.reshape(q1,(2**l1,2**l2,2**l,2**(n-l1-l2-l)))\n",
    "  q_out=np.zeros((2**l1,2**l2,2**l,2**(n-l1-l2-l)))\n",
    "  for i in range(2**l1):\n",
    "    for j in range(2**l2):\n",
    "      q_out[i,j,:]=bitnode_gen_unitary(dg_list1[i][0],dg_list2[j][0],dg_list1[i][1],dg_list2[j][1])@q2[i,j,:]\n",
    "  q_out=np.reshape(q_out,2*np.ones(n,dtype=int))\n",
    "  q_out=np.moveaxis(q_out,np.array(range(n)),pi)\n",
    "  q_out=np.reshape(q_out,np.shape(q_state))\n",
    "  return q_out\n",
    "def apply_conditional_unitary_inverse(qubit_list,conditional_list1,conditional_list2,dg_list1,dg_list2,q_state,number_qubits):\n",
    "  n=number_qubits\n",
    "  l=len(qubit_list)\n",
    "  l1=len(conditional_list1)\n",
    "  l2=len(conditional_list2)\n",
    "  q=np.reshape(q_state,2*np.ones(n,dtype=int))\n",
    "  p=np.zeros(n,dtype=int)\n",
    "  p,pi=apply_permutation(qubit_list,n,conditional_list1,conditional_list2)\n",
    "  q1=np.moveaxis(q,np.array(range(n)),p)\n",
    "  q2=np.reshape(q1,(2**l1,2**l2,2**l,2**(n-l1-l2-l)))\n",
    "  q_out=np.zeros((2**l1,2**l2,2**l,2**(n-l1-l2-l)))\n",
    "  for i in range(2**l1):\n",
    "    for j in range(2**l2):\n",
    "      q_out[i,j,:]=np.transpose(bitnode_gen_unitary(dg_list1[i][0],dg_list2[j][0],dg_list1[i][1],dg_list2[j][1]))@q2[i,j,:]\n",
    "  q_out=np.reshape(q_out,2*np.ones(n,dtype=int))\n",
    "  q_out=np.moveaxis(q_out,np.array(range(n)),pi)\n",
    "  q_out=np.reshape(q_out,np.shape(q_state))\n",
    "  return q_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8ffa87e1-39a6-4560-b417-8eeb3122f09a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_state(l,v):\n",
    "  o=int(np.random.random()>l[0])\n",
    "  if o==0:\n",
    "    return v[:,0]\n",
    "  else:\n",
    "    return v[:,1]\n",
    "def sample_state_one(l,v):\n",
    "  o=int\n",
    "  pass\n",
    "def sample_joint_state(l,v,number_qubits):\n",
    "  if number_qubits==1:\n",
    "    return sample_state(l,v)\n",
    "  else:\n",
    "    return np.kron(sample_joint_state(l,v,number_qubits-1),sample_state(l,v))\n",
    "\n",
    "\n",
    "def sample_joint_random_state(l,v,v1,number_qubits,bit_sequence):\n",
    "  b=bit_sequence\n",
    "  if b[0]==0:\n",
    "    s=sample_state(l,v)\n",
    "  else:\n",
    "    s=sample_state(l,v1)\n",
    "  for i in range(1,number_qubits):\n",
    "    if b[i]==0:\n",
    "      s=np.kron(s,sample_state(l,v))\n",
    "    else:\n",
    "      s=np.kron(s,sample_state(l,v1))\n",
    "  return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "653cb7b2-4a1a-4c0c-b927-4f5bcd39ad5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cnop(quantum_state,informtion_qubit1,information_qubit2,conditional_list1,conditional_list2,number_qubits):\n",
    "  qubit_list=np.array([informtion_qubit1,information_qubit2])\n",
    "  q_out=apply_unitary(qubit_list,CU,quantum_state,number_qubits)\n",
    "  l1=len(conditional_list1[0])\n",
    "  l2=len(conditional_list2[0])\n",
    "  conditional_list=[[information_qubit2],[]]\n",
    "  c1=[]\n",
    "  c2=[]\n",
    "  for i in range(2**l1):\n",
    "    for j in range(2**l2):\n",
    "      rc=checknode(conditional_list1[1][i][0],conditional_list2[1][j][0],conditional_list1[1][i][1],conditional_list2[1][j][1])[0]\n",
    "      c1.append(rc[0])\n",
    "      c2.append(rc[1])\n",
    "  conditional_list[1]=c1+c2\n",
    "\n",
    "  for i in range(l1):\n",
    "    conditional_list[0].append(conditional_list1[0][i])\n",
    "  for j in range(l2):\n",
    "    conditional_list[0].append(conditional_list2[0][j])\n",
    "  return q_out, conditional_list\n",
    "\n",
    "def cnop_inverse(quantum_state,informtion_qubit1,information_qubit2,conditional_list1,conditional_list2,number_qubits):\n",
    "  qubit_list=np.array([informtion_qubit1,information_qubit2])\n",
    "  q_out=apply_unitary_inverse(qubit_list,CU,quantum_state,number_qubits)\n",
    "  return q_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3e41d1dc-2ce9-4809-bc65-8105ce2ca5e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def vnop(quantum_state,informtion_qubit1,information_qubit2,conditional_list1,conditional_list2,number_qubits):\n",
    "  qubit_list=np.array([informtion_qubit1,information_qubit2])\n",
    "  l1=len(conditional_list1[0])\n",
    "  l2=len(conditional_list2[0])\n",
    "  if l1==0 and l2==0:\n",
    "    q_out=apply_unitary(qubit_list,bitnode_gen_unitary(conditional_list1[1][0][0],conditional_list2[1][0][0],conditional_list1[1][0][1],conditional_list2[1][0][1]),quantum_state,number_qubits)\n",
    "  else:\n",
    "    q_out=apply_conditional_unitary(qubit_list,conditional_list1[0],conditional_list2[0],conditional_list1[1],conditional_list2[1],quantum_state,number_qubits)\n",
    "  conditional_list=[[information_qubit2],[]]\n",
    "  c1=[]\n",
    "  c2=[]\n",
    "  for i in range(2**l1):\n",
    "    for j in range(2**l2):\n",
    "      rc=bitnode(conditional_list1[1][i][0],conditional_list2[1][j][0],conditional_list1[1][i][1],conditional_list2[1][j][1])[0]\n",
    "      c1.append(rc[0])\n",
    "      c2.append(rc[1])\n",
    "  conditional_list[1]=c1+c2\n",
    "\n",
    "  for i in range(l1):\n",
    "    conditional_list[0].append(conditional_list1[0][i])\n",
    "  for j in range(l2):\n",
    "    conditional_list[0].append(conditional_list2[0][j])\n",
    "  return q_out, conditional_list\n",
    "\n",
    "def vnop_inverse(quantum_state,informtion_qubit1,information_qubit2,conditional_list1,conditional_list2,number_qubits):\n",
    "  qubit_list=np.array([informtion_qubit1,information_qubit2])\n",
    "  l1=len(conditional_list1[0])\n",
    "  l2=len(conditional_list2[0])\n",
    "  if l1==0 and l2==0:\n",
    "    q_out=apply_unitary_inverse(qubit_list,bitnode_gen_unitary(conditional_list1[1][0][0],conditional_list2[1][0][0],conditional_list1[1][0][1],conditional_list2[1][0][1]),quantum_state,number_qubits)\n",
    "  else:\n",
    "    q_out=apply_conditional_unitary_inverse(qubit_list,conditional_list1[0],conditional_list2[0],conditional_list1[1],conditional_list2[1],quantum_state,number_qubits)\n",
    "  return q_out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "54ba7906-314d-4f4d-932a-48b8a6310b52",
   "metadata": {},
   "outputs": [],
   "source": [
    "#returns Helstrom error probability after measuring first qubit of joint qubit state in some arbitrary basis\n",
    "def measure_first_qubit(quantum_state,number_qubits,projection_operator):\n",
    "  q=quantum_state\n",
    "  n=number_qubits\n",
    "  pi=projection_operator\n",
    "  q1=np.reshape(q,(2,2**(n-1)))\n",
    "  q2=pi@q1\n",
    "  q3=np.reshape(q2,2**n)\n",
    "  return q3@q\n",
    "\n",
    "\n",
    "def apply_projection_first_qubit(quantum_state,number_qubits,projection_operator):\n",
    "  q=quantum_state\n",
    "  n=number_qubits\n",
    "  pi=projection_operator\n",
    "  q1=np.reshape(q,(2,2**(n-1)))\n",
    "  q2=pi@q1\n",
    "  q3=np.reshape(q2,2**n)\n",
    "  return q3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6da2f608-60ca-4b01-b660-d22eca927892",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_flip(q_state,number_qubits,index):\n",
    "  n=number_qubits\n",
    "  x=np.array([[0,1],[1,0]])\n",
    "  p=np.array(range(n))\n",
    "  p[0]=index\n",
    "  p[index]=0\n",
    "  q=np.reshape(q_state,2*np.ones(n,dtype=int))\n",
    "  q1=np.moveaxis(q,np.array(range(n)),p)\n",
    "  q2=np.reshape(q1,(2,2**(n-1)))\n",
    "  q_out=x@q2\n",
    "  q_out=np.reshape(q_out,2*np.ones(n,dtype=int))\n",
    "  q_out=np.moveaxis(q_out,np.array(range(n)),p)\n",
    "  q_out=np.reshape(q_out,np.shape(q_state))\n",
    "  return q_out\n",
    "\n",
    "def hard_decision_flip(q_state,number_qubits,hard_decision,index):\n",
    "  if hard_decision==1:\n",
    "    return apply_flip(q_state,number_qubits,index)\n",
    "  else:\n",
    "    return q_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "3ccd76c9-97a9-4ca4-b596-3dc339703457",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def polar_transform(u):\n",
    "#   if len(u)==1:\n",
    "#     x=u\n",
    "#   else:\n",
    "#     u1=np.mod(u[0::2]+u[1::2],2)\n",
    "#     u2=u[1::2]\n",
    "#     x=np.concatenate([polar_transform(u1),polar_transform(u2)])\n",
    "#   return x\n",
    "\n",
    "@njit('(int64[:])(int64[:])') # Input/output specifications to make Numba work\n",
    "def polar_transform(u):\n",
    "    '''\n",
    "    Encode polar information vector u\n",
    "\n",
    "          Arguments:\n",
    "                  u (int64[:]): Numpy array of input bits\n",
    "\n",
    "          Returns:\n",
    "                  x (int64[:]): Numpy array of encoded bits\n",
    "    '''\n",
    "    # Recurse down to length 1\n",
    "    if (len(u)==1):\n",
    "        return u;\n",
    "    else:\n",
    "        # R_N maps odd/even indices (i.e., u1u2/u2) to first/second half\n",
    "        # Compute odd/even outputs of (I_{N/2} \\otimes G_2) transform\n",
    "        x = np.zeros(len(u), dtype=np.int64)\n",
    "        x[:len(u)//2] = polar_transform((u[::2]+u[1::2])%2)\n",
    "        x[len(u)//2:] = polar_transform(u[1::2])\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2b1f8109-d19a-4f26-b639-5a5e33831941",
   "metadata": {},
   "outputs": [],
   "source": [
    "def polar_decoder_cq_output(yi,yd,quantum_state,number_qubits,frozen_set):\n",
    "  q=quantum_state\n",
    "  N=len(yi)\n",
    "  f=frozen_set\n",
    "  #pi=np.kron(np.outer(qubit(0),qubit(0)),np.eye(2**number_qubits//2))\n",
    "  pi=np.outer(qubit(0),qubit(0))\n",
    "  if N==1:\n",
    "    if f[0]==1:\n",
    "      p=measure_first_qubit(q,number_qubits,pi)\n",
    "      o=int(np.random.random()>p)\n",
    "      if o==0:\n",
    "        q=apply_projection_first_qubit(q,number_qubits,pi)\n",
    "        return o\n",
    "      else:\n",
    "        q=apply_projection_first_qubit(q,number_qubits,np.eye(2)-pi)\n",
    "        return o\n",
    "    else:\n",
    "      return f[0]\n",
    "\n",
    "  yd1=[]\n",
    "  for i in range(0,N,2):\n",
    "    q,list1=cnop(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "    yd1.append(list1)\n",
    "  out1=polar_decoder_cq_output(yi[::2],yd1,q,number_qubits,f[:(N//2)])\n",
    "  for i in range(0,N,2):\n",
    "    q=cnop_inverse(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "\n",
    "  yd2=[]\n",
    "  for i in range(0,N,2):\n",
    "    q,list2=vnop(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "    yd2.append(list2)\n",
    "  out2=polar_decoder_cq_output(yi[::2],yd2,q,number_qubits,f[N//2:])\n",
    "\n",
    "  for i in range(0,N,2):\n",
    "    q=vnop_inverse(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "\n",
    "  out=np.hstack([out1,out2])\n",
    "  return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "10da27b4-e034-4550-bf3f-76bf05502f2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def polar_decoder_cq_out_avg_error(dp,gp,number_samples,number_qubits,frozen_set):\n",
    "  r=rhom(dp,gp,0)\n",
    "  l,v=LA.eigh(r)\n",
    "  list1=[[],[[dp,gp]]]\n",
    "  yi=np.array(range(number_qubits),dtype=int)\n",
    "  yd=[]\n",
    "  f=frozen_set\n",
    "  for i in range(number_qubits):\n",
    "    yd.append(list1)\n",
    "\n",
    "  h=np.zeros(number_qubits)\n",
    "\n",
    "  for i in range(number_samples):\n",
    "    q=sample_joint_state(l,v,number_qubits)\n",
    "    h=h+polar_decoder_cq_output(yi,yd,q,number_qubits,f)\n",
    "  return h/number_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 305,
   "id": "be7c3073-dc2a-4e69-9f06-658a0f6f3241",
   "metadata": {},
   "outputs": [],
   "source": [
    "def polar_decoder_cq(yi,yd,quantum_state,number_qubits,frozen_set):\n",
    "  q=quantum_state\n",
    "  N=len(yi)\n",
    "  #pi=np.kron(np.outer(qubit(0),qubit(0)),np.eye(2**number_qubits//2))\n",
    "  pi=np.outer(qubit(0),qubit(0))\n",
    "  \n",
    "  f=frozen_set\n",
    "  if N==1:\n",
    "    p1=measure_first_qubit(q,number_qubits,pi)\n",
    "    o=int(np.random.random()>p1)\n",
    "    q=hard_decision_flip(q,number_qubits,f[0],0)\n",
    "    # print(q,'flip')\n",
    "    # print(q)\n",
    "    #if o==1:\n",
    "     # q=apply_flip(q,number_qubits,0)\n",
    "    #o=np.random.choice([0,1],p=[p1,1-p1])\n",
    "    # if f[0]==1:\n",
    "      \n",
    "    #   q=apply_flip(q,number_qubits,0) \n",
    "    # else:\n",
    "    #   print(1)\n",
    "    return o,q\n",
    "  yd1=[]\n",
    "  for i in range(0,N,2):\n",
    "    q,list1=cnop(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "    yd1.append(list1)\n",
    "  # print(q,'cnop')\n",
    "  out1,q=polar_decoder_cq(yi[::2],yd1,q,number_qubits,f[:(N//2)])\n",
    "\n",
    "  \n",
    "  for i in range(0,N,2):\n",
    "    q=cnop_inverse(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "  # print(q,'cnop_inverse')  \n",
    "#   hard1=f[:(N//2)]\n",
    "#   for i in range(0,N,2):\n",
    "#     q=hard_decision_flip(q,number_qubits,hard1[i//2],yi[i])\n",
    "    \n",
    "  yd2=[]\n",
    "  for i in range(0,N,2):\n",
    "    q,list2=vnop(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "    yd2.append(list2)\n",
    "  # print(q,'vnop')\n",
    "  out2,q=polar_decoder_cq(yi[::2],yd2,q,number_qubits,f[(N//2):])\n",
    "  \n",
    "  for i in range(0,N,2):\n",
    "    q=vnop_inverse(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "  # print(q,'vnop_inverse')\n",
    "  # for i in range(0,N,2):\n",
    "  #   q=hard_decision_flip(q,number_qubits,hard1[i//2],yi[i])\n",
    "  out=np.hstack([out1,out2])\n",
    "  return out,q\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 306,
   "id": "2cf7639c-d0e0-489f-866f-54d230059ad5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def polar_decoder_cq_avg_error(dp,gp,number_samples,number_qubits,codeword):\n",
    "  r=rhom(dp,gp,0)\n",
    "  encoded_bits=polar_transform(np.array(codeword))\n",
    "  x=np.array([[0,1],[1,0]])\n",
    "  l,v=LA.eigh(r)\n",
    "  l1,v1=LA.eigh(x@r@x)\n",
    "  list1=[[],[[dp,gp]]]\n",
    "  \n",
    "  yi=np.array(range(number_qubits),dtype=int)\n",
    "  yd=[]\n",
    "  for i in range(number_qubits):\n",
    "    yd.append(list1)\n",
    "    \n",
    "\n",
    "  h=np.zeros(number_qubits)\n",
    "\n",
    "  for i in range(number_samples):\n",
    "    q=sample_joint_random_state(l,v,v1,number_qubits,encoded_bits)\n",
    "    h1,q1=polar_decoder_cq(yi,yd,q,number_qubits,codeword)\n",
    "    h=h+np.mod(h1+codeword,2)\n",
    "               \n",
    "  return h/number_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 307,
   "id": "5e25d555-8801-4403-8666-1bcf3467fb71",
   "metadata": {},
   "outputs": [],
   "source": [
    "h=polar_decoder_cq_avg_error(0.1,0.1,1000,8,[0,0,0,0,0,0,0,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 308,
   "id": "2f849f80-56f6-4404-93ec-efe81a689ef9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.423 0.28  0.296 0.068 0.271 0.053 0.052 0.004]\n"
     ]
    }
   ],
   "source": [
    "print(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 267,
   "id": "80aa671e-b1b8-4098-b843-91bdbc216139",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error rate for channel 1: 0.41604365268962873\n",
      "Error rate for channel 2: 0.27872502680798883\n",
      "Error rate for channel 3: 0.2796547846448365\n",
      "Error rate for channel 4: 0.08216633626864445\n",
      "Error rate for channel 5: 0.27888671816715527\n",
      "Error rate for channel 6: 0.07190103587830066\n",
      "Error rate for channel 7: 0.058029348294409755\n",
      "Error rate for channel 8: 0.0033242216626361287\n"
     ]
    }
   ],
   "source": [
    "dp=0.1\n",
    "gp=0.1\n",
    "n=3\n",
    "d=np.ones(20000)*dp\n",
    "g=np.ones(20000)*gp\n",
    "p=polar(n,d,g)\n",
    "#print(p.shape)\n",
    "ww = np.mean(p,axis=2)[:,0]\n",
    "#print(ww)\n",
    "#f=polar_design(ww,0.001)\n",
    "# print('rate for code with given design constraint:',np.mean(f==0.5))\n",
    "for i in range(2**n):\n",
    "  d=np.mean(p[i][0])\n",
    "  g=np.mean(p[i][1])\n",
    "  r1=rhom(d,g,0)\n",
    "  h=helstrom_success(r1,pauli(1))\n",
    "  print(f'Error rate for channel {i+1}:', 1-h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 244,
   "id": "a39977a4-6d72-42ea-b361-e5c9534d6f37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.014928749927334057\n",
      "0.9850712500726662\n",
      "[0 0 0 1]\n"
     ]
    }
   ],
   "source": [
    "dp,gp=0.1,0.1\n",
    "r=rhom(dp,gp,0)\n",
    "x=np.array([[0,1],[1,0]])\n",
    "l,v=LA.eigh(r)\n",
    "l1,v1=LA.eigh(x@r@x)\n",
    "pi=np.outer(qubit(0),qubit(0))\n",
    "print(v[:,1]@pi@v[:,1])\n",
    "print(v1[:,1]@pi@v1[:,1])\n",
    "x=polar_transform(np.array([1,1,1,1]))\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "id": "d22861ae-1cbd-4757-bed8-e5d2407dd512",
   "metadata": {},
   "outputs": [],
   "source": [
    "q=np.kron(np.kron(v[:,1],v[:,1]),np.kron(v[:,1],v1[:,1]))\n",
    "pi1=np.kron(pi,np.eye(8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 246,
   "id": "8333b3d9-1fdd-4a54-8f53-c7a83486ecd9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.05709342560553635\n"
     ]
    }
   ],
   "source": [
    "q1=apply_unitary([0,1],CU,q,4)\n",
    "q2=apply_unitary([2,3],CU,q1,4)\n",
    "q3=apply_unitary([0,2],CU,q2,4)\n",
    "print(q3@pi1@q3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 247,
   "id": "646269ca-713c-4e69-8549-6e5900678e28",
   "metadata": {},
   "outputs": [],
   "source": [
    "flip1=np.kron(pauli(1),np.eye(8))\n",
    "\n",
    "#q4=flip1@q3\n",
    "q4=apply_flip(q3,4,0)\n",
    "q5=apply_unitary_inverse([0,2],CU,q4,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 248,
   "id": "e746b69a-1ab9-4d88-aa07-33563b278120",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.004206649156410568\n"
     ]
    }
   ],
   "source": [
    "rc,pc=checknode(dp,dp,gp,gp)\n",
    "dc1=rc[0][0]\n",
    "dc2=rc[1][0]\n",
    "gc1=rc[0][1]\n",
    "gc2=rc[1][1]\n",
    "\n",
    "q6=apply_conditional_unitary([0,2],[1],[3],np.array([[dc1,gc1],[dc2,gc2]]),np.array([[dc1,gc1],[dc2,gc2]]),q5,4)\n",
    "print(q6@pi1@q6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 249,
   "id": "2dc8c941-3696-4875-9d32-008aaebe55ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.002548782266514117\n"
     ]
    }
   ],
   "source": [
    "q7=flip1@q6\n",
    "q7=apply_conditional_unitary_inverse([0,2],[1],[3],np.array([[dc1,gc1],[dc2,gc2]]),np.array([[dc1,gc1],[dc2,gc2]]),q7,4)\n",
    "\n",
    "\n",
    "q8=apply_unitary_inverse([2,3],CU,q7,4)\n",
    "q9=apply_unitary_inverse([0,1],CU,q8,4)\n",
    "V=bitnode_gen_unitary(dp,dp,gp,gp)\n",
    "q10=apply_unitary([0,1],V,q9,4)\n",
    "q11=apply_unitary([2,3],V,q10,4)\n",
    "q12=apply_unitary([0,2],CU,q11,4)\n",
    "print(q12@pi1@q12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 250,
   "id": "99e20e6b-2daa-4fbb-86d4-5be7b9c83a5a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.633160611898096e-06\n"
     ]
    }
   ],
   "source": [
    "q12=flip1@q12\n",
    "q13=apply_unitary_inverse([0,2],CU,q12,4)\n",
    "rb,pb=bitnode(dp,dp,gp,gp)\n",
    "db1=rb[0][0]\n",
    "db2=rb[1][0]\n",
    "gb1=rb[0][1]\n",
    "gb2=rb[1][1]\n",
    "q14=apply_conditional_unitary([0,2],[1],[3],np.array([[db1,gb1],[db2,gb2]]),np.array([[db1,gb1],[db2,gb2]]),q13,4)\n",
    "print(q14@pi1@q14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 252,
   "id": "42079a7a-9646-4d9d-a4db-9dd35d41fe6c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[5.70934256e-02 4.20664916e-03 2.54878227e-03 7.63316061e-06]\n"
     ]
    }
   ],
   "source": [
    "list1=[[],[[dp,gp]]]\n",
    "yi=np.array(range(4),dtype=int)\n",
    "yd=[]\n",
    "for i in range(4):\n",
    "  yd.append(list1)\n",
    "print(polar_decoder_cq(yi,yd,q,4,[1,1,1,1])[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 273,
   "id": "615e7126-3105-42ba-b12c-847fee53dbee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def polar_decoder_cq_output(yi,yd,quantum_state,number_qubits,info_set,codeword):\n",
    "  q=quantum_state\n",
    "  N=len(yi)\n",
    "  #pi=np.kron(np.outer(qubit(0),qubit(0)),np.eye(2**number_qubits//2))\n",
    "  pi=np.outer(qubit(0),qubit(0))\n",
    "  \n",
    "  f=info_set\n",
    "  if N==1:\n",
    "    if f[0]==1:\n",
    "        p1=measure_first_qubit(q,number_qubits,pi)\n",
    "        o=int(np.random.random()>p1)\n",
    "        if o==0:\n",
    "            q=apply_projection_first_qubit(q,number_qubits,pi)\n",
    "            return o,q\n",
    "        else:\n",
    "            q=apply_projection_first_qubit(q,number_qubits,np.eye(2)-pi)\n",
    "            q=apply_flip(q,number_qubits,0)\n",
    "            return o,q\n",
    "    else:\n",
    "        q=hard_decision_flip(q,number_qubits,codeword[0],0)\n",
    "        return codeword[0],q\n",
    "    # print(q,'flip')\n",
    "    # print(q)\n",
    "    #if o==1:\n",
    "     # q=apply_flip(q,number_qubits,0)\n",
    "    #o=np.random.choice([0,1],p=[p1,1-p1])\n",
    "    # if f[0]==1:\n",
    "      \n",
    "    #   q=apply_flip(q,number_qubits,0) \n",
    "    # else:\n",
    "    #   print(1)\n",
    "    return o,q\n",
    "  yd1=[]\n",
    "  for i in range(0,N,2):\n",
    "    q,list1=cnop(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "    yd1.append(list1)\n",
    "  # print(q,'cnop')\n",
    "  out1,q=polar_decoder_cq_output(yi[::2],yd1,q,number_qubits,f[:(N//2)],codeword[:(N//2)])\n",
    "\n",
    "  \n",
    "  for i in range(0,N,2):\n",
    "    q=cnop_inverse(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "  # print(q,'cnop_inverse')  \n",
    "#   hard1=f[:(N//2)]\n",
    "#   for i in range(0,N,2):\n",
    "#     q=hard_decision_flip(q,number_qubits,hard1[i//2],yi[i])\n",
    "    \n",
    "  yd2=[]\n",
    "  for i in range(0,N,2):\n",
    "    q,list2=vnop(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "    yd2.append(list2)\n",
    "  # print(q,'vnop')\n",
    "  out2,q=polar_decoder_cq_output(yi[::2],yd2,q,number_qubits,f[(N//2):],codeword[(N//2):])\n",
    "  \n",
    "  for i in range(0,N,2):\n",
    "    q=vnop_inverse(q,yi[i],yi[i+1],yd[i],yd[i+1],number_qubits)\n",
    "  # print(q,'vnop_inverse')\n",
    "  # for i in range(0,N,2):\n",
    "  #   q=hard_decision_flip(q,number_qubits,hard1[i//2],yi[i])\n",
    "  out=np.hstack([out1,out2])\n",
    "  return out,q\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 279,
   "id": "cbf3eae8-6a2a-406a-8730-26ebdafb02e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def polar_decoder_cq_output_avg_error(dp,gp,number_samples,number_qubits,info_set,codeword):\n",
    "  r=rhom(dp,gp,0)\n",
    "  encoded_bits=polar_transform(np.array(codeword))\n",
    "  x=np.array([[0,1],[1,0]])\n",
    "  l,v=LA.eigh(r)\n",
    "  l1,v1=LA.eigh(x@r@x)\n",
    "  list1=[[],[[dp,gp]]]\n",
    "  \n",
    "  yi=np.array(range(number_qubits),dtype=int)\n",
    "  yd=[]\n",
    "  for i in range(number_qubits):\n",
    "    yd.append(list1)\n",
    "    \n",
    "\n",
    "  h=np.zeros(number_qubits)\n",
    "  block=0\n",
    "\n",
    "  for i in range(number_samples):\n",
    "    q=sample_joint_random_state(l,v,v1,number_qubits,encoded_bits)\n",
    "    h1,q1=polar_decoder_cq_output(yi,yd,q,number_qubits,info_set,codeword)\n",
    "    h=h+np.mod(h1+codeword,2)\n",
    "    if np.sum(np.mod(h1+codeword,2))>=1:\n",
    "        block=block+1\n",
    "               \n",
    "  return h/number_samples,block/number_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 342,
   "id": "5d2a53f2-1e6a-4548-9933-a5ca8b5c531e",
   "metadata": {},
   "outputs": [],
   "source": [
    "h,block=polar_decoder_cq_output_avg_error(0.05,0.15,500,8,[0,0,0,1,0,1,1,1],[1,1,1,1,1,1,0,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 343,
   "id": "fdedff89-d643-43ed-ad4b-014ddc3ce996",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.    0.    0.    0.012 0.    0.014 0.06  0.002]\n"
     ]
    }
   ],
   "source": [
    "print(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 344,
   "id": "c6da3e4b-c06d-4865-a691-0140e0396883",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.064\n"
     ]
    }
   ],
   "source": [
    "print(block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 336,
   "id": "cecb93f9-4d8f-4049-b2c2-b93a6ea4a60d",
   "metadata": {},
   "outputs": [],
   "source": [
    "h,block=polar_decoder_cq_output_avg_error(0.05,0.15,1000,4,[0,0,1,1],[1,0,1,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 337,
   "id": "f8ed1eef-16bb-435e-a22c-a8f450064539",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.    0.    0.061 0.049]\n"
     ]
    }
   ],
   "source": [
    "print(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 338,
   "id": "6cda6634-43a7-4dba-9558-50208d0b10db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.075\n"
     ]
    }
   ],
   "source": [
    "print(block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60f54e09-0dc7-40cd-9496-6fdacfb4a8dd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}