added closest support for both qubit and cavity

squadds/calcs/__init__.py

@@ -0,0 +1,2 @@
+from .qubit import *
+from .transmon_cross import *

squadds/calcs/qubit.py

@@ -0,0 +1,72 @@
+from abc import ABC, abstractmethod
+import pandas as pd
+
+class QubitHamiltonian(ABC):
+    #TODO: make method names more general
+    def __init__(self, analysis):
+        self.analysis = analysis
+        self.db = self.analysis.db
+        self.df = self.analysis.df
+        self.qubit_type = self.db.selected_qubit
+        self.cavity_type = self.db.selected_cavity
+        self.target_param_keys = self.analysis.H_param_keys
+        self.target_params = self.analysis.target_params
+        self.f_q = None
+        self.alpha = None
+        self.Lj = None
+        self.EJEC = None
+        self.g = None
+
+    @abstractmethod
+    def plot_data(self, data_frame):
+        pass
+
+    @abstractmethod
+    def E01(self, EJ, EC):
+        pass
+
+    @abstractmethod
+    def E01_and_anharmonicity(self, EJ, EC):
+        pass
+
+    @abstractmethod
+    def EJ_and_LJ(self, f_q, alpha):
+        pass
+
+    @abstractmethod
+    def EJ(self, f_q, alpha):
+        pass
+
+    @abstractmethod
+    def EC(self, cross_to_claw, cross_to_ground):
+        pass
+
+    @abstractmethod
+    def calculate_target_quantities(self, f_res, alpha, g, w_q, N, Z_0=50):
+        pass
+
+    @abstractmethod
+    def g_and_alpha(self, C, C_c, f_q, EJ, f_r, res_type, Z0=50):
+        pass
+
+    @abstractmethod
+    def g_alpha_freq(self, C, C_c, EJ, f_r, res_type, Z0=50):
+        pass
+
+    @abstractmethod
+    def get_freq_alpha_fixed_LJ(self, fig4_df, LJ_target):
+        pass
+
+    @abstractmethod
+    def g_from_cap_matrix(self, C, C_c, EJ, f_r, res_type, Z0=50):
+        pass
+
+    @abstractmethod
+    def add_qubit_H_params(self):
+        pass
+
+    @abstractmethod
+    def add_cavity_coupled_H_params(self):
+        pass
+
+

squadds/calcs/transmon_cross.py

