# sweep_plotting.py
#
# This file is part of scqubits.
#
# Copyright (c) 2019 and later, Jens Koch and Peter Groszkowski
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
############################################################################
from typing import TYPE_CHECKING, List, Tuple, Union
import numpy as np
from matplotlib.axes import Axes
from matplotlib.figure import Figure
import scqubits.core.sweep_generators as sweep_gen
import scqubits.utils.plot_defaults as defaults
import scqubits.utils.plotting as plot
if TYPE_CHECKING:
from scqubits import DataStore, Grid1d, Oscillator, ParameterSweep, SpectrumData
from scqubits.core.qubit_base import QuantumSystem, QubitBaseClass, QubitBaseClass1d
QuantumSys = Union[QubitBaseClass, Oscillator]
[docs]def bare_spectrum(
sweep: "ParameterSweep",
subsys: "QuantumSys",
which: Union[int, List[int]] = -1,
**kwargs
) -> Tuple[Figure, Axes]:
"""
Plots energy spectrum of bare system `subsys` for given ParameterSweep `sweep`.
Parameters
----------
sweep:
subsys:
which:
default: -1, signals to plot all wavefunctions within the truncated Hilbert space;
int>0: plot wavefunctions 0..int-1; list(int) plot specific wavefunctions
**kwargs:
standard plotting option (see separate documentation)
"""
subsys_index = sweep.get_subsys_index(subsys)
specdata = sweep.bare_specdata_list[subsys_index]
if which is None:
which = subsys.truncated_dim
return specdata.plot_evals_vs_paramvals(which=which, **kwargs)
[docs]def dressed_spectrum(sweep: "ParameterSweep", **kwargs) -> Tuple[Figure, Axes]:
"""
Plots energy spectrum of dressed system
Parameters
----------
sweep:
**kwargs:
standard plotting option (see separate documentation)
"""
return sweep.dressed_specdata.plot_evals_vs_paramvals(
subtract_ground=True, **defaults.dressed_spectrum(sweep, **kwargs)
)
[docs]def difference_spectrum(
sweep: "ParameterSweep", initial_state_ind: int = 0, **kwargs
) -> Tuple[Figure, Axes]:
"""
Plots a transition energy spectrum with reference to the given initial_state_ind, obtained by taking energy
differences of the eigenenergy spectrum.
Parameters
----------
sweep:
initial_state_ind:
**kwargs:
standard plotting option (see separate documentation)
"""
return sweep_gen.generate_diffspec_sweep(
sweep, initial_state_ind
).plot_evals_vs_paramvals(**kwargs)
[docs]def n_photon_qubit_spectrum(
sweep: "ParameterSweep",
photonnumber: int,
initial_state_labels: Tuple[int, ...],
**kwargs
) -> Tuple[Figure, Axes]:
"""
Plots the n-photon qubit transition spectrum.
Parameters
----------
sweep:
photonnumber:
number of photons used in the transition
initial_state_labels:
bare state index of the initial state for the transitions
**kwargs:
standard plotting option (see separate documentation)
"""
label_list, specdata = sweep_gen.generate_qubit_transitions_sweep(
sweep, photonnumber, initial_state_labels
)
return specdata.plot_evals_vs_paramvals(label_list=label_list, **kwargs)
[docs]def bare_wavefunction(
sweep: "ParameterSweep",
param_val: float,
subsys: "QubitBaseClass1d",
which: Union[int, List[int]] = -1,
phi_grid: "Grid1d" = None,
**kwargs
) -> Tuple[Figure, Axes]:
"""
Plot bare wavefunctions for given parameter value and subsystem.
Parameters
----------
sweep:
param_val:
value of the external parameter
subsys:
which:
default: -1, signals to plot all wavefunctions; int>0: plot wavefunctions 0..int-1; list(int) plot specific
wavefunctions
phi_grid:
used for setting a custom grid for phi; if None use self._default_grid
**kwargs:
standard plotting option (see separate documentation)
"""
subsys_index = sweep.get_subsys_index(subsys)
sweep.update_hilbertspace(param_val)
param_index = np.searchsorted(sweep.param_vals, param_val)
evals = sweep.bare_specdata_list[subsys_index].energy_table[param_index]
evecs = sweep.bare_specdata_list[subsys_index].state_table[param_index]
return subsys.plot_wavefunction(
esys=(evals, evecs), which=which, mode="real", phi_grid=phi_grid, **kwargs
)
[docs]def chi(datastore: "DataStore", **kwargs) -> Tuple[Figure, Axes]:
"""
Plot dispersive shifts chi_j for a given pair of qubit and oscillator.
Parameters
----------
datastore:
contains sweep data for the dispersive shift, stored as specdata.chi
**kwargs:
standard plotting option (see separate documentation)
"""
ydata = datastore.chi
xdata = datastore.param_vals
state_count = ydata.shape[1]
label_list = list(range(state_count))
return plot.data_vs_paramvals(
xdata,
ydata,
label_list=label_list,
**defaults.chi(datastore.param_name, **kwargs)
)
[docs]def kerr(datastore: "DataStore", qubit_level=None, **kwargs) -> Tuple[Figure, Axes]:
"""
Plot dispersive Kerr energy for a given pair of qubit and oscillator.
Parameters
----------
datastore:
contains sweep data for the Kerr shift, stored as specdata.kerr
**kwargs:
standard plotting option (see separate documentation)
"""
ydata = datastore.kerr if qubit_level is None else datastore.kerr[:, qubit_level]
xdata = datastore.param_vals
state_count = len(ydata)
label_list = list(range(state_count)) if qubit_level is None else None
return plot.data_vs_paramvals(
xdata, ydata.T, label_list=label_list
) # , **defaults.chi(datastore.param_name, **kwargs))
[docs]def chi_01(
datastore: "DataStore", param_index: int = 0, **kwargs
) -> Tuple[Figure, Axes]:
"""
Plot the dispersive shift chi01 for a given pair of qubit and oscillator.
Parameters
----------
datastore: DataStore
param_index: int, optional
index of the external parameter to be used
**kwargs: dict
standard plotting option (see separate documentation)
"""
ydata = datastore.chi
xdata = datastore.param_vals
yval = ydata[param_index]
return plot.data_vs_paramvals(
xdata,
ydata,
label_list=None,
**defaults.chi01(datastore.param_name, yval, **kwargs)
)
[docs]def charge_matrixelem(
specdata: "SpectrumData",
qbt_index_subsys: Tuple[int, "QuantumSystem"],
initial_state_idx: int = 0,
**kwargs
) -> Tuple[Figure, Axes]:
"""
Parameters
----------
specdata:
qbt_index_subsys:
index of the qubit system within the underlying HilbertSpace, and qubit object
initial_state_idx:
index of initial state
**kwargs:
standard plotting option (see separate documentation)
"""
(qbt_index, qbt_subsys) = qbt_index_subsys
label_list = [
(initial_state_idx, final_idx) for final_idx in range(qbt_subsys.truncated_dim)
]
return plot.matelem_vs_paramvals(
specdata,
select_elems=label_list,
mode="abs",
**defaults.charge_matrixelem(specdata.param_name or "", **kwargs)
)