Utils

inspeqtor.utils

inspeqtor.utils.SyntheticDataModel dataclass

A utility dataclass holding objects necessary for simulating single qubit quantum device.

Source code in src/inspeqtor/v2/utils.py

@dataclass
class SyntheticDataModel:
    """A utility dataclass holding objects necessary for simulating single qubit quantum device."""

    control_sequence: ControlSequence
    qubit_information: QubitInformation
    dt: float
    ideal_hamiltonian: typing.Callable[..., jnp.ndarray]
    total_hamiltonian: typing.Callable[..., jnp.ndarray]
    solver: typing.Callable[..., jnp.ndarray]
    quantum_device: typing.Callable[..., jnp.ndarray] | None
    whitebox: typing.Callable[..., jnp.ndarray] | None

inspeqtor.utils.single_qubit_shot_quantum_device

single_qubit_shot_quantum_device(
    key: ndarray,
    control_parameters: ndarray,
    solver: Callable[[ndarray], ndarray],
    SHOTS: int,
    expectation_value_receipt: Sequence[
        ExpectationValue
    ] = get_complete_expectation_values(1),
) -> ndarray

This is the shot estimate expectation value quantum device

Parameters:

Name	Type	Description	Default
control_parameters	ndarray	The control parameter to be feed to simlulator	required
key	ndarray	Random key	required
solver	Callable[[ndarray], ndarray]	The ODE solver for propagator	required
SHOTS	int	The number of shots used to estimate expectation values	required

Returns:

Type	Description
ndarray	jnp.ndarray: The expectation value of shape (control_parameters.shape[0], 18)

Source code in src/inspeqtor/v2/utils.py

def single_qubit_shot_quantum_device(
    key: jnp.ndarray,
    control_parameters: jnp.ndarray,
    solver: typing.Callable[[jnp.ndarray], jnp.ndarray],
    SHOTS: int,
    expectation_value_receipt: typing.Sequence[
        ExpectationValue
    ] = get_complete_expectation_values(1),
) -> jnp.ndarray:
    """This is the shot estimate expectation value quantum device

    Args:
        control_parameters (jnp.ndarray): The control parameter to be feed to simlulator
        key (jnp.ndarray): Random key
        solver (typing.Callable[[jnp.ndarray], jnp.ndarray]): The ODE solver for propagator
        SHOTS (int): The number of shots used to estimate expectation values

    Returns:
        jnp.ndarray: The expectation value of shape (control_parameters.shape[0], 18)
    """

    expectation_values = jnp.zeros(
        (control_parameters.shape[0], len(expectation_value_receipt))
    )
    unitaries = jax.vmap(solver)(control_parameters)[:, -1, :, :]

    for idx, exp in enumerate(expectation_value_receipt):
        key, sample_key = jax.random.split(key)
        sample_keys = jax.random.split(sample_key, num=unitaries.shape[0])

        expectation_value = jax.vmap(
            calculate_shots_expectation_value,
            in_axes=(0, None, 0, None, None),
        )(
            sample_keys,
            get_initial_state(exp.initial_state, dm=True),
            unitaries,
            get_observable_operator(exp.observable),
            SHOTS,
        )

        expectation_values = expectation_values.at[..., idx].set(expectation_value)

    return expectation_values

inspeqtor.utils.finite_shot_quantum_device

finite_shot_quantum_device(
    key: ndarray,
    param: Param,
    solver: Callable[[Param, ndarray], ndarray],
    shots: int,
    expval: ExpectationValue,
)

Source code in src/inspeqtor/v2/utils.py

def finite_shot_quantum_device(
    key: jnp.ndarray,
    param: Param,
    solver: typing.Callable[[Param, jnp.ndarray], jnp.ndarray],
    shots: int,
    expval: ExpectationValue,
):
    initial_state = get_initial_state(expval.initial_state, dm=True)

    state = solver(param, initial_state)
    prob = get_measurement_probability(state, expval.observable)

    return finite_shot_expectation_value(key, prob, shots)

inspeqtor.utils.get_measurement_probability