@@ -0,0 +1,117 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import scqubits as scq
+from squadds.calcs.qubit import QubitHamiltonian
+from scqubits.core.transmon import Transmon
+from pyEPR.calcs import Convert
+from scipy.constants import e, h, hbar
+
+class TransmonCrossHamiltonian(QubitHamiltonian):
+    #TODO: make method handling more general
+    def __init__(self, analysis):
+        super().__init__(analysis)
+        scq.set_units("GHz")
+
+    def plot_data(self, data_frame):
+        data_frame.plot(kind='box', subplots=True, layout=(1, 3), sharex=False, sharey=False)
+        plt.show()
+
+    def EC(self, cross_to_claw, cross_to_ground):
+        C_eff_fF = abs(cross_to_ground) + abs(cross_to_claw)
+        EC = Convert.Ec_from_Cs(C_eff_fF, units_in='fF',units_out='GHz')
+        return EC
+
+    def _calculate_target_qubit_params(self, w_q, alpha, Z_0=50):
+        EJ, EC = Transmon.find_EJ_EC(w_q, alpha)
+        EJEC = EJ / EC
+        Lj = Convert.Lj_from_Ej(EJ, units_in='GHz', units_out='nH')
+        self.EJ = EJ
+        self.EC = EC
+        self.EJEC = EJEC
+        self.Lj = Lj
+        return EJ, EC, EJEC, Lj
+
+    def EJ_and_LJ(self, w_q, alpha, *args, **kwargs):
+        EJ, EC = Transmon.find_EJ_EC(w_q, alpha)
+        Lj = Convert.Lj_from_Ej(EJ, units_in='GHz', units_out='nH')
+        self.EJ = EJ
+        self.Lj = Lj
+        return EJ, Lj
+
+    def EJ(self, w_q, alpha):
+        EJ, EC = Transmon.find_EJ_EC(w_q, alpha)
+        self.EJ = EJ
+        return EJ
+
+    def calculate_target_quantities(self, f_res, alpha, g, w_q, N, Z_0=50):
+        EJ, EC = Transmon.find_EJ_EC(w_q, alpha)
+        C_q = Convert.Cs_from_Ec(EC, units_in='GHz', units_out='fF')
+        omega_r = 2 * np.pi * f_res
+        prefactor = np.sqrt(N * Z_0 * e**2 / (hbar * np.pi)) * (EJ / (8 * EC))**(1/4)
+        denominator = omega_r * prefactor
+        numerator = g * C_q
+        C_c = numerator / denominator
+        EJ_EC_ratio = EJ / EC
+        return C_q, C_c, EJ, EC, EJ_EC_ratio
+
+    def g_and_alpha(self, C, C_c, f_q, EJ, f_r, res_type, Z0=50):
+        C, C_c = abs(C) * 1e-15, abs(C_c) * 1e-15
+        C_q = C + C_c
+        g = self.calculate_g_with_C(C, C_c, EJ, f_r, res_type, Z0)
+        EC = Convert.Ec_from_Cs(C_q, units_in='F', units_out='GHz')
+        transmon = Transmon(EJ=EJ, EC=EC, ng=0, ncut=30)
+        alpha = transmon.anharmonicity() * 1E3  # MHz
+        return g, alpha
+
+    def g_alpha_freq(self, C, C_c, EJ, f_r, res_type, Z0=50):
+        C, C_c = abs(C) * 1e-15, abs(C_c) * 1e-15
+        C_q = C + C_c
+        g = self.calculate_g_with_C(C, C_c, EJ, f_r, res_type, Z0)
+        EC = Convert.Ec_from_Cs(C_q, units_in='F', units_out='GHz')
+        transmon = Transmon(EJ=EJ, EC=EC, ng=0, ncut=30)
+        alpha = transmon.anharmonicity() * 1E3  # MHz
+        freq = transmon.E01()
+        return g, alpha, freq
+
+    def get_freq_alpha_fixed_LJ(self, fig4_df, LJ_target):
+        EJ = Convert.Ej_from_Lj(LJ_target, units_in='nH', units_out='GHz')
+        EC = fig4_df["EC"].values
+        transmons = [Transmon(EJ=EJ, EC=EC[i], ng=0, ncut=30) for i in range(len(EC))]
+        alpha = [transmon.anharmonicity() * 1E3 for transmon in transmons]
+        freq = [transmon.E01() for transmon in transmons]
+        return freq, alpha
+
+    def g_from_cap_matrix(self, C, C_c, EJ, f_r, res_type, Z0=50):
+        C = abs(C) * 1e-15  # F
+        C_c = abs(C_c) * 1e-15  # F
+        C_q = C_c + C
+        omega_r = 2 * np.pi * f_r * 1e9
+        EC = Convert.Ec_from_Cs(C_q, units_in='F', units_out='GHz')
+        g = (abs(C_c) / C_q) * omega_r * np.sqrt(res_type * Z0 * e**2 / (hbar * np.pi)) * (EJ / (8 * EC))**(1/4)
+        return (g * 1E-6) / (2 * np.pi)  # MHz
+
+    def E01_and_anharmonicity(self, EJ, EC, ng=0, ncut=30):
+        transmon = Transmon(EJ=EJ, EC=EC, ng=ng, ncut=ncut)
+        E01 = transmon.E01()
+        alpha = transmon.anharmonicity() * 1E3  # MHz
+        return E01, alpha
+
+    def E01(self, EJ, EC, ng=0, ncut=30):
+        transmon = Transmon(EJ=EJ, EC=EC, ng=ng, ncut=ncut)
+        E01 = transmon.E01()
+        return E01
+
+    def add_qubit_H_params(self):
+        EJ_target = self.EJ(self.target_params["qubit_frequency_GHz"], self.target_params["anharmonicity_MHz"]*1e-3)
+
+        self.df["EC"] = self.df.apply(lambda row: self.EC(row["cross_to_claw"], row["cross_to_ground"]), axis=1)
+
+        self.df['EJ'] = EJ_target
+        self.df["EJEC"] = self.df.apply(lambda row: row["EJ"] / row["EC"], axis=1)
+
+        self.df['qubit_frequency_GHz'], self.df['anharmonicity_MHz'] = zip(*self.df.apply(lambda row: self.E01_and_anharmonicity(row['EJ'], row['EC']), axis=1))
+
+
+    def add_cavity_coupled_H_params(self):
+        self.add_qubit_H_params()
+        self.df['g_MHz'] = self.df.apply(lambda row: self.compute_g_alpha_freq(row['C'], row['C_c'], row['EJ'], row['f_r'], row['res_type'], row['Z0'])[0], axis=1)

squadds/core/analysis.py

@@ -7,6 +7,7 @@
 from pandasai import SmartDataframe
 from pandasai.llm import OpenAI, Starcoder, Falcon
 from dotenv import load_dotenv
+from squadds.calcs import *
 
 from squadds.core.metrics import *
 from squadds.core.db import SQuADDS_DB
@@ -59,23 +60,61 @@ def __init__(self, db):
         self.selected_coupler = self.db.selected_coupler
         self.selected_system = self.db.selected_system
         self.df = self.db.selected_df
-        self.closest_designs = None
+        self.closest_design = None
         self.closest_df_entry = None
