{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autograd experiments\n",
    "\n",
    "**References**\n",
    "* https://github.com/HIPS/autograd/blob/master/docs/tutorial.md\n",
    "* https://github.com/HIPS/autograd/blob/master/docs/updateguide.md\n",
    "\n",
    "abbreviations:\n",
    "* ax= axis\n",
    "* dwrt= derivative with respect to"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example of @primitive usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import autograd.numpy as anp\n",
    "from autograd import grad, jacobian\n",
    "from autograd.extend import primitive, defvjp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fun(x, y, z):\n",
    "    \"\"\"\n",
    "    x : float\n",
    "    y: float\n",
    "    z : bool\n",
    "    \n",
    "    \"\"\"\n",
    "    return x * y**2 if z else y**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def d_fun(x, y, z, dwrt):\n",
    "    if z:\n",
    "        if dwrt == 0:\n",
    "            return y**2\n",
    "        elif dwrt == 1:\n",
    "            return 2*x*y\n",
    "        else:\n",
    "            assert False, 'z not differentiable'\n",
    "    else:\n",
    "        if dwrt == 0:\n",
    "            return 0.\n",
    "        elif dwrt == 1:\n",
    "            return 2*y\n",
    "        else:\n",
    "            assert False, 'z not differentiable'  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "fun_x = grad(fun, 0)\n",
    "fun_y = grad(fun, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rrtucci/anaconda3/envs/qubiter/lib/python3.7/site-packages/autograd/tracer.py:14: UserWarning: Output seems independent of input.\n",
      "  warnings.warn(\"Output seems independent of input.\")\n"
     ]
    }
   ],
   "source": [
    "for z in [True,False]:\n",
    "    print(fun_x(.5, .1, z) - d_fun(.5, .1, z, dwrt=0))\n",
    "    print(fun_y(.5, .1, z) - d_fun(.5, .1, z, dwrt=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@primitive\n",
    "def pfun(x, y, z):\n",
    "    return fun(x, y, z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "defvjp(pfun,\n",
    "       lambda ans, x, y, z: lambda g: g*d_fun(x, y, z, 0),  # g and pfun scalars\n",
    "       lambda ans, x, y, z: lambda g: g*d_fun(x, y, z, 1),\n",
    "       argnums=[0, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "pfun_x = grad(pfun, 0)\n",
    "pfun_y = grad(pfun, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n"
     ]
    }
   ],
   "source": [
    "for z in [True,False]:\n",
    "    print(pfun_x(.5, .1, z) - d_fun(.5, .1, z, 0))\n",
    "    print(pfun_y(.5, .1, z) - d_fun(.5, .1, z, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Storing Pauli matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1],\n",
       "       [1, 0]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sigx = anp.array([[0, 1],[1,0]])\n",
    "sigx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.+0.j, -0.-1.j],\n",
       "       [ 0.+1.j,  0.+0.j]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sigy = anp.array([[0, -1j],[1j,0]])\n",
    "sigy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1,  0],\n",
       "       [ 0, -1]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sigz = anp.array([[1, 0],[0, -1]])\n",
    "sigz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.+0.j,  1.+0.j],\n",
       "       [ 1.+0.j,  0.+0.j],\n",
       "       [ 0.+0.j, -0.-1.j],\n",
       "       [ 0.+1.j,  0.+0.j],\n",
       "       [ 1.+0.j,  0.+0.j],\n",
       "       [ 0.+0.j, -1.+0.j]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig_all = anp.vstack([sigx, sigy, sigz])\n",
    "sig_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0.+0.j,  0.+0.j,  1.+0.j],\n",
       "        [ 1.+0.j, -0.-1.j,  0.+0.j]],\n",
       "\n",
       "       [[ 1.+0.j,  0.+1.j,  0.+0.j],\n",
       "        [ 0.+0.j,  0.+0.j, -1.+0.j]]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig_all = anp.reshape(sig_all, (3, 2, 2)).transpose(1, 2, 0)\n",
    "sig_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sigx_=\n",
      " [[0.+0.j 1.+0.j]\n",
      " [1.+0.j 0.+0.j]]\n",
      "sigy_=\n",
      " [[0.+0.j 0.-1.j]\n",
      " [0.+1.j 0.+0.j]]\n",
      "sigz_=\n",
      " [[ 1.+0.j  0.+0.j]\n",
      " [ 0.+0.j -1.+0.j]]\n"
     ]
    }
   ],
   "source": [
    "ex = anp.array([1,0,0])\n",
    "ey = anp.array([0,1,0])\n",
    "ez = anp.array([0,0,1])\n",
    "sigx_ = anp.dot(sig_all, ex)\n",
    "sigy_ = anp.dot(sig_all, ey)\n",
    "sigz_ = anp.dot(sig_all, ez)\n",
    "print('sigx_=\\n', sigx_)\n",
    "print('sigy_=\\n', sigy_)\n",
    "print('sigz_=\\n', sigz_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## rot1, single parameter rotation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$ U = e^{i\\sigma_3\\theta_3} = C + i\\sigma_3 S$\n",
    "\n",
    "$S = \\sin\\theta_3, C = \\cos \\theta_3$\n",
    "\n",
    "$\\frac{dU}{dt} = \\dot{\\theta}_3(-S + i\\sigma_3 C)$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rot1(t, ax):\n",
    "    assert ax in [1, 2, 3]\n",
    "    return anp.eye(2)*anp.cos(t) + 1j*sig_all[:, :, ax-1]*anp.sin(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rot1_t(t, ax):\n",
    "    assert  ax in [1, 2, 3]\n",
    "    return -anp.eye(2)*anp.sin(t) + 1j*sig_all[:, :, ax-1]*anp.cos(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def d_auto_rot1(t, ax):\n",
    "    assert  ax in [1, 2, 3]\n",
    "    def rot1r(t, ax):\n",
    "        return anp.real(rot1(t, ax))\n",
    "    def rot1i(t, ax):\n",
    "        return anp.imag(rot1(t, ax))\n",
    "    return jacobian(rot1r, 0)(t, ax) + 1j*jacobian(rot1i, 0)(t, ax) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.+0.j 0.+0.j]\n",
      " [0.+0.j 0.+0.j]]\n",
      "[[0.+0.j 0.+0.j]\n",
      " [0.+0.j 0.+0.j]]\n",
      "[[0.+0.j 0.+0.j]\n",
      " [0.+0.j 0.+0.j]]\n"
     ]
    }
   ],
   "source": [
    "for ax in range(1, 4):\n",
    "    print(d_auto_rot1(.5, 1) - rot1_t(.5, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arbitrary 2-dim unitary (4 parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$ U = e^{i\\sigma_k\\theta_k} = C +\n",
    "i\\sigma_k \\frac{\\theta_k}{\\theta} S$\n",
    "\n",
    "$\\theta = \\sqrt{\\theta_k\\theta_k},\n",
    "S = \\sin\\theta, C = \\cos \\theta$\n",
    "\n",
    "$\\frac{dU}{dt}=-S \\frac{\\theta_k}{\\theta}\n",
    "\\dot{\\theta_k}+ i\\sigma_k\\dot{\\theta_r}\n",
    "\\left[\\frac{\\theta_k\\theta_r}{\\theta^2} C+\n",
    "\\frac{S}{\\theta}(-\\frac{\\theta_k\\theta_r}{\\theta^2}+\\delta_{k, r})\\right]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u2(*tlist):\n",
    "    assert len(tlist) == 4\n",
    "    t = anp.sqrt(tlist[1]**2 + tlist[2]**2 + tlist[3]**2)\n",
    "    tvec = anp.array([tlist[1]/t, tlist[2]/t, tlist[3]/t])\n",
    "    out = anp.eye(2)*anp.cos(t) + 1j*anp.dot(sig_all, tvec)*anp.sin(t) \n",
    "    return anp.exp(1j*tlist[0])*out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def d_u2(dwrt, *tlist):\n",
    "    assert dwrt in range(4)\n",
    "    assert len(tlist) == 4\n",
    "    if dwrt == 0:\n",
    "        return 1j*u2(*tlist)\n",
    "    dwrt -= 1\n",
    "    t = anp.sqrt(tlist[1]**2 + tlist[2]**2 + tlist[3]**2)\n",
    "    tvec = anp.array([tlist[1]/t, tlist[2]/t, tlist[3]/t])\n",
    "    dotted_vec = tvec*tvec[dwrt]*anp.cos(t) + (anp.sin(t)/t)*(-tvec*tvec[dwrt] + anp.eye(3)[dwrt, :])\n",
    "    out = -anp.sin(t)*tvec[dwrt]*anp.eye(2) + 1j*anp.dot(sig_all, dotted_vec)\n",
    "    return anp.exp(1j*tlist[0])*out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.81615064+0.46474597j,  0.26526593+0.21804425j],\n",
       "       [-0.30329697+0.16099768j,  0.89221274-0.29333789j]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u2(.1, .2, .3, .4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.+0.j, 0.+0.j],\n",
       "       [0.+0.j, 0.+0.j]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u2(0., 0., 0., .4) - rot1(.4, ax=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def d_auto_u2(dwrt, *tlist):\n",
    "    def u2r(*tlist):\n",
    "        return anp.real(u2(*tlist))\n",
    "    def u2i(*tlist):\n",
    "        return anp.imag(u2(*tlist))\n",
    "    return jacobian(u2r, dwrt)(*tlist) + 1j*jacobian(u2i, dwrt)(*tlist) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.47942554+0.87758256j  0.        +0.j        ]\n",
      " [ 0.        +0.j          0.47942554+0.87758256j]]\n",
      "[[0.+0.j         0.+0.95885108j]\n",
      " [0.+0.95885108j 0.+0.j        ]]\n",
      "[[ 0.        +0.j  0.95885108+0.j]\n",
      " [-0.95885108+0.j  0.        +0.j]]\n",
      "[[-0.47942554+0.87758256j  0.        +0.j        ]\n",
      " [ 0.        +0.j         -0.47942554-0.87758256j]]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rrtucci/anaconda3/envs/qubiter/lib/python3.7/site-packages/numpy/core/_asarray.py:85: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return array(a, dtype, copy=False, order=order)\n"
     ]
    }
   ],
   "source": [
    "for dwrt in range(4):\n",
    "    print(d_auto_u2(dwrt, 0., 0., 0., .5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.47942554+0.87758256j,  0.        +0.j        ],\n",
       "       [ 0.        +0.j        , -0.47942554-0.87758256j]])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rot1_t(.5, ax=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0\n",
      "2.5814177514652363e-16\n",
      "1.4833660915169986e-16\n",
      "2.786810452962585e-16\n"
     ]
    }
   ],
   "source": [
    "tlist = [.3, 1.1, .7, .5]\n",
    "for dwrt in range(4):\n",
    "    print(anp.linalg.norm(d_auto_u2(dwrt, *tlist) - d_u2(dwrt, *tlist)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## u2 as primitive\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "@primitive\n",
    "def pu2r(*tlist):\n",
    "    return anp.real(u2(*tlist))\n",
    "@primitive\n",
    "def pu2i(*tlist):\n",
    "    return anp.imag(u2(*tlist))\n",
    "def pu2(*tlist):\n",
    "    return pu2r(*tlist) + 1j*pu2i(*tlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "defvjp(pu2r,\n",
    "       lambda ans, *tlist: lambda g: anp.sum(g*anp.real(d_u2(0, *tlist))), # g.shape==pu2r.shape\n",
    "       lambda ans, *tlist: lambda g: anp.sum(g*anp.real(d_u2(1, *tlist))),\n",
    "       lambda ans, *tlist: lambda g: anp.sum(g*anp.real(d_u2(2, *tlist))),\n",
    "       lambda ans, *tlist: lambda g: anp.sum(g*anp.real(d_u2(3, *tlist))),\n",
    "       argnums=range(4))\n",
    "defvjp(pu2i,\n",
    "       lambda ans, *tlist: lambda g: anp.sum(g*anp.imag(d_u2(0, *tlist))), # g.shape==pu2i.shape\n",
    "       lambda ans, *tlist: lambda g: anp.sum(g*anp.imag(d_u2(1, *tlist))),\n",
    "       lambda ans, *tlist: lambda g: anp.sum(g*anp.imag(d_u2(2, *tlist))),\n",
    "       lambda ans, *tlist: lambda g: anp.sum(g*anp.imag(d_u2(3, *tlist))),\n",
    "       argnums=range(4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def d_auto_pu2(dwrt, *tlist):\n",
    "    assert dwrt in range(4)\n",
    "    return jacobian(pu2r, dwrt)(*tlist) + 1j*jacobian(pu2i, dwrt)(*tlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0\n",
      "0.0\n",
      "0.0\n",
      "0.0\n"
     ]
    }
   ],
   "source": [
    "tlist = [.1, .2, .3, .4]\n",
    "for dwrt in range(4):\n",
    "    print(anp.linalg.norm(d_auto_pu2(dwrt, *tlist) - d_u2(dwrt, *tlist)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10.9"
  },
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "12px",
    "width": "252px"
   },
   "navigate_menu": true,
   "number_sections": true,
   "sideBar": true,
   "threshold": 4,
   "toc_cell": false,
   "toc_section_display": "block",
   "toc_window_display": false,
   "widenNotebook": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}