sqdtoolz 0.9.0.dev1__py3-none-any.whl

sqdtoolz/Experiments/Experimental/ExpZIRandomisedBenchmarking.py

@@ -0,0 +1,126 @@
+import numpy as np
+from sqdtoolz.Experiments.Experimental.ExpZIqubit import ExpZIqubit
+from sqdtoolz.Variable import VariablePropertyTransient
+from sqdtoolz.HAL.WaveformGeneric import *
+from sqdtoolz.HAL.WaveformSegments import *
+from sqdtoolz.Utilities.DataFitting import *
+from sqdtoolz.Experiments.Experimental.ExpCalibGE import *
+from sqdtoolz.Utilities.QubitGates import QubitGatesBase
+import json
+from sqdtoolz.Experiments.Experimental.ZI import single_qubit_gates_sweep
+from sqdtoolz.Experiments.Experimental.ZI import sinlge_qubit_gates_sweep_chunking
+from sqdtoolz.Utilities.QubitGates import QubitGatesBase
+from sqdtoolz.Variable import VariableInternalTransient
+
+class ExpZIBenchmarkRandomised(ExpZIqubit):   
+    def __init__(self, name, expt_config, hal_QPU, qubit_ids, **kwargs):
+        self._qubit_datasets = qubit_ids
+
+        self._hal_QPU = hal_QPU
+
+        self._dont_show_plot = kwargs.pop('dont_show_plot', False)
+        assert (not 'update' in kwargs) or ('update' in kwargs and not kwargs['update']), "Don't set 'update=True'. The updates shall be done by calling update_qubit after running the experiment."
+        kwargs['update'] = False
+
+        kwargs['coordinate_system'] = kwargs.get('coordinate_system', 'RH')
+        assert kwargs['coordinate_system'] in ['LH', 'RH'], "The 'coordinate_system' must be either LH or RH for left/right handed."
+
+        self._fit_vals = []
+
+        #extra_gate_seqs = kwargs.get('extra_gate_seqs', [])
+        self._rb_seed = kwargs.get('rb_seed', 88)
+        self._rng = np.random.default_rng(seed = self._rb_seed)
+
+        self._sequence_lengths = kwargs.pop('sequence_lengths', [3,4,5,6,7,8,9,10,11,12])
+        self._num_trials = kwargs.pop('num_trials',5)
+
+
+        self._gate_set = ['X', 'X/2', '-X/2', 'Y', 'Y/2', '-Y/2']
+
+    
+        self._all_seqs = {}
+        kwargs['gate_lists'] = []
+        print("Generating gate sequences...")
+        for seq_len in self._sequence_lengths:
+            self._all_seqs[int(seq_len)] = {}
+            for trial in range(self._num_trials):
+                cur_seq = self._generate_sequence(seq_len)
+                kwargs['gate_lists'].append([cur_seq]*len(qubit_ids))
+                self._all_seqs[int(seq_len)][int(trial)] = [cur_seq]*len(qubit_ids)
+        self._all_seqs_list = kwargs['gate_lists']
+        print("Done.")
+        super().__init__(name, expt_config, sinlge_qubit_gates_sweep_chunking, hal_QPU, qubit_ids, **kwargs)
+
+    def _generate_sequence(self, seq_len):
+        #Keep trying to find random sequences that enable the final gate to be within the gate-set - e.g. X/2, Y will give a 45° rotation...
+        while(True):
+            rand_int_array = self._rng.integers(low=0, high=len(self._gate_set)-1, size=seq_len-1)
+            gate_seq = [self._gate_set[i] for i in rand_int_array]
+            gate_mats = [QubitGatesBase.get_rotation_from_Pauli_Matrix(g) for g in gate_seq]
+            gate_mats.reverse()
+            effective_gate = np.linalg.multi_dot(gate_mats)
+            #Final state should be excited-state...
+            XReInv = QubitGatesBase.get_rotation_from_Pauli_Matrix('X') @ np.linalg.inv(effective_gate)
+
+            try:
+                final_gate = QubitGatesBase.convert_Pauli_rotation_to_natural( *QubitGatesBase.compute_rotation_Pauli_Matrices(XReInv))
+            except:
+                continue
+            if final_gate == 'Z' or final_gate == 'Z/2' or final_gate == '-Z/2':
+                final_gate = 'I'
+            gate_seq.append(final_gate)
+
+            return gate_seq
+        
+    def _post_process(self, data):
+        for ind_qubit, qubit_dataset in enumerate(self._qubit_datasets):          
+            leData = self.retrieve_last_dataset(qubit_dataset)
+            arr = leData.get_numpy_array()
+            seqs = leData.param_vals[0]
+            seq_lens = np.array(self._sequence_lengths)
+            norm = ExpZIqubit.normalise_qubit_data(self.retrieve_last_dataset(qubit_dataset+'_calib'), self._transition)
+            probs = norm.normalise_data(arr)
+            mean_vals = []
+            std_vals = []
+
+            fig, axs = plt.subplots(ncols=2); fig.set_figwidth(12); axs[0].grid(); axs[1].grid()
+
+            for m in range(len(self._sequence_lengths)):
+                trial_probs = []
+                xvals = []
+                for k in range(self._num_trials):
+                    cur_seq_index = m*self._num_trials + k
+                    if probs[cur_seq_index] < 0.8: 
+                        continue
+                    trial_probs.append(probs[cur_seq_index])
+                    xvals.append(self._sequence_lengths[m])
+                cur_x = xvals
+                cur_y = trial_probs
+                axs[0].plot(cur_x, cur_y, 'kx', alpha = 0.3)
+                mean_vals.append(np.mean(cur_y))
+                std_vals.append(np.std(cur_y))
+            mean_vals = np.array(mean_vals)
+            std_vals = np.array(std_vals)
+
+            
+
+            axs[0].plot(seq_lens, mean_vals)
+            axs[0].fill_between(seq_lens, mean_vals-std_vals, mean_vals+std_vals, alpha=0.5)
+            axs[0].set_xlabel('Sequence Length')
+            axs[0].set_ylabel('Excited State Probability')
+            axs[0].set_title(f'Number of Trials per Sequence Length: {self._num_trials}')
+
+            zProj = mean_vals - 0.5
+            slope, intercept = np.polyfit(seq_lens, np.log(zProj), deg=1)
+            error_per_gate = np.exp(slope)
+            axs[1].plot(seq_lens, np.log(zProj), 'kx')
+            axs[1].plot(seq_lens, seq_lens*slope+intercept, 'r-')
+            axs[1].set_xlabel('Sequence Length')
+            axs[1].set_ylabel(r'$\ell n(Z_{proj})$')
+            axs[1].set_title(f"Error per Gate: {(error_per_gate*100):.6g}%")
+            axs[1].legend(['Raw Data', 'Fitted Error line'])
+
+            fig.show()
+            fig.savefig(self._file_path + 'Summary.png')
+
+        return data

