decoder_demo.py

import json
from pathlib import Path
from typing import Literal, Optional

import numpy as np
import numpy.typing as npt
from matplotlib import pyplot as plt

LOCAL_DATA_DIRECTORY_ROOT = Path("downloaded_data")


def read_file(remote_path: str | Path) -> npt.NDArray | list:
    extension = Path(remote_path).suffix

    if extension == ".npy":
        file = np.load(remote_path)
    elif extension == ".json":
        with open(remote_path, "r") as file_:
            file = json.load(file_)
    else:
        raise ValueError(f"Unknown file extension {extension}")

    return file


def sift_and_group_quadratures(
    batch_dir: str | Path, use_decoder_decision: bool
) -> dict[str, npt.NDArray]:
    """Starting from a time array of quadratures, acquired under active decoding,
    performs the following steps:

        - keep quadratures acquired in the decoder decision time bin (and omit
            the others)
        - discard time bins in which pairs 2 or 3 included at least
            one heralded state
        - group remaining quadratures according to either
            - decoding result: "keep" or "cut", if use_decoder_decision=True
            - measurement basis: p (X-measurement) or q (Z-measurement),
                if use_decoder_decision=False

    Args:
        batch_dir: directory of the acquisition batch
        use_decoder_decision: whether to separate the acquired quadratures
            into "keep" or "cut" decisions of the decoder (as opposed to
            separating them by the measured LO phase setting)

    Returns:
        dictionary of sifted quadratures, grouped by
            - either decoder decision (if use_decoder_decision is True)
            - or measurement basis (if use_decoder_decision is False)
    """
    # the number or time bins in one decoding cycle
    decoder_cycles = 5

    # index of a decoding cycle's decision time bin (i.e. the time bin in which
    # the adaptive measurement takes place)
    decision_time_bin = 4

    # obtain quadrature array, of shape (12, total time bins)
    quadratures = read_file(
        remote_path=LOCAL_DATA_DIRECTORY_ROOT / batch_dir / "quadratures.npy"
    )[:, decision_time_bin::decoder_cycles]

    # the QPU switch settings corresponding to the batch's decision time bins
    if not use_decoder_decision:
        # use the acquired switch setting of QPU 2 (M_0) to inform when
        # entanglement was kept or cut; just focus on the decision time bins
        switch_settings = read_file(
            remote_path=LOCAL_DATA_DIRECTORY_ROOT
            / batch_dir
            / "switch_settings_qpu_2.npy"
        )[decision_time_bin::decoder_cycles]

        # the integer values in switch_settings correspond to the following
        # measurement bases:
        # 0 -- Z-measurement || cut  || q || 0 deg  || (keep_or_cut=False)
        # 3 -- X-measurement || keep || p || 90 deg || (keep_or_cut=True)
        # (numbers 1 and 2 are reserved for additional measurement angles that
        # were, however, not used in this experiment)
        keep_or_cut = np.array(switch_settings) == 3
    else:
        # use decoder decisions to inform when entanglement was kept
        # or cut. decoder_bits can be:
        # - -1: non-decision time bin (so time bins 0--3 before the decision),
        # - 0: cut,
        # - 1: keep,
        # - -2: timeout -- no recovery found in 5 clock cycles (cut)
        decisions = read_file(
            remote_path=LOCAL_DATA_DIRECTORY_ROOT / batch_dir / "decisions.npy"
        )

        # a decoder decision in time bin N will only affect the adaptive
        # measurement in time bin N+1; so in order to later compare decisions
        # and measurements we shift the array for one clock cycle
        decisions = np.roll(decisions, 1)

        # just focus on the decision time bins
        keep_or_cut = decisions[decision_time_bin::decoder_cycles]

        # as a sanity check, make sure there is no -1 left in the keep_or_cut
        # values (a decoding cycle has five time bins: four decoding bins
        # followed by a decision bin; the value -1 is only allowed in decoding
        # bins which, after the above sifting, should be discarded)
        if -1 in keep_or_cut:
            raise ValueError(
                "Unexpected value -1 in keep_or_cut. Only values 0, 1, and -2 are valid in a decision time bin."
            )

    # boolean arrays indicating "keep" and "cut" decisions as well as timeout
    # events (also corresponding to "cut")
    keep_event = keep_or_cut == 1
    cut_event = keep_or_cut == 0
    timeout_event = keep_or_cut == -2

    # load the switch settings of multiplexers 2 and 3
    switch_settings_mux_2_top = read_file(
        remote_path=LOCAL_DATA_DIRECTORY_ROOT
        / batch_dir
        / "switch_settings_mux_2_top.npy"
    )
    switch_settings_mux_2_bot = read_file(
        remote_path=LOCAL_DATA_DIRECTORY_ROOT
        / batch_dir
        / "switch_settings_mux_2_bot.npy"
    )
    switch_settings_mux_3_top = read_file(
        remote_path=LOCAL_DATA_DIRECTORY_ROOT
        / batch_dir
        / "switch_settings_mux_3_top.npy"
    )
    switch_settings_mux_3_bot = read_file(
        remote_path=LOCAL_DATA_DIRECTORY_ROOT
        / batch_dir
        / "switch_settings_mux_3_bot.npy"
    )

    # just focus on the decision time bins
    switch_settings_mux_2_top = switch_settings_mux_2_top[
        decision_time_bin::decoder_cycles
    ]
    switch_settings_mux_2_bot = switch_settings_mux_2_bot[
        decision_time_bin::decoder_cycles
    ]
    switch_settings_mux_3_top = switch_settings_mux_3_top[
        decision_time_bin::decoder_cycles
    ]
    switch_settings_mux_3_bot = switch_settings_mux_3_bot[
        decision_time_bin::decoder_cycles
    ]

    # the elements in these arrays correspond to the index of the mux input
    # routed to the entangled-pair generation at each individual time bin; input
    # 0 always corresponds to a squeezed-light mode, inputs 1--3 correspond to a
    # heralded state; therefore, if the sum of all four inputs is zero, no
    # heralded state is involved in any of the four multiplexers at that
    # particular time bin
    switch_settings_sum = np.sum(
        [
            switch_settings_mux_2_top,
            switch_settings_mux_2_bot,
            switch_settings_mux_3_top,
            switch_settings_mux_3_bot,
        ],
        axis=0,
    )
    # boolean array indicating events of all-gaussian pairs
    no_heralded_event = switch_settings_sum == 0

    # finally, group quadratures acquired at the decision time bins by
    # measurement basis, discarding decision time bins that involve non-gaussian
    # states in pairs 2 or 3
    quadratures_p = quadratures[:, keep_event * no_heralded_event]
    quadratures_q = quadratures[:, cut_event * no_heralded_event]
    quadratures_timeout = quadratures[:, timeout_event * no_heralded_event]

    return {
        "quadratures_p": quadratures_p,
        "quadratures_q": quadratures_q,
        "quadratures_timeout": quadratures_timeout,
    }


