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
def plot_wavefunction(
self,
which: Union[int, Iterable[int]] = 0,
mode: str = "real",
esys: Tuple[ndarray, ndarray] = None,
phi_grid: Grid1d = None,
scaling: float = None,
**kwargs,
) -> Tuple[Figure, Axes]:
"""Plot 1d phase-basis wave function(s). Must be overwritten by
higher-dimensional qubits like FluxQubits and ZeroPi.
Parameters
----------
which:
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.
mode:
choices as specified in `constants.MODE_FUNC_DICT`
(default value = 'abs_sqr')
esys:
eigenvalues, eigenvectors
phi_grid:
used for setting a custom grid for phi; if None use self._default_grid
scaling:
custom scaling of wave function amplitude/modulus
**kwargs:
standard plotting option (see separate documentation)
"""
wavefunc_indices = process_which(which, self.truncated_dim)
if esys is None:
evals_count = max(wavefunc_indices) + 1
evals = self.eigenvals(evals_count=evals_count)
else:
evals, _ = esys
energies = evals[list(wavefunc_indices)]
phi_grid = phi_grid or self._default_grid
potential_vals = self.potential(phi_grid.make_linspace())
amplitude_modifier = constants.MODE_FUNC_DICT[mode]
wavefunctions = []
for wavefunc_index in wavefunc_indices:
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)
wavefunctions.append(phi_wavefunc)
fig_ax = kwargs.get("fig_ax") or plt.subplots()
kwargs["fig_ax"] = fig_ax
kwargs = {
**self.wavefunction1d_defaults(mode,
evals,
wavefunc_count=len(wavefunc_indices)),
**kwargs,
}
# in merging the dictionaries in the previous line: if any duplicates,
# later ones survive
plot.wavefunction1d(
wavefunctions,
potential_vals=potential_vals,
offset=energies,
scaling=scaling,
**kwargs,
)
return fig_ax