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
def plot_wavefunction(self, esys, which=0, phi_range=None, phi_count=None, mode='real', scaling=None, xlabel=r'$\varphi$', ylabel=r'$\psi_j(\varphi),\, V(\varphi)$', y_range=None, title=None, filename=None, fig_ax=None): """Plot 1d phase-basis wave function(s). Must be overwritten by higher-dimensional qubits like FluxQubits and ZeroPi. Parameters ---------- esys: ndarray, ndarray 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_range: tuple(float, float), optional used for setting a custom plot range for phi phi_count: int, optional number of points on the x-axis (resolution) (default value = 251) 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 xlabel, ylabel: str, optional axes labels y_range: tuple(float, float), optional used to set custom y range for plot title: str, optional plot title filename: str, optional file path and name (not including suffix) for output fig_ax: Figure, Axes Returns ------- Figure, Axes """ modefunction = constants.MODE_FUNC_DICT[mode] index_list = process_which(which, self.truncated_dim) if fig_ax is None: fig = plt.figure() ax = fig.add_subplot(111) else: fig, ax = fig_ax phi_wavefunc = self.wavefunction(esys, which=index_list[-1], phi_range=phi_range, phi_count=phi_count) potential_vals = self.potential(phi_wavefunc.basis_labels) if scaling is None: if isinstance(self, qubits.Transmon): scale = 0.2 * self.EJ elif isinstance(self, qubits.Fluxonium): scale = 0.125 * (np.max(potential_vals) - np.min(potential_vals)) else: scale = scaling for wavefunc_index in index_list: phi_wavefunc = self.wavefunction(esys, which=wavefunc_index, phi_range=phi_range, phi_count=phi_count) if np.sum(phi_wavefunc.amplitudes) < 0: phi_wavefunc.amplitudes *= -1.0 phi_wavefunc.amplitudes = modefunction(phi_wavefunc.amplitudes) plot.wavefunction1d(phi_wavefunc, potential_vals=potential_vals, offset=phi_wavefunc.energy, scaling=scale, xlabel=xlabel, ylabel=ylabel, y_range=y_range, title=title, fig_ax=(fig, ax), filename=filename) return fig, ax