get_measurement_probability(
    state: ndarray, operator: str
) -> ndarray

Calculate the probability of measuring each projector of tensor product of Pauli operators

Parameters:

Name	Type	Description	Default
state	ndarray	The quantum state to measure	required
operator	str	The string representation of the measurement operator, e.g., 'XY'	required

Returns:

Type	Description
ndarray	jnp.ndarray: An array of probability where each index is a base 10 representation of base 2 measurement result.

Source code in src/inspeqtor/v2/utils.py

def get_measurement_probability(state: jnp.ndarray, operator: str) -> jnp.ndarray:
    """Calculate the probability of measuring each projector of tensor product of Pauli operators

    Args:
        state (jnp.ndarray): The quantum state to measure
        operator (str): The string representation of the measurement operator, e.g., 'XY'

    Returns:
        jnp.ndarray: An array of probability where each index is a base 10 representation of base 2 measurement result.
    """

    return jnp.array(
        [
            jnp.trace(state @ tensor_product(*g_projector))
            for g_projector in product(
                *[(projectors[op][0], projectors[op][1]) for op in operator]
            )
        ]
    ).real

inspeqtor.utils.finite_shot_expectation_value

finite_shot_expectation_value(
    key: ndarray, prob: ndarray, shots: int
)

Source code in src/inspeqtor/v2/utils.py

def finite_shot_expectation_value(key: jnp.ndarray, prob: jnp.ndarray, shots: int):
    return jnp.mean(
        jax.random.choice(
            key,
            jax.vmap(check_parity)(jnp.arange(0, prob.size, dtype=jnp.int_)),
            shape=(shots,),
            p=prob,
        )
    )

inspeqtor.utils.calculate_expectation_values

calculate_expectation_values(
    unitaries: ndarray,
    expectation_value_order: list[
        ExpectationValue
    ] = get_complete_expectation_values(1),
) -> ndarray

Source code in src/inspeqtor/v2/utils.py

def calculate_expectation_values(
    unitaries: jnp.ndarray,
    expectation_value_order: list[ExpectationValue] = get_complete_expectation_values(
        1
    ),
) -> jnp.ndarray:
    # Calculate the ideal expectation values of the original pulse
    ideal_expectation_values = jnp.zeros(tuple(unitaries.shape[:-2]) + (18,))
    for idx, exp in enumerate(expectation_value_order):
        expvals = calculate_exp(
            unitaries,
            get_observable_operator(exp.observable),
            get_initial_state(exp.initial_state, dm=True),
        )
        ideal_expectation_values = ideal_expectation_values.at[..., idx].set(expvals)

    return ideal_expectation_values

inspeqtor.utils.dictorization

dictorization(
    expvals: ndarray, order: list[ExpectationValue]
)

This function formats expectation values of shape (18, N) to a dictionary with the initial state as outer key and the observable as inner key.

Parameters:

Name	Type	Description	Default
expvals	ndarray	Expectation values of shape (18, N). Assumes that order is as in default_expectation_values_order.	required

Returns:

Type	Description
	dict[str, dict[str, jnp.ndarray]]: A dictionary with the initial state as outer key and the observable as inner key.

Source code in src/inspeqtor/v2/utils.py

def dictorization(expvals: jnp.ndarray, order: list[ExpectationValue]):
    """This function formats expectation values of shape (18, N) to a dictionary
    with the initial state as outer key and the observable as inner key.

    Args:
        expvals (jnp.ndarray): Expectation values of shape (18, N). Assumes that order is as in default_expectation_values_order.

    Returns:
        dict[str, dict[str, jnp.ndarray]]: A dictionary with the initial state as outer key and the observable as inner key.
    """
    expvals_dict: dict[str, dict[str, jnp.ndarray]] = {}
    for idx, exp in enumerate(order):
        if exp.initial_state not in expvals_dict:
            expvals_dict[exp.initial_state] = {}

        expvals_dict[exp.initial_state][exp.observable] = expvals[idx]

    return expvals_dict

inspeqtor.utils.random_split

