Source code for HybridSuperQubits.ferbo

from collections.abc import Iterable
from typing import Any, Literal, Optional, Union

import matplotlib.pyplot as plt
import numpy as np
from qutip import Bloch, Qobj, wigner
from scipy.linalg import cosm, eigh, sinm

from .operators import (
    creation,
    destroy,
    ptrace,
    sigma_x,
    sigma_y,
    sigma_z,
    state_to_density_matrix,
)
from .qubit_base import QubitBase


[docs] class Ferbo(QubitBase): PARAM_LABELS = { "Ec": r"$E_C$", "El": r"$E_L$", "Ej": r"$E_J$", "Gamma": r"$\Gamma$", "delta_Gamma": r"$\delta \Gamma$", "er": r"$\epsilon_r$", "phase": r"$\Phi_{\mathrm{ext}} / \Phi_0$", } OPERATOR_LABELS = { "n_operator": r"\hat{n}", "phase_operator": r"\hat{\phi}", "d_hamiltonian_d_ng": r"\partial \hat{H} / \partial n_g", "d_hamiltonian_d_phase": r"\partial \hat{H} / \partial \phi_{{ext}}", "d_hamiltonian_d_EL": r"\partial \hat{H} / \partial E_L", "d_hamiltonian_d_deltaGamma": r"\partial \hat{H} / \partial \delta \Gamma", "d_hamiltonian_d_er": r"\partial \hat{H} / \partial \epsilon_r", }
[docs] def __init__( self, Ec, El, Ej, Gamma, delta_Gamma, er, phase, dimension, flux_grouping: Literal["EL", "ABS"] = "ABS", Delta=40, ): """ Initializes the Ferbo class with the given parameters. Parameters ---------- Ec : float Charging energy. El : float Inductive energy. Ej : float Josephson energy. Gamma : float Coupling strength. delta_Gamma : float Coupling strength difference. er : float Energy relaxation rate. phase : float External magnetic phase. dimension : int Dimension of the Hilbert space. flux_grouping : Literal["EL", "ABS"], optional Flux grouping ('EL' or 'ABS') (default is 'ABS'). Delta : float Superconducting gap. """ if flux_grouping not in ["EL", "ABS"]: raise ValueError("Invalid flux grouping; must be 'EL' or 'ABS'.") self.Ec = Ec self.El = El self.Ej = Ej self.Gamma = Gamma self.delta_Gamma = delta_Gamma self.er = er self.phase = phase self.dimension = dimension // 2 * 2 self.flux_grouping = flux_grouping self.Delta = Delta super().__init__(self.dimension)
@property def phase_zpf(self) -> float: """ Returns the zero-point fluctuation of the phase. Returns ------- float Zero-point fluctuation of the phase. """ return (2 * self.Ec / self.El) ** 0.25 @property def n_zpf(self) -> float: """ Returns the zero-point fluctuation of the charge number. Returns ------- float Zero-point fluctuation of the charge number. """ return 1 / 2 * (self.El / 2 / self.Ec) ** 0.25 @property def lc_energy(self) -> float: """ Returns the plasma energy. Returns ------- float Plasma energy. """ return np.sqrt(8 * self.Ec * self.El) @property def transparency(self) -> float: """ Return the transparency of the weak link. Returns ------- float Transparency of the weak link. """ return (self.Gamma**2 - self.delta_Gamma**2) / (self.Gamma**2 + self.er**2)
[docs] def phi_osc(self) -> float: """ Returns the oscillator length for the LC oscillator composed of the inductance and capacitance. Returns ------- float Oscillator length. """ return (8.0 * self.Ec / self.El) ** 0.25
[docs] def n_operator(self) -> np.ndarray: """ Returns the charge number operator. Returns ------- np.ndarray The charge number operator. """ single_mode_n_operator = ( 1j * self.n_zpf * (creation(self.dimension // 2) - destroy(self.dimension // 2)) ) return np.kron(np.eye(2), single_mode_n_operator)
[docs] def phase_operator(self) -> np.ndarray: """ Returns the total phase operator. Returns ------- np.ndarray The total phase operator. """ single_mode_phase_operator = self.phase_zpf * ( creation(self.dimension // 2) + destroy(self.dimension // 2) ) return np.kron(np.eye(2), single_mode_phase_operator)
[docs] def jrl_potential(self) -> np.ndarray: """ Returns the Josephson Resonance Level potential. Returns ------- np.ndarray The Josephson Resonance Level potential. """ phase_op = self.phase_operator()[: self.dimension // 2, : self.dimension // 2] if self.flux_grouping == "ABS": phase_op -= self.phase * np.eye(self.dimension // 2) return ( -self.Gamma * np.kron(sigma_z(), cosm(phase_op / 2)) - self.delta_Gamma * np.kron(sigma_y(), sinm(phase_op / 2)) + self.er * np.kron(sigma_x(), np.eye(self.dimension // 2)) )
[docs] def hamiltonian(self) -> np.ndarray: """ Returns the Hamiltonian of the system. Returns ------- np.ndarray The Hamiltonian of the system. """ # n_x = self.delta_Gamma / 4 / (self.Gamma + self.Delta) n_op = ( self.n_operator() ) # + n_x * np.kron(sigma_x(), np.eye(self.dimension // 2)) charge_term = 4 * self.Ec * n_op @ n_op phase_op = self.phase_operator() if self.flux_grouping == "ABS": inductive_term = 0.5 * self.El * phase_op @ phase_op phase_op -= self.phase * np.eye(self.dimension) josephson_term = -self.Ej * cosm(phase_op) else: josephson_term = -self.Ej * cosm(phase_op) phase_op += self.phase * np.eye(self.dimension) inductive_term = 0.5 * self.El * phase_op @ phase_op potential = self.jrl_potential() return charge_term + inductive_term + potential + josephson_term
[docs] def d_hamiltonian_d_EC(self) -> np.ndarray: """ Returns the derivative of the Hamiltonian with respect to the charging energy. Returns ------- np.ndarray The derivative of the Hamiltonian with respect to the charging energy. """ n_x = self.delta_Gamma / 4 / (self.Gamma + self.Delta) n_op = self.n_operator() + n_x * np.kron(sigma_x(), np.eye(self.dimension // 2)) return 8 * n_op @ n_op
[docs] def d_hamiltonian_d_EL(self) -> np.ndarray: """ Returns the derivative of the Hamiltonian with respect to the inductive energy. Returns ------- np.ndarray The derivative of the Hamiltonian with respect to the inductive energy. """ if self.flux_grouping == "EL": phase_op = self.phase_operator() elif self.flux_grouping == "ABS": phase_op = self.phase_operator() - self.phase * np.eye(self.dimension) return 1 / 2 * np.dot(phase_op, phase_op)
[docs] def d_hamiltonian_d_EJ(self) -> np.ndarray: """ Returns the derivative of the Hamiltonian with respect to the Josephson energy. Returns ------- np.ndarray The derivative of the Hamiltonian with respect to the Josephson energy. """ phase_op = self.phase_operator() if self.flux_grouping == "ABS": phase_op -= self.phase * np.eye(self.dimension) return -cosm(phase_op)
[docs] def d_hamiltonian_d_Gamma(self) -> np.ndarray: """ Returns the derivative of the Hamiltonian with respect to Gamma. Returns ------- np.ndarray The derivative of the Hamiltonian with respect to Gamma. """ phase_op = self.phase_operator()[: self.dimension // 2, : self.dimension // 2] if self.flux_grouping == "ABS": phase_op -= self.phase * np.eye(self.dimension // 2) return -np.kron(sigma_z(), sinm(phase_op / 2))
[docs] def d_hamiltonian_d_er(self) -> np.ndarray: """ Returns the derivative of the Hamiltonian with respect to the energy relaxation rate. Returns ------- np.ndarray The derivative of the Hamiltonian with respect to the energy relaxation rate. """ return +np.kron(sigma_x(), np.eye(self.dimension // 2))
[docs] def d_hamiltonian_d_deltaGamma(self) -> np.ndarray: """ Returns the derivative of the Hamiltonian with respect to the coupling strength difference. Returns ------- np.ndarray The derivative of the Hamiltonian with respect to the coupling strength difference. """ phase_op = self.phase_operator()[: self.dimension // 2, : self.dimension // 2] if self.flux_grouping == "ABS": phase_op -= self.phase * np.eye(self.dimension // 2) return -np.kron(sigma_y(), sinm(phase_op / 2))
[docs] def d2_hamiltonian_d_Gamma2(self) -> np.ndarray: """ Returns the second derivative of the Hamiltonian with respect to Gamma. Returns ------- np.ndarray The second derivative of the Hamiltonian with respect to Gamma. """ return np.zeros((self.dimension, self.dimension))
[docs] def d2_hamiltonian_d_deltaGamma2(self) -> np.ndarray: """ Returns the second derivative of the Hamiltonian with respect to the coupling strength difference. Returns ------- np.ndarray The second derivative of the Hamiltonian with respect to the coupling strength difference. """ return np.zeros((self.dimension, self.dimension))
[docs] def d2_hamiltonian_d_er2(self) -> np.ndarray: """ Returns the second derivative of the Hamiltonian with respect to the energy relaxation rate. Returns ------- np.ndarray The second derivative of the Hamiltonian with respect to the energy relaxation rate. """ return np.zeros((self.dimension, self.dimension))
[docs] def d_hamiltonian_d_ng(self) -> np.ndarray: """ Returns the derivative of the Hamiltonian with respect to the number of charge offset. Returns ------- np.ndarray The derivative of the Hamiltonian with respect to the number of charge offset. """ n_x = self.delta_Gamma / 4 / (self.Gamma + self.Delta) return ( -8 * self.Ec * ( self.n_operator() + n_x * np.kron(sigma_x(), np.eye(self.dimension // 2)) ) )
[docs] def d2_hamiltonian_d_ng2(self) -> np.ndarray: """ Returns the second derivative of the Hamiltonian with respect to the number of charge offset. Returns ------- np.ndarray The second derivative of the Hamiltonian with respect to the number of charge offset. """ return 8 * self.Ec * np.eye(self.dimension)
[docs] def d_hamiltonian_d_phase(self) -> np.ndarray: """ Returns the derivative of the Hamiltonian with respect to the external magnetic phase. Returns ------- np.ndarray The derivative of the Hamiltonian with respect to the external magnetic phase. """ if self.flux_grouping == "EL": return self.El * ( self.phase_operator() + self.phase * np.eye(self.dimension) ) elif self.flux_grouping == "ABS": phase_op = self.phase_operator()[ : self.dimension // 2, : self.dimension // 2 ] - self.phase * np.eye(self.dimension // 2) return -self.Gamma / 2 * np.kron( sigma_z(), sinm(phase_op / 2) ) + self.delta_Gamma / 2 * np.kron(sigma_y(), cosm(phase_op / 2)) else: raise ValueError(f"Unknown flux_grouping: {self.flux_grouping}")
[docs] def d2_hamiltonian_d_phase2(self) -> np.ndarray: """ Returns the second derivative of the Hamiltonian with respect to the external magnetic phase. Returns ------- np.ndarray The second derivative of the Hamiltonian with respect to the external magnetic phase. """ if self.flux_grouping == "EL": return self.El * np.eye(self.dimension) elif self.flux_grouping == "ABS": phase_op = self.phase_operator()[ : self.dimension // 2, : self.dimension // 2 ] - self.phase * np.eye(self.dimension // 2) return self.Gamma / 4 * np.kron( sigma_z(), cosm(phase_op / 2) ) + self.delta_Gamma / 4 * np.kron(sigma_y(), sinm(phase_op / 2)) else: raise ValueError(f"Unknown flux_grouping: {self.flux_grouping}")
[docs] def wigner( self, which: int = 0, phi_grid: Optional[np.ndarray] = None, n_grid: Optional[np.ndarray] = None, esys: Optional[tuple[np.ndarray, np.ndarray]] = None, ): """ Computes the Wigner function for a given wavefunction. Parameters ---------- which : int, optional Index of desired wavefunction (default is 0). phi_grid : np.ndarray, optional Custom grid for phi; if None, a default grid is used. n_grid : np.ndarray, optional Custom grid for n; if None, a default grid is used. esys : Tuple[np.ndarray, np.ndarray], optional Precomputed eigenvalues and eigenvectors. Returns ------- np.ndarray The Wigner function. """ rho_reduced = self.reduced_density_matrix(which=which, esys=esys, subsys=0) if phi_grid is None: phi_grid = np.linspace(-5, 5, 151) if n_grid is None: n_grid = np.linspace(-5, 5, 151) rho_reduced_qobj = Qobj(rho_reduced) wigner_func = wigner(rho_reduced_qobj, phi_grid, n_grid) return wigner_func
[docs] def reduced_density_matrix( self, which: int = 0, esys: tuple[np.ndarray, np.ndarray] = None, subsys: int = 0, ) -> np.ndarray: """ Computes the reduced density matrix for a given wavefunction. Parameters ---------- which : int, optional Index of desired wavefunction (default is 0). esys : Tuple[np.ndarray, np.ndarray], optional Precomputed eigenvalues and eigenvectors. subsys : int, optional Subsystem to compute the reduced density matrix for (default is 0). 0 for the tracing out the Fock states, 1 for the Andreev states. Returns ------- np.ndarray The reduced density matrix. """ if esys is None: evals_count = max(which + 1, 3) _, evecs = self.eigensys(evals_count) else: _, evecs = esys rho = state_to_density_matrix(evecs[:, which]) rho_reduced = ptrace(rho, dims=(2, self.dimension // 2), subsys=subsys) return rho_reduced
[docs] def wavefunction( self, which: int = 0, phi_grid: Optional[np.ndarray] = None, esys: Optional[tuple[np.ndarray, np.ndarray]] = None, basis: Literal["phase", "charge"] = "phase", representation: Literal["ballistic", "Andreev"] = "ballistic", ) -> dict[str, Any]: """ Returns a wave function in the phi basis. Parameters ---------- which : int, optional Index of desired wave function (default is 0). phi_grid : np.ndarray, optional Custom grid for phi; if None, a default grid is used. basis : Literal["phase", "charge"], optional Basis in which to return the wavefunction ('phase' or 'charge') (default is 'phase'). representation : Literal["ballistic", "Andreev"], optional Representation used for the two-component wavefunction. Use 'ballistic' for the basis used to define the Hamiltonian matrix, or 'Andreev' for the local phi-dependent eigenbasis of the Andreev sector. Returns ------- Dict[str, Any] Wave function data containing basis labels, amplitudes, and energy. """ if representation not in {"ballistic", "Andreev"}: raise ValueError("Invalid representation; must be 'ballistic' or 'Andreev'.") if esys is None: evals_count = max(which + 1, 3) evals, evecs = self.eigensys(evals_count) else: evals, evecs = esys dim = self.dimension // 2 evecs = evecs.T if basis == "phase": l_osc = self.phase_zpf elif basis == "charge": l_osc = self.n_zpf if phi_grid is None: phi_grid = np.linspace(-5 * np.pi, 5 * np.pi, 151) phi_basis_labels = phi_grid wavefunc_osc_basis_amplitudes = evecs[which, :] phi_wavefunc_amplitudes = np.zeros((2, len(phi_grid)), dtype=np.complex128) # Compute the wavefunction in the ballistic Andreev basis used in the Hamiltonian. for n in range(dim): phi_wavefunc_amplitudes[0] += wavefunc_osc_basis_amplitudes[ n ] * self.harm_osc_wavefunction(n, phi_basis_labels, l_osc) phi_wavefunc_amplitudes[1] += wavefunc_osc_basis_amplitudes[ self.dimension // 2 + n ] * self.harm_osc_wavefunction(n, phi_basis_labels, l_osc) # Optionally rotate from the ballistic Andreev basis to the # phi-dependent Andreev eigenbasis. if representation == "Andreev": # Get rotation matrices for each phi point _, rotation_matrices = self.potential(phi_grid, return_evecs=True) # Rotate wavefunction at each phi point rotated_amplitudes = np.zeros_like(phi_wavefunc_amplitudes) for i in range(len(phi_grid)): # rotation_matrices[i] transforms from computational to Andreev eigenbasis # We need to apply the adjoint to transform the wavefunction state_vector = np.array( [phi_wavefunc_amplitudes[0, i], phi_wavefunc_amplitudes[1, i]] ) rotated_state = rotation_matrices[i].conj().T @ state_vector rotated_amplitudes[0, i] = rotated_state[0] rotated_amplitudes[1, i] = rotated_state[1] phi_wavefunc_amplitudes = rotated_amplitudes if basis == "charge": phi_wavefunc_amplitudes[0] /= np.sqrt(self.n_zpf) phi_wavefunc_amplitudes[1] /= np.sqrt(self.n_zpf) elif basis == "phase": phi_wavefunc_amplitudes[0] /= np.sqrt(self.phase_zpf) phi_wavefunc_amplitudes[1] /= np.sqrt(self.phase_zpf) # Fix global phase: ensure the wavefunction is mostly positive # For each Andreev component, choose phase so that the sum of real parts is positive for j in range(2): real_sum = np.sum(np.real(phi_wavefunc_amplitudes[j])) if real_sum < 0: phi_wavefunc_amplitudes[j] *= -1 return { "basis_labels": phi_basis_labels, "amplitudes": phi_wavefunc_amplitudes, "energy": evals[which], }
[docs] def potential( self, phi: Union[float, np.ndarray], return_evecs: bool = False, include_berry: bool = True, ) -> Union[np.ndarray, tuple[np.ndarray, np.ndarray]]: """ Calculates the potential energy for given values of phi. Parameters ---------- phi : Union[float, np.ndarray] The phase values at which to calculate the potential. return_evecs : bool, optional If True, returns both eigenvalues and eigenvectors (default is False). include_berry : bool, optional If True, includes the scalar Berry correction in the potential eigenvalues (default is True). Returns ------- Union[np.ndarray, Tuple[np.ndarray, np.ndarray]] If return_evecs is False: The potential energy values (eigenvalues). If return_evecs is True: A tuple of (eigenvalues, eigenvectors). Eigenvalues shape: (len(phi), 2) Eigenvectors shape: (len(phi), 2, 2) where each [i, :, :] is the rotation matrix at phi[i] """ phi_array = np.atleast_1d(phi) evals_array = np.zeros((len(phi_array), 2)) if return_evecs: evecs_array = np.zeros((len(phi_array), 2, 2), dtype=complex) for i, phi_val in enumerate(phi_array): # Construct only the Andreev part if self.flux_grouping == "ABS": andreev_phase = phi_val - self.phase andreev_term = ( -self.Gamma * np.cos(andreev_phase / 2) * sigma_z() - self.delta_Gamma * np.sin(andreev_phase / 2) * sigma_y() + self.er * sigma_x() ) # Inductive and Josephson contributions (diagonal, add to eigenvalues) inductive_contribution = 0.5 * self.El * phi_val**2 josephson_contribution = -self.Ej * np.cos(phi_val - self.phase) elif self.flux_grouping == "EL": andreev_phase = phi_val andreev_term = ( -self.Gamma * np.cos(andreev_phase / 2) * sigma_z() - self.delta_Gamma * np.sin(andreev_phase / 2) * sigma_y() + self.er * sigma_x() ) # Inductive and Josephson contributions (diagonal, add to eigenvalues) inductive_contribution = 0.5 * self.El * (phi_val + self.phase) ** 2 josephson_contribution = -self.Ej * np.cos(phi_val) else: raise ValueError(f"Unknown flux_grouping: {self.flux_grouping}") berry_contribution = ( self._berry_contribution(andreev_phase) if include_berry else 0.0 ) # Diagonalize only the Andreev part if return_evecs: andreev_evals, andreev_evecs = eigh( andreev_term, check_finite=False, ) # Add diagonal contributions to eigenvalues evals_array[i] = ( andreev_evals + inductive_contribution + josephson_contribution + berry_contribution ) # Gauge fixing: ensure phase continuity with previous point if i > 0: for j in range(2): # For each eigenstate # Compute overlap with previous eigenstate overlap = np.vdot(evecs_array[i - 1, :, j], andreev_evecs[:, j]) # If overlap is negative or has large imaginary part, fix phase if np.real(overlap) < 0: andreev_evecs[:, j] *= -1 evecs_array[i] = andreev_evecs else: andreev_evals = eigh( andreev_term, eigvals_only=True, check_finite=False, ) # Add diagonal contributions to eigenvalues evals_array[i] = ( andreev_evals + inductive_contribution + josephson_contribution + berry_contribution ) if return_evecs: return evals_array, evecs_array return evals_array
def _berry_contribution( self, phi: Union[float, np.ndarray] ) -> Union[float, np.ndarray]: """ Returns the scalar Berry contribution in the Andreev basis. The correction is 4 * self.Ec * 1/4 * [ (d chi / d phi)^2 + (d theta / d phi)^2 ], multiplied by the identity in the Andreev subspace. """ phi_array = np.asarray(phi, dtype=float) b_perp_sq = ( self.Gamma**2 * np.cos(phi_array / 2) ** 2 + self.delta_Gamma**2 * np.sin(phi_array / 2) ** 2 ) b_sq = b_perp_sq + self.er**2 b_perp = np.sqrt(b_perp_sq) with np.errstate(divide="ignore", invalid="ignore"): dchi_dphi = self.Gamma * self.delta_Gamma / (2 * b_perp_sq) dtheta_dphi = ( self.er * (self.delta_Gamma**2 - self.Gamma**2) * np.sin(phi_array) / (4 * b_perp * b_sq) ) berry_contribution = 4 * self.Ec * 0.25 * (dchi_dphi**2 + dtheta_dphi**2) if np.isscalar(phi): return float(berry_contribution) return berry_contribution
[docs] def tphi_1_over_f_flux( self, A_noise: float = 1e-6, esys: tuple[np.ndarray, np.ndarray] = None, get_rate: bool = False, **kwargs, ) -> float: return self.tphi_1_over_f( A_noise, ["d_hamiltonian_d_phase", "d2_hamiltonian_d_phase2"], esys=esys, get_rate=get_rate, **kwargs, )
[docs] def plot_wavefunction( self, which: Union[int, Iterable[int]] = 0, phi_grid: np.ndarray = None, esys: tuple[np.ndarray, np.ndarray] = None, scaling: Optional[float] = 1, plot_potential: bool = False, basis: Literal["phase", "charge"] = "phase", andreev_basis: Literal["ballistic", "Andreev"] = "ballistic", mode: Literal["abs", "abs2", "real", "imag"] = "abs", **kwargs, ) -> tuple[plt.Figure, plt.Axes]: """ Plot the wave function in the phi basis. Parameters ---------- which : Union[int, Iterable[int]], optional Index or indices of desired wave function(s) (default is 0). phi_grid : np.ndarray, optional Custom grid for phi; if None, a default grid is used. esys : Tuple[np.ndarray, np.ndarray], optional Precomputed eigenvalues and eigenvectors. scaling : float, optional Scaling factor for the wavefunction (default is 1). plot_potential : bool, optional Whether to plot the potential (default is False). basis : Literal["phase", "charge"], optional Basis in which to return the wavefunction ('phase' or 'charge') (default is 'phase'). andreev_basis : Literal["ballistic", "Andreev"], optional Andreev basis used for plotting. Use 'ballistic' for the basis used to define the Hamiltonian matrix, or 'Andreev' for the local phi-dependent Andreev eigenbasis. mode : Literal["abs", "abs2", "real", "imag"], optional Mode of the wavefunction ('abs', 'abs2', 'real', or 'imag') (default is 'abs'). **kwargs Additional arguments for plotting. Can include: - fig_ax: Tuple[plt.Figure, plt.Axes], optional Figure and axes to use for plotting. If not provided, a new figure and axes are created. Returns ------- Tuple[plt.Figure, plt.Axes] The figure and axes of the plot. """ if isinstance(which, int): which = [which] if andreev_basis not in {"ballistic", "Andreev"}: raise ValueError("Invalid andreev_basis; must be 'ballistic' or 'Andreev'.") if phi_grid is None: phi_grid = np.linspace(-5 * np.pi, 5 * np.pi, 151) fig_ax = kwargs.get("fig_ax") if fig_ax is None: fig, ax = plt.subplots() fig.suptitle(self._generate_suptitle()) else: fig, ax = fig_ax if plot_potential: potential = self.potential(phi=phi_grid) ax.plot( phi_grid / 2 / np.pi, potential[:, 0], color="black", label="Potential" ) ax.plot(phi_grid / 2 / np.pi, potential[:, 1], color="black") for idx in which: wavefunc_data = self.wavefunction( which=idx, phi_grid=phi_grid, esys=esys, basis=basis, representation=andreev_basis, ) phi_basis_labels = wavefunc_data["basis_labels"] wavefunc_amplitudes = wavefunc_data["amplitudes"] wavefunc_energy = wavefunc_data["energy"] if mode == "abs": y_values = np.abs(wavefunc_amplitudes[0]) y_values_down = np.abs(wavefunc_amplitudes[1]) elif mode == "abs2": y_values = np.abs(wavefunc_amplitudes[0]) ** 2 y_values_down = np.abs(wavefunc_amplitudes[1]) ** 2 elif mode == "real": y_values = wavefunc_amplitudes[0].real y_values_down = wavefunc_amplitudes[1].real elif mode == "imag": y_values = wavefunc_amplitudes[0].imag y_values_down = wavefunc_amplitudes[1].imag else: raise ValueError("Invalid mode; must be 'abs', 'real', or 'imag'.") ax.plot( phi_basis_labels / 2 / np.pi, wavefunc_energy + scaling * y_values, label=rf"$\Psi_{idx} \uparrow $", ) ax.plot( phi_basis_labels / 2 / np.pi, wavefunc_energy + scaling * y_values_down, label=rf"$\Psi_{idx} \downarrow $", ) if basis == "phase": ax.set_xlabel(r"$\Phi / \Phi_0$") ax.set_ylabel(r"$\psi(\varphi)$, Energy [GHz]") elif basis == "charge": ax.set_xlabel(r"$n$") ax.set_ylabel(r"$\psi(n)$, Energy [GHz]") ax.legend() ax.grid(True) return fig, ax
[docs] def plot_state( self, which: int = 0, phi_grid: np.ndarray = None, n_grid: np.ndarray = None, wigner_func: bool = False, esys: tuple[np.ndarray, np.ndarray] = None, plot_bloch: bool = False, **kwargs, ) -> tuple[plt.Figure, plt.Axes]: """ Plot the Wigner function of the state and the Bloch sphere. Parameters ---------- which : int, optional Index of desired wavefunction (default is 0). phi_grid : np.ndarray, optional Custom grid for phi; if None, a default grid is used. n_grid : np.ndarray, optional Custom grid for n; if None, a default grid is used. wigner_func : bool, optional Precomputed wigner_func function (default is False). esys : Tuple[np.ndarray, np.ndarray], optional Precomputed eigenvalues and eigenvectors. plot_bloch : bool, optional Whether to plot the Bloch sphere (default is False). **kwargs : dict, optional Additional arguments for plotting. Can include: - fig_ax: Tuple[plt.Figure, plt.Axes], optional Figure and axes to use for plotting. If not provided, a new figure and axes are created. - cmap: str, optional Colormap to use for the Wigner function (default is 'seismic'). - bloch_view: Tuple[float, float], optional Tuple with (elevation, azimuth) for Bloch sphere view (default is (-30, 60)). - bloch_position: Tuple[float, float, float, float], optional Position of the Bloch sphere inset in figure coordinates (left, bottom, width, height). If not provided, a default position is calculated. Returns ------- Tuple[plt.Figure, plt.Axes] The figure and axes of the plot. """ if phi_grid is None: phi_grid = np.linspace(-5, 5, 151) if n_grid is None: n_grid = np.linspace(-5, 5, 151) if wigner_func is False: wigner_func = self.wigner( which=which, phi_grid=phi_grid, n_grid=n_grid, esys=esys ) fig_ax = kwargs.get("fig_ax") if fig_ax is None: fig, ax = plt.subplots() fig.suptitle(self._generate_suptitle()) else: fig, ax = fig_ax cmap = kwargs.get("cmap", "seismic") map = ax.imshow( wigner_func, aspect="auto", origin="lower", extent=[phi_grid[0], phi_grid[-1], n_grid[0], n_grid[-1]], cmap=cmap, ) vmax = np.max(np.abs(wigner_func)) vmin = -vmax map.set_clim(vmin, vmax) ax.set_xlabel(r"$\varphi/2\pi$") ax.set_ylabel(r"$n$") # ax.set_aspect('equal') if plot_bloch: bloch_view = kwargs.get("bloch_view", (30, -60)) bloch_position = kwargs.get( "bloch_position" ) # Custom position for Bloch sphere bbox = ( ax.get_position() ) # posición del Axes principal en coordenadas de figura if bloch_position is None: inset_width = 0.3 * bbox.width inset_height = 0.3 * bbox.height inset_left = bbox.width - inset_width * 1.01 inset_bottom = bbox.height - inset_height * 1.01 bloch_position = [inset_left, inset_bottom, inset_width, inset_height] bloch_position[0] += bbox.x0 bloch_position[1] += bbox.y0 inset_ax = fig.add_axes(bloch_position, projection="3d") inset_ax.view_init(elev=bloch_view[0], azim=bloch_view[1]) rho_reduced = self.reduced_density_matrix(which=which, esys=esys, subsys=1) rho_reduced_qobj = Qobj(rho_reduced) b = Bloch(fig=fig, axes=inset_ax) b.vector_width = 1.5 b.xlabel = ["", ""] b.ylabel = ["", ""] b.zlabel = ["", ""] b.add_states(state=rho_reduced_qobj, colors="blue") b.render() # fig.colorbar(map, ax=ax, label=r"$W(\Phi,n)$", shrink=0.8) # fig.tight_layout() return fig, ax