Skip to content

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
37
38
39
40
41
42
43
44
45
46
47
48
@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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
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
200
201
202
203
204
205
206
207
208
209
210
211
212
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
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
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
189
190
191
192
193
194
195
196
197
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
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
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
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
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
235
236
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
269
270
271
272
273
274
275
276
277
278
279
280
281
282
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
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
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
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
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
339
340
341
342
343
344
345
346
347
348
349
350
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
353
354
355
356
357
358
359
360
361
362
363
364
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
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
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
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
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
539
540
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
543
544
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
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
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

inspeqtor.utils.Y module-attribute

inspeqtor.utils.Z module-attribute

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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
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
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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