random_split(
    key: ndarray, test_size: int, *data_arrays: ndarray
)

The random_split function splits the data into training and testing sets.

Examples:

>>> key = jax.random.key(0)
>>> x = jnp.arange(10)
>>> y = jnp.arange(10)
>>> x_train, y_train, x_test, y_test = random_split(key, 2, x, y)
>>> assert x_train.shape[0] == 8 and y_train.shape[0] == 8
>>> assert x_test.shape[0] == 2 and y_test.shape[0] == 2

Parameters:

Name	Type	Description	Default
key	ndarray	Random key.	required
test_size	int	The size of the test set. Must be less than the size of the data.	required

Returns:

Type	Description
	typing.Sequence[jnp.ndarray]: The training and testing sets in the same order.

Source code in src/inspeqtor/v1/utils.py

def random_split(key: jnp.ndarray, test_size: int, *data_arrays: jnp.ndarray):
    """The random_split function splits the data into training and testing sets.

    Examples:
        >>> key = jax.random.key(0)
        >>> x = jnp.arange(10)
        >>> y = jnp.arange(10)
        >>> x_train, y_train, x_test, y_test = random_split(key, 2, x, y)
        >>> assert x_train.shape[0] == 8 and y_train.shape[0] == 8
        >>> assert x_test.shape[0] == 2 and y_test.shape[0] == 2

    Args:
        key (jnp.ndarray): Random key.
        test_size (int): The size of the test set. Must be less than the size of the data.

    Returns:
        typing.Sequence[jnp.ndarray]: The training and testing sets in the same order.
    """
    # * General random split
    idx = jax.random.permutation(key, data_arrays[0].shape[0])
    train_data = []
    test_data = []

    for data in data_arrays:
        train_data.append(data[idx][test_size:])
        test_data.append(data[idx][:test_size])

    return (*train_data, *test_data)

inspeqtor.utils.dataloader

dataloader(
    arrays: Sequence[ndarray],
    batch_size: int,
    num_epochs: int,
    *,
    key: ndarray,
)

The dataloader function creates a generator that yields batches of data.

Parameters:

Name	Type	Description	Default
arrays	Sequence[ndarray]	The list or tuple of arrays to be batched.	required
batch_size	int	The size of the batch.	required
num_epochs	int	The number of epochs. If set to -1, the generator will run indefinitely.	required
key	ndarray	The random key.	required

Returns:

Name	Type	Description
None		stop the generator.

Yields:

Type	Description
	typing.Any: (step, batch_idx, is_last_batch, epoch_idx), (array_batch, ...)

Source code in src/inspeqtor/v1/utils.py

def dataloader(
    arrays: typing.Sequence[jnp.ndarray],
    batch_size: int,
    num_epochs: int,
    *,
    key: jnp.ndarray,
):
    """The dataloader function creates a generator that yields batches of data.

    Args:
        arrays (typing.Sequence[jnp.ndarray]): The list or tuple of arrays to be batched.
        batch_size (int): The size of the batch.
        num_epochs (int): The number of epochs. If set to -1, the generator will run indefinitely.
        key (jnp.ndarray): The random key.

    Returns:
        None: stop the generator.

    Yields:
        typing.Any: (step, batch_idx, is_last_batch, epoch_idx), (array_batch, ...)
    """
    # * General dataloader
    # Check that all arrays have the same size in the first dimension
    dataset_size = arrays[0].shape[0]
    # assert all(array.shape[0] == dataset_size for array in arrays)
    # Generate random indices
    indices = jnp.arange(dataset_size)
    step = 0
    epoch_idx = 0
    while True:
        if epoch_idx == num_epochs:
            return None
        perm = jax.random.permutation(key, indices)
        (key,) = jax.random.split(key, 1)
        batch_idx = 0
        start = 0
        end = batch_size
        is_last_batch = False
        while not is_last_batch:
            batch_perm = perm[start:end]
            # Check if this is the last batch
            is_last_batch = end >= dataset_size
            yield (
                (step, batch_idx, is_last_batch, epoch_idx),
                tuple(array[batch_perm] for array in arrays),
            )
            start = end
            end = start + batch_size
            step += 1
            batch_idx += 1

        epoch_idx += 1