sqdtoolz/Experiments/Experimental/ExpZIRes.py

@@ -0,0 +1,123 @@
+from sqdtoolz.Experiments.Experimental.ExpZIqubit import ExpZIqubit
+from sqdtoolz.Utilities.DataFitting import*
+from sqdtoolz.Utilities.ResonatorTools import ResonatorPowerSweep
+from laboneq_applications.experiments import (
+    resonator_spectroscopy,
+    resonator_spectroscopy_amplitude,
+)
+
+class ExpZIRes(ExpZIqubit):
+    def __init__(self, name, expt_config, hal_QPU, qubit_id, **kwargs):
+        assert isinstance(qubit_id, str), "Supply qubit_id as the solitary string ID here (i.e. not a list)."
+        self._qubit_dataset = qubit_id
+
+        self._update_qubit = kwargs.pop('update_qubit_params', True)
+
+        self._dont_show_plot = kwargs.pop('dont_show_plot', False)
+
+        self._iq_indices = kwargs.pop('iq_indices', [0,1])
+        self._is_trough = kwargs.pop('is_trough', False)
+        self._fit_type = kwargs.pop('fit_type', 'Default')  #Default, Fano, Full
+        assert self._fit_type in ["Default", "Fano", "Full"], "Choose fit_type as either 'Default', 'Fano' or 'Full'."
+        assert self._is_trough or (not self._is_trough and not self._fit_type=='Fano'), "Fano resonance fitting only supports troughs at the moment."
+        self._param_centre = kwargs.pop('param_freq', None)
+        self._param_width = kwargs.pop('param_width', None)
+        self._param_amplitude = kwargs.pop('param_amplitude', None)
+        self._param_offset = kwargs.pop('param_offset', None)
+        self._param_fano = kwargs.pop('param_fano', None)
+        self._dont_plot = kwargs.pop('dont_plot', False)
+        self._xUnits = kwargs.pop('plot_x_units', 'Hz')
+        self._hal_QPU = hal_QPU
+
+        self._frequencies = kwargs.pop('frequencies', None)
+        self._span = kwargs.pop('span', 5e6)
+        self._num_freq_pts = kwargs.pop('num_frequency_points', 401)
+
+        if self._frequencies is None:
+            self._frequencies = np.linspace(-self._span, self._span, self._num_freq_pts) + self._hal_QPU.get_qubit_obj(self._qubit_dataset).ReadoutFrequency
+
+        kwargs['frequencies'] = self._frequencies
+
+        super().__init__(name, expt_config, resonator_spectroscopy, hal_QPU, [qubit_id], **kwargs)
+    
+    def _post_process(self, data):
+        data = self.retrieve_last_dataset(self._qubit_dataset)
+        assert len(data.param_names), "The sweep should only be 1D."
+
+        arr = data.get_numpy_array()
+        data_x = data.param_vals[0]
+        data_i = arr[:,self._iq_indices[0]]
+        data_q = arr[:,self._iq_indices[1]]
+        data_y = np.sqrt(arr[:,self._iq_indices[0]]**2 + arr[:,self._iq_indices[1]]**2)
+
+        if self._fit_type == "Default":
+            dfit = DFitPeakLorentzian()
+            dpkt = dfit.get_fitted_plot(data_x, data_y**2, xLabel="Frequency (Hz)", dip=self._is_trough, dontplot=self._dont_plot, xUnits=self._xUnits)
+            #Commit to parameters...
+            if self._update_qubit:
+                cur_qubit = self._hal_QPU.get_qubit_obj(self._qubit_dataset)
+                cur_qubit.ReadoutFrequency = dpkt['centre']
+            if self._param_centre:
+                self._param_centre.Value = dpkt['centre']
+            if self._param_width:
+                self._param_width.Value = dpkt['width']
+            if self._param_amplitude:
+                self._param_amplitude.Value = dpkt['amplitude']
+            if self._param_offset:
+                self._param_offset.Value = dpkt['offset']
+            dpkt['fit_data'] = {'squared_amplitude': dpkt['fit_data']}
+        elif self._fit_type == "Fano":
+            dfit = DFitFanoResonance()
+            dpkt = dfit.get_fitted_plot(data_x, data_y**2, xLabel="Frequency (Hz)", yLabel="Squared IQ Amplitude", dontplot=self._dont_plot, xUnits=self._xUnits) #, dip=self._is_trough)
+            #Commit to parameters...
+            if self._update_qubit:
+                cur_qubit = self._hal_QPU.get_qubit_obj(self._qubit_dataset)
+                cur_qubit.ReadoutFrequency = dpkt['xMinimum']
+            if self._param_centre:
+                self._param_centre.Value = dpkt['xMinimum']
+            if self._param_width:
+                self._param_width.Value = dpkt['width']
+            if self._param_amplitude:
+                self._param_amplitude.Value = dpkt['amplitude']
+            if self._param_offset:
+                self._param_offset.Value = dpkt['offset']
+            if self._param_fano:
+                self._param_fano.Value = dpkt['FanoFac']
+            dpkt['fit_data'] = {'squared_amplitude': dpkt['fit_data']}
+        elif self._fit_type == "Full":
+            if self._hal_QPU.get_qubit_obj(self._qubit_dataset).ReadoutLineAttenuation_dB:
+                atten = self._hal_QPU.get_qubit_obj(self._qubit_dataset).ReadoutLineAttenuation_dB
+                readout_pwr = self._hal_QPU.get_qubit_obj(self._qubit_dataset).ReadoutPower + 20*np.log10(self._hal_QPU.get_qubit_obj(self._qubit_dataset).ReadoutAmplitude)
+                dpkt = ResonatorPowerSweep.single_circlefit(data_x, data_i, data_q, power_dBm=atten + readout_pwr, dont_plot=self._dont_plot)
+            else:
+                dpkt = ResonatorPowerSweep.single_circlefit(data_x, data_i, data_q, power_dBm=-100, dont_plot=self._dont_plot)
+            #Commit to parameters...
+            if self._update_qubit:
+                cur_qubit = self._hal_QPU.get_qubit_obj(self._qubit_dataset)
+                if dpkt and dpkt['fr']:
+                    cur_qubit.ReadoutFrequency = dpkt['fr']
+                    cur_qubit.ReadoutQi = dpkt['Qi_dia_corr']
+                    cur_qubit.ReadoutQc = dpkt['Qc_dia_corr']
+                    cur_qubit.ReadoutQl = dpkt['Ql']
+                    dpkt['fit_data'] = {'real': np.real(dpkt['fit_data']), 'imag': np.imag(dpkt['fit_data'])}
+                    cur_qubit.ReadoutKappa = dpkt['fr'] / dpkt['Qc_dia_corr']
+                else:
+                    print('Qubit parameters were not updated: bad fit.')
+            if self._param_centre:
+                self._param_centre.Value = dpkt['fr']
+        
+        if dpkt:
+            if not self._dont_plot:
+                dpkt['fig'].savefig(self._file_path + 'fitted_plot.png')
+                if not self._dont_show_plot:
+                    dpkt['fig'].show()
+                else:
+                    plt.close(dpkt['fig'])
+            if 'fit_data' in dpkt:
+                np.save(self._file_path + 'fitted_data.npy', dpkt['fit_data'])
+        else:
+            if not self._dont_plot:
+                plt.plot(data_x, data_y)
+                plt.xlabel("Frequency (Hz)")
+                plt.ylabel("|S21|")
+                plt.show()