def covariance_matrix(quadratures: npt.NDArray) -> npt.NDArray:
    """Computes a covariance matrix of quadrature values, sorted by macronode.

    Args:
        quadratures: quadratures of shape (12, time bins)

    Returns: covariance matrix
    """

    def array_to_qpu_dict(arr: npt.NDArray) -> dict:
        """Reorganizes a nested array of length 12 (number of FTB spatial modes) into a
        dict with one key per QPU.

        Args:
            arr: array of raw data or quadratures per spatial mode

        Returns: dict of raw data or quadratures, organized by QPU
        """
        return {
            "qpu_0": np.array([arr[0], arr[2]]),
            "qpu_1": np.array([arr[1], arr[3], arr[4]]),
            "qpu_2": np.array([arr[5], arr[6]]),
            "qpu_3": np.array([arr[10], arr[8], arr[7]]),
            "qpu_4": np.array([arr[11], arr[9]]),
        }

    # group the spatial modes by QPU (macronode)
    qpu_dict = array_to_qpu_dict(quadratures)
    x_2 = np.array(
        [
            qpu_dict[f"qpu_{i}"][j]
            for i in range(5)
            for j in range(2 if i % 2 == 0 else 3)
        ]
    )

    return np.cov(x_2)


def get_batch_cov(quadratures_dict: dict[str, npt.NDArray]) -> dict[str, npt.NDArray]:
    """Computes covariance matrices from quadrature arrays:

    Args:
        quadratures_dict: dict of quadratures with keys quadratures_p
            and quadratures_q, each one corresponding to an array of shape
            (12, shots)

    Returns:
        dict of covariance matrices with keys cov_p and cov_p, each one
        corresponding to an array of shape (24, 24)
    """
    return {
        "cov_p": covariance_matrix(quadratures_dict["quadratures_p"]),
        "cov_q": covariance_matrix(quadratures_dict["quadratures_q"]),
    }