inspeqtor.utils.variance_of_observable

variance_of_observable(expval: ndarray, shots: int = 1)

Source code in src/inspeqtor/v1/utils.py

def variance_of_observable(expval: jnp.ndarray, shots: int = 1):
    return (1 - expval**2) / shots

inspeqtor.utils.expectation_value_to_prob_plus

expectation_value_to_prob_plus(
    expectation_value: ndarray,
) -> ndarray

Calculate the probability of -1 and 1 for the given expectation value E[O] = -1 * P[O = -1] + 1 * P[O = 1], where P[O = -1] + P[O = 1] = 1 Thus, E[O] = -1 * (1 - P[O = 1]) + 1 * P[O = 1] E[O] = 2 * P[O = 1] - 1 -> P[O = 1] = (E[O] + 1) / 2 Args: expectation_value (jnp.ndarray): Expectation value of quantum observable

Returns:

Type	Description
ndarray	jnp.ndarray: Probability of measuring plus eigenvector

Source code in src/inspeqtor/v1/utils.py

def expectation_value_to_prob_plus(expectation_value: jnp.ndarray) -> jnp.ndarray:
    """
    Calculate the probability of -1 and 1 for the given expectation value
    E[O] = -1 * P[O = -1] + 1 * P[O = 1], where P[O = -1] + P[O = 1] = 1
    Thus, E[O] = -1 * (1 - P[O = 1]) + 1 * P[O = 1]
    E[O] = 2 * P[O = 1] - 1 -> P[O = 1] = (E[O] + 1) / 2
    Args:
        expectation_value (jnp.ndarray): Expectation value of quantum observable

    Returns:
        jnp.ndarray: Probability of measuring plus eigenvector
    """

    return (expectation_value + 1) / 2

inspeqtor.utils.expectation_value_to_prob_minus

expectation_value_to_prob_minus(
    expectation_value: ndarray,
) -> ndarray

Convert quantum observable expectation value to probability of measuring -1.

For a binary quantum observable \(\hat{O}\) with eigenvalues \(b = \{-1, 1\}\), this function calculates the probability of measuring the eigenvalue -1 given its expectation value.

Derivation: $$ \langle \hat{O} \rangle = -1 \cdot \Pr(b=-1) + 1 \cdot \Pr(b = 1) $$ With the constraint \(\Pr(b = -1) + \Pr(b = 1) = 1\):

\[ \langle \hat{O} \rangle = -1 \cdot \Pr(b=-1) + 1 \cdot (1 - \Pr(b=-1)) \ \langle \hat{O} \rangle = -\Pr(b=-1) + 1 - \Pr(b=-1) \ \langle \hat{O} \rangle = 1 - 2\Pr(b=-1) \ \Pr(b=-1) = \frac{1 - \langle \hat{O} \rangle}{2} \]

Parameters:

Name	Type	Description	Default
expectation_value	ndarray	Expectation value of the quantum observable, must be in range [-1, 1].	required

Returns:

Type	Description
ndarray	jnp.ndarray: Probability of measuring the -1 eigenvalue.

Source code in src/inspeqtor/v1/utils.py

def expectation_value_to_prob_minus(expectation_value: jnp.ndarray) -> jnp.ndarray:
    """Convert quantum observable expectation value to probability of measuring -1.

    For a binary quantum observable $\\hat{O}$ with eigenvalues $b = \\{-1, 1\\}$, this function
    calculates the probability of measuring the eigenvalue -1 given its expectation value.

    Derivation:
    $$
        \\langle \\hat{O} \\rangle = -1 \\cdot \\Pr(b=-1) + 1 \\cdot \\Pr(b = 1)
    $$
        With the constraint $\\Pr(b = -1) + \\Pr(b = 1) = 1$:

    $$
        \\langle \\hat{O} \\rangle = -1 \\cdot \\Pr(b=-1) + 1 \\cdot (1 - \\Pr(b=-1)) \\
        \\langle \\hat{O} \\rangle = -\\Pr(b=-1) + 1 - \\Pr(b=-1) \\
        \\langle \\hat{O} \\rangle = 1 - 2\\Pr(b=-1) \\
        \\Pr(b=-1) = \\frac{1 - \\langle \\hat{O} \\rangle}{2}
    $$

    Args:
        expectation_value (jnp.ndarray): Expectation value of the quantum observable,
            must be in range [-1, 1].

    Returns:
        jnp.ndarray: Probability of measuring the -1 eigenvalue.
    """
    return (1 - expectation_value) / 2