sqdtoolz/Experiments/Experimental/ExpZIResFluxSweep.py

@@ -0,0 +1,101 @@
+from sqdtoolz.Experiments.Experimental.ExpZIqubit import ExpZIqubit
+from laboneq_applications.experiments import resonator_spectroscopy
+from sqdtoolz.Variable import VariablePropertyTransient
+import numpy as np
+import matplotlib.pyplot as plt
+import scipy.interpolate
+import scipy.optimize
+from sqdtoolz.Utilities.DataFitting import DFitPeakLorentzian
+
+class ExpZIResFluxSweep(ExpZIqubit):
+    def __init__(self, name, expt_config, hal_QPU, qubit_id, **kwargs):
+        assert isinstance(qubit_id, str), "Supply qubit_id as the solitary string ID here (i.e. not a list)."
+        self._qubit_id = qubit_id
+
+        self._update_qubit = kwargs.pop('update_qubit_params', True)
+
+        self._dont_show_plot = kwargs.pop('dont_show_plot', False)
+
+        self._is_trough = kwargs.pop('is_trough', True)
+        self._dont_plot = kwargs.pop('dont_plot', False)
+        self._xUnits = kwargs.pop('plot_x_units', 'Hz')
+        self._flux_range = kwargs.pop('flux_range', np.linspace(-1, 1, 10))
+        self._frequencies = kwargs.pop('frequencies', None)
+        self._span = kwargs.pop('span', 10e6)
+        self._num_freq_pts = kwargs.pop('num_frequency_points', 401)
+        self._hal_QPU = hal_QPU
+
+        self._opt_flux_val = None
+
+        if self._frequencies is None:
+            self._frequencies = np.linspace(-self._span, self._span, self._num_freq_pts) + self._hal_QPU.get_qubit_obj(self._qubit_id).ReadoutFrequency
+
+        kwargs['frequencies'] = self._frequencies
+        super().__init__(name, expt_config, resonator_spectroscopy, hal_QPU, [qubit_id], **kwargs)
+    
+    def _run(self, file_path, sweep_vars=[], **kwargs):
+        assert len(sweep_vars) == 0, "Supply the frequency/flux when defining the Experiment object."
+
+        var_flux = VariablePropertyTransient('Flux', self._hal_QPU.get_qubit_obj(self._qubit_id),'FluxDC')
+        super()._run(file_path, sweep_vars=[(var_flux, self._flux_range)], **kwargs)
+
+    def _post_process(self, data):
+        leData = self.retrieve_last_dataset(self._qubit_id)
+        flux_vals,freq_vals = leData.param_vals
+        arr = leData.get_numpy_array()
+
+        ampl = np.sqrt(arr[:,:,0]**2 + arr[:,:,1]**2)
+
+        is_trough = self._is_trough
+
+        dfit = DFitPeakLorentzian()
+        res_freqs = []
+        for cur_slice in ampl:
+            dpkt = dfit.get_fitted_plot(freq_vals, cur_slice, dip=is_trough, dontplot=True)
+            res_freqs.append(dpkt['centre'])
+        res_freqs = np.array(res_freqs)
+
+        norm_fac = 10**np.round(np.log10(res_freqs).mean())
+        if flux_vals[1]-flux_vals[0] < 0:
+            spl = scipy.interpolate.UnivariateSpline(flux_vals[::-1], res_freqs[::-1]/norm_fac)
+        else:
+            spl = scipy.interpolate.UnivariateSpline(flux_vals, res_freqs/norm_fac)
+        spl.set_smoothing_factor(0.5)
+        soln = scipy.optimize.minimize_scalar(lambda x:-spl(x), bounds=(np.min(flux_vals), np.max(flux_vals)))
+
+        opt_flux = soln.x
+        sweet_freq = spl(opt_flux)*norm_fac
+
+        fitted_data = {
+            'raw_fit_freqs': res_freqs,
+            'smoothed_fit_freqs': spl(flux_vals)*norm_fac,
+            'sweet_freq_flux': np.array([sweet_freq, opt_flux])
+        }
+        np.save(self._file_path + 'fitted_data.npy', fitted_data)
+
+        if not self._dont_plot:
+            fig, ax = plt.subplots(1); fig.set_figwidth(15)
+            ExpZIResFluxSweep.plot_fitted_results(ax, self._qubit_id, freq_vals, flux_vals, ampl, fitted_data)
+            if not self._dont_show_plot:
+                fig.show()
+            else:
+                plt.close(fig)
+            fig.savefig(self._file_path + 'fitted_plot.png')
+
+        self._opt_flux_val = opt_flux
+    
+    @staticmethod
+    def plot_fitted_results(ax, qubit_id, freq_vals, flux_vals, ampl, fitted_data):
+        ax.pcolor(freq_vals, flux_vals, ampl)
+        ax.plot(fitted_data['smoothed_fit_freqs'], flux_vals, 'w-')
+        ax.plot(fitted_data['raw_fit_freqs'], flux_vals, 'wo')
+        ax.plot([fitted_data['sweet_freq_flux'][0]], [fitted_data['sweet_freq_flux'][1]], 'ro')
+        ax.set_title(f"{qubit_id} resonator flux sweep: $f_r$={fitted_data['sweet_freq_flux'][0]*1e-9:.4f} GHz, $V_\Phi$={fitted_data['sweet_freq_flux'][1]:.4f} V")
+        ax.set_xlabel('Frequency (Hz)')
+        ax.set_ylabel('Flux (V)')
+
+    def update_qubit(self):
+        assert self._opt_flux_val != None, "Must run an experiment before updating qubit parameters."
+        self._hal_QPU.get_qubit_obj(self._qubit_id).FluxDC = self._opt_flux_val
+        # TODO: update flux range to minimum
+        self._opt_flux_val = None

