From 8c3af6644a8b16aedfa3af3b3dc85211cc58b8ef Mon Sep 17 00:00:00 2001
From: Boxi Li <etamin1201@gmail.com>
Date: Sun, 5 Feb 2023 18:21:53 +0100
Subject: [PATCH] Add a helper for qobj compatibility

In qutip-v5, a bra ket multiplication returns a number instead of a Qobj. This makes it hard to keep the code clean and compatible for both version. This helper function changes the returned value so that it is identical to the one returned by qutip-v4.
---
 doc/pulse-paper/customize.py    |  3 ++-
 doc/pulse-paper/dj_algorithm.py |  6 +++---
 src/qutip_qip/utilities.py      | 16 ++++++++++++++++
 src/qutip_qip/vqa.py            |  7 +++----
 4 files changed, 24 insertions(+), 8 deletions(-)
 create mode 100644 src/qutip_qip/utilities.py

diff --git a/doc/pulse-paper/customize.py b/doc/pulse-paper/customize.py
index 90c3dcd3..2002f6f4 100644
--- a/doc/pulse-paper/customize.py
+++ b/doc/pulse-paper/customize.py
@@ -15,7 +15,8 @@
 import numpy as np
 from qutip import (
     fidelity, sigmax, sigmay, sigmaz, basis, qeye, tensor, Qobj, fock_dm)
-from qutip_qip.circuit import QubitCircuit, Gate
+from qutip_qip.circuit import QubitCircuit
+from qutip_qip.operations import Gate
 from qutip_qip.device import ModelProcessor, Model
 from qutip_qip.compiler import GateCompiler, Instruction
 import qutip
diff --git a/doc/pulse-paper/dj_algorithm.py b/doc/pulse-paper/dj_algorithm.py
index 190e1c3c..77a46c9d 100644
--- a/doc/pulse-paper/dj_algorithm.py
+++ b/doc/pulse-paper/dj_algorithm.py
@@ -18,6 +18,8 @@
     OptPulseProcessor, LinearSpinChain, SCQubits, SpinChainModel)
 from qutip_qip.circuit import QubitCircuit
 from qutip import sigmaz, sigmax, identity, tensor, basis
+from qutip_qip.utilities import _qobj_compatibility_helper
+
 
 # Deutsch-Josza algorithm
 dj_circuit = QubitCircuit(num_qubits)
@@ -79,7 +81,6 @@
     ax[i].fill_between([full_tlist[point2], full_tlist[point3]], [vmin ,vmin],  [vmax, vmax], color="lightgray", alpha=0.5)
     ax[i].vlines([full_tlist[point2], full_tlist[point3]], vmin, vmax, "gray", "--", linewidth=0.8, alpha=0.5)
 
-fig.tight_layout()
 fig.savefig("optimal_control_pulse.pdf")
 fig.show()
 
@@ -124,7 +125,6 @@
 
 full_tlist = scqubits_processor.get_full_tlist()
 
-fig3.tight_layout()
 fig3.savefig("transmon_pulse.pdf")
 fig3.show()
 
@@ -146,7 +146,7 @@
 for state in result1.states:
     tmp = state.ptrace([0,1])
     tmp = basis([2,2], [0,0]).dag() * tmp * basis([2,2], [0,0])
-    expect.append(np.real(tmp[0, 0]))
+    expect.append(np.real(_qobj_compatibility_helper(tmp)[0, 0]))
 
 fig5, ax5 = plt.subplots(figsize=(LINEWIDTH, LINEWIDTH*0.7), dpi=200)
 ax5.plot(t_record, expect, color="slategrey")
diff --git a/src/qutip_qip/utilities.py b/src/qutip_qip/utilities.py
new file mode 100644
index 00000000..fd32ac58
--- /dev/null
+++ b/src/qutip_qip/utilities.py
@@ -0,0 +1,16 @@
+from itertools import chain
+from functools import reduce
+from packaging.version import parse as parse_version
+import numpy as np
+import qutip
+
+
+def _qobj_compatibility_helper(qobj):
+    """
+    Make qobj generated by operations in qutip-v5 identical to the one by qutip-v4.
+    Here, if `qobj` is a number (e.g. generated by a bra-ket multiplication), it is replaced by a Qobj. 
+    """
+    if parse_version(qutip.__version__) >= parse_version("5.dev"):
+        if np.isscalar(qobj):
+            return qutip.Qobj([[qobj]], dims=[[1], [1]])
+    return qobj
diff --git a/src/qutip_qip/vqa.py b/src/qutip_qip/vqa.py
index 5b0fae57..95f2c260 100644
--- a/src/qutip_qip/vqa.py
+++ b/src/qutip_qip/vqa.py
@@ -7,6 +7,7 @@
 from scipy.optimize import minimize
 from scipy.linalg import expm_frechet
 from .operations import gate_sequence_product
+from .utilities import _qobj_compatibility_helper
 
 
 class VQA:
@@ -429,10 +430,8 @@ def cost_derivative(self, U, dU):
         dCost = (init.dag() * dU.dag()) * obs * (U * init) + (
             init.dag() * U.dag()
         ) * obs * (dU * init)
-        if isinstance(dCost, Qobj):  # qutip version < 5
-            return np.real(dCost.full().item())
-        else:  # qutip version >= 5
-            return np.real(dCost)
+        dCost = _qobj_compatibility_helper(dCost)
+        return np.real(dCost.full().item())
 
     def compute_jac(self, angles, indices_to_compute=None):
         """