sqil-core 0.1.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

sqil_core/utils/_plot.py

@@ -1,8 +1,26 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import matplotlib.pyplot as plt
 import numpy as np
+from matplotlib.gridspec import GridSpec
+
+from sqil_core.fit import transform_data
+
+from ._analysis import remove_linear_background, remove_offset, soft_normalize
+from ._const import PARAM_METADATA
+from ._formatter import (
+    ParamInfo,
+    format_number,
+    get_relevant_exp_parameters,
+    param_info_from_schema,
+)
+from ._read import extract_h5_data, get_data_and_info, map_data_dict, read_json
 
-from ._const import _PARAM_METADATA
-from ._formatter import format_number
-from ._read import read_json
+if TYPE_CHECKING:
+    from sqil_core.fit._core import FitResult
+    from sqil_core.utils import ParamDict
 
 
 def set_plot_style(plt):
@@ -12,7 +30,7 @@ def set_plot_style(plt):
         "xtick.labelsize": 18,  # X-axis tick labels
         "ytick.labelsize": 18,  # Y-axis tick labels
         "lines.linewidth": 2.5,  # Line width
-        "lines.marker": "o",
+        # "lines.marker": "o",
         "lines.markersize": 7,  # Marker size
         "lines.markeredgewidth": 1.5,  # Marker line width
         "lines.markerfacecolor": "none",
@@ -24,6 +42,7 @@ def set_plot_style(plt):
         "ytick.major.width": 1.5,
         "figure.figsize": (20, 7),
     }
+    reset_plot_style(plt)
     return plt.rcParams.update(style)
 
 
@@ -67,10 +86,10 @@ def build_title(title: str, path: str, params: list[str]) -> str:
     dic = read_json(f"{path}/param_dict.json")
     title += " with "
     for idx, param in enumerate(params):
-        if not (param in _PARAM_METADATA.keys()) or not (param in dic):
+        if not (param in PARAM_METADATA.keys()) or not (param in dic):
             title += f"{param} = ? & "
             continue
-        meta = _PARAM_METADATA[param]
+        meta = PARAM_METADATA[param]
         value = format_number(dic[param], 3, meta["unit"])
         title += f"${meta['symbol']} =${value} & "
         if idx % 2 == 0 and idx != 0:
@@ -105,3 +124,250 @@ def guess_plot_dimension(
         return "1.5"
     else:
         return "1"
+
+
+def finalize_plot(
+    fig,
+    title,
+    fit_res: FitResult = None,
+    qubit_params: ParamDict = {},
+    updated_params: dict = {},
+    sweep_info={},
+    relevant_params=[],
+):
+    """
+    Annotates a matplotlib figure with experiment parameters, fit quality, and title.
+
+    Parameters
+    ----------
+    fig : matplotlib.figure.Figure
+        The figure object to annotate.
+    title : str
+        Title text to use for the plot.
+    fit_res : FitResult, optional
+        Fit result object containing model name and quality summary.
+    qubit_params : ParamDict, optional
+        Dictionary of experimental qubit parameters, indexed by parameter ID.
+    updated_params : dict, optional
+        Dictionary of updated parameters (e.g., from fitting), where keys are param IDs
+        and values are numeric or symbolic parameter values.
+    sweep_info : dict, optional
+        Information about sweep parameters (e.g., their IDs and labels).
+    relevant_params : list, optional
+        List of parameter IDs considered relevant for display under "Experiment".
+    """
+    # Make a summary of relevant experimental parameters
+    exp_params_keys = get_relevant_exp_parameters(
+        qubit_params, relevant_params, [info.id for info in sweep_info]
+    )
+    params_str = ",   ".join(
+        [qubit_params[id].symbol_and_value for id in exp_params_keys]
+    )
+    # Make a summary of the updated qubit parameters
+    updated_params_info = {k: ParamInfo(k, v) for k, v in updated_params.items()}
+    update_params_str = ",   ".join(
+        [updated_params_info[id].symbol_and_value for id in updated_params_info.keys()]
+    )
+
+    # Find appropriate y_position to print text
+    bbox = fig.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+    fig_height_inches = bbox.height
+    if fig_height_inches < 8:
+        y_pos = -0.05
+    elif fig_height_inches < 10:
+        y_pos = -0.03
+    elif fig_height_inches < 13:
+        y_pos = -0.02
+    else:
+        y_pos = -0.01
+
+    # Add text to the plot
+    fig.suptitle(f"{title}\n" + update_params_str)
+    if fit_res:
+        fig.text(0.02, y_pos, f"Model: {fit_res.model_name} - {fit_res.quality()}")
+    if params_str:
+        fig.text(0.4, y_pos, "Experiment:   " + params_str, ha="left")
+
+
+def plot_mag_phase(path=None, datadict=None, raw=False):
+    """
+    Plot the magnitude and phase of complex measurement data from an db path or in-memory dictionary.
+
+    This function generates either a 1D or 2D plot of the magnitude and phase of complex data,
+    depending on the presence of sweep parameters. It supports normalization and background
+    subtraction.
+
+    Parameters
+    ----------
+    path : str or None, optional
+        Path to the folder containing measurement data. Required if `datadict` is not provided.
+    datadict : dict or None, optional
+        Pre-loaded data dictionary with schema, typically extracted using `extract_h5_data`.
+        Required if `path` is not provided.
+    raw : bool, default False
+        If True, skip normalization and background subtraction for 2D plots. Useful for viewing raw data.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The matplotlib Figure object containing the plot.
+    axs : matplotlib.axes.Axes or ndarray of Axes
+        The Axes object(s) used for the subplot(s).
+
+    Raises
+    ------
+    Exception
+        If neither `path` nor `datadict` is provided.
+
+    Notes
+    -----
+    - Axes and units are automatically inferred from the schema in the dataset.
+    """
+
+    all_data, all_info, _ = get_data_and_info(path=path, datadict=datadict)
+    x_data, y_data, sweeps = all_data
+    x_info, y_info, sweep_info = all_info
+
+    # Rescale data
+    x_data_scaled = x_data * x_info.scale
+    y_data_scaled = y_data * y_info.scale
+    y_unit = f" [{y_info.rescaled_unit}]" if y_info.unit else ""
+
+    set_plot_style(plt)
+
+    if len(sweeps) == 0:  # 1D plot
+        fig, axs = plt.subplots(2, 1, figsize=(20, 12), sharex=True)
+
+        axs[0].plot(x_data_scaled, np.abs(y_data_scaled), "o")
+        axs[0].set_ylabel("Magnitude" + y_unit)
+        axs[0].tick_params(labelbottom=True)
+        axs[0].xaxis.set_tick_params(
+            which="both", labelbottom=True
+        )  # Redundant for safety
+
+        axs[1].plot(x_data_scaled, np.unwrap(np.angle(y_data_scaled)), "o")
+        axs[1].set_xlabel(x_info.name_and_unit)
+        axs[1].set_ylabel("Phase [rad]")
+    else:  # 2D plot
+        fig, axs = plt.subplots(1, 2, figsize=(24, 12), sharex=True, sharey=True)
+
+        # Process mag and phase
+        mag, phase = np.abs(y_data), np.unwrap(np.angle(y_data))
+        if not raw:
+            mag = soft_normalize(remove_offset(mag))
+            flat_phase = remove_linear_background(x_data, phase, points_cut=1)
+            phase = soft_normalize(flat_phase)
+        # Load sweep parameter
+        sweep0_info = sweep_info[0]
+        sweep0_scaled = sweeps[0] * sweep0_info.scale
+
+        c0 = axs[0].pcolormesh(
+            x_data_scaled,
+            sweep0_scaled,
+            mag,
+            shading="auto",
+            cmap="PuBu",
+        )
+        if raw:
+            fig.colorbar(c0, ax=axs[0])
+            axs[0].set_title("Magnitude" + y_unit)
+        else:
+            axs[0].set_title("Magnitude (normalized)")
+        axs[0].set_xlabel(x_info.name_and_unit)
+        axs[0].set_ylabel(sweep0_info.name_and_unit)
+
+        c1 = axs[1].pcolormesh(
+            x_data_scaled,
+            sweep0_scaled,
+            phase,
+            shading="auto",
+            cmap="PuBu",
+        )
+        if raw:
+            fig.colorbar(c1, ax=axs[1])
+            axs[1].set_title("Phase [rad]")
+        else:
+            axs[1].set_title("Phase (normalized)")
+        axs[1].set_xlabel(x_info.name_and_unit)
+        axs[1].tick_params(labelleft=True)
+        axs[1].xaxis.set_tick_params(
+            which="both", labelleft=True
+        )  # Redundant for safety
+
+    fig.tight_layout()
+    return fig, axs
+
+
+def plot_projection_IQ(path=None, datadict=None, proj_data=None, full_output=False):
+    """
+    Plots the real projection of complex I/Q data versus the x-axis and the full IQ plane.
+
+    Parameters
+    ----------
+    path : str, optional
+        Path to the HDF5 file containing the data. Required if `datadict` is not provided.
+    datadict : dict, optional
+        Pre-loaded data dictionary with schema, typically extracted using `extract_h5_data`.
+        Required if `path` is not provided.
+    proj_data : np.ndarray, optional
+        Precomputed projected data (real part of transformed complex values).
+        If not provided, it will be computed using `transform_data`.
+    full_output : bool, default False
+        Whether to return projected data and the inverse transformation function.
+
+    Returns
+    -------
+    res : tuple
+        If `full_output` is False:
+            (fig, [ax_proj, ax_iq])
+        If `full_output` is True:
+            (fig, [ax_proj, ax_iq], proj_data, inv)
+        - `fig`: matplotlib Figure object.
+        - `ax_proj`: Axis for projection vs x-axis.
+        - `ax_iq`: Axis for I/Q scatter plot.
+        - `proj_data`: The real projection of the complex I/Q data.
+        - `inv`: The inverse transformation function used during projection.
+
+    Notes
+    -----
+    This function supports only 1D datasets. If sweep dimensions are detected, no plot is created.
+    The projection is performed using a data transformation routine (e.g., PCA or rotation).
+    """
+
+    all_data, all_info, _ = get_data_and_info(path=path, datadict=datadict)
+    x_data, y_data, sweeps = all_data
+    x_info, y_info, sweep_info = all_info
+
+    # Get y_unit
+    y_unit = f" [{y_info.rescaled_unit}]" if y_info.unit else ""
+
+    set_plot_style(plt)
+
+    if len(sweeps) == 0:
+        # Project data
+        if proj_data is None:
+            proj_data, inv = transform_data(y_data, inv_transform=True)
+
+        set_plot_style(plt)
+        fig = plt.figure(figsize=(20, 7), constrained_layout=True)
+        gs = GridSpec(nrows=1, ncols=10, figure=fig, wspace=0.2)
+
+        # Plot the projection
+        ax_proj = fig.add_subplot(gs[:, :6])  # 6/10 width
+        ax_proj.plot(x_data * x_info.scale, proj_data.real * y_info.scale, "o")
+        ax_proj.set_xlabel(x_info.name_and_unit)
+        ax_proj.set_ylabel("Projected" + y_unit)
+
+        # Plot IQ data
+        ax_iq = fig.add_subplot(gs[:, 6:])  # 4/10 width
+        ax_iq.scatter(0, 0, marker="+", color="black", s=150)
+        ax_iq.plot(y_data.real * y_info.scale, y_data.imag * y_info.scale, "o")
+        ax_iq.set_xlabel("In-Phase" + y_unit)
+        ax_iq.set_ylabel("Quadrature" + y_unit)
+        ax_iq.set_aspect(aspect="equal", adjustable="datalim")
+
+    if full_output:
+        res = (fig, [ax_proj, ax_iq], proj_data, inv)
+    else:
+        res = (fig, [ax_proj, ax_iq])
+    return res

sqil_core/utils/_read.py

@@ -1,14 +1,26 @@
+from __future__ import annotations
+
 import json
 import os
+import shutil
+from typing import TYPE_CHECKING
 
 import h5py
 import numpy as np
+import yaml
+from laboneq import serializers
+
+from sqil_core.utils._formatter import param_info_from_schema
+
+from ._const import _EXP_UNIT_MAP, PARAM_METADATA
 
-from ._const import _EXP_UNIT_MAP, _PARAM_METADATA
+if TYPE_CHECKING:
+    from laboneq.dsl.quantum.qpu import QPU
 
 
+# TODO: add tests for schema
 def extract_h5_data(
-    path: str, keys: list[str] | None = None
+    path: str, keys: list[str] | None = None, schema=False
 ) -> dict | tuple[np.ndarray, ...]:
     """Extract data at the given keys from an HDF5 file. If no keys are
     given (None) returns the data field of the object.
@@ -42,6 +54,11 @@ def extract_h5_data(
     with h5py.File(path, "r") as h5file:
         data = h5file["data"]
         data_keys = data.keys()
+
+        db_schema = None
+        if schema:
+            db_schema = json.loads(data.attrs.get("__schema__"))
+
         # Extract only the requested keys
         if bool(keys) and (len(keys) > 0):
             res = []
@@ -51,9 +68,13 @@ def extract_h5_data(
                     res.append([])
                     continue
                 res.append(np.array(data[key][:]))
-            return tuple(res)
+            if not schema and len(res) == 1:
+                return res[0]
+            return tuple(res) if not schema else (*tuple(res), db_schema)
         # Extract the whole data dictionary
-        return _h5_to_dict(data)
+        h5_dict = _h5_to_dict(data)
+        return h5_dict if not schema else {**h5_dict, "schema": db_schema}
+    #
 
 
 def _h5_to_dict(obj) -> dict:
@@ -68,112 +89,174 @@ def _h5_to_dict(obj) -> dict:
     return data_dict
 
 
-def read_json(path: str) -> dict:
-    """Reads a json file and returns the data as a dictionary."""
-    with open(path) as f:
-        dictionary = json.load(f)
-    return dictionary
+def map_data_dict(data_dict: dict):
+    """
+    Maps experimental data to standardized arrays using a provided schema.
 
+    This function interprets the structure of a measurement data dictionary
+    (obtained using extract_h5_data) by extracting relevant data fields according
+    to roles specified in the database schema. It returns the x-axis values, y-axis data,
+    any additional sweep parameters, and a mapping of keys used for each role.
 
-class ParamInfo:
-    """Parameter information for items of param_dict
+    Parameters
+    ----------
+    data_dict : dict
+        Dictionary containing measurement data and an associated 'schema' key
+        that defines the role of each field (e.g., "x-axis", "data", "axis").
 
-    Attributes:
-        id (str): param_dict key
-        value (any): the value of the parameter
-        name (str): full name of the parameter (e.g. Readout frequency)
-        symbol (str): symbol of the parameter in Latex notation (e.g. f_{RO})
-        unit (str): base unit of measurement (e.g. Hz)
-        scale (int): the scale that should be generally applied to raw data (e.g. 1e-9 to take raw Hz to GHz)
+    Returns
+    -------
+    x_data : np.ndarray
+        Array containing the x-axis values.
+    y_data : np.ndarray
+        Array containing the y-axis (measured) data.
+    sweeps : list[np.ndarray]
+        List of additional swept parameter arrays (if any).
+    key_map : dict
+        Dictionary with keys `"x_data"`, `"y_data"`, and `"sweeps"` indicating
+        the corresponding keys used in the original `data_dict`.
+
+    Notes
+    -----
+    - If the schema is missing, the function prints a warning and returns empty arrays.
+    - Each item in the schema must be a dictionary with a `"role"` key.
+
+    Examples
+    --------
+    >>> x, y, sweeps, mapping = map_data_dict(experiment_data)
+    >>> print(f"x-axis data from key: {mapping['x_data']}")
     """
 
-    def __init__(self, id, value):
-        self.id = id
-        self.value = value
-        if id in _PARAM_METADATA:
-            meta = _PARAM_METADATA[id]
-        else:
-            meta = {}
-        self.name = meta["name"] if "name" in meta else id
-        self.symbol = meta["symbol"] if "symbol" in meta else id
-        self.unit = meta["unit"] if "unit" in meta else ""
-        self.scale = meta["scale"] if "scale" in meta else 1
-
-    def get_name_and_unit(self):
-        res = self.name
-        if self.unit != "":
-            exponent = -(int(f"{self.scale:.0e}".split("e")[1]) // 3) * 3
-            unit = f" [{_EXP_UNIT_MAP[exponent]}{self.unit}]"
-            res += unit
-        return res
-
-    def to_dict(self):
-        """Convert ParamInfo to a dictionary."""
-        return {
-            "id": self.id,
-            "value": self.value,
-            "name": self.name,
-            "symbol": self.symbol,
-            "unit": self.unit,
-            "scale": self.scale,
-        }
-
-    def __str__(self):
-        """Return a JSON-formatted string of the object."""
-        return json.dumps(self.to_dict())
-
-    def __eq__(self, other):
-        if isinstance(other, ParamInfo):
-            return (self.id == other.id) & (self.value == other.value)
-        if isinstance(other, (int, float, complex, str)):
-            return self.value == other
-        return False
-
-    def __bool__(self):
-        return bool(self.id)
-
-
-ParamDict = dict[str, ParamInfo | dict[str, ParamInfo]]
-
-
-def _enrich_param_dict(param_dict: dict) -> ParamDict:
-    """Add metadata to param_dict entries."""
-    res = {}
-    for key, value in param_dict.items():
-        if isinstance(value, dict):
-            # Recursive step for nested dictionaries
-            res[key] = _enrich_param_dict(value)
-        else:
-            res[key] = ParamInfo(key, value)
-    return res
-
-
-def read_param_dict(path: str) -> ParamDict:
-    """Read param_dict and include additional information for each entry.
+    schema = data_dict.get("schema", None)
+    if schema is None:
+        print(
+            "Cannot automatically read data: no database schema was provided by the experiment."
+        )
+
+    x_data, y_data, sweeps = np.array([]), np.array([]), []
+    key_map = {"x_data": "", "y_data": "", "sweeps": []}
+
+    for key, value in schema.items():
+        if type(value) is not dict:
+            continue
+        role = value.get("role", None)
+        if role == "data":
+            key_map["y_data"] = key
+            y_data = data_dict[key]
+        elif role == "x-axis":
+            key_map["x_data"] = key
+            x_data = data_dict[key]
+        elif role == "axis":
+            key_map["sweeps"].append(key)
+            sweeps.append(data_dict[key])
+
+    return x_data, y_data, sweeps, key_map
+
+
+def extract_mapped_data(path: str):
+    """
+    Loads measurement data from an HDF5 file and maps it into x_data, y_data and sweeps.
+    The map and the database schema on which it relies are also returned.
 
     Parameters
     ----------
-    path : str
-        Path to the file or a folder in which is contained a param_dict.json file
+    path : str or Path
+        Path to the HDF5 file containing experimental data and schema definitions.
 
     Returns
     -------
-    ParamDict
-        The param_dict with additional metadata
+    x_data : np.ndarray
+        Array of x-axis values extracted according to the schema.
+    y_data : np.ndarray
+        Array of measured data values (y-axis).
+    sweeps : list[np.ndarray]
+        List of arrays for any additional swept parameters defined in the schema.
+    datadict_map : dict
+        Mapping of keys used for `"x_data"`, `"y_data"`, and `"sweeps"` in the original file.
+    schema : dict
+        The schema used to interpret the data structure and field roles.
+
+    Notes
+    -----
+    - This function expects the file to contain a top-level "schema" key that defines the
+      role of each dataset (e.g., "data", "x-axis", "axis").
+    - Uses `extract_h5_data` and `map_data_dict` internally for loading and interpretation.
+
+    Examples
+    --------
+    >>> x, y, sweeps, datadict_map, schema = extract_mapped_data(path)
     """
-    # If the path is to a folder open /param_dict.json
-    if os.path.isdir(path):
-        path = os.path.join(path, "param_dict.json")
-    return _enrich_param_dict(read_json(path))
+
+    datadict = extract_h5_data(path, schema=True)
+    schema = datadict.get("schema")
+    x_data, y_data, sweeps, datadict_map = map_data_dict(datadict)
+    return x_data, y_data, sweeps, datadict_map, schema
 
 
-def get_sweep_param(path: str, exp_id: str):
-    params = read_param_dict(path)
-    sweep_id = params[exp_id]["sweep"].value
-    if sweep_id:
-        return params[sweep_id]
-    return ParamInfo("", "")
+def get_data_and_info(path=None, datadict=None):
+    if path is None and datadict is None:
+        raise Exception("At least one of `path` and `datadict` must be specified.")
+
+    if path is not None:
+        datadict = extract_h5_data(path, schema=True)
+
+    # Get schema and map data
+    schema = datadict.get("schema")
+    x_data, y_data, sweeps, datadict_map = map_data_dict(datadict)
+
+    # Get metadata on x_data and y_data
+    x_info = param_info_from_schema(
+        datadict_map["x_data"], schema[datadict_map["x_data"]]
+    )
+    y_info = param_info_from_schema(
+        datadict_map["y_data"], schema[datadict_map["y_data"]]
+    )
+
+    sweep_info = []
+    for sweep_key in datadict_map["sweeps"]:
+        sweep_info.append(param_info_from_schema(sweep_key, schema[sweep_key]))
+
+    return (x_data, y_data, sweeps), (x_info, y_info, sweep_info), datadict
+
+
+def read_json(path: str) -> dict:
+    """Reads a json file and returns the data as a dictionary."""
+    with open(path) as f:
+        dictionary = json.load(f)
+    return dictionary
+
+
+def read_yaml(path: str) -> dict:
+    with open(path) as stream:
+        try:
+            return yaml.safe_load(stream)
+        except yaml.YAMLError as exc:
+            print(exc)
+
+
+def read_qpu(dir_path: str, filename: str) -> QPU:
+    """Reads QPU file stored in dir_path/filename using laboneq serializers."""
+    qpu = serializers.load(os.path.join(dir_path, filename))
+    return qpu
 
 
 def get_measurement_id(path):
     return os.path.basename(path)[0:5]
+
+
+def copy_folder(src: str, dst: str):
+    # Ensure destination exists
+    os.makedirs(dst, exist_ok=True)
+
+    # Copy files recursively
+    for root, dirs, files in os.walk(src):
+        for dir_name in dirs:
+            os.makedirs(
+                os.path.join(dst, os.path.relpath(os.path.join(root, dir_name), src)),
+                exist_ok=True,
+            )
+        for file_name in files:
+            shutil.copy2(
+                os.path.join(root, file_name),
+                os.path.join(dst, os.path.relpath(os.path.join(root, file_name), src)),
+            )