sqdtoolz/Experiments/Experimental/ExpZIResOptimal.py

@@ -0,0 +1,224 @@
+from sqdtoolz.Experiments.Experimental.ExpZIqubit import ExpZIqubit
+import matplotlib.pyplot as plt
+from sqdtoolz.Utilities.Miscellaneous import Miscellaneous
+from laboneq_applications.experiments import dispersive_shift
+import numpy as np
+from sqdtoolz.Utilities.DataIQDiscriminate import DataIQDiscriminate
+from sqdtoolz.Experiments.Experimental.ExpZIBlobs import ExpZIBlobs
+
+class ExpZIResOptimal(ExpZIqubit):
+    def __init__(self, name, expt_config, hal_QPU, qubit_ids, **kwargs):
+        self._dont_show_plot = kwargs.pop('dont_show_plot', False)
+        assert (not 'update' in kwargs) or ('update' in kwargs and not kwargs['update']), "Don't set 'update=True'. This is just a diagnostic experiment."
+        kwargs['update'] = False
+        # self._fit_vals = []
+        assert 'states' in kwargs, "Must provide the states to measure traces over; e.g. 'ge' or 'gef'"
+        self._states = kwargs['states']
+        self._calc_single_shot_fidelities = kwargs.pop('calc_single_shot_fidelities', False)
+        if self._calc_single_shot_fidelities:
+            kwargs['do_analysis'] = False   #The default ZI analysis will fail in this mode...
+        self._fit_data = {}
+        self._readout_fidelity = 0.0
+        super().__init__(name, expt_config, dispersive_shift, hal_QPU, qubit_ids, **kwargs)
+    
+    def _run(self, file_path, sweep_vars=[], **kwargs):
+        if self._calc_single_shot_fidelities:
+            kwargs['override_ACQ_params'] = {'AveragingOrder': "SingleShot"}
+        return super()._run(file_path, sweep_vars, **kwargs)
+
+    def _post_process(self, data):
+        #This experiment only supports one qubit at a time...
+        leDatasets = [self.retrieve_last_dataset(self._qubit_ids[0] + '_' + x) for x in self._states]
+        arrs = [leDatasets[m].get_numpy_array() for m in range((len(self._states)))]
+
+        leCols = plt.rcParams['axes.prop_cycle'].by_key()['color']
+        if self._calc_single_shot_fidelities:
+            freqs = leDatasets[0].param_vals[1]
+            norm_fac, norm_prefix = Miscellaneous.get_metric_multiplier(freqs)  #Should all be the same anyway...
+
+            fig = plt.figure(figsize=(8, 10))
+            gs = fig.add_gridspec(5, 4)
+
+            # First three rows each span all four columns
+            ax_mag = fig.add_subplot(gs[0, :])
+            axSeps = fig.add_subplot(gs[1, :], sharex=ax_mag)
+            axFids = fig.add_subplot(gs[2, :], sharex=ax_mag)
+            
+            ax_mag.grid(); axSeps.grid(); axFids.grid()
+
+            axFids.set_ylabel('State Fidelities')
+
+            axFids.set_xlabel(f'Readout Frequency ({norm_prefix}Hz)')
+        else:
+            freqs = leDatasets[0].param_vals[0]
+            norm_fac, norm_prefix = Miscellaneous.get_metric_multiplier(freqs)  #Should all be the same anyway...
+            fig, axs = plt.subplots(nrows=2)
+            ax_mag = axs[0]
+            axSeps = axs[1]
+            ax_mag.grid(); axSeps.grid()
+        ax_mag.set_xlabel(f'Readout Frequency ({norm_prefix}Hz)')
+        axSeps.set_xlabel(f'Readout Frequency ({norm_prefix}Hz)')
+        axSeps.set_ylabel('IQ separation')
+
+        ax_phase = ax_mag.twinx()
+
+        handles = []
+        for m in range(len(self._states)):
+            if self._calc_single_shot_fidelities:
+                arr = np.mean(arrs[m],axis=0)
+            else:
+                arr = arrs[m]
+            mag = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2)
+            phase = np.unwrap(np.arctan2(arr[:, 1], arr[:, 0]))
+            # Plot magnitude and capture the line object
+            line, = ax_mag.plot(freqs/norm_fac, mag, label=self._states[m])
+            handles.append(line)
+            # Use the same colour for phase, but dashed
+            ax_phase.plot(freqs/norm_fac, phase,':',color=line.get_color())
+
+        ax_mag.set_ylabel('Magnitude')
+        ax_phase.set_ylabel('Phase (rad) (dotted)')
+        ax_mag.legend(handles=handles)
+
+
+        if len(self._states) == 2:
+            combs = [[0,1]]
+        else:
+            combs = [(0,1), (1,2), (0,2)]
+
+        all_diffs = []
+        maxSepInds = []
+        for m in range(len(combs)):
+            if self._calc_single_shot_fidelities:
+                arr1 = np.mean(arrs[combs[m][0]],axis=0)
+                arr2 = np.mean(arrs[combs[m][1]],axis=0)
+            else:
+                arr1 = arrs[combs[m][0]]
+                arr2 = arrs[combs[m][1]]
+            cur_diff = np.linalg.norm(arr1 - arr2, axis=1)
+            axSeps.plot(freqs/norm_fac, cur_diff, color=leCols[m])
+            maxSepInds.append(np.argmax(cur_diff))
+            axSeps.plot([freqs[maxSepInds[-1]]/norm_fac], [cur_diff[maxSepInds[-1]]], 'o', color=leCols[m], label='_nolegend_')
+            all_diffs.append(cur_diff)
+        if len(self._states) == 3:
+            cur_diff = np.sum(all_diffs, axis=0)
+            axSeps.plot(freqs/norm_fac, cur_diff, color=leCols[m+1])
+            maxSepInds.append(np.argmax(cur_diff))
+            axSeps.plot([freqs[maxSepInds[-1]]/norm_fac], [cur_diff[maxSepInds[-1]]], 'o', color=leCols[m+1], label='_nolegend_')
+        if len(self._states) == 2:
+            axSeps.legend(['GE'])
+            leTitlesSS = ['G','E','Mean']
+        else:
+            axSeps.legend(['GE', 'EF', 'GF', 'Total'])
+            leTitlesSS = ['G','E','F','Mean']
+
+        if self._calc_single_shot_fidelities:
+            leIQDiscs = [DataIQDiscriminate([arrs[x][:,m,:] for x in range(len(self._states))]) for m in range(freqs.size)] 
+            leFids = [leIQDiscs[m].get_fidelities() for m in range(freqs.size)]
+            leFids = np.array(leFids)
+            #Gather fidelity data for each state the total
+            fidData = []
+            for m in range(len(self._states)):
+                fidData.append(leFids[:,m])
+            fidData.append(np.mean(leFids, axis=1))
+
+            maxFidInds = []
+            axFirst = None
+            for m in range(len(fidData)):
+                axFids.plot(freqs/norm_fac, fidData[m], color = leCols[m])
+                maxFidInds.append(np.argmax(fidData[m]))
+                axFids.plot(freqs[maxFidInds[-1]]/norm_fac, [fidData[m][maxFidInds[-1]]], 'o', color=leCols[m], label='_nolegend_')
+
+                if axFirst == None:
+                    ax = fig.add_subplot(gs[3, m])
+                    axFirst = ax
+                else:
+                    ax = fig.add_subplot(gs[3, m], sharey=axFirst)
+                leIQDiscs[maxFidInds[-1]].plot_points(ax)
+                ax.set_title(leTitlesSS[m])
+                ax.set_xticklabels([])
+                ax.set_yticklabels([])
+                #
+                axA = fig.add_subplot(gs[4, m])
+                leIQDiscs[maxFidInds[-1]].plot_assignment_matrix(axA, sigFigs=2)
+                self._readout_fidelity = (leIQDiscs[maxFidInds[-1]].get_average_fidelity()*100)
+                axA.set_title(f"Mean: {self._readout_fidelity:.4g}%")
+                if m > 0:
+                    ax.set_ylabel('')
+                    axA.set_ylabel('')
+                    axA.set_yticklabels([])
+
+            if len(self._states) == 2:
+                axFids.legend(['G','E','Mean'])
+            else:
+                axFids.legend(['G','E','F','Mean'])
+        
+        self._fit_data = {'freqs':freqs, 'maxSepIndices':maxSepInds, 'maxFidIndices':maxFidInds, 'discriminators':leIQDiscs}
+
+        fig.subplots_adjust(hspace=0.1)
+        fig.tight_layout()
+
+        fig.savefig(self._file_path + f'fitted_plot_{self._qubit_ids[0]}.png')
+        if not self._dont_show_plot:
+            fig.show()
+        else:
+            plt.close(fig)
+
+    def update_qubits_by_separation(self, transition:str='Total'):
+        """
+        transition is given as 'ge', 'ef', 'gf' or 'Total' (can capitalise etc.)
+        """
+        assert len(self._fit_data) > 0, "Must run experiment first."
+        qubit_obj = self._hal_QPU.get_qubit_obj(self._qubit_ids[0])
+        transition = transition.lower()
+        if len(self._states) == 2:
+            assert transition in ['ge', 'total'], "Invalid transition (must be either 'ge' or 'total')"
+            ind = 0
+        else:
+            assert transition in ['ge', 'ef', 'gf', 'total'], "Invalid transition (must be either 'ge', 'ef', 'gf' or 'total')"
+            ind = ['ge', 'ef', 'gf', 'total'].index(transition)
+        qubit_obj.ReadoutFrequency = float( self._fit_data['freqs'][ self._fit_data['maxSepIndices'][ind] ] )
+        qubit_obj.FidelityReadout = self._readout_fidelity
+
+    def update_qubits_by_fidelity(self, state_fidelity:str='Mean'):
+        """
+        state_fidelity is given as 'g', 'e', 'f' or 'Mean' (can capitalise etc.) to take the highest of the respective fidelities as the point.
+        """
+        assert len(self._fit_data) > 0, "Must run experiment first."
+        assert self._calc_single_shot_fidelities, "Must run the experiment in single-shot by setting calc_single_shot_fidelities to True."
+        qubit_obj = self._hal_QPU.get_qubit_obj(self._qubit_ids[0])
+        state_fidelity = state_fidelity.lower()
+        if len(self._states) == 2:
+            assert state_fidelity in ['g', 'e', 'mean'], "Invalid state (must be either 'g', 'e' or 'mean')"
+            ind = ['g', 'e', 'mean'].index(state_fidelity)
+        else:
+            assert state_fidelity in ['g', 'e', 'f', 'mean'], "Invalid state (must be either 'g', 'e', 'f' or 'mean')"
+            ind = ['g', 'e', 'f', 'mean'].index(state_fidelity)
+        qubit_obj.ReadoutFrequency = float( self._fit_data['freqs'][ self._fit_data['maxFidIndices'][ind] ] )
+        qubit_obj.FidelityReadout = self._readout_fidelity
+
+    def print_best_frequencies_by_separation(self):
+        assert len(self._fit_data) > 0, "Must run experiment first."
+        if len(self._states) == 2:
+            print(f"Max GE: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxSepIndices'][0] ])}Hz")
+        else:
+            print(f"Max GE: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxSepIndices'][0] ])}Hz")
+            print(f"Max EF: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxSepIndices'][1] ])}Hz")
+            print(f"Max GF: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxSepIndices'][2] ])}Hz")
+
+    def print_best_frequencies_by_fidelity(self):
+        assert len(self._fit_data) > 0, "Must run experiment first."
+        assert self._calc_single_shot_fidelities, "Must run the experiment in single-shot by setting calc_single_shot_fidelities to True."
+        print(f"Max G: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxFidIndices'][0] ])}Hz")
+        print(f"Max E: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxFidIndices'][1] ])}Hz")
+        if len(self._states) == 2:
+            print(f"Max Mean: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxFidIndices'][2] ])}Hz")
+        else:
+            print(f"Max F: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxFidIndices'][2] ])}Hz")
+            print(f"Max Mean: {Miscellaneous.get_units(self._fit_data['freqs'][ self._fit_data['maxFidIndices'][3] ])}Hz")
+
+    def plot_blobs(self, frequency):
+        assert len(self._fit_data) > 0, "Must run experiment first."
+        assert self._calc_single_shot_fidelities, "Must run the experiment in single-shot by setting calc_single_shot_fidelities to True."
+        ind = np.argmin(np.abs(frequency-self._fit_data['freqs']))
+        ExpZIBlobs.plot_fitted_results(self._fit_data['discriminators'][ind], f"(Frequency: {Miscellaneous.get_units(self._fit_data['freqs'][ind])}Hz)")

