Merge pull request #27 from stanfordLINQS/dev-tq-phi

Add phi derivative support

SQcircuit/circuit.py

@@ -136,6 +136,9 @@ def update_circuit_loop_from_element(
             loop.add_index(B_idx)
             loop.addK1(self.w)
 
+            if get_optim_mode():
+                self.circ.add_to_parameters(loop)
+
     def process_edge_and_update_circ(
         self,
         B_idx: int,
@@ -412,18 +415,29 @@ def safecopy(self, save_eigs=False):
         # Instantiate new container
         new_circuit = copy(self)
 
-        # Explicitly copy any non-leaf tensors
-        # (these don't implement a __deepcopy__ method)
+        # Explicitly copy all elements to new circuit, in particular those with
+        # non-leaf `.internal_value`s ((these don't implement a `._deepcopy__`
+        # method)
         if get_optim_mode():
             new_circuit.C = new_circuit.C.detach()
             new_circuit.L = new_circuit.L.detach()
 
+            new_loops: List[Loop] = []
+            replacement_dict: Dict[Union[Loop, Element], Union[Loop, Element]] = {}
+            for loop in self.loops:
+                new_loop = copy(loop)
+                new_loop.internal_value = loop.internal_value.detach().clone()
+                new_loops.append(new_loop)
+                replacement_dict[loop] = new_loop
+            new_circuit.loops = new_loops
+
             new_elements = defaultdict(list)
-            replacement_dict = dict()
             for edge in self.elements:
                 for el in self.elements[edge]:
                     new_el = copy(el)
                     new_el.internal_value = el.internal_value.detach().clone()
+                    if hasattr(el, 'loop'):
+                        new_el.loop = replacement_dict[el.loop]
                     new_elements[edge].append(new_el)
 
                     replacement_dict[el] = new_el
@@ -544,7 +558,7 @@ def add_to_parameters(self, el: Element) -> None:
         """Add elements with ``requires_grad=True`` to parameters.
         """
 
-        if el.requires_grad:
+        if el.requires_grad and el not in self._parameters:
             self._parameters[el] = el.internal_value
 
     def add_loop(self, loop: Loop) -> None:
@@ -1059,7 +1073,7 @@ def description(
 
         loopTxt = txt.loops() + txt.tab()
         for i in range(len(self.loops)):
-            phiExt = self.loops[i].value() / 2 / np.pi
+            phiExt = sqf.numpy(self.loops[i].value()) / 2 / np.pi
             loopTxt += txt.phiExt(i + 1) + txt.tPi() + txt.eq() + str(
                 phiExt) + txt.tab()
 
@@ -1473,7 +1487,11 @@ def _get_LC_hamil(self) -> Qobj:
 
         return LC_hamil
 
-    def _get_external_flux_at_element(self, B_idx: int) -> float:
+    def _get_external_flux_at_element(
+            self, 
+            B_idx: int, 
+            torch = False
+    ) -> Union[float, Tensor]:
         """
         Return the external flux at an inductive element.
 
@@ -1483,11 +1501,15 @@ def _get_external_flux_at_element(self, B_idx: int) -> float:
                 An integer point to each row of B matrix (external flux
                 distribution of that element)
         """
-        phi_ext = 0.0
+        phi_ext = sqf.zero()
         for i, loop in enumerate(self.loops):
+            print(loop.value())
             phi_ext += loop.value() * self.B[B_idx, i]
 
-        return phi_ext
+        if isinstance(phi_ext, Tensor) and not torch:
+            return sqf.numpy(phi_ext)
+        else:
+            return phi_ext
 
     def _get_inductive_hamil(self) -> Qobj:
 
@@ -1572,6 +1594,37 @@ def charge_op(self, mode: int, basis: str = 'FC') -> Qobj:
 
             return qt.Qobj(Q_eig)
 
+    def _get_W_idx(self, my_el: Junction, my_B_idx: int) -> Optional[int]:
+        for _, el, B_idx, W_idx in self.elem_keys[Junction]:
+            if el == my_el and B_idx == my_B_idx:
+                return W_idx
+
+        return None
+
+    def op(self, name: str, keywords: Dict):
+        """
+        Returns the `name` operator of the circuit, specified by `keywords`,
+        as a sparse torch matrix.
+        """
+        if name == 'sin_half':
+            B_idx = keywords['B_idx']
+            el = keywords['el']
+            if get_optim_mode():
+                W_idx = self._get_W_idx(el, B_idx)
+
+                phi = self._get_external_flux_at_element(B_idx, torch=True)
+                root_exp = (
+                    torch.exp(1j * phi / 2)
+                    * sqf.qobj_to_tensor(self._memory_ops["root_exp"][W_idx])
+                )
+
+                sin_half = (root_exp - sqf.dag(root_exp)) / 2j
+                return sin_half ## TO CHECK: need to squeeze?
+            else:
+                return self._memory_ops['sin_half'][el, B_idx]
+        else:
+            raise NotImplementedError
+
     def diag_np(
         self,
         n_eig: int
@@ -2118,8 +2171,8 @@ def dec_rate(
                             sqf.operator_inner_product(state1, op, state2)) ** 2
 
         if dec_type == "quasiparticle":
-            for el, _ in self._memory_ops['sin_half']:
-                op = self._memory_ops['sin_half'][(el, _)]
+            for el, B_idx in self._memory_ops['sin_half']:
+                op = self.op('sin_half', {'el': el, 'B_idx': B_idx})
                 if np.isnan(sqf.numpy(el.Y(omega, ENV["T"]))):
                     decay = decay + 0
                 else:
@@ -2334,7 +2387,7 @@ def _get_partial_omega_mn(
         partial_omega_n = sqf.operator_inner_product(state_n, partial_H, state_n)
 
         partial_omega_mn = partial_omega_m -  partial_omega_n
-        assert sqf.imag(partial_omega_mn)/sqf.real(partial_omega_mn) < 1e-6
+        # assert sqf.imag(partial_omega_mn)/sqf.real(partial_omega_mn) < 1e-6
 
         return sqf.real(partial_omega_mn)
 
@@ -2502,4 +2555,4 @@ def get_all_circuit_elements(self):
         def flatten(l):
             return [item for sublist in l for item in sublist]
         elements = flatten(list(self.elements.values()))
-        return elements
+        return elements

SQcircuit/elements.py

@@ -626,10 +626,15 @@ def __init__(
         self,
         value: float = 0,
         A: float = 1e-6,
+        requires_grad: bool = False,
         id_str: Optional[str] = None
     ) -> None:
 
-        self.lpValue = value * 2 * np.pi
+        self.set_flux(value)
+
+        if requires_grad:
+            self.requires_grad = requires_grad
+
         self.A = A * 2 * np.pi
         # indices of inductive elements.
         self.indices = []
@@ -669,6 +674,9 @@ def set_flux(self, value: float) -> None:
             value:
                 The external flux value
         """
+        if get_optim_mode():
+            value = torch.as_tensor(value)
+
         self.lpValue = value * 2 * np.pi
 
     def add_index(self, index):
@@ -690,6 +698,27 @@ def getP(self):
         p = np.linalg.inv(np.concatenate((b, K1.T), axis=0)) @ b.T
         return p.T
 
+    @property
+    def requires_grad(self) -> bool:
+        raise_optim_error_if_needed()
+
+        return self.lpValue.requires_grad
+
+    @requires_grad.setter
+    def requires_grad(self, f: bool) -> None:
+
+        raise_optim_error_if_needed()
+
+        self.lpValue.requires_grad = f
+
+    @property
+    def internal_value(self) -> Union[float, Tensor]:
+        return self.lpValue
+
+    @internal_value.setter
+    def internal_value(self, v: Union[float, Tensor]) -> None:
+        self.lpValue = v
+
 
 class Charge:
     """Class that contains the charge island properties.

SQcircuit/functions.py

@@ -105,6 +105,11 @@ def init_op(size):
         # return torch.sparse_coo_tensor(size = size, dtype=torch.complex128)
     return qt.Qobj()
 
+def zero(dtype=torch.complex128):
+    if get_optim_mode():
+        return torch.tensor(0, dtype=dtype)
+    return 0
+
 
 def zeros(shape, dtype=torch.complex128):
     if get_optim_mode():

SQcircuit/tests/conftest.py

@@ -25,6 +25,7 @@ def setup_class(cls):
 
     def test_transform_process(self):
         # load the data
+        set_optim_mode(False)
         data = SQdata.load(DATADIR + "/" + self.fileName)
 
         # build the new circuit based on data circuit parameters
@@ -103,6 +104,36 @@ def create_transmon_torch(trunc_num):
     return circuit_torch
 
 
+def create_flux_transmon_numpy(trunc_num, phi_ext):
+    cap_value, ind_value, Q = 7.746, 5, 1e6
+    cap_unit, ind_unit = 'fF', 'GHz'
+    disorder = 1.1
+
+    set_optim_mode(False)
+    loop = Loop(phi_ext)
+    C_numpy = Capacitor(cap_value, cap_unit, Q=Q)
+    J1_numpy = Junction(ind_value, ind_unit, loops=[loop])
+    J2_numpy = Junction(ind_value * disorder, ind_unit, loops=[loop])
+    circuit_numpy = Circuit({(0, 1): [C_numpy, J1_numpy, J2_numpy], })
+    circuit_numpy.set_trunc_nums([trunc_num, ])
+    return circuit_numpy
+
+
+def create_flux_transmon_torch(trunc_num, phi_ext):
+    cap_value, ind_value, Q = 7.746, 5, 1e6
+    cap_unit, ind_unit = 'fF', 'GHz'
+    disorder = 1.1
+
+    set_optim_mode(True)
+    loop_torch = Loop(phi_ext, requires_grad=True)
+    C_torch = Capacitor(cap_value, cap_unit, Q=Q, requires_grad=True)
+    J1_torch = Junction(ind_value, ind_unit, loops=[loop_torch], requires_grad=True)
+    J2_torch = Junction(ind_value * disorder, ind_unit, loops=[loop_torch], requires_grad=True)
+    circuit_torch = Circuit({(0, 1): [C_torch, J1_torch, J2_torch], })
+    circuit_torch.set_trunc_nums([trunc_num, ])
+    return circuit_torch
+
+
 def create_fluxonium_numpy(trunc_num, phi_ext=0):
     set_optim_mode(False)
     loop = Loop(phi_ext)
@@ -123,3 +154,60 @@ def create_fluxonium_torch(trunc_num, phi_ext=0):
     circuit_torch = Circuit({(0, 1): [C_torch, L_torch, JJ_torch], }, flux_dist='junctions')
     circuit_torch.set_trunc_nums([trunc_num, ])
     return circuit_torch
+
+
+def create_fluxonium_torch_flux(trunc_num, phi_ext=0):
+    set_optim_mode(True)
+    loop = Loop(phi_ext, requires_grad=True)
+    C_torch = Capacitor(3.6, 'GHz', Q=1e6, requires_grad=True)
+    L_torch = Inductor(0.46, 'GHz', Q=500e6, loops=[loop], requires_grad=True)
+    JJ_torch = Junction(10.2, 'GHz', A=1e-7, x=3e-06, loops=[loop], requires_grad=True)
+    circuit_torch = Circuit({(0, 1): [C_torch, L_torch, JJ_torch], }, flux_dist='junctions')
+    circuit_torch.set_trunc_nums([trunc_num, ])
+    return circuit_torch
+
+def create_JJL_numpy(trunc_num, phi_ext):
+    set_optim_mode(False)
+    Cs = [3.6, 1.2, 2.0]
+    Js = [10.2, 6.8]
+    Ls = [0.46]
+    loop = Loop(phi_ext)
+    C1 = Capacitor(Cs[0], 'GHz', requires_grad=False)
+    C2 = Capacitor(Cs[1], 'GHz', requires_grad=False)
+    C3 = Capacitor(Cs[2], 'GHz', requires_grad=False)
+    L = Inductor(Ls[0], 'GHz', loops=[loop], requires_grad=False)
+    J1 = Junction(Js[0], 'GHz',loops=[loop], requires_grad=False)
+    J2 = Junction(Js[1], 'GHz',loops=[loop], requires_grad=False)
+
+    circuit_numpy = Circuit(
+        {(0, 1): [C1, J1],
+         (1, 2): [C2, J2],
+         (2, 0): [L, C3]
+         },
+        flux_dist='junctions'
+    )
+    circuit_numpy.set_trunc_nums([trunc_num, trunc_num])
+    return circuit_numpy
+
+def create_JJL_torch(trunc_num, phi_ext):
+    set_optim_mode(True)
+    Cs = [3.6, 1.2, 2.0]
+    Js = [10.2, 6.8]
+    Ls = [0.46]
+    loop = Loop(phi_ext, requires_grad=False)
+    C1 = Capacitor(Cs[0], 'GHz', requires_grad=True)
+    C2 = Capacitor(Cs[1], 'GHz', requires_grad=True)
+    C3 = Capacitor(Cs[2], 'GHz', requires_grad=True)
+    L = Inductor(Ls[0], 'GHz', loops=[loop], requires_grad=True)
+    J1 = Junction(Js[0], 'GHz',loops=[loop], requires_grad=True)
+    J2 = Junction(Js[1], 'GHz',loops=[loop], requires_grad=True)
+
+    circuit_torch = Circuit(
+        {(0, 1): [C1, J1],
+         (1, 2): [C2, J2],
+         (2, 0): [L, C3]
+         },
+        flux_dist='junctions'
+    )
+    circuit_torch.set_trunc_nums([trunc_num, trunc_num])
+    return circuit_torch