import time
from typing import Any
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from squadds.calcs.transmon_cross import TransmonCrossHamiltonian
from squadds.core.analysis_enrichment import (
extract_coupler_options,
extract_cpw_options,
extract_qubit_options,
fix_cavity_claw_dataframe,
)
from squadds.core.analysis_plotting import build_closest_design_hspace_plot
from squadds.core.analysis_search import (
SUPPORTED_METRICS,
filter_df_by_target_params,
get_H_param_keys_for_system,
outside_bounds,
rank_closest_indices,
remove_resonator_type_from_target_params,
resolve_metric_strategy,
)
from squadds.core.metrics import MetricStrategy
from squadds.core.processing import merge_dfs
from squadds.core.utils import create_unified_design_options
"""
=====================================================================================
HELPER FUNCTIONS
=====================================================================================
"""
# Helper function to scale values with 'um' in them
[docs]
def scale_value(value, ratio):
"""
Scales the given value by the specified ratio.
Args:
- value (str): The value to be scaled, in the format 'Xum' where X is a number.
- ratio (float): The scaling ratio.
Returns:
scaled_value (str): The scaled value in the format 'Xum' where X is the scaled number.
"""
scaled_value = str(float(value.replace("um", "")) * ratio) + "um"
return scaled_value
"""
=====================================================================================
Analyzer
=====================================================================================
"""
[docs]
class Analyzer:
"""
The Analyzer class is responsible for analyzing designs and finding the closest designs based on target parameters.
Methods:
_add_target_params_columns(): Adds target parameter columns to the dataframe based on the selected system.
_fix_cavity_claw_df(): Fixes the cavity claw DataFrame by renaming columns and updating values.
_get_H_param_keys(): Gets the parameter keys for the Hamiltonian based on the selected system.
target_param_keys(): Returns the target parameter keys.
set_metric_strategy(strategy: MetricStrategy): Sets the metric strategy to use for calculating the distance metric.
_outside_bounds(df: pd.DataFrame, params: dict, display=True) -> bool: Checks if entered parameters are outside the bounds of a dataframe.
find_closest(target_params: dict, num_top: int, metric: str = 'Euclidean', display: bool = True): Finds the closest designs in the library based on the target parameters.
get_interpolated_design(target_params: dict, metric: str = 'Euclidean', display: bool = True): Gets the interpolated design based on the target parameters.
get_design(df): Extracts the design parameters from the dataframe and returns a dict.
"""
__supported_metrics__ = SUPPORTED_METRICS
__supported_estimation_methods__ = ["Interpolation"]
def __init__(self, db=None):
"""
Initializes an instance of the Analysis class.
Parameters:
- db: The database object.
Attributes:
- db: The database object.
- selected_component_name: The name of the selected component.
- selected_component: The selected component.
- selected_data_type: The selected data type.
- selected_confg: The selected configuration.
- selected_qubit: The selected qubit.
- selected_cavity: The selected cavity.
- selected_coupler: The selected coupler.
- selected_system: The selected system.
- df: The selected dataframe.
- closest_df_entry: The closest dataframe entry.
- closest_design: The closest design.
- presimmed_closest_cpw_design: The presimmed closest CPW design.
- presimmed_closest_qubit_design: The presimmed closest qubit design.
- presimmed_closest_coupler_design: The presimmed closest coupler design.
- interpolated_design: The interpolated design.
- metric_strategy: The metric strategy (will be set dynamically).
- custom_metric_func: The custom metric function.
- metric_weights: The metric weights.
- target_params: The target parameters.
- H_param_keys: The H parameter keys.
"""
from squadds.core.db import SQuADDS_DB
self.db = db if db is not None else SQuADDS_DB()
self.reload_db()
def _initialize_attributes(self):
self.selected_component_name = self.db.selected_component_name
self.selected_component = self.db.selected_component
self.selected_data_type = self.db.selected_data_type
self.selected_confg = self.db.selected_confg
self.selected_qubit = self.db.selected_qubit
self.selected_cavity = self.db.selected_cavity
self.selected_resonator_type = self.db.selected_resonator_type
self.selected_coupler = self.db.selected_coupler
self.selected_system = self.db.selected_system
self.df = self.db.selected_df
self.qubit_df = self.db.qubit_df
self.cavity_df = self.db.cavity_df
self.coupler_df = self.db.coupler_df
self.closest_df_entry = None
self.closest_design = None
self.closest_df = None
self.presimmed_closest_cpw_design = None
self.presimmed_closest_qubit_design = None
self.presimmed_closest_coupler_design = None
self.interpolated_design = None
self.closest_design_found = False
self.params_computed = False
self.metric_strategy = None # Will be set dynamically
self.custom_metric_func = None
self.metric_weights = None
self.target_params = None
self.H_param_keys = self._get_H_param_keys()
[docs]
def reload_db(self):
"""
Reload the Analyzer with the current singleton SQuADDS_DB object.
"""
self._initialize_attributes()
[docs]
def _add_target_params_columns(self):
"""
Adds target parameter columns to the dataframe based on the selected system.
If the selected system is "qubit", it adds qubit Hamiltonian parameters to the dataframe.
If the selected system is "cavity_claw", it fixes the dataframe for the cavity_claw system.
If the selected system is "coupler", it does nothing.
If the selected system is ["qubit", "cavity_claw"] or ["cavity_claw", "qubit"], it fixes the dataframe for the cavity_claw system and adds cavity-coupled Hamiltonian parameters to the dataframe.
Raises:
a ValueError if the selected system is invalid.
"""
self.params_computed = True
#! 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":
self._fix_cavity_claw_df()
elif self.selected_system == "coupler":
pass
elif (self.selected_system == ["qubit", "cavity_claw"]) or (self.selected_system == ["cavity_claw", "qubit"]):
self._fix_cavity_claw_df()
qubit_H = TransmonCrossHamiltonian(self)
start = time.time()
qubit_H.add_cavity_coupled_H_params()
end = time.time()
print(f"Time taken to add the coupled H params: {end - start} seconds")
self.df = qubit_H.df
else:
raise ValueError("Invalid system.")
[docs]
def _fix_cavity_claw_df(self):
"""
Fix the cavity claw DataFrame by renaming columns and updating values.
If the columns 'cavity_frequency' or 'kappa' exist in the DataFrame, they will be renamed to
'cavity_frequency_GHz' and 'kappa_kHz' respectively. The values in these columns will also be
updated by multiplying them with appropriate conversion factors.
Args:
None
Returns:
None
"""
self.df = fix_cavity_claw_dataframe(self.df)
[docs]
def _get_H_param_keys(self):
"""
Get the parameter keys for the Hamiltonian (H) based on the selected system.
Returns:
list: A list of parameter keys for the Hamiltonian.
Raises:
ValueError: If the selected system is invalid.
"""
self.H_param_keys = get_H_param_keys_for_system(self.selected_system)
return self.H_param_keys
[docs]
def target_param_keys(self):
"""
Returns:
list: The target parameter keys.
"""
return self.H_param_keys
[docs]
def set_metric_strategy(self, strategy: MetricStrategy):
"""
Sets the metric strategy to use for calculating the distance metric.
Args:
strategy (MetricStrategy): The strategy to use for calculating the distance metric.
Raises:
ValueError: If the specified metric is not supported.
"""
self.metric_strategy = strategy
def _outside_bounds(self, df: pd.DataFrame, params: dict, display=True) -> bool:
"""
Check if entered parameters are outside the bounds of a dataframe.
Args:
df (pd.DataFrame): Dataframe to give warning.
params (dict): Keys are column names of `df`. Values are values to check for bounds.
Returns:
bool: True if any value is outside of bounds. False if all values are inside bounds.
"""
return outside_bounds(df, params, display=display)
[docs]
def get_complete_df(self, target_params: dict, metric: str = "Euclidean", display: bool = True):
"""
Returns the complete DataFrame (design + Hamiltonian parameters) sourced using the target parameters.
Args:
- target_params (dict): A dictionary containing the target parameters.
- metric (str, optional): The distance metric to use for calculating distances. Defaults to 'Euclidean'.
- display (bool, optional): Whether to display warnings for parameters outside of the library bounds. Defaults to True.
Returns:
- complete_df (DataFrame): A DataFrame containing all designs and Hamiltonian parameters.
Raises:
- ValueError: If the specified metric is not supported or if num_top is bigger than the size of the library.
- ValueError: If the metric is invalid.
"""
### Checks
# Check for supported metric
if metric not in self.__supported_metrics__:
raise ValueError(f"`metric` must be one of the following: {self.__supported_metrics__}")
self.target_params = target_params
remove_resonator_type_from_target_params(
self.target_params,
self.selected_resonator_type,
missing_ok=False,
)
if not self.params_computed:
self._add_target_params_columns()
else:
print("Target parameters have already been computed.")
return self.df
[docs]
def find_closest(
self,
target_params: dict,
num_top: int,
metric: str = "Euclidean",
display: bool = True,
parallel: bool = False,
num_cpu: str = "auto",
skip_df_gen: bool = False,
):
"""
Find the closest designs in the library based on the target parameters.
Args:
- target_params (dict): A dictionary containing the target parameters.
- num_top (int): The number of closest designs to retrieve.
- metric (str, optional): The distance metric to use for calculating distances. Defaults to 'Euclidean'.
- display (bool, optional): Whether to display warnings for parameters outside of the library bounds. Defaults to True.
- parallell (bool, optional): Whether to run metric calculation in a parallelized way
- num_cpu (str/int, optional): The number of CPUs to run a job over
- skip_df_gen (bool, optional): Whether to generate the df or run from memory
Returns:
- closest_df (DataFrame): A DataFrame containing the closest designs.
Raises:
- ValueError: If the specified metric is not supported or if num_top is bigger than the size of the library.
- ValueError: If the metric is invalid.
"""
### Checks
# Check for supported metric
if metric not in self.__supported_metrics__:
raise ValueError(f"`metric` must be one of the following: {self.__supported_metrics__}")
self.target_params = target_params
remove_resonator_type_from_target_params(
self.target_params,
self.selected_resonator_type,
missing_ok=True,
)
if (skip_df_gen) or (not self.params_computed):
self._add_target_params_columns()
elif self.selected_resonator_type == "quarter":
pass
else:
print(
"Either `skip_df_gen` flag is set to True or all target params have been precomputed at an earlier step. Using `df` from memory.\nPlease set this to False if `target_parameters` have changed."
)
target_params_list = list(self.target_params.keys())
filtered_df = self.df[target_params_list]
self._outside_bounds(df=filtered_df, params=target_params, display=display)
self.set_metric_strategy(resolve_metric_strategy(metric, self.metric_weights, self.custom_metric_func))
# Main logic
# Filter DataFrame based on target parameters that are string
filtered_df = filter_df_by_target_params(filtered_df, self.target_params)
# if the filtered_df is empty, raise a User input error
if filtered_df.empty:
raise ValueError(
f"No geometries found with the specified parameters:\n{target_params}\nPlease double-check your targets (especially ``resonator_type``) and try again."
)
# Calculate distances
# Use vectorized calculation (much faster and avoids joblib overhead)
if parallel:
print(
"Using vectorized calculation for speed (ignoring num_cpu parameter as it's not needed for vectorization)."
)
sorted_indices = rank_closest_indices(filtered_df, self.target_params, self.metric_strategy, num_top)
# Sort distances and get the closest ones
self.closest_df = self.df.loc[sorted_indices]
# set the closest design found flag
self.closest_design_found = True
if self.selected_resonator_type == "quarter":
# store the best design
self.closest_df_entry = self.closest_df.iloc[0]
self.closest_design = self.closest_df.iloc[0]["design_options"]
if len(self.selected_system) == 2: #! TODO: make this more general
self.presimmed_closest_cpw_design = self.closest_df_entry["design_options_cavity_claw"]
self.presimmed_closest_qubit_design = self.closest_df_entry["design_options_qubit"]
elif self.selected_resonator_type == "half":
# retrieve the best designs
self.closest_qubit = self.qubit_df.iloc[self.closest_df.index_qc].copy()
self.closest_coupler = self.coupler_df.iloc[self.closest_df.index_cplr]
self.closest_cavity = self.get_closest_cavity()
for merger_term in self.db.claw_merger_terms:
self.closest_qubit[merger_term] = self.closest_qubit["design_options"].map(
lambda x, mt=merger_term: x["connection_pads"]["readout"].get(mt)
)
# Create a unified design options column
merged_df = merge_dfs(self.closest_qubit, self.closest_cavity, self.db.claw_merger_terms)
# Add a temporary key column for cross join
self.closest_df["_temp_key"] = 1
merged_df["_temp_key"] = 1
# Perform the cross join
self.closest_df = pd.merge(
self.closest_df, merged_df, on="_temp_key", how="inner", suffixes=("_closest", "_merged")
).drop("_temp_key", axis=1)
# Create the unified design options column
self.closest_df["design_options"] = self.closest_df.apply(create_unified_design_options, axis=1)
self.closest_df_entry = self.closest_df.iloc[0]
return self.closest_df
[docs]
def get_closest_cavity(self):
"""
Returns the closest cavity design.
Returns:
pd.Series: The closest cavity design.
"""
# Extract the values you're looking for
closest_index_cc = self.closest_df.index_cc.values[0]
closest_index_cplr = self.closest_df.index_cplr.values[0]
# Use np.where for fast boolean indexing
mask = np.where(
(self.cavity_df["index_cc"].values == closest_index_cc)
& (self.cavity_df["index_cplr"].values == closest_index_cplr)
)[0]
# Get the index from the DataFrame
index = self.cavity_df.index[mask]
return self.cavity_df.loc[index]
[docs]
def compute_metric_distances(self, row):
return self.metric_strategy.calculate(self.target_params, row)
[docs]
def get_interpolated_design(self, target_params: dict, metric: str = "Euclidean", display: bool = True):
""" """
raise NotImplementedError
[docs]
def get_design(self, df):
"""
Extracts the design parameters from the dataframe and returns a dict.
Returns:
dict: A dict containing the design parameters.
"""
return df["design_options"].to_dict()[0]
[docs]
def get_param(self, design, param):
"""
Extracts a specific parameter from the design dict.
"""
raise NotImplementedError
[docs]
def closest_design_in_H_space(self):
"""Plots a scatter plot of the closest design in the H-space.
This method creates a scatter plot with two subplots. The first subplot shows the relationship between 'cavity_frequency_GHz' and 'kappa_kHz', while the second subplot shows the relationship between 'anharmonicity_MHz' and 'g_MHz'. The scatter plot includes pre-simulated data, target data, and the closest design entry from the database.
Returns:
None
"""
build_closest_design_hspace_plot(
self.df,
self.target_params,
self.closest_df_entry,
self.selected_resonator_type,
)
plt.show()
[docs]
def get_qubit_options(self, df: pd.DataFrame) -> dict[str, list[Any]]:
"""
Extracts qubit design options from the dataframe.
Parameters:
df (pd.DataFrame): The dataframe containing design options.
Returns:
Dict[str, List[Any]]: A dictionary containing lists of the extracted qubit options.
"""
return extract_qubit_options(df)
[docs]
def get_cpw_options(self, df: pd.DataFrame) -> dict[str, list[Any]]:
"""
Extracts CPW options from the dataframe.
Parameters:
df (pd.DataFrame): The dataframe containing design options.
Returns:
Dict[str, List[Any]]: A dictionary containing lists of the extracted CPW options.
"""
return extract_cpw_options(df)
[docs]
def get_coupler_options(self, df: pd.DataFrame) -> dict[str, list[Any]]:
"""
Extracts coupler options from the dataframe.
Parameters:
df (pd.DataFrame): The dataframe containing design options.
Returns:
Dict[str, List[Any]]: A dictionary containing lists of the extracted coupler options.
"""
return extract_coupler_options(df)
[docs]
def get_Ljs(self, df: pd.DataFrame):
"""
Extracts the EJ values from the dataframe. Converts them to Josephson inductance values using pyEPR
Parameters:
df (pd.DataFrame): The dataframe containing design options.
Returns:
np.array: An array of Josephson inductance values.
"""
from pyEPR.calcs import Convert
# EJ values are stored in a column named 'EJ' in GHz
EJ_values = df["EJ"].values
# Convert EJ (in MHz) to Josephson inductance (Lj) in nH
Ljs = np.array([Convert.Lj_from_Ej(Ej, units_in="GHz", units_out="nH") for Ej in EJ_values])
return Ljs