def get_covs(
    type_: Literal["signal", "random", "vacuum"],
    use_decoder_decision: bool,
    start_batch: int = 0,
    stop_batch: Optional[int] = None,
) -> dict[str, npt.NDArray]:
    """Processes a sequence of decoding jobs and returns a dictionary with covariance
    matrix per measurement setting and job.

    Args:
        type_: specifies the type of measurement sequence to be processed:
            regular decoding experiment ("signal"), control experiment
            ("random) or vacuum measurement ("vacuum")
        use_decoder_decision: whether to separate the acquired quadratures
            into "keep" or "cut" decisions of the decoder (as opposed to
            separating them by the measured LO phase setting)
        start_batch: first batch of the measurement sequence to be processed;
            defaults to zero
        stop_batch: last batch of the measurement sequence to be processed;
            defaults to the total available number of batches

    Returns:
        dictionary with covariance matrix per measurement setting and job
    """
    if stop_batch is None:
        if type_ == "signal":
            # a total of 72 batches were acquired in the signal acquisition
            stop_batch = 73
        elif type_ == "vacuum":
            # a total of 62 batches were acquired in the vacuum acquisition
            stop_batch = 63
        elif type_ == "random":
            # a total of 73 batches were acquired in the random acquisition
            stop_batch = 74
        else:
            raise ValueError

    num_batches = stop_batch - start_batch

    # pre-allocate arrays of covariance matrices
    all_covs = {
        "cov_p": np.full((num_batches, 12, 12), np.nan),
        "cov_q": np.full((num_batches, 12, 12), np.nan),
    }

    # iterate through the individual jobs of the measurement sequence
    for i in range(start_batch, stop_batch):
        batch_dir = Path(f"decoder_demo/{type_}/batch_{i}")
        quadratures_dict = sift_and_group_quadratures(
            batch_dir=batch_dir, use_decoder_decision=use_decoder_decision
        )
        covs = get_batch_cov(quadratures_dict=quadratures_dict)
        all_covs["cov_p"][i, ...] = covs["cov_p"]
        all_covs["cov_q"][i, ...] = covs["cov_q"]

    return {
        "cov_p": np.nanmean(all_covs["cov_p"], axis=0),
        "cov_q": np.nanmean(all_covs["cov_q"], axis=0),
    }


def plot_covariance(
    covariance: npt.NDArray,
    range_: Optional[float] = None,
    show_values: bool = True,
    cmap: str = "RdBu_r",
    file_name: Optional[Path] = None,
    show_plot: bool = False,
) -> None:
    """Creates a colour plot of the quadrature covariance matrix.

    Args:
        covariance: array of shape (12, 12)
        range_: the range of values shown in the colour plot (-range_ , range_)
        show_values: whether to display the matrix values in the colour plot
        cmap: matplotlib colour map
        file_name: file path and name under which generated figure is saved
        show_plot: whether to show the plot or close the figure instead
            trace
    """
    # plot width and height
    width = 7
    height = 5
    fig, ax = plt.subplots(figsize=(width, height))

    # the diagonal of the covariance matrix just shows the single-mode variance
    # of each mode; it is usually much larger than any off-diagonal elements and
    # therefore dominates any colour plot; in order to enhance visibility of the
    # inter-mode correlations, we set the diagonal to zero
    covariance = np.copy(covariance)
    np.fill_diagonal(covariance, 0.0)

    if range_ is None:
        range_ = np.nanmax(np.abs(covariance))

    cov_x_plot = ax.pcolor(
        covariance,
        vmin=-range_,
        vmax=range_,
        cmap=cmap,
    )
    ax.invert_yaxis()
    ax.set_aspect("equal", adjustable="box")

    separation_a_0_b_0 = 2
    separation_b_0_c = 5
    separation_c_b_1 = 7
    separation_b_1_a_1 = 10
    ax.hlines(
        [separation_a_0_b_0, separation_b_0_c, separation_c_b_1, separation_b_1_a_1],
        0,
        12,
        linestyle="solid",
        linewidth=0.5,
        color="grey",
    )
    ax.vlines(
        [separation_a_0_b_0, separation_b_0_c, separation_c_b_1, separation_b_1_a_1],
        0,
        12,
        linestyle="solid",
        linewidth=0.5,
        color="grey",
    )

    ticks = np.array([1, 3.5, 6, 8.5, 11])
    labels = [r"$M_{3}$", r"$M_{1}$", r"$M_{0}$", r"$M_{2}$", r"$M_{4}$"]
    ax.set_xticks(ticks, labels, fontsize=15)
    ax.set_yticks(ticks, labels, fontsize=15)

    if show_values:
        for (j, i), value in np.ndenumerate(covariance):
            plt.text(
                i + 0.5, j + 0.5, f"{value:.3f}", ha="center", va="center", fontsize=6
            )

    color_bar = fig.colorbar(cov_x_plot, ax=ax)
    color_bar.ax.tick_params(labelsize=15)
    if file_name is not None:
        fig.savefig(file_name)
    plt.show()
    if not show_plot:
        plt.close(fig)