inspeqtor.utils.expectation_value_to_eigenvalue

expectation_value_to_eigenvalue(
    expectation_value: ndarray, SHOTS: int
) -> ndarray

Convert expectation value to eigenvalue

Parameters:

Name	Type	Description	Default
expectation_value	ndarray	Expectation value of quantum observable	required
SHOTS	int	The number of shots used to produce expectation value	required

Returns:

Type	Description
ndarray	jnp.ndarray: Array of eigenvalues

Source code in src/inspeqtor/v1/utils.py

def expectation_value_to_eigenvalue(
    expectation_value: jnp.ndarray, SHOTS: int
) -> jnp.ndarray:
    """Convert expectation value to eigenvalue

    Args:
        expectation_value (jnp.ndarray): Expectation value of quantum observable
        SHOTS (int): The number of shots used to produce expectation value

    Returns:
        jnp.ndarray: Array of eigenvalues
    """
    return jnp.where(
        jnp.broadcast_to(jnp.arange(SHOTS), expectation_value.shape + (SHOTS,))
        < jnp.around(
            expectation_value_to_prob_plus(
                jnp.reshape(expectation_value, expectation_value.shape + (1,))
            )
            * SHOTS
        ).astype(jnp.int32),
        1,
        -1,
    ).astype(jnp.int32)

inspeqtor.utils.eigenvalue_to_binary

eigenvalue_to_binary(eigenvalue: ndarray) -> ndarray

Convert -1 to 1, and 0 to 1 This implementation should be differentiable

Parameters:

Name	Type	Description	Default
eigenvalue	ndarray	Eigenvalue to convert to bit value	required

Returns:

Type	Description
ndarray	jnp.ndarray: Binary array

Source code in src/inspeqtor/v1/utils.py

def eigenvalue_to_binary(eigenvalue: jnp.ndarray) -> jnp.ndarray:
    """Convert -1 to 1, and 0 to 1
    This implementation should be differentiable

    Args:
        eigenvalue (jnp.ndarray): Eigenvalue to convert to bit value

    Returns:
        jnp.ndarray: Binary array
    """

    return (-1 * eigenvalue + 1) / 2

inspeqtor.utils.binary_to_eigenvalue

binary_to_eigenvalue(binary: ndarray) -> ndarray

Convert 1 to -1, and 0 to 1 This implementation should be differentiable

Parameters:

Name	Type	Description	Default
binary	ndarray	Bit value to convert to eigenvalue	required

Returns:

Type	Description
ndarray	jnp.ndarray: Eigenvalue array

Source code in src/inspeqtor/v1/utils.py

def binary_to_eigenvalue(binary: jnp.ndarray) -> jnp.ndarray:
    """Convert 1 to -1, and 0 to 1
    This implementation should be differentiable

    Args:
        binary (jnp.ndarray): Bit value to convert to eigenvalue

    Returns:
        jnp.ndarray: Eigenvalue array
    """

    return -1 * (binary * 2 - 1)

inspeqtor.utils.recursive_vmap

recursive_vmap(func, in_axes)

Perform recursive vmap on the given axis

Note

def func(x):
    assert x.ndim == 1
    return x ** 2
x = jnp.arange(10)
x_test = jnp.broadcast_to(x, (2, 3, 4,) + x.shape)
x_test.shape, recursive_vmap(func, (0,) * (x_test.ndim - 1))(x_test).shape
((2, 3, 4, 10), (2, 3, 4, 10))

Examples:

>>> def func(x):
...     assert x.ndim == 1
...     return x ** 2
>>> x = jnp.arange(10)
>>> x_test = jnp.broadcast_to(x, (2, 3, 4,) + x.shape)
>>> x_test.shape, recursive_vmap(func, (0,) * (x_test.ndim - 1))(x_test).shape
((2, 3, 4, 10), (2, 3, 4, 10))

Parameters:

Name	Type	Description	Default
func	Any	The function for vmap	required
in_axes	Any	The axes for vmap	required

Returns:

Type	Description
	typing.Any: description

Source code in src/inspeqtor/v1/utils.py

def recursive_vmap(func, in_axes):
    """Perform recursive vmap on the given axis

    Note:
        ```python
        def func(x):
            assert x.ndim == 1
            return x ** 2
        x = jnp.arange(10)
        x_test = jnp.broadcast_to(x, (2, 3, 4,) + x.shape)
        x_test.shape, recursive_vmap(func, (0,) * (x_test.ndim - 1))(x_test).shape
        ((2, 3, 4, 10), (2, 3, 4, 10))
        ```

    Examples:
        >>> def func(x):
        ...     assert x.ndim == 1
        ...     return x ** 2
        >>> x = jnp.arange(10)
        >>> x_test = jnp.broadcast_to(x, (2, 3, 4,) + x.shape)
        >>> x_test.shape, recursive_vmap(func, (0,) * (x_test.ndim - 1))(x_test).shape
        ((2, 3, 4, 10), (2, 3, 4, 10))

    Args:
        func (typing.Any): The function for vmap
        in_axes (typing.Any): The axes for vmap

    Returns:
        typing.Any: _description_
    """
    if not in_axes:
        # Base case: no more axes to vectorize over
        return func

    # Apply vmap over the first axis specified in in_axes
    vmap_func = jax.vmap(func, in_axes=in_axes[0])

    # Recursively apply vmap over the remaining axes
    return recursive_vmap(vmap_func, in_axes[1:])

inspeqtor.utils.calculate_shots_expectation_value

calculate_shots_expectation_value(
    key: ndarray,
    initial_state: ndarray,
    unitary: ndarray,
    operator: ndarray,
    shots: int,
) -> ndarray

Calculate finite-shots estimate of expectation value

Parameters:

Name	Type	Description	Default
key	ndarray	Random key	required
initial_state	ndarray	Inital state	required
unitary	ndarray	Unitary operator	required
plus_projector	ndarray	The eigenvector corresponded to +1 eigenvalue of Pauli observable.	required
shots	int	Number of shot to be used in estimation of expectation value	required

Returns:

Type	Description
ndarray	jnp.ndarray: Finite-shot estimate expectation value

Source code in src/inspeqtor/v1/utils.py

def calculate_shots_expectation_value(
    key: jnp.ndarray,
    initial_state: jnp.ndarray,
    unitary: jnp.ndarray,
    operator: jnp.ndarray,
    shots: int,
) -> jnp.ndarray:
    """Calculate finite-shots estimate of expectation value

    Args:
        key (jnp.ndarray): Random key
        initial_state (jnp.ndarray): Inital state
        unitary (jnp.ndarray): Unitary operator
        plus_projector (jnp.ndarray): The eigenvector corresponded to +1 eigenvalue of Pauli observable.
        shots (int): Number of shot to be used in estimation of expectation value

    Returns:
        jnp.ndarray: Finite-shot estimate expectation value
    """
    expval = jnp.trace(unitary @ initial_state @ unitary.conj().T @ operator).real
    prob = expectation_value_to_prob_plus(expval)

    return jax.random.choice(
        key, jnp.array([1, -1]), shape=(shots,), p=jnp.array([prob, 1 - prob])
    ).mean()

inspeqtor.utils.enable_jax_x64

enable_jax_x64()

Source code in src/inspeqtor/v1/utils.py

def enable_jax_x64():
    jax.config.update("jax_enable_x64", True)

inspeqtor.utils.disable_jax_x64

disable_jax_x64()

