def plot_n_wavefunction(self,
esys: Tuple[ndarray, ndarray] = None,
mode: str = 'real',
which: int = 0,
nrange: Tuple[int, int] = None,
**kwargs) -> Tuple[Figure, Axes]:
"""Plots transmon wave function in charge basis
Parameters
----------
esys:
eigenvalues, eigenvectors
mode:
`'abs_sqr', 'abs', 'real', 'imag'`
which:
index or indices of wave functions to plot (default value = 0)
nrange:
range of `n` to be included on the x-axis (default value = (-5,6))
**kwargs:
plotting parameters
"""
if nrange is None:
nrange = self._default_n_range
n_wavefunc = self.numberbasis_wavefunction(esys, which=which)
amplitude_modifier = constants.MODE_FUNC_DICT[mode]
n_wavefunc.amplitudes = amplitude_modifier(n_wavefunc.amplitudes)
kwargs = {
**defaults.wavefunction1d_discrete(mode),
**kwargs
} # if any duplicates, later ones survive
return plot.wavefunction1d_discrete(n_wavefunc, xlim=nrange, **kwargs)
def wavefunction1d_discrete(wavefunc, **kwargs):
"""
Plots the amplitude of a real-valued 1d wave function in a discrete basis. (Example: transmon in the charge basis.)
Parameters
----------
wavefunc: WaveFunction object
basis and amplitude data of wave function to be plotted
**kwargs: dict
standard plotting option (see separate documentation)
Returns
-------
tuple(Figure, Axes)
matplotlib objects for further editing
"""
fig, axes = kwargs.get('fig_ax') or plt.subplots()
x_vals = wavefunc.basis_labels
width = .75
axes.bar(x_vals, wavefunc.amplitudes, width=width)
axes.set_xticks(x_vals)
axes.set_xticklabels(x_vals)
_process_options(fig, axes, defaults.wavefunction1d_discrete(), **kwargs)
return fig, axes
def plot_wavefunction(self,
which=0,
mode='real',
esys=None,
phi_grid=None,
scaling=None,
**kwargs):
"""Plot 1d phase-basis wave function(s). Must be overwritten by higher-dimensional qubits like FluxQubits and
ZeroPi.
Parameters
----------
esys: (ndarray, ndarray), optional
eigenvalues, eigenvectors
which: int or tuple or list, optional
single index or tuple/list of integers indexing the wave function(s) to be plotted.
If which is -1, all wavefunctions up to the truncation limit are plotted.
phi_grid: Grid1d, optional
used for setting a custom grid for phi; if None use self._default_grid
mode: str, optional
choices as specified in `constants.MODE_FUNC_DICT` (default value = 'abs_sqr')
scaling: float or None, optional
custom scaling of wave function amplitude/modulus
**kwargs: dict
standard plotting option (see separate documentation)
Returns
-------
Figure, Axes
"""
fig_ax = kwargs.get('fig_ax') or plt.subplots()
kwargs['fig_ax'] = fig_ax
index_list = process_which(which, self.truncated_dim)
phi_wavefunc = self.wavefunction(esys,
which=index_list[-1],
phi_grid=phi_grid)
potential_vals = self.potential(phi_wavefunc.basis_labels)
scale = set_scaling(self, scaling, potential_vals)
amplitude_modifier = constants.MODE_FUNC_DICT[mode]
kwargs = {
**defaults.wavefunction1d_discrete(mode),
**kwargs
} # if any duplicates, later ones survive
for wavefunc_index in index_list:
phi_wavefunc = self.wavefunction(esys,
which=wavefunc_index,
phi_grid=phi_grid)
phi_wavefunc.amplitudes = standardize_sign(phi_wavefunc.amplitudes)
phi_wavefunc.amplitudes = amplitude_modifier(
phi_wavefunc.amplitudes)
plot.wavefunction1d(phi_wavefunc,
potential_vals=potential_vals,
offset=phi_wavefunc.energy,
scaling=scale,
**kwargs)
return fig_ax
