separated out the nonlinearity

EQuS · Jul 31, 2024 · 195a675 · 195a675
1 parent dbfc5d0
commit 195a675
Show file tree

Hide file tree

Showing 5 changed files with 175 additions and 45 deletions.
diff --git a/qcsys/devices/ats.py b/qcsys/devices/ats.py
@@ -45,24 +45,51 @@ def get_H_linear(self):
         return w*(self.linear_ops["a_dag"]@self.linear_ops["a"] + 0.5 * self.linear_ops["id"])
 
 
-    def get_H_nonlinear(self, phi_op, id_op):
+    def get_H_nonlinear(self, phi_op):
         """Return nonlinear terms in H."""
 
         Ej = self.params["Ej"]
         dEj = self.params["dEj"]
         Ej2 = self.params["Ej2"]
+
+        phi_sum = self.params["phi_sum_ext"]
+        phi_delta = self.params["phi_delta_ext"]
+
+        cos_phi_op = jqt.cosm(phi_op)
+        sin_phi_op = jqt.sinm(phi_op)
+
+        cos_2phi_op = cos_phi_op @ cos_phi_op - sin_phi_op @ sin_phi_op
+        sin_2phi_op = 2 * cos_phi_op @ sin_phi_op
+
+        H_nl_Ej = - 2 * Ej * (
+            cos_phi_op * jnp.cos(2 * jnp.pi * phi_delta) 
+            - sin_phi_op * jnp.sin(2 * jnp.pi * phi_delta)
+        ) * jnp.cos(2 * jnp.pi * phi_sum)
+        H_nl_dEj = 2 * dEj * (
+            sin_phi_op * jnp.cos(2 * jnp.pi * phi_delta)
+            + cos_phi_op * jnp.sin(2 * jnp.pi * phi_delta)
+        ) * jnp.sin(2 * jnp.pi * phi_sum) 
+        H_nl_Ej2 = 2 * Ej2 * (
+            cos_2phi_op * jnp.cos(2 * 2 * jnp.pi * phi_delta)
+            - sin_2phi_op * jnp.sin(2 * 2 * jnp.pi * phi_delta)
+        ) * jnp.cos(2 * 2 * jnp.pi * phi_sum)
+
+        H_nl = H_nl_Ej + H_nl_dEj + H_nl_Ej2
+
+
+        # id_op = jqt.identity_like(phi_op)
+        # phi_delta_ext_op = self.params["phi_delta_ext"] * id_op
+        # H_nl_old = - 2 * Ej * jqt.cosm(phi_op + 2 * jnp.pi * phi_delta_ext_op) * jnp.cos(2 * jnp.pi * self.params["phi_sum_ext"])
+        # H_nl_old += 2 * dEj * jqt.sinm(phi_op + 2 * jnp.pi * phi_delta_ext_op) * jnp.sin(2 * jnp.pi * self.params["phi_sum_ext"]) 
+        # H_nl_old += 2 * Ej2 * jqt.cosm(2*phi_op + 2 * 2 * jnp.pi * phi_delta_ext_op) * jnp.cos(2 * 2 * jnp.pi * self.params["phi_sum_ext"])
 
-        phi_delta_ext_op = self.params["phi_delta_ext"] * id_op
-        H_nl = - 2 * Ej * jqt.cosm(phi_op + 2 * jnp.pi * phi_delta_ext_op) * jnp.cos(2 * jnp.pi * self.params["phi_sum_ext"])
-        H_nl += 2 * dEj * jqt.sinm(phi_op + 2 * jnp.pi * phi_delta_ext_op) * jnp.sin(2 * jnp.pi * self.params["phi_sum_ext"]) 
-        H_nl += 2 * Ej2 * jqt.cosm(2*phi_op + 2 * 2 * jnp.pi * phi_delta_ext_op) * jnp.cos(2 * 2 * jnp.pi * self.params["phi_sum_ext"])
         return H_nl
 
     def get_H_full(self):
         """Return full H in linear basis."""
         id_op = self.linear_ops["id"]
         phi_b = self.linear_ops["phi"]
-        H_nl = self.get_H_nonlinear(phi_b, id_op)
+        H_nl = self.get_H_nonlinear(phi_b)
         H = self.get_H_linear() + H_nl
         return H
 

diff --git a/qcsys/devices/fluxonium.py b/qcsys/devices/fluxonium.py
@@ -54,16 +54,20 @@ def get_H_linear(self):
 
     def get_H_full(self):
         """Return full H in linear basis."""
-        op_cos_phi = jqt.cosm(self.linear_ops["phi"])
-        op_sin_phi = jqt.sinm(self.linear_ops["phi"])
+
+        phi_op = self.linear_ops["phi"]
+        return self.get_H_linear() + self.get_H_nonlinear(phi_op)
+
+    def get_H_nonlinear(self, phi_op):
+        op_cos_phi = jqt.cosm(phi_op)
+        op_sin_phi = jqt.sinm(phi_op)
 
         phi_ext = self.params["phi_ext"]
         Hcos = op_cos_phi * jnp.cos(2.0 * jnp.pi * phi_ext) + op_sin_phi * jnp.sin(
             2.0 * jnp.pi * phi_ext
         )
