iqm-benchmarks 1.7__py3-none-any.whl → 1.9__py3-none-any.whl

iqm/benchmarks/compressive_gst/gst_analysis.py

@@ -0,0 +1,851 @@
+"""
+Data analysis code for compressive gate set tomography
+"""
+
+from itertools import product
+from time import perf_counter
+from typing import Any, List, Tuple, Union
+
+from matplotlib.figure import Figure
+import matplotlib.pyplot as plt
+from matplotlib.transforms import Bbox
+from numpy import ndarray
+import numpy as np
+from pandas import DataFrame
+from pygsti.models.model import Model
+from pygsti.tools import change_basis
+import xarray as xr
+
+from iqm.benchmarks.benchmark_definition import (
+    BenchmarkAnalysisResult,
+    BenchmarkObservation,
+    BenchmarkObservationIdentifier,
+    BenchmarkRunResult,
+)
+from mGST import additional_fns, algorithm, compatibility
+from mGST.low_level_jit import contract
+from mGST.qiskit_interface import qiskit_gate_to_operator
+from mGST.reporting import figure_gen, reporting
+
+
+def dataframe_to_figure(
+    df: DataFrame, row_labels: Union[List[str], None] = None, col_width: float = 2, fontsize: int = 12
+) -> Figure:
+    """Turns a pandas DataFrame into a figure
+    This is needed to conform with the standard file saving routine of QCVV.
+
+    Args:
+        df: Pandas DataFrame
+            A dataframe table containing GST results
+        row_labels: List[str]
+            The row labels for the dataframe
+        col_width: int
+            Used to control cell width in the table
+        fontsize: int
+            Font size of text/numbers in table cells
+
+    Returns:
+        figure: Matplotlib figure object
+            A figure representing the dataframe.
+    """
+
+    if row_labels is None:
+        row_labels = list(np.arange(df.shape[0]))
+
+    row_height = fontsize / 70 * 2
+    n_cols = df.shape[1]
+    n_rows = df.shape[0]
+    figsize = np.array([n_cols + 1, n_rows + 1]) * np.array([col_width, row_height])
+
+    fig, ax = plt.subplots(figsize=figsize)
+
+    fig.patch.set_visible(False)
+    ax.axis("off")
+    ax.axis("tight")
+    data_array = (df.to_numpy(dtype="str")).copy()
+    column_names = df.columns.tolist()
+    table = ax.table(
+        cellText=data_array,
+        colLabels=column_names,
+        rowLabels=row_labels,
+        cellLoc="center",
+        colColours=["#7FA1C3" for _ in range(n_cols)],
+        bbox=Bbox([[0, 0], [1, 1]]),
+    )
+    table.set_fontsize(fontsize)
+    table.set_figure(fig)
+    return fig
+
+
+def bootstrap_errors(
+    dataset: xr.Dataset,
+    y: ndarray,
+    K: ndarray,
+    X: ndarray,
+    E: ndarray,
+    rho: ndarray,
+    target_mdl: Model,
+    parametric: bool = False,
+) -> tuple[Any, Any, Any, Any, Any]:
+    """Resamples circuit outcomes a number of times and computes GST estimates for each repetition
+    All results are then returned in order to compute bootstrap-error bars for GST estimates.
+    Parametric bootstrapping uses the estimated gate set to create a newly sampled data set.
+    Non-parametric bootstrapping uses the initial dataset and resamples according to the
+    corresp. outcome probabilities.
+    Each bootstrap run is initialized with the estimated gate set in order to save processing time.
+
+    Parameters
+    ----------
+    dataset: xarray.Dataset
+        A dataset containing counts from the experiment and configurations
+    qubit_layout: List[int]
+        The list of qubits for the current GST experiment
+    y: ndarray
+        The circuit outcome probabilities as a num_povm x num_circuits array
+    K : ndarray
+        Each subarray along the first axis contains a set of Kraus operators.
+        The second axis enumerates Kraus operators for a gate specified by the first axis.
+    X : 3D ndarray
+        Array where reconstructed CPT superoperators in standard basis are stacked along the first axis.
+    E : ndarray
+        Current POVM estimate
+    rho : ndarray
+        Current initial state estimate
+    target_mdl : pygsti model object
+        The target gate set
+    parametric : bool
+        If set to True, parametric bootstrapping is used, else non-parametric bootstrapping. Default: False
+
+    Returns
+    -------
+    X_array : ndarray
+        Array containing all estimated gate tensors of different bootstrapping repetitions along first axis
+    E_array : ndarray
+        Array containing all estimated POVM tensors of different bootstrapping repetitions along first axis
+    rho_array : ndarray
+        Array containing all estimated initial states of different bootstrapping repetitions along first axis
+    df_g_array : ndarray
+        Contains gate quality measures of bootstrapping repetitions
+    df_o_array : ndarray
+        Contains SPAM and other quality measures of bootstrapping repetitions
+
+    """
+    if parametric:
+        y = np.real(
+            np.array(
+                [
+                    [E[i].conj() @ contract(X, j) @ rho for j in dataset.attrs["J"]]
+                    for i in range(dataset.attrs["num_povm"])
+                ]
+            )
+        )
+    X_array = np.zeros((dataset.attrs["bootstrap_samples"], *X.shape)).astype(complex)
+    E_array = np.zeros((dataset.attrs["bootstrap_samples"], *E.shape)).astype(complex)
+    rho_array = np.zeros((dataset.attrs["bootstrap_samples"], *rho.shape)).astype(complex)
+    df_g_list = []
+    df_o_list = []
+
+    for i in range(dataset.attrs["bootstrap_samples"]):
+        y_sampled = additional_fns.sampled_measurements(y, dataset.attrs["shots"]).copy()
+        _, X_, E_, rho_, _ = algorithm.run_mGST(
+            y_sampled,
+            dataset.attrs["J"],
+            dataset.attrs["seq_len_list"][-1],
+            dataset.attrs["num_gates"],
+            dataset.attrs["pdim"] ** 2,
+            dataset.attrs["rank"],
+            dataset.attrs["num_povm"],
+            dataset.attrs["batch_size"],
+            dataset.attrs["shots"],
+            method=dataset.attrs["opt_method"],
+            max_inits=dataset.attrs["max_inits"],
+            max_iter=0,
+            final_iter=dataset.attrs["max_iterations"][1],
+            threshold_multiplier=dataset.attrs["convergence_criteria"][0],
+            target_rel_prec=dataset.attrs["convergence_criteria"][1],
+            init=[K, E, rho],
+            testing=False,
+        )
+
+        X_opt, E_opt, rho_opt = reporting.gauge_opt(X_, E_, rho_, target_mdl, dataset.attrs[f"gauge_weights"])
+        df_g, df_o = reporting.report(
+            X_opt,
+            E_opt,
+            rho_opt,
+            dataset.attrs["J"],
+            y_sampled,
+            target_mdl,
+            dataset.attrs["gate_labels"],
+        )
+        df_g_list.append(df_g.values)
+        df_o_list.append(df_o.values)
+
+        X_opt_pp, E_opt_pp, rho_opt_pp = compatibility.std2pp(X_opt, E_opt, rho_opt)
+
+        X_array[i] = X_opt_pp
+        E_array[i] = E_opt_pp
+        rho_array[i] = rho_opt_pp
+
+    return X_array, E_array, rho_array, np.array(df_g_list), np.array(df_o_list)
+
+
+def generate_non_gate_results(
+    dataset: xr.Dataset, qubit_layout: List[int], df_o: DataFrame
+) -> tuple[DataFrame, Figure]:
+    """
+    Creates error bars (if bootstrapping was used) and formats results for non-gate errors.
+    The resulting tables are also turned into figures, so that they can be saved automatically.
+
+    Args:
+        dataset: xr.Dataset
+            A dataset containing counts from the experiment and configurations
+        qubit_layout: List[int]
+                The list of qubits for the current GST experiment
+        df_o: Pandas DataFrame
+            A dataframe containing the non-gate quality metrics (SPAM errors and fit quality)
+
+    Returns:
+        df_o_final: Pandas DataFrame
+            The final formated results
+    """
+    identifier = BenchmarkObservationIdentifier(qubit_layout).string_identifier
+    if dataset.attrs["bootstrap_samples"] > 0:
+        _, _, _, _, df_o_array = dataset.attrs["results_layout_" + identifier]["bootstrap_data"]
+        df_o_array[df_o_array == -1] = np.nan
+        percentiles_o_low, percentiles_o_high = np.nanpercentile(df_o_array, [2.5, 97.5], axis=0)
+        df_o_final = DataFrame(
+            {
+                f"Mean TVD: estimate - data": reporting.number_to_str(
+                    df_o.values[0, 1].copy(), [percentiles_o_high[0, 1], percentiles_o_low[0, 1]], precision=5
+                ),
+                f"Mean TVD: target - data": reporting.number_to_str(
+                    df_o.values[0, 2].copy(), [percentiles_o_high[0, 2], percentiles_o_low[0, 2]], precision=5
+                ),
+                f"POVM - diamond dist.": reporting.number_to_str(
+                    df_o.values[0, 3].copy(), [percentiles_o_high[0, 3], percentiles_o_low[0, 3]], precision=5
+                ),
+                f"State - trace dist.": reporting.number_to_str(
+                    df_o.values[0, 4].copy(), [percentiles_o_high[0, 4], percentiles_o_low[0, 4]], precision=5
+                ),
+            },
+            index=[""],
+        )
+    else:
+        df_o_final = DataFrame(
+            {
+                f"Mean TVD: estimate - data": reporting.number_to_str(df_o.values[0, 1].copy(), precision=5),
+                f"Mean TVD: target - data": reporting.number_to_str(df_o.values[0, 2].copy(), precision=5),
+                f"POVM - diamond dist.": reporting.number_to_str(df_o.values[0, 3].copy(), precision=5),
+                f"State - trace dist.": reporting.number_to_str(df_o.values[0, 4].copy(), precision=5),
+            },
+            index=[""],
+        )
+    fig = dataframe_to_figure(df_o_final, [""])  # dataframe_to_figure(df_o_final, [""])
+    return df_o_final, fig
+
+
+def generate_unit_rank_gate_results(
+    dataset: xr.Dataset, qubit_layout: List[int], df_g: DataFrame, X_opt: ndarray, K_target: ndarray
+) -> Tuple[DataFrame, DataFrame, Figure, Figure]:
+    """
+    Produces all result tables for Kraus rank 1 estimates and turns them into figures.
+
+    This includes parameters of the Hamiltonian generators in the Pauli basis for all gates,
+    as well as the usual performance metrics (Fidelities and Diamond distances). If bootstrapping
+    data is available, error bars will also be generated.
+
+    Args:
+        dataset: xarray.Dataset
+            A dataset containing counts from the experiment and configurations
+        qubit_layout: List[int]
+            The list of qubits for the current GST experiment
+        df_g: Pandas DataFrame
+            The dataframe with properly formated results
+        X_opt: 3D numpy array
+            The gate set after gauge optimization
+        K_target: 4D numpy array
+            The Kraus operators of all target gates, used to compute distance measures.
+
+    Returns:
+        df_g_final: Pandas DataFrame
+            The dataframe with properly formated results of standard gate errors
+        df_g_rotation Pandas DataFrame
+            A dataframe containing Hamiltonian (rotation) parameters
+        fig_g: Figure
+            A table in Figure format of gate results (fidelities etc.)
+        fig_rotation: Figure
+            A table in Figure format of gate Hamiltonian parameters
+
+
+    """
+    identifier = BenchmarkObservationIdentifier(qubit_layout).string_identifier
+    pauli_labels = generate_basis_labels(dataset.attrs["pdim"], basis="Pauli")
+    if dataset.attrs["bootstrap_samples"] > 0:
+        X_array, E_array, rho_array, df_g_array, _ = dataset.attrs["results_layout_" + identifier]["bootstrap_data"]
+        df_g_array[df_g_array == -1] = np.nan
+        percentiles_g_low, percentiles_g_high = np.nanpercentile(df_g_array, [2.5, 97.5], axis=0)
+
+        df_g_final = DataFrame(
+            {
+                r"Avg. gate fidelity": [
+                    reporting.number_to_str(
+                        df_g.values[i, 0], [percentiles_g_high[i, 0], percentiles_g_low[i, 0]], precision=5
+                    )
+                    for i in range(len(dataset.attrs["gate_labels"]))
+                ],
+                r"Diamond distance": [
+                    reporting.number_to_str(
+                        df_g.values[i, 1], [percentiles_g_high[i, 1], percentiles_g_low[i, 1]], precision=5
+                    )
+                    for i in range(dataset.attrs["num_gates"])
+                ],
+            }
+        )
+
+        U_opt = reporting.phase_opt(X_opt, K_target)
+        pauli_coeffs = reporting.compute_sparsest_Pauli_Hamiltonian(U_opt)
+
+        bootstrap_pauli_coeffs = np.zeros((len(X_array), dataset.attrs["num_gates"], dataset.attrs["pdim"] ** 2))
+        for i, X_ in enumerate(X_array):
+            X_std, _, _ = compatibility.pp2std(X_, E_array[i], rho_array[i])
+            U_opt_ = reporting.phase_opt(
+                np.array([change_basis(X_std[j], "pp", "std") for j in range(dataset.attrs["num_gates"])]), K_target
+            )
+            pauli_coeffs_ = reporting.compute_sparsest_Pauli_Hamiltonian(U_opt_)
+            bootstrap_pauli_coeffs[i, :, :] = pauli_coeffs_
+        pauli_coeffs_low, pauli_coeffs_high = np.nanpercentile(bootstrap_pauli_coeffs, [2.5, 97.5], axis=0)
+
+        df_g_rotation = DataFrame(
+            np.array(
+                [
+                    [
+                        reporting.number_to_str(
+                            pauli_coeffs[i, j], [pauli_coeffs_high[i, j], pauli_coeffs_low[i, j]], precision=5
+                        )
+                        for i in range(dataset.attrs["num_gates"])
+                    ]
+                    for j in range(dataset.attrs["pdim"] ** 2)
+                ]
+            ).T
+        )
+
+        df_g_rotation.columns = [f"h_%s" % label for label in pauli_labels]
+        df_g_rotation.rename(index=dataset.attrs["gate_labels"], inplace=True)
+
+    else:
+        df_g_final = DataFrame(
+            {
+                "Avg. gate fidelity": [
+                    reporting.number_to_str(df_g.values[i, 0], precision=5) for i in range(dataset.attrs["num_gates"])
+                ],
+                "Diamond distance": [
+                    reporting.number_to_str(df_g.values[i, 1], precision=5) for i in range(dataset.attrs["num_gates"])
+                ],
+            }
+        )
+        U_opt = reporting.phase_opt(X_opt, K_target)
+        pauli_coeffs = reporting.compute_sparsest_Pauli_Hamiltonian(U_opt)
+
+        df_g_rotation = DataFrame(
+            np.array(
+                [
+                    [
+                        reporting.number_to_str(pauli_coeffs[i, j], precision=5)
+                        for i in range(dataset.attrs["num_gates"])
+                    ]
+                    for j in range(dataset.attrs["pdim"] ** 2)
+                ]
+            ).T
+        )
+        df_g_rotation.columns = [f"h_%s" % label for label in pauli_labels]
+        df_g_final.rename(index=dataset.attrs["gate_labels"], inplace=True)
+
+    fig_g = dataframe_to_figure(df_g_final, dataset.attrs["gate_labels"])
+    fig_rotation = dataframe_to_figure(df_g_rotation, dataset.attrs["gate_labels"])
+    return df_g_final, df_g_rotation, fig_g, fig_rotation
+
+
+def generate_gate_results(
+    dataset: xr.Dataset,
+    qubit_layout: List[int],
+    df_g: DataFrame,
+    X_opt: ndarray,
+    E_opt: ndarray,
+    rho_opt: ndarray,
+    max_evals: int = 6,
+) -> Tuple[DataFrame, DataFrame, Figure, Figure]:
+    """
+    Produces all result tables for arbitrary Kraus rank estimates and turns them into figures.
+
+    Args:
+        df_g: Pandas DataFrame
+            The dataframe with properly formated results
+        X_opt: 3D numpy array
+            The gate set after gauge optimization
+        E_opt: 3D numpy array
+            An array containg all the POVM elements as matrices after gauge optimization
+        rho_opt: 2D numpy array
+            The density matrix after gauge optmization
+        max_evals: int
+            The maximum number of eigenvalues of the Choi matrices which are returned.
+
+    Returns:
+        df_g_final: Pandas DataFrame
+            The dataframe with properly formated results of standard gate errors
+        df_g_evals_final Pandas DataFrame
+            A dataframe containing eigenvalues of the Choi matrices for all gates
+        fig_g: Figure
+            A table in Figure format of gate results (fidelities etc.)
+        fig_choi: Figure
+            A table in Figure format of eigenvalues of the Choi matrices of all gates
+
+    """
+    identifier = BenchmarkObservationIdentifier(qubit_layout).string_identifier
+    n_evals = np.min([max_evals, dataset.attrs["pdim"] ** 2])
+    X_opt_pp, _, _ = compatibility.std2pp(X_opt, E_opt, rho_opt)
+    df_g_evals = reporting.generate_Choi_EV_table(X_opt, n_evals, dataset.attrs["gate_labels"])
+
+    if dataset.attrs["bootstrap_samples"] > 0:
+        X_array, E_array, rho_array, df_g_array, _ = dataset.attrs["results_layout_" + identifier]["bootstrap_data"]
+        df_g_array[df_g_array == -1] = np.nan
+        percentiles_g_low, percentiles_g_high = np.nanpercentile(df_g_array, [2.5, 97.5], axis=0)
+        bootstrap_unitarities = np.array(
+            [reporting.unitarities(X_array[i]) for i in range(dataset.attrs["bootstrap_samples"])]
+        )
+        percentiles_u_low, percentiles_u_high = np.nanpercentile(bootstrap_unitarities, [2.5, 97.5], axis=0)
+        X_array_std = [
+            compatibility.pp2std(X_array[i], E_array[i], rho_array[i])[0]
+            for i in range(dataset.attrs["bootstrap_samples"])
+        ]
+        bootstrap_evals = np.array(
+            [
+                reporting.generate_Choi_EV_table(X_array_std[i], n_evals, dataset.attrs["gate_labels"])
+                for i in range(dataset.attrs["bootstrap_samples"])
+            ]
+        )
+        percentiles_evals_low, percentiles_evals_high = np.nanpercentile(bootstrap_evals, [2.5, 97.5], axis=0)
+        eval_strs = [
+            [
+                reporting.number_to_str(
+                    df_g_evals.values[i, j],
+                    [percentiles_evals_high[i, j], percentiles_evals_low[i, j]],
+                    precision=5,
+                )
+                for i in range(dataset.attrs["num_gates"])
+            ]
+            for j in range(n_evals)
+        ]
+
+        df_g_final = DataFrame(
+            {
+                r"Avg. gate fidelity": [
+                    reporting.number_to_str(
+                        df_g.values[i, 0], [percentiles_g_high[i, 0], percentiles_g_low[i, 0]], precision=5
+                    )
+                    for i in range(dataset.attrs["num_gates"])
+                ],
+                r"Diamond distance": [
+                    reporting.number_to_str(
+                        df_g.values[i, 1], [percentiles_g_high[i, 1], percentiles_g_low[i, 1]], precision=5
+                    )
+                    for i in range(dataset.attrs["num_gates"])
+                ],
+                r"Unitarity": [
+                    reporting.number_to_str(
+                        reporting.unitarities(X_opt_pp)[i],
+                        [percentiles_u_high[i], percentiles_u_low[i]],
+                        precision=5,
+                    )
+                    for i in range(dataset.attrs["num_gates"])
+                ],
+            }
+        )
+
+    else:
+        df_g_final = DataFrame(
+            {
+                "Avg. gate fidelity": [
+                    reporting.number_to_str(df_g.values[i, 0].copy(), precision=5)
+                    for i in range(len(dataset.attrs["gate_labels"]))
+                ],
+                "Diamond distance": [
+                    reporting.number_to_str(df_g.values[i, 1].copy(), precision=5)
+                    for i in range(len(dataset.attrs["gate_labels"]))
+                ],
+                "Unitarity": [
+                    reporting.number_to_str(reporting.unitarities(X_opt_pp)[i], precision=5)
+                    for i in range(len(dataset.attrs["gate_labels"]))
+                ],
+                # "Entanglemen fidelity to depol. channel": [reporting.number_to_str(reporting.eff_depol_params(X_opt_pp)[i], precision=5)
+                #                                            for i in range(len(gate_labels))],
+                # "Min. spectral distances": [number_to_str(df_g.values[i, 2], precision=5) for i in range(len(gate_labels))]
+            }
+        )
+        eval_strs = [
+            [
+                reporting.number_to_str(df_g_evals.values[i, j].copy(), precision=5)
+                for i in range(dataset.attrs["num_gates"])
+            ]
+            for j in range(n_evals)
+        ]
+
+    df_g_evals_final = DataFrame(eval_strs).T
+    df_g_evals_final.rename(index=dataset.attrs["gate_labels"], inplace=True)
+
+    fig_g = dataframe_to_figure(df_g_final, dataset.attrs["gate_labels"])
+    fig_choi = dataframe_to_figure(df_g_evals_final, dataset.attrs["gate_labels"])
+    return df_g_final, df_g_evals_final, fig_g, fig_choi
+
+
+def generate_basis_labels(pdim: int, basis: Union[str, None] = None) -> List[str]:
+    """Generate a list of labels for the Pauli basis or the standard basis
+
+    Args:
+        pdim: int
+            Physical dimension
+        basis: str
+            Which basis the labels correspond to, currently default is standard basis and "Pauli" can be choose
+            for Pauli basis labels like "II", "IX", "XX", ...
+
+    Returns:
+        labels: List[str]
+            A list of all string combinations for the given dimension and basis
+    """
+    separator = ""
+    if basis == "Pauli":
+        pauli_labels_loc = ["I", "X", "Y", "Z"]
+        pauli_labels_rep = [pauli_labels_loc for _ in range(int(np.log2(pdim)))]
+        labels = [separator.join(map(str, x)) for x in product(*pauli_labels_rep)]
+    else:
+        std_labels_loc = ["0", "1"]
+        std_labels_rep = [std_labels_loc for _ in range(int(np.log2(pdim)))]
+        labels = [separator.join(map(str, x)) for x in product(*std_labels_rep)]
+
+    return labels
+
+
+def result_str_to_floats(result_str: str, err: str) -> Tuple[float, float]:
+    """Converts formated string results from mgst to float (value, uncertainty) pairs
+
+    Args:
+        result_str: str
+            The value of a result parameter formated as str
+        err: str
+            The error interval of the parameters
+
+    Returns:
+        value: float
+            The parameter value as floar
+        uncertainty: float
+            A single uncertainty value
+    """
+    if err:
+        value = float(result_str.split("[")[0])
+        rest = result_str.split("[")[1].split(",")
+        uncertainty = float(rest[1][:-1]) - float(rest[0])
+        return value, uncertainty
+    return float(result_str), np.NaN
+
+
+def pandas_results_to_observations(
+    dataset: xr.Dataset, df_g: DataFrame, df_o: DataFrame, identifier: BenchmarkObservationIdentifier
+) -> List[BenchmarkObservation]:
+    """Converts high level GST results from a pandas Dataframe to a simple observation dictionary
+
+    Args:
+    dataset: xarray.Dataset
+        A dataset containing counts from the experiment and configurations
+    qubit_layout: List[int]
+        The list of qubits for the current GST experiment
+    df_g: Pandas DataFrame
+        The dataframe with properly formated gate results
+    df_o: Pandas DataFrame
+        The dataframe with properly formated non-gate results like SPAM error measures or fit quality.
+    identifier: BenchmarkObservationIdentifier
+        An identifier object for the current GST run
+
+    Returns:
+    observation_list: List[BenchmarkObservation]
+        List of observations converted from the pandas dataframes
+    """
+    observation_list: list[BenchmarkObservation] = []
+    err = dataset.attrs["bootstrap_samples"] > 0
+    for idx, gate_label in enumerate(dataset.attrs["gate_labels"].values()):
+        observation_list.extend(
+            [
+                BenchmarkObservation(
+                    name=f"{name} - {gate_label}",
+                    identifier=identifier,
+                    value=result_str_to_floats(df_g[name].iloc[idx], err)[0],
+                    uncertainty=result_str_to_floats(df_g[name].iloc[idx], err)[1],
+                )
+                for name in df_g.columns.tolist()
+            ]
+        )
+    observation_list.extend(
+        [
+            BenchmarkObservation(
+                name=f"{name}",
+                identifier=identifier,
+                value=result_str_to_floats(df_o[name].iloc[0], err)[0],
+                uncertainty=result_str_to_floats(df_o[name].iloc[0], err)[1],
+            )
+            for name in df_o.columns.tolist()
+        ]
+    )
+    return observation_list
+
+
+def dataset_counts_to_mgst_format(dataset: xr.Dataset, qubit_layout: List[int]) -> ndarray:
+    """Turns the dictionary of outcomes obtained from qiskit backend
+        into the format which is used in mGST
+
+    Args:
+        dataset: xarray.Dataset
+            A dataset containing counts from the experiment and configurations
+        qubit_layout: List[int]
+            The list of qubits for the current GST experiment
+
+    Returns
+    -------
+    y : numpy array
+        2D array of measurement outcomes for sequences in J;
+        Each column contains the outcome probabilities for a fixed sequence
+
+    """
+    num_qubits = len(qubit_layout)
+    num_povm = dataset.attrs["num_povm"]
+    y_list = []
+    for run_index in range(dataset.attrs["num_circuits"]):
+        # result = dataset[f"{qubit_layout}_counts_{run_index}"].to_dict()
+        result = dataset[f"{qubit_layout}_state_{run_index}"].data.tolist()
+        basis_dict = {entry: int("".join([entry[::-1][i] for i in range(num_qubits)][::-1]), 2) for entry in result}
+        # Sort by index:
+        basis_dict = dict(sorted(basis_dict.items(), key=lambda item: item[1]))
+
+        counts_normalized = (
+            dataset[f"{qubit_layout}_counts_{run_index}"] / dataset[f"{qubit_layout}_counts_{run_index}"].sum()
+        )
+        row = [counts_normalized.loc[key].data for key in basis_dict]
+        # row = [result[key] for key in basis_dict]
+        if len(row) < num_povm:
+            missing_entries = list(np.arange(num_povm))
+            for given_entry in basis_dict.values():
+                missing_entries.remove(given_entry)
+            for missing_entry in missing_entries:
+                row.insert(missing_entry, 0)  # 0 measurement outcomes in not recorded entry
+        y_list.append(row)
+    y = np.array(y_list).T
+    return y
+
+
+def run_mGST(
+    dataset: xr.Dataset, y: ndarray
+) -> tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray, ndarray, ndarray]:
+    """Wrapper function for mGST algorithm execution which prepares an initialization and sets the alg. parameters
+
+    Args:
+        dataset: xarray.Dataset
+            A dataset containing counts from the experiment and configurations
+        y: ndarray
+            The circuit outcome probabilities as a num_povm x num_circuits array
+
+    Returns:
+        K : ndarray
+            Kraus estimate array where each subarray along the first axis contains a set of Kraus operators.
+            The second axis enumerates Kraus operators for a gate specified by the first axis.
+        X : ndarray
+            Superoperator estimate array where reconstructed CPT superoperators in
+            standard basis are stacked along the first axis.
+        E : ndarray
+            Current POVM estimate
+        rho : ndarray
+            Current initial state estimate
+        K_target : ndarray
+            Target gate Kraus array where each subarray along the first axis contains a set of Kraus operators.
+            The second axis enumerates Kraus operators for a gate specified by the first axis.
+        X_target : ndarray
+            Target gate superoperator estimate array where reconstructed CPT superoperators in
+            standard basis are stacked along the first axis.
+        E_target : ndarray
+            Target POVM
+        rho_target : ndarray
+            Target initial state
+    """
+
+    K_target = qiskit_gate_to_operator(dataset.attrs["gate_set"])
+    X_target = np.einsum("ijkl,ijnm -> iknlm", K_target, K_target.conj()).reshape(
+        (dataset.attrs["num_gates"], dataset.attrs["pdim"] ** 2, dataset.attrs["pdim"] ** 2)
+    )  # tensor of superoperators
+
+    rho_target = (
+        np.kron(additional_fns.basis(dataset.attrs["pdim"], 0).T.conj(), additional_fns.basis(dataset.attrs["pdim"], 0))
+        .reshape(-1)
+        .astype(np.complex128)
+    )
+
+    # Computational basis measurement:
+    E_target = np.array(
+        [
+            np.kron(
+                additional_fns.basis(dataset.attrs["pdim"], i).T.conj(), additional_fns.basis(dataset.attrs["pdim"], i)
+            ).reshape(-1)
+            for i in range(dataset.attrs["pdim"])
+        ]
+    ).astype(np.complex128)
+
+    # Run mGST
+    if dataset.attrs["from_init"]:
+        K_init = additional_fns.perturbed_target_init(X_target, dataset.attrs["rank"])
+        init_params = [K_init, E_target, rho_target]
+    else:
+        init_params = None
+
+    K, X, E, rho, _ = algorithm.run_mGST(
+        y,
+        dataset.attrs["J"],
+        dataset.attrs["seq_len_list"][-1],
+        dataset.attrs["num_gates"],
+        dataset.attrs["pdim"] ** 2,
+        dataset.attrs["rank"],
+        dataset.attrs["num_povm"],
+        dataset.attrs["batch_size"],
+        dataset.attrs["shots"],
+        method=dataset.attrs["opt_method"],
+        max_inits=dataset.attrs["max_inits"],
+        max_iter=dataset.attrs["max_iterations"][0],
+        final_iter=dataset.attrs["max_iterations"][1],
+        threshold_multiplier=dataset.attrs["convergence_criteria"][0],
+        target_rel_prec=dataset.attrs["convergence_criteria"][1],
+        init=init_params,
+        testing=False,
+    )
+
+    return K, X, E, rho, K_target, X_target, E_target, rho_target
+
+
+def mgst_analysis(run: BenchmarkRunResult) -> BenchmarkAnalysisResult:
+    """Analysis function for compressive GST
+
+    Args:
+        run: BenchmarkRunResult
+            A BenchmarkRunResult instance storing the dataset
+    Returns:
+        result: BenchmarkAnalysisResult
+            An BenchmarkAnalysisResult instance with the updated dataset, as well as plots and observations
+    """
+    dataset = run.dataset
+    pdim = dataset.attrs["pdim"]
+    plots = {}
+    for i, qubit_layout in enumerate(dataset.attrs["qubit_layouts"]):
+        identifier = BenchmarkObservationIdentifier(qubit_layout).string_identifier
+
+        # Computing circuit outcome probabilities from counts
+        y = dataset_counts_to_mgst_format(dataset, qubit_layout)
+
+        # Main GST reconstruction
+        start_timer = perf_counter()
+        K, X, E, rho, K_target, X_target, E_target, rho_target = run_mGST(dataset, y)
+        main_gst_time = perf_counter() - start_timer
+
+        # Gauge optimization
+        start_timer = perf_counter()
+        target_mdl = compatibility.arrays_to_pygsti_model(X_target, E_target, rho_target, basis="std")
+        X_opt, E_opt, rho_opt = reporting.gauge_opt(X, E, rho, target_mdl, dataset.attrs[f"gauge_weights"])
+        gauge_optimization_time = perf_counter() - start_timer
+
+        # Quick report
+        df_g, _ = reporting.quick_report(
+            X_opt, E_opt, rho_opt, dataset.attrs["J"], y, target_mdl, dataset.attrs["gate_labels"]
+        )
+
+        # Gate set in the Pauli basis
+        X_opt_pp, _, _ = compatibility.std2pp(X_opt, E_opt, rho_opt)
+        X_target_pp, _, _ = compatibility.std2pp(X_target, E_target, rho_target)
+
+        # Saving
+        dataset.attrs["results_layout_" + identifier] = {
+            "raw_Kraus_operators": K,
+            "raw_gates": X,
+            "raw_POVM": E.reshape((dataset.attrs["num_povm"], pdim, pdim)),
+            "raw_state": rho.reshape((pdim, pdim)),
+            "gauge_opt_gates": X_opt,
+            "gauge_opt_gates_Pauli_basis": X_opt_pp,
+            "gauge_opt_POVM": E_opt.reshape((dataset.attrs["num_povm"], pdim, pdim)),
+            "gauge_opt_state": rho_opt.reshape((pdim, pdim)),
+            "main_mGST_time": main_gst_time,
+            "gauge_optimization_time": gauge_optimization_time,
+        }
+
+        ### Bootstrap
+        if dataset.attrs["bootstrap_samples"] > 0:
+            bootstrap_results = bootstrap_errors(dataset, y, K, X, E, rho, target_mdl)
+            dataset.attrs["results_layout_" + identifier].update({"bootstrap_data": bootstrap_results})
+
+        _, df_o_full = reporting.report(
+            X_opt, E_opt, rho_opt, dataset.attrs["J"], y, target_mdl, dataset.attrs["gate_labels"]
+        )
+        df_o_final, fig_o = generate_non_gate_results(dataset, qubit_layout, df_o_full)
+
+        ### Result table generation and full report
+        if dataset.attrs["rank"] == 1:
+            df_g_final, df_g_rotation, fig_g, fig_rotation = generate_unit_rank_gate_results(
+                dataset, qubit_layout, df_g, X_opt, K_target
+            )
+            dataset.attrs["results_layout_" + identifier].update({"hamiltonian_parameters": df_g_rotation.to_dict()})
+            plots[f"layout_{qubit_layout}_hamiltonian_parameters"] = fig_rotation
+        else:
+            df_g_final, df_g_evals, fig_g, fig_choi = generate_gate_results(
+                dataset, qubit_layout, df_g, X_opt, E_opt, rho_opt
+            )
+            dataset.attrs["results_layout_" + identifier].update({"choi_evals": df_g_evals.to_dict()})
+            plots[f"layout_{qubit_layout}_choi_eigenvalues"] = fig_choi
+        plots[f"layout_{qubit_layout}_gate_metrics"] = fig_g
+        plots[f"layout_{qubit_layout}_other_metrics"] = fig_o
+
+        observation_list = pandas_results_to_observations(
+            dataset, df_g_final, df_o_final, BenchmarkObservationIdentifier(qubit_layout)
+        )
+
+        dataset.attrs["results_layout_" + identifier].update(
+            {"full_metrics": {"Gates": df_g_final.to_dict(), "Outcomes and SPAM": df_o_final.to_dict()}}
+        )
+
+        ### Process matrix plots
+        pauli_labels = generate_basis_labels(pdim, basis="Pauli")
+        std_labels = generate_basis_labels(pdim)
+
+        figures = figure_gen.generate_gate_err_pdf(
+            "",
+            X_opt_pp,
+            X_target_pp,
+            basis_labels=pauli_labels,
+            gate_labels=dataset.attrs["gate_labels"],
+            return_fig=True,
+        )
+        for i, figure in enumerate(figures):
+            plots[f"layout_{qubit_layout}_process_matrix_{i}"] = figure
+
+        plots[f"layout_{qubit_layout}_SPAM_matrices_real"] = figure_gen.generate_spam_err_std_pdf(
+            "",
+            E_opt.real,
+            rho_opt.real,
+            E_target.real,
+            rho_target.real,
+            basis_labels=std_labels,
+            title=f"Real part of state and measurement effects in the standard basis",
+            return_fig=True,
+        )
+        plots[f"layout_{qubit_layout}_SPAM_matrices_imag"] = figure_gen.generate_spam_err_std_pdf(
+            "",
+            E_opt.imag,
+            rho_opt.imag,
+            E_target.imag,
+            rho_target.imag,
+            basis_labels=std_labels,
+            title=f"Imaginary part of state and measurement effects in the standard basis",
+            return_fig=True,
+        )
+        plt.close("all")
+
+    return BenchmarkAnalysisResult(dataset=dataset, observations=observation_list, plots=plots)