{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QobjEvo \"advanced\" methods\n",
    "\n",
    "Test and present diverse feature of QobjEvo.\n",
    "\n",
    "Made by Eric Giguere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qutip as qt\n",
    "import numpy as np\n",
    "from qutip import QobjEvo\n",
    "%load_ext cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 5\n",
    "destroy, create, Id = qt.destroy(N), qt.create(N), qt.qeye(N)\n",
    "def exp_i(t,args):\n",
    "    return np.exp(-1j*t)\n",
    "def cos_w(t,args):\n",
    "    return np.cos(args[\"w\"]*t)\n",
    "tlist = np.linspace(0,10,10000)\n",
    "tlistlog = np.logspace(-3,1,10000)\n",
    "\n",
    "# state vector as np array\n",
    "vec = np.arange(N)*.5+.5j\n",
    "vec_super = np.arange(N**2)*.5+.5j\n",
    "mat_c = (np.arange(N**2)*.5+.5j).reshape((5,5))\n",
    "mat_f = np.asfortranarray(mat_c*1.)\n",
    "\n",
    "# Construct QobjEvo of all type\n",
    "td_cte1 = QobjEvo(Id)\n",
    "td_cte2 = QobjEvo([Id])\n",
    "\n",
    "td_func = QobjEvo([Id,[create,exp_i],[destroy,cos_w]],args={\"w\":2})\n",
    "td_str = QobjEvo([Id,[create,\"exp(-1j*t)\"],[destroy,\"cos(w*t)\"]],args={'w':2.})\n",
    "td_array = QobjEvo([Id,[create,np.exp(-1j*tlist)],[destroy,np.cos(2*tlist)]],tlist=tlist)\n",
    "td_array_log = QobjEvo([Id,[create,np.exp(-1j*tlistlog)],[destroy,np.cos(2*tlistlog)]],tlist=tlistlog)\n",
    "\n",
    "td_super = qt.liouvillian(td_func, c_ops=td_cte1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "## Product and expectation value\n",
    "\n",
    "QobjEvo.mul_vec(t,state) = spmv(QobjEvo(t), state)  \n",
    "QobjEvo.expect(t, state, real) = cy_expect_psi/cy_expect_rho_vec  (QobjEvo(t), state, real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qutip.cy.spmatfuncs import spmv, cy_expect_rho_vec, cy_expect_psi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ True,  True,  True,  True,  True])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spmv(td_func(2).data, vec) == td_func.mul_vec(2,vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ True  True  True  True  True]\n",
      " [ True  True  True  True  True]\n",
      " [ True  True  True  True  True]\n",
      " [ True  True  True  True  True]\n",
      " [ True  True  True  True  True]]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "  C_CONTIGUOUS : True\n",
       "  F_CONTIGUOUS : False\n",
       "  OWNDATA : False\n",
       "  WRITEABLE : True\n",
       "  ALIGNED : True\n",
       "  WRITEBACKIFCOPY : False\n",
       "  UPDATEIFCOPY : False"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(td_func(2).data * mat_c == td_func.mul_mat(2,mat_c))\n",
    "mat_c.flags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ True  True  True  True  True]\n",
      " [ True  True  True  True  True]\n",
      " [ True  True  True  True  True]\n",
      " [ True  True  True  True  True]\n",
      " [ True  True  True  True  True]]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "  C_CONTIGUOUS : False\n",
       "  F_CONTIGUOUS : True\n",
       "  OWNDATA : True\n",
       "  WRITEABLE : True\n",
       "  ALIGNED : True\n",
       "  WRITEBACKIFCOPY : False\n",
       "  UPDATEIFCOPY : False"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(td_func(2).data * mat_f == td_func.mul_mat(2,mat_f))\n",
    "mat_f.flags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cy_expect_psi(td_str(2).data, vec, 0) == td_str.expect(2, vec, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cy_expect_rho_vec(td_super(2).data, vec_super, 0) == td_super.expect(2, vec_super, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function with state\n",
    "\n",
    "For the solvers which have the option \"rhs_with_state\"\n",
    "\n",
    "the QobjEvo can take coefficient function with the signature (for backward compatibility):\n",
    "\n",
    "**def coeff(t, psi. args)**\n",
    "\n",
    "or use advanced args: **args={\"psi=vec\":psi0}**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def coeff_state(t, args):\n",
    "    return np.max(args[\"psi\"])*args[\"w\"]\n",
    "td_state = QobjEvo([Id, [destroy, coeff_state]],args={\"w\":1, \"psi=vec\":qt.basis(N,0)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "Quantum object: dims = [[5], [5]], shape = (5, 5), type = oper, isherm = False\\begin{equation*}\\left(\\begin{array}{*{11}c}1.0 & (2.0+0.500j) & 0.0 & 0.0 & 0.0\\\\0.0 & 1.0 & (2.828+0.707j) & 0.0 & 0.0\\\\0.0 & 0.0 & 1.0 & (3.464+0.866j) & 0.0\\\\0.0 & 0.0 & 0.0 & 1.0 & (4.0+1.0j)\\\\0.0 & 0.0 & 0.0 & 0.0 & 1.0\\\\\\end{array}\\right)\\end{equation*}"
      ],
      "text/plain": [
       "Quantum object: dims = [[5], [5]], shape = (5, 5), type = oper, isherm = False\n",
       "Qobj data =\n",
       "[[1.        +0.j         2.        +0.5j        0.        +0.j\n",
       "  0.        +0.j         0.        +0.j        ]\n",
       " [0.        +0.j         1.        +0.j         2.82842712+0.70710678j\n",
       "  0.        +0.j         0.        +0.j        ]\n",
       " [0.        +0.j         0.        +0.j         1.        +0.j\n",
       "  3.46410162+0.8660254j  0.        +0.j        ]\n",
       " [0.        +0.j         0.        +0.j         0.        +0.j\n",
       "  1.        +0.j         4.        +1.j        ]\n",
       " [0.        +0.j         0.        +0.j         0.        +0.j\n",
       "  0.        +0.j         1.        +0.j        ]]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "td_state(2,state=vec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cython object and \"Compiling\"\n",
    "\n",
    "There is a cython version of the qobjevo for fast call:  \n",
    "- qutip.cy.cqobjevo.CQobjEvo\n",
    "\n",
    "The cython is created when the \"compile\" method is created.\n",
    "The cython object contain fast version of the call, expect and rhs (spmv) methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before compilation\n",
      "195 µs ± 8.68 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n",
      "205 µs ± 1.98 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n",
      "219 µs ± 15.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n",
      "268 µs ± 41.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n",
      "206 µs ± 1.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n",
      "After compilation\n",
      "13.2 µs ± 477 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n",
      "8.53 µs ± 52.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n",
      "8.66 µs ± 48.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n",
      "9.18 µs ± 49.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n",
      "10.7 µs ± 288 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "td_str.compiled = False\n",
    "print(\"Before compilation\")\n",
    "%timeit td_str(2, data=True)\n",
    "%timeit td_str.mul_vec(2,vec)\n",
    "%timeit td_str.mul_mat(2,mat_c)\n",
    "%timeit td_str.mul_mat(2,mat_f)\n",
    "%timeit td_str.expect(2,vec,0)\n",
    "td_str.compile()\n",
    "print(\"After compilation\")\n",
    "%timeit td_str(2, data=True)\n",
    "%timeit td_str.mul_vec(2,vec)\n",
    "%timeit td_str.mul_mat(2,mat_c)\n",
    "%timeit td_str.mul_mat(2,mat_f)\n",
    "%timeit td_str.expect(2,vec,0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## apply\n",
    "Pass a function ( Qobj, \\*args, \\*\\*kwargs) -> Qobj to act on each component of the qobjevo.\n",
    "\n",
    "Will only be mathematicaly valid if the transformation is linear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "def multiply(qobj, b, factor = 3.):\n",
    "    return qobj*b*factor\n",
    "\n",
    "print(td_func.apply(multiply,2)(2) == td_func(2)*6)\n",
    "print(td_func.apply(multiply,2,factor=2)(2) == td_func(2)*4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## apply_decorator\n",
    "Transform the functions containing the time dependence using a decorator.\n",
    "\n",
    "The decorator must return a function of (t, **kwargs).\n",
    "\n",
    "Do not modify the constant part (the contant part do not have a function f(t) = 1).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "def rescale_time_and_scale(f_original, time_scale, factor=2.):\n",
    "    def f(t, *args, **kwargs):\n",
    "        return f_original(time_scale*t, *args, **kwargs)*factor\n",
    "    return f\n",
    "\n",
    "print(td_func.apply_decorator(rescale_time_and_scale, 2)(2) == td_func(4)*2-Id)\n",
    "print(td_func.apply_decorator(rescale_time_and_scale, 3, factor=3)(2) == \n",
    "      td_func(6)*3.0 - 2*Id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "QobjEvo based on string and np.array are changed to a function then the decorator is applied. There are option so that the type of coefficient stay unchanged: \n",
    "\n",
    "    str_mod : change the string -> str_mod[0] + str + str_mod[1]\n",
    "    \n",
    "    inplace_np : modify the array  (array[i] = decorator(lambda v: v)(array[i]))\n",
    "         *any modification that rescale the time will not work properly\n",
    "        \n",
    "Decorator can cause problem when used in parallel. (function cannot be pickled error)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "<function exp_i at 0x7f1c1e7bb158>\n",
      "True\n",
      "str not updated: <function square_f.<locals>.f at 0x7f1c1dee0c80>\n",
      "True\n",
      "str  updated: (exp(-1j*t))**2\n",
      "True\n",
      "array not updated: <function square_f.<locals>.f at 0x7f1c1dee09d8>\n",
      "True\n",
      "array  updated: [1.        -0.j         0.999998  -0.0020002j  0.999992  -0.00400039j ...\n",
      " 0.41173093-0.91130546j 0.40990732-0.91212718j 0.40808206-0.91294525j]\n"
     ]
    }
   ],
   "source": [
    "td_func_1 = QobjEvo([[create,exp_i]],args={\"w\":2})\n",
    "td_str_1 = QobjEvo([[create,\"exp(-1j*t)\"]],args={'w':2.})\n",
    "td_array_1 = QobjEvo([[create,np.exp(-1j*tlist)]],tlist=tlist)\n",
    "\n",
    "def square_qobj(qobj):\n",
    "    return qobj*qobj\n",
    "\n",
    "def square_f(f_original):\n",
    "    def f(t, *args, **kwargs):\n",
    "        return f_original(t, *args, **kwargs)**2\n",
    "    return f\n",
    "\n",
    "t1 = td_func_1.apply(square_qobj).apply_decorator(square_f)\n",
    "print(t1(2) == td_func_1(2)*td_func_1(2))\n",
    "print((t1.ops[0][2]))\n",
    "\n",
    "t1 = td_str_1.apply(square_qobj).apply_decorator(square_f)\n",
    "print(t1(2) == td_str_1(2)*td_str_1(2))\n",
    "print(\"str not updated:\", (t1.ops[0][2]))\n",
    "\n",
    "t1 = td_str_1.apply(square_qobj).apply_decorator(square_f, str_mod=[\"(\",\")**2\"])\n",
    "print(t1(2) == td_str_1(2)*td_str_1(2))\n",
    "print(\"str  updated:\",(t1.ops[0][2]))\n",
    "\n",
    "t1 = td_array_1.apply(square_qobj).apply_decorator(square_f)\n",
    "print(t1(2) == td_array_1(2)*td_array_1(2))\n",
    "print(\"array not updated:\",(t1.ops[0][2]))\n",
    "\n",
    "t1 = td_array_1.apply(square_qobj).apply_decorator(square_f, inplace_np=1)\n",
    "print(t1(2) == td_array_1(2)*td_array_1(2))\n",
    "print(\"array  updated:\",(t1.ops[0][2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Removing redundant Qobj\n",
    "\n",
    "When multiple components of the QobjEvo are made from the same Qobj, you can unite them with the \"compress\" method.\n",
    "It is only done with the same form of time dependance:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "small = qt.destroy(2)\n",
    "def f1(t,args):\n",
    "    return np.sin(t)\n",
    "def f2(t,args):\n",
    "    return np.cos(args[\"w\"]*t)\n",
    "def f3(t,args):\n",
    "    return np.sin(args[\"w\"]*t)\n",
    "def f4(t,args):\n",
    "    return np.cos(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = True\n",
      "Qobj data =\n",
      "[[1. 0.]\n",
      " [0. 1.]], [Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = False\n",
      "Qobj data =\n",
      "[[0. 1.]\n",
      " [0. 0.]], '(sin(t)) + (cos(w*t)) + (sin(w*t)) + (cos(t))']]\n",
      "1\n",
      "True\n",
      "1\n",
      "True\n",
      "1\n",
      "True\n",
      "3\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "td_redoundance = QobjEvo([qt.qeye(2),[small,\"sin(t)\"],[small,\"cos(w*t)\"],[small,\"sin(w*t)\"],\n",
    "                                  [small,\"cos(t)\"]],args={'w':2.})\n",
    "td_redoundance1 = QobjEvo([qt.qeye(2),[small,\"sin(t)\"],[small,\"cos(w*t)\"],[small,\"sin(w*t)\"],\n",
    "                                   [small,\"cos(t)\"]],args={'w':2.})\n",
    "td_redoundance2 = QobjEvo([qt.qeye(2),[small,f1],[small,f2],[small,f3],[small,f4]],args={'w':2.})\n",
    "td_redoundance3 = QobjEvo([qt.qeye(2),[small,np.sin(tlist)],[small,np.cos(2*tlist)],\n",
    "                                   [small,np.sin(2*tlist)],[small,np.cos(tlist)]],tlist=tlist)\n",
    "td_redoundance4 = QobjEvo([qt.qeye(2),[small,f1],[small,\"cos(w*t)\"],\n",
    "                                   [small,np.sin(2*tlist)],[small,\"cos(t)\"]],args={'w':2.},tlist=tlist)\n",
    "\n",
    "td_redoundance1.compress()\n",
    "print(td_redoundance1.to_list())\n",
    "print(len(td_redoundance1.ops))\n",
    "print(td_redoundance(1.) == td_redoundance1(1.))\n",
    "td_redoundance2.compress()\n",
    "print(len(td_redoundance2.ops))\n",
    "print(td_redoundance(1.) == td_redoundance2(1.))\n",
    "td_redoundance3.compress()\n",
    "print(len(td_redoundance3.ops))\n",
    "print(td_redoundance(1.) == td_redoundance3(1.))\n",
    "td_redoundance4.compress()\n",
    "print(len(td_redoundance4.ops))\n",
    "print(td_redoundance(1.) == td_redoundance4(1.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = True\n",
       " Qobj data =\n",
       " [[1. 0.]\n",
       "  [0. 1.]],\n",
       " [Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = True\n",
       "  Qobj data =\n",
       "  [[1. 1.]\n",
       "   [1. 1.]], 'sin(t)']]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "td_redoundance_v2 = QobjEvo([qt.qeye(2),[qt.qeye(2),\"sin(t)\"],[small,\"sin(t)\"],[qt.create(2),\"sin(t)\"]])\n",
    "td_redoundance_v2.compress()\n",
    "td_redoundance_v2.to_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## cimport cython object\n",
    "\n",
    "The cython object can be 'cimported'. \n",
    "\n",
    "```\n",
    "cdef class CQobjEvo:\n",
    "    cdef void _mul_vec(self, double t, complex* vec, complex* out)\n",
    "    cdef void _mul_matf(self, double t, complex* mat, complex* out,int nrow, int ncols)\n",
    "    cdef void _mul_matc(self, double t, complex* mat, complex* out,int nrow, int ncols)\n",
    "    cdef complex _expect(self, double t, complex* vec, int isherm)\n",
    "    cdef complex _expect_super(self, double t, complex* rho, int isherm)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from qutip.cy.cqobjevo cimport CQobjEvo\n",
    "cimport numpy as np\n",
    "\n",
    "\n",
    "def rhs_call_from_cy(CQobjEvo qobj, double t, np.ndarray[complex, ndim=1] vec, np.ndarray[complex, ndim=1] out):\n",
    "    qobj._mul_vec(t,&vec[0],&out[0])\n",
    "    \n",
    "    \n",
    "def expect_call_from_cy(CQobjEvo qobj, double t, np.ndarray[complex, ndim=1] vec, int isherm):\n",
    "    return qobj._expect(t,&vec[0])\n",
    "    \n",
    "    \n",
    "def rhs_cdef_timing(CQobjEvo qobj, double t, np.ndarray[complex, ndim=1] vec, np.ndarray[complex, ndim=1] out):\n",
    "    cdef int i\n",
    "    for i in range(10000):\n",
    "        qobj._mul_vec(t,&vec[0],&out[0])\n",
    "\n",
    "        \n",
    "def expect_cdef_timing(CQobjEvo qobj, double t, np.ndarray[complex, ndim=1] vec, int isherm):\n",
    "    cdef complex aa = 0.\n",
    "    cdef int i\n",
    "    for i in range(10000):\n",
    "        aa = qobj._expect(t, &vec[0])\n",
    "    return aa\n",
    "\n",
    "def rhs_def_timing(qobj, double t, np.ndarray[complex, ndim=1] vec, complex[::1] out):\n",
    "    cdef int i\n",
    "    for i in range(10000):\n",
    "        out = qobj.mul_vec(t,vec)\n",
    "        \n",
    "def expect_def_timing(qobj, double t, np.ndarray[complex, ndim=1] vec, int isherm):\n",
    "    cdef complex aa = 0.\n",
    "    cdef int i\n",
    "    for i in range(10000):\n",
    "        aa = qobj.expect(t, vec)\n",
    "    return aa\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0j\n",
      "56.1 ms ± 447 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
      "43.4 ms ± 365 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "td_str.compile()\n",
    "print(expect_call_from_cy(td_str.compiled_qobjevo, 2, vec, 0) - td_str.expect(2,vec,0))\n",
    "%timeit expect_def_timing(td_str.compiled_qobjevo, 2, vec, 0)\n",
    "%timeit expect_cdef_timing(td_str.compiled_qobjevo, 2, vec, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0j, 0j, 0j, 0j, 0j]\n",
      "64.6 ms ± 2.26 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
      "11 ms ± 24.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "out = np.zeros(N,dtype=np.complex128)\n",
    "rhs_call_from_cy(td_str.compiled_qobjevo, 2, vec, out)\n",
    "print( [a - b for a,b in zip(out, td_str.mul_vec(2,vec))])\n",
    "%timeit rhs_def_timing(td_str.compiled_qobjevo, 2, vec, out)\n",
    "%timeit rhs_cdef_timing(td_str.compiled_qobjevo, 2, vec, out)\n",
    "# Most of the time gained is from allocating the out vector, not the "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0j, 0j, 0j, 0j, 0j]\n",
      "39.7 ms ± 209 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
      "635 µs ± 1.07 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "td_cte = QobjEvo([Id])\n",
    "td_cte.compile()\n",
    "out = np.zeros(N,dtype=np.complex128)\n",
    "rhs_call_from_cy(td_cte.compiled_qobjevo, 2, vec, out)\n",
    "print( [a - b for a,b in zip(out, td_cte.mul_vec(2,vec))])\n",
    "%timeit rhs_def_timing(td_cte.compiled_qobjevo, 2, vec, out)\n",
    "%timeit rhs_cdef_timing(td_cte.compiled_qobjevo, 2, vec, out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compiled string code\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#!python\n",
      "#cython: language_level=3\n",
      "# This file is generated automatically by QuTiP.\n",
      "\n",
      "import numpy as np\n",
      "cimport numpy as np\n",
      "import scipy.special as spe\n",
      "cimport cython\n",
      "np.import_array()\n",
      "cdef extern from \"numpy/arrayobject.h\" nogil:\n",
      "    void PyDataMem_NEW_ZEROED(size_t size, size_t elsize)\n",
      "    void PyArray_ENABLEFLAGS(np.ndarray arr, int flags)\n",
      "from qutip.cy.spmatfuncs cimport spmvpy\n",
      "from qutip.cy.inter cimport _spline_complex_t_second, _spline_complex_cte_second\n",
      "from qutip.cy.inter cimport _spline_float_t_second, _spline_float_cte_second\n",
      "from qutip.cy.interpolate cimport (interp, zinterp)\n",
      "from qutip.cy.cqobjevo_factor cimport StrCoeff\n",
      "from qutip.cy.cqobjevo cimport CQobjEvo\n",
      "from qutip.cy.math cimport erf, zerf\n",
      "from qutip.qobj import Qobj\n",
      "cdef double pi = 3.14159265358979323\n",
      "\n",
      "include '/home/eric/anaconda3/lib/python3.7/site-packages/qutip-4.4.0.dev0+8b414dd-py3.7-linux-x86_64.egg/qutip/cy/complex_math.pxi'\n",
      "\n",
      "cdef class CompiledStrCoeff(StrCoeff):\n",
      "    cdef double w\n",
      "\n",
      "    def set_args(self, args):\n",
      "        self.w=args['w']\n",
      "\n",
      "    cdef void _call_core(self, double t, complex * coeff):\n",
      "        cdef double w = self.w\n",
      "\n",
      "        coeff[0] = exp(-1j*t)\n",
      "        coeff[1] = cos(w*t)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "td_str.compiled = False\n",
    "print(td_str.compile(code=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "QuTiP: Quantum Toolbox in Python\n",
      "Copyright (c) 2011 and later.\n",
      "A. J. Pitchford, P. D. Nation, R. J. Johansson, A. Grimsmo, and C. Granade\n",
      "\n",
      "QuTiP Version:      4.4.0.dev0+8b414dd\n",
      "Numpy Version:      1.16.2\n",
      "Scipy Version:      1.2.1\n",
      "Cython Version:     0.29.6\n",
      "Matplotlib Version: 3.0.1\n",
      "Python Version:     3.7.0\n",
      "Number of CPUs:     4\n",
      "BLAS Info:          OPENBLAS\n",
      "OPENMP Installed:   False\n",
      "INTEL MKL Ext:      False\n",
      "Platform Info:      Linux (x86_64)\n",
      "Installation path:  /home/eric/anaconda3/lib/python3.7/site-packages/qutip-4.4.0.dev0+8b414dd-py3.7-linux-x86_64.egg/qutip\n",
      "==============================================================================\n",
      "Please cite QuTiP in your publication.\n",
      "==============================================================================\n",
      "For your convenience a bibtex reference can be easily generated using `qutip.cite()`\n"
     ]
    }
   ],
   "source": [
    "qt.about()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}