sqdtoolz/Experiments/Experimental/ExpZIResPowerSweep.py

@@ -0,0 +1,66 @@
+from sqdtoolz.Experiments.Experimental.ExpZIqubit import ExpZIqubit
+from laboneq_applications.experiments import resonator_spectroscopy
+from sqdtoolz.Variable import VariablePropertyTransient
+import numpy as np
+import matplotlib.pyplot as plt
+import scipy.interpolate
+import scipy.optimize
+from sqdtoolz.Utilities.DataFitting import DFitPeakLorentzian
+
+class ExpZIResPowerSweep(ExpZIqubit):
+    def __init__(self, name, expt_config, hal_QPU, qubit_id, frequencies, **kwargs):
+        assert isinstance(qubit_id, str), "Supply qubit_id as the solitary string ID here (i.e. not a list)."
+        self._qubit_id = qubit_id
+
+        self._amplitude_range = kwargs.pop('amplitude_range', np.linspace(0.001, 0.9, 20))
+
+        self._dont_show_plot = kwargs.pop('dont_show_plot', False)
+
+        self._is_trough = kwargs.pop('is_trough', True)
+        self._dont_plot = kwargs.pop('dont_plot', False)
+        self._frequencies = frequencies
+        self._hal_QPU = hal_QPU
+
+        kwargs['frequencies'] = frequencies
+        super().__init__(name, expt_config, resonator_spectroscopy, hal_QPU, [qubit_id], **kwargs)
+    
+    def _run(self, file_path, sweep_vars=[], **kwargs):
+        var_ampl = VariablePropertyTransient('Amplitude', self._hal_QPU.get_qubit_obj(self._qubit_id),'ReadoutAmplitude')
+        super()._run(file_path, sweep_vars=[(var_ampl, self._amplitude_range)], **kwargs)
+
+    def _post_process(self, data):
+        leData = self.retrieve_last_dataset(self._qubit_id)
+        ampl_vals,freq_vals = leData.param_vals
+        arr = leData.get_numpy_array()
+
+        ampl = np.sqrt(arr[:,:,0]**2 + arr[:,:,1]**2)
+
+        is_trough = self._is_trough
+
+        dfit = DFitPeakLorentzian()
+        res_freqs = []
+        for cur_slice in ampl:
+            dpkt = dfit.get_fitted_plot(freq_vals, cur_slice, dip=is_trough, dontplot=True)
+            res_freqs.append(dpkt['centre'])
+        res_freqs = np.array(res_freqs)
+
+        fitted_data = {'raw_fit_freqs': res_freqs,}
+
+        if not self._dont_plot:
+            fig, ax = plt.subplots(1); fig.set_figwidth(15)
+            ExpZIResPowerSweep.plot_fitted_results(ax, self._hal_QPU, self._qubit_id, freq_vals, ampl_vals, ampl, fitted_data)
+            if not self._dont_show_plot:
+                fig.show()
+            else:
+                plt.close(fig)
+            fig.savefig(self._file_path + 'fitted_plot.png')
+    
+    @staticmethod
+    def plot_fitted_results(ax, hal_QPU, qubit_id, freq_vals, ampl_vals, ampl, fitted_data):
+        row_means = np.nanmean(ampl, axis=1)
+        ampl_corrected = (ampl.T - row_means).T
+        ax.pcolor(freq_vals, ampl_vals, ampl_corrected)
+        ax.plot(fitted_data['raw_fit_freqs'], ampl_vals, 'wo')
+        ax.set_title(f"{qubit_id} resonator power sweep at {hal_QPU.get_qubit_obj(qubit_id).ReadoutPower} dBm")
+        ax.set_xlabel('Frequency (Hz)')
+        ax.set_ylabel('Readout amplitude')