-        self.interpolate_design = None
+        self.interpolated_design = None
 
-        self.H_param_keys = self.db.target_param_keys
         self.metric_strategy = None  # Will be set dynamically
         self.custom_metric_func = None
         self.metric_weights = None
-        self.smart_df = None
-        self.coupling_type = None
+        self.target_params = None
+
+        self.H_param_keys = self._get_H_param_keys()
+
+    def _add_target_params_columns(self):
+        #TODO: make this more general and read the param keys from the database
+        if self.selected_system == "qubit":
+            qubit_H = TransmonCrossHamiltonian(self)
+            qubit_H.add_qubit_H_params()
+            self.df = qubit_H.df 
+        elif self.selected_system == "cavity_claw":
+            # rename the columns cavity_frequency_GHz and kappa_kHz and update the values
+            if ("cavity_frequency" in self.df.columns) or ("kappa" in self.df.columns):
+                self.df = self.df.rename(columns={"cavity_frequency": "cavity_frequency_GHz", "kappa": "kappa_kHz"})
+                self.df["cavity_frequency_GHz"] = self.df["cavity_frequency_GHz"] * 1e-9
+                self.df["kappa_kHz"] = self.df["kappa_kHz"] * 1e-3
+            else:
+                pass
+        elif self.selected_system == "coupler":
+            pass
+        elif (self.selected_system == ["qubit","cavity_claw"]) or (self.selected_system == ["cavity_claw","qubit"]):
+            self.db = self.add_qubit_params(self.db)
+            self.db = self.add_cavity_claw_params(self.db)
+            self.db = self.add_coupling_params(self.db) # adds g
+        else:
+            raise ValueError("Invalid system.")
+
+    def _get_H_param_keys(self):
+        #TODO: make this more general and read the param keys from the database
+        self.H_param_keys = None
+        if self.selected_system == "qubit":
+            self.H_param_keys = ["qubit_frequency_GHz", "anharmonicity_MHz"]
+        elif self.selected_system == "cavity_claw":
+            self.H_param_keys = ["resonator_type", "cavity_frequency_GHz", "kappa_kHz"]
+        elif self.selected_system == "coupler":
+            pass
+        elif (self.selected_system == ["qubit","cavity_claw"]) or (self.selected_system == ["cavity_claw","qubit"]):
+            self.H_param_keys = ["qubit_frequency_GHz", "anharmonicity_MHz", "resonator_type", "cavity_frequency_GHz", "kappa_kHz", "g_MHz"]
+        else:
+            raise ValueError("Invalid system.")
+        return self.H_param_keys
 
     def target_param_keys(self):
         """
         Returns:
             list: The target parameter keys.
         """
-        return self.db.target_param_keys
+        return self.H_param_keys
 
     def set_metric_strategy(self, strategy: MetricStrategy):
         """
@@ -102,6 +141,7 @@ def _outside_bounds(self, df: pd.DataFrame, params: dict, display=True) -> bool:
             bool: True if any value is outside of bounds. False if all values are inside bounds.
         """
         outside_bounds = False
+
         filtered_df = df.copy()
 
         for param, value in params.items():
@@ -144,6 +184,9 @@ def find_closest(self,
         if (num_top > len(self.df)):
             raise ValueError('`num_top` cannot be bigger than size of read-in library.')
 
+        self.target_params = target_params
+        self._add_target_params_columns()
+
         # Log if parameters outside of library
         filtered_df = self.df[self.H_param_keys]  # Filter DataFrame based on H_param_keys
         self._outside_bounds(df=filtered_df, params=target_params, display=display)
@@ -164,12 +207,23 @@ def find_closest(self,
             raise ValueError("Invalid metric.")
 
         # Main logic
+
+        # Filter DataFrame based on target parameters that are string
+        for param, value in target_params.items():
+            if isinstance(value, str):
+                filtered_df = filtered_df[filtered_df[param] == value]
+
+        # Calculate distances
         distances = filtered_df.apply(lambda row: self.metric_strategy.calculate(target_params, row), axis=1)
 
         # Sort distances and get the closest ones
         sorted_indices = distances.nsmallest(num_top).index
         closest_df = self.df.loc[sorted_indices]
 
+        # store the best design 
+        self.closest_df_entry = closest_df.iloc[0]
+        self.closest_design = closest_df.iloc[0]["design_options"]
+
         return closest_df
 
     def get_interpolated_design(self,

squadds/core/metrics.py

@@ -40,7 +40,7 @@ def calculate(self, target_params, df_row):
         distance = 0.0
         for column, target_value in target_params.items():
             if isinstance(target_value, (int, float)):  # Only numerical columns
-                distance += ((df_row[column] - target_value)**2 / target_value)
+                distance += ((df_row[column] - target_value)**2 / target_value**2)
         return np.sqrt(distance)
 
 class ManhattanMetric(MetricStrategy):