Source code in src/inspeqtor/v1/utils.py

def disable_jax_x64():
    jax.config.update("jax_enable_x64", False)

inspeqtor.utils.default_expectation_values_order module-attribute

default_expectation_values_order = (
    get_default_expectation_values_order()
)

inspeqtor.utils.get_default_expectation_values_order

get_default_expectation_values_order()

Source code in src/inspeqtor/v2/constant.py

def get_default_expectation_values_order():
    return [
        ExpectationValue(observable="X", initial_state="+"),
        ExpectationValue(observable="X", initial_state="-"),
        ExpectationValue(observable="X", initial_state="r"),
        ExpectationValue(observable="X", initial_state="l"),
        ExpectationValue(observable="X", initial_state="0"),
        ExpectationValue(observable="X", initial_state="1"),
        ExpectationValue(observable="Y", initial_state="+"),
        ExpectationValue(observable="Y", initial_state="-"),
        ExpectationValue(observable="Y", initial_state="r"),
        ExpectationValue(observable="Y", initial_state="l"),
        ExpectationValue(observable="Y", initial_state="0"),
        ExpectationValue(observable="Y", initial_state="1"),
        ExpectationValue(observable="Z", initial_state="+"),
        ExpectationValue(observable="Z", initial_state="-"),
        ExpectationValue(observable="Z", initial_state="r"),
        ExpectationValue(observable="Z", initial_state="l"),
        ExpectationValue(observable="Z", initial_state="0"),
        ExpectationValue(observable="Z", initial_state="1"),
    ]

inspeqtor.utils.SX module-attribute

SX = (
    exp(1j * pi / 4) / sqrt(2) * array([[1, -1j], [-1j, 1]])
)

inspeqtor.utils.X module-attribute

X = get_observable_operator('X')

inspeqtor.utils.Y module-attribute

Y = get_observable_operator('Y')

inspeqtor.utils.Z module-attribute

Z = get_observable_operator('Z')

inspeqtor.utils.plot_control_envelope

plot_control_envelope(
    waveform: ndarray,
    x_axis: ndarray,
    ax: Axes,
    font_size: int = 12,
)

Source code in src/inspeqtor/v1/visualization.py

def plot_control_envelope(
    waveform: jnp.ndarray,
    x_axis: jnp.ndarray,
    ax: Axes,
    font_size: int = 12,
):
    ax.bar(
        x_axis, jnp.real(waveform), label="Real component", color="orange", alpha=0.5
    )
    ax.bar(x_axis, jnp.imag(waveform), label="Imag component", color="blue", alpha=0.5)

    ax.set_xlabel("Time", fontsize=font_size)

    # Text size
    ax.tick_params(axis="both", labelsize=font_size)

    ax.set_xlabel("Time (dt)", fontsize=font_size)
    ax.set_ylabel("Amplitude", fontsize=font_size)

    ax.legend(fontsize=font_size, loc="upper right")

inspeqtor.utils.plot_expectation_values

plot_expectation_values(
    expvals_dict: dict[str, dict[str, ndarray]], title: str
)

Source code in src/inspeqtor/v1/visualization.py

def plot_expectation_values(
    expvals_dict: dict[str, dict[str, jnp.ndarray]],
    title: str,
):
    fig, axes = plt.subplot_mosaic(
        """
        +r0
        -l1
        """,
        figsize=(10, 5),
        sharex=True,
        sharey=True,
    )

    colormap = {
        "X": "#ef4444",
        "Y": "#6366f1",
        "Z": "#10b981",
    }

    for idx, (initial_state, expvals) in enumerate(expvals_dict.items()):
        ax = axes[initial_state]
        for observable, expval in expvals.items():
            ax.plot(expval, "-", label=observable, color=colormap[observable])
        ax.set_title(f"Initial state: {initial_state}")
        ax.set_ylim(-1.05, 1.05)
        ax.legend(loc="upper left")

    # Set title for the figure
    fig.suptitle(title)

    fig.tight_layout()
    return fig, axes

inspeqtor.utils.assert_list_of_axes

assert_list_of_axes(axes) -> list[Axes]