-
-        H = self.get_H_linear() - self.params["Ej"] * Hcos
-        return H
+        H_nl = - self.params["Ej"] * Hcos
+        return H_nl
 
     def potential(self, phi):
         """Return potential energy for a given phi."""

diff --git a/tutorials/3-fluxonium.ipynb b/tutorials/3-fluxonium.ipynb
diff --git a/tutorials/ats-coupler/0-resonator-ats-resonator.ipynb b/tutorials/ats-coupler/0-resonator-ats-resonator.ipynb
@@ -14,7 +14,16 @@
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/shanj/miniconda3/envs/jax-new/lib/python3.12/site-packages/qutip/__init__.py:66: UserWarning: The new version of Cython, (>= 3.0.0) is not supported.\n",
+      "  warnings.warn(\n"
+     ]
+    }
+   ],
    "source": [
     "from jax import jit, grad\n",
     "import qcsys as qs\n",
@@ -33,21 +42,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "{'g_3': Array(1.28654004e-05+0.j, dtype=complex128),\n",
-       " 'g_cd': Array(2.57308009e-05+0.j, dtype=complex128),\n",
+       "{'g_3': Array(-1.28654004e-05+0.j, dtype=complex128),\n",
+       " 'g_cd': Array(-2.57308009e-05+0.j, dtype=complex128),\n",
        " 'g_ex': Array(-0.00212182+0.j, dtype=complex128),\n",
        " 'ω_ATS': Array(6.02181564, dtype=float64),\n",
        " 'ω_ResonatorA': Array(4.99584201, dtype=float64),\n",
        " 'ω_ResonatorB': Array(7.00151547, dtype=float64)}"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -58,6 +67,29 @@
     "metrics0"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Array([ 0.00000000e+00, -1.77635684e-06,  0.00000000e+00,  1.06581410e-05,\n",
+       "       -1.77635684e-05,  1.42108547e-05], dtype=float64)"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Calculate Kerr in full system\n",
+    "metrics, system, ϕ0, metrics0, system0 = get_metrics_normal_rar({})#{\"ResonatorA_frequency\": 3})\n",
+    "metrics[\"K_a\"]*1e9"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 6,
@@ -66,8 +98,8 @@
     {
      "data": {
       "text/plain": [
-       "Array([ 1.77635684e-06, -3.55271368e-06,  3.55271368e-06,  3.55271368e-06,\n",
-       "       -1.77635684e-05,  2.13162821e-05], dtype=float64)"
+       "Array([ 0.00000000e+00, -1.77635684e-06,  0.00000000e+00,  1.06581410e-05,\n",
+       "       -1.77635684e-05,  1.42108547e-05], dtype=float64)"
       ]
      },
      "execution_count": 6,
@@ -89,8 +121,8 @@
     {
      "data": {
       "text/plain": [
-       "Array([ 9.07340869e+00,  9.07284026e+00,  9.07218478e+00,  9.07082054e+00,\n",
-       "       -1.22036532e+01,  4.27481506e+05], dtype=float64)"
+       "Array([ 9.07343711e+00,  9.07287578e+00,  9.07212083e+00,  9.07078501e+00,\n",
+       "       -1.22035537e+01,  4.27481506e+05], dtype=float64)"
       ]
      },
      "execution_count": 7,
@@ -106,6 +138,31 @@
     "metrics[\"K_a\"]*1e9"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Array([ 9.07343711e+00,  9.07287578e+00,  9.07212083e+00,  9.07078501e+00,\n",
+       "       -1.22035537e+01,  4.27481506e+05], dtype=float64)"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Can override parameters!\n",
+    "metrics, system, ϕ0, metrics0, system0 = get_metrics_normal_rar({\n",
+    "    \"ATS__dE_J\": 37.0 * 0.01\n",
+    "})\n",
+    "metrics[\"K_a\"]*1e9"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -115,7 +172,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -124,7 +181,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -137,7 +194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -148,33 +205,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/latex": [
-       "Quantum object: dims=[[8, 10, 8], [8, 10, 8]], shape=(640, 640), type='oper', dtype=Dense, isherm=True$$\\left(\\begin{array}{cc}0 & -0.697 & 0 & 2.651\\times10^{ -5 } & 0 & \\cdots & 0 & 0 & 0 & 0 & 0\\\\-0.697 & 0 & -0.986 & 0 & 5.302\\times10^{ -5 } & \\cdots & 0 & 0 & 0 & 0 & 0\\\\0 & -0.986 & 0 & -1.207 & 0 & \\cdots & 0 & 0 & 0 & 0 & 0\\\\2.651\\times10^{ -5 } & 0 & -1.207 & 0 & -1.394 & \\cdots & 0 & 0 & 0 & 0 & 0\\\\0 & 5.302\\times10^{ -5 } & 0 & -1.394 & 0 & \\cdots & 0 & 0 & 0 & 0 & 0\\\\\\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\ddots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots\\\\0 & 0 & 0 & 0 & 0 & \\cdots & 0 & -1.139 & 0 & 9.688\\times10^{ -5 } & 0\\\\0 & 0 & 0 & 0 & 0 & \\cdots & -1.139 & 0 & -1.274 & 0 & 1.282\\times10^{ -4 }\\\\0 & 0 & 0 & 0 & 0 & \\cdots & 0 & -1.274 & 0 & -1.395 & 0\\\\0 & 0 & 0 & 0 & 0 & \\cdots & 9.688\\times10^{ -5 } & 0 & -1.395 & 0 & -1.507\\\\0 & 0 & 0 & 0 & 0 & \\cdots & 0 & 1.282\\times10^{ -4 } & 0 & -1.507 & 0\\end{array}\\right)$$"
+       "Quantum object: dims = [[8, 10, 8], [8, 10, 8]], shape = (640, 640), type = oper, isherm = True $ \\\\ \\left(\\begin{matrix}0.0 & 0.697 & 0.0 & -2.651\\times10^{-05} & 0.0 & \\cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\\\0.697 & 0.0 & 0.986 & 0.0 & -5.302\\times10^{-05} & \\cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\\\0.0 & 0.986 & 0.0 & 1.207 & 0.0 & \\cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\\\-2.651\\times10^{-05} & 0.0 & 1.207 & 0.0 & 1.394 & \\cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\\\0.0 & -5.302\\times10^{-05} & 0.0 & 1.394 & 0.0 & \\cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\\\\\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\ddots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots\\\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \\cdots & 0.0 & 1.139 & 0.0 & -9.688\\times10^{-05} & 0.0\\\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \\cdots & 1.139 & 0.0 & 1.274 & 0.0 & -1.282\\times10^{-04}\\\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \\cdots & 0.0 & 1.274 & 0.0 & 1.395 & 0.0\\\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \\cdots & -9.688\\times10^{-05} & 0.0 & 1.395 & 0.0 & 1.507\\\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \\cdots & 0.0 & -1.282\\times10^{-04} & 0.0 & 1.507 & 0.0\\\\\\end{matrix}\\right)$"
       ],
       "text/plain": [
-       "Quantum object: dims=[[8, 10, 8], [8, 10, 8]], shape=(640, 640), type='oper', dtype=Dense, isherm=True\n",
+       "Quantum object: dims = [[8, 10, 8], [8, 10, 8]], shape = (640, 640), type = oper, isherm = True\n",
        "Qobj data =\n",
-       "[[ 0.         -0.69714856  0.         ...  0.          0.\n",
-       "   0.        ]\n",
-       " [-0.69714856  0.         -0.98587103 ...  0.          0.\n",
-       "   0.        ]\n",
-       " [ 0.         -0.98587103  0.         ...  0.          0.\n",
-       "   0.        ]\n",
+       "[[0.         0.69714856 0.         ... 0.         0.         0.        ]\n",
+       " [0.69714856 0.         0.98587103 ... 0.         0.         0.        ]\n",
+       " [0.         0.98587103 0.         ... 0.         0.         0.        ]\n",
        " ...\n",
-       " [ 0.          0.          0.         ...  0.         -1.39529043\n",
-       "   0.        ]\n",
-       " [ 0.          0.          0.         ... -1.39529043  0.\n",
-       "  -1.50720291]\n",
-       " [ 0.          0.          0.         ...  0.         -1.50720291\n",
-       "   0.        ]]"
+       " [0.         0.         0.         ... 0.         1.39529043 0.        ]\n",
+       " [0.         0.         0.         ... 1.39529043 0.         1.50720291]\n",
+       " [0.         0.         0.         ... 0.         1.50720291 0.        ]]"
       ]
      },
-     "execution_count": 16,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -200,7 +251,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,

diff --git a/tutorials/ats-coupler/rar_helpers.py b/tutorials/ats-coupler/rar_helpers.py
@@ -249,8 +249,8 @@ def get_system_normal_rar(params):
 
     couplings = []
     coupling_term = 0 
-    coupling_term += ats.get_H_nonlinear(phi, id_op_c)
-    coupling_term -= ats.get_H_nonlinear(phi_c, id_op_c)
+    coupling_term += ats.get_H_nonlinear(phi)
+    coupling_term -= ats.get_H_nonlinear(phi_c)
     couplings.append(coupling_term)
 
     system = qs.System.create(devices, couplings=couplings)