Assert the provide object that they are a list of Axes

Parameters:

Name	Type	Description	Default
axes	Any	Expected to be numpy array of Axes	required

Returns:

Type	Description
list[Axes]	list[Axes]: The list of Axes

Source code in src/inspeqtor/v1/visualization.py

def assert_list_of_axes(axes) -> list[Axes]:
    """Assert the provide object that they are a list of Axes

    Args:
        axes (typing.Any): Expected to be numpy array of Axes

    Returns:
        list[Axes]: The list of Axes
    """
    assert isinstance(axes, np.ndarray)
    axes = axes.flatten()

    for ax in axes:
        assert isinstance(ax, Axes)
    return axes.tolist()

inspeqtor.utils.set_fontsize

set_fontsize(ax: Axes, fontsize: float | int)

Set all fontsize of the Axes object

Parameters:

Name	Type	Description	Default
ax	Axes	The Axes object which fontsize to be changed.	required
fontsize	float | int	The fontsize.	required

Source code in src/inspeqtor/v1/visualization.py

def set_fontsize(ax: Axes, fontsize: float | int):
    """Set all fontsize of the Axes object

    Args:
        ax (Axes): The Axes object which fontsize to be changed.
        fontsize (float | int): The fontsize.
    """
    for item in (
        [ax.title, ax.xaxis.label, ax.yaxis.label]
        + ax.get_xticklabels()
        + ax.get_yticklabels()
    ):
        item.set_fontsize(fontsize)

    legend, handles = ax.get_legend_handles_labels()

    ax.legend(legend, handles, fontsize=fontsize)

inspeqtor.utils.plot_loss_with_moving_average

plot_loss_with_moving_average(
    x: ndarray | ndarray,
    y: ndarray | ndarray,
    ax: Axes,
    window: int = 50,
    annotate_at: list[float] = [0.2, 0.4, 0.6, 0.8, 1.0],
    **kwargs,
) -> Axes

Plot the moving average of the given argument y

Parameters:

Name	Type	Description	Default
x	ndarray | ndarray	The horizontal axis	required
y	ndarray | ndarray	The vertical axis	required
ax	Axes	Axes object	required
window	int	The moving average window. Defaults to 50.	50
annotate_at	list[int]	The list of x positions to annotate the y value. Defaults to [2000, 4000, 6000, 8000, 10000].	[0.2, 0.4, 0.6, 0.8, 1.0]

Returns:

Name	Type	Description
Axes	Axes	Axes object.

Source code in src/inspeqtor/v1/visualization.py

def plot_loss_with_moving_average(
    x: jnp.ndarray | np.ndarray,
    y: jnp.ndarray | np.ndarray,
    ax: Axes,
    window: int = 50,
    annotate_at: list[float] = [0.2, 0.4, 0.6, 0.8, 1.0],
    **kwargs,
) -> Axes:
    """Plot the moving average of the given argument y

    Args:
        x (jnp.ndarray | np.ndarray): The horizontal axis
        y (jnp.ndarray | np.ndarray): The vertical axis
        ax (Axes): Axes object
        window (int, optional): The moving average window. Defaults to 50.
        annotate_at (list[int], optional): The list of x positions to annotate the y value. Defaults to [2000, 4000, 6000, 8000, 10000].

    Returns:
        Axes: Axes object.
    """
    moving_average = pd.Series(np.asarray(y)).rolling(window=window).mean()

    ax.plot(
        x,
        moving_average,
        **kwargs,
    )

    for percentile in annotate_at:
        # Calculate the data index that corresponds to the percentile
        idx = int(percentile * (len(x) - 1))

        loss_value = moving_average[idx]

        # Skip annotation if the moving average value is not available (e.g., at the beginning)
        if pd.isna(loss_value):
            continue

        ax.annotate(
            f"{loss_value:.3g}",
            xy=(float(x[idx].item()), float(loss_value)),
            xytext=(-10, 10),  # Offset the text for better readability
            textcoords="offset points",
            ha="center",
            va="bottom",
        )

    return ax