Spectrum Analysis

def matrix_element(state1: Union[np.ndarray, qt.Qobj], operator: Union[np.ndarray, csc_matrix, qt.Qobj], 
                   state2: Union[np.ndarray, qt.Qobj]) -> Union[float, complex]:
    """
    Calculate matrix element <state1|operator|state2>.
    
    Parameters:
    - state1 (Union[np.ndarray, qt.Qobj]): Bra state vector  
    - operator (Union[np.ndarray, csc_matrix, qt.Qobj]): Operator matrix
    - state2 (Union[np.ndarray, qt.Qobj]): Ket state vector
    
    Returns:
    - Union[float, complex]: Matrix element value
    """

def get_matrixelement_table(operator: Union[np.ndarray, csc_matrix, qt.Qobj], 
                           state_table: Union[np.ndarray, qt.Qobj]) -> np.ndarray:
    """
    Generate comprehensive matrix element table.
    
    Parameters:
    - operator (Union[np.ndarray, csc_matrix, qt.Qobj]): Operator for matrix elements
    - state_table (Union[np.ndarray, qt.Qobj]): Table of eigenvectors in scipy eigsh format
    
    Returns:
    - np.ndarray: Matrix element table [i,j] = <i|operator|j>
    """

def absorption_spectrum(transitions: dict, initial_state_labels: list, 
                       lineshape_func, **kwargs) -> tuple:
    """
    Calculate absorption spectrum from transition data.
    
    Parameters:
    - transitions (dict): Dictionary of transition frequencies and strengths
    - initial_state_labels (list): Labels for initial states
    - lineshape_func: Function defining spectral lineshape
    - **kwargs: Additional parameters for lineshape function
    
    Returns:
    - tuple: (frequencies, spectrum) arrays
    """

def emission_spectrum(transitions: dict, initial_state_labels: list,
                     lineshape_func, **kwargs) -> tuple:
    """
    Calculate emission spectrum from transition data.
    
    Parameters:
    - transitions (dict): Dictionary of transition frequencies and strengths  
    - initial_state_labels (list): Labels for initial states
    - lineshape_func: Function defining spectral lineshape
    - **kwargs: Additional parameters for lineshape function
    
    Returns:
    - tuple: (frequencies, spectrum) arrays
    """

def order_eigensystem(evals: np.ndarray, evecs: np.ndarray, 
                     previous_evecs: np.ndarray = None) -> tuple:
    """
    Order eigenvalues and eigenvectors consistently.
    
    Parameters:
    - evals (np.ndarray): Eigenvalues to order
    - evecs (np.ndarray): Corresponding eigenvectors
    - previous_evecs (np.ndarray): Reference eigenvectors for consistent ordering
    
    Returns:
    - tuple: (ordered_evals, ordered_evecs)
    """

def standardize_phases(evecs: np.ndarray) -> np.ndarray:
    """
    Standardize phase conventions for eigenvectors.
    
    Parameters:
    - evecs (np.ndarray): Eigenvector matrix
    
    Returns:
    - np.ndarray: Eigenvectors with standardized phases
    """

def identity_wrap(operator, subsystem_index: int, subsystem_list: list) -> np.ndarray:
    """
    Wrap operator with identity matrices for composite systems.
    
    Parameters:
    - operator: Operator matrix to wrap
    - subsystem_index (int): Index of subsystem for the operator
    - subsystem_list (list): List of all subsystems
    
    Returns:
    - np.ndarray: Wrapped operator for composite Hilbert space
    """

def transition_table(system, operator, evals_count: int = None) -> np.ndarray:
    """
    Generate table of all transition matrix elements.
    
    Parameters:
    - system: Quantum system object
    - operator: Transition operator (e.g., charge, flux)
    - evals_count (int): Number of energy levels to include
    
    Returns:
    - np.ndarray: Transition matrix elements [initial, final]
    """

def plot_transitions(system, operator, initial_state: int = 0, 
                    final_states: list = None, **kwargs):
    """
    Plot transition strengths as function of energy.
    
    Parameters:
    - system: Quantum system object
    - operator: Transition operator
    - initial_state (int): Index of initial state
    - final_states (list): Indices of final states to plot
    - **kwargs: Plotting parameters
    """

def oscillator_strength(system, operator, initial_state: int, 
                       final_state: int) -> float:
    """
    Calculate oscillator strength for transition.
    
    Parameters:
    - system: Quantum system object
    - operator: Transition operator
    - initial_state (int): Initial state index
    - final_state (int): Final state index
    
    Returns:
    - float: Oscillator strength
    """

def lorentzian_lineshape(frequency: float, center_freq: float, 
                        gamma: float) -> float:
    """
    Lorentzian lineshape function.
    
    Parameters:
    - frequency (float): Frequency at which to evaluate lineshape
    - center_freq (float): Center frequency of transition
    - gamma (float): FWHM linewidth
    
    Returns:
    - float: Lineshape amplitude at given frequency
    """

def gaussian_lineshape(frequency: float, center_freq: float, 
                      sigma: float) -> float:
    """
    Gaussian lineshape function.
    
    Parameters:
    - frequency (float): Frequency at which to evaluate lineshape
    - center_freq (float): Center frequency of transition
    - sigma (float): Standard deviation of Gaussian
    
    Returns:
    - float: Lineshape amplitude at given frequency
    """

import scqubits as scq
import numpy as np
from scqubits.utils import spectrum_utils

# Create transmon and calculate eigensystem
transmon = scq.Transmon(EJ=25.0, EC=0.2, ng=0.0, ncut=30)
evals, evecs = transmon.eigensys(evals_count=6)

# Calculate charge matrix elements
n_matrix = spectrum_utils.get_matrixelement_table(
    transmon, 
    transmon.n_operator, 
    evecs=evecs
)

print("Charge matrix elements:")
print(n_matrix)

# Calculate specific transition matrix element
transition_01 = spectrum_utils.matrix_element(
    evecs[:, 0],           # Ground state
    transmon.n_operator,   # Charge operator
    evecs[:, 1]            # First excited state
)
print(f"0->1 charge matrix element: {transition_01}")

# Calculate oscillator strengths
osc_strength_01 = spectrum_utils.oscillator_strength(
    transmon, transmon.n_operator, 0, 1
)
print(f"0->1 oscillator strength: {osc_strength_01}")

# Analyze all transitions from ground state
ground_state = 0
excited_states = [1, 2, 3, 4, 5]

print("Transitions from ground state:")
for final_state in excited_states:
    # Transition energy
    transition_energy = evals[final_state] - evals[ground_state]
    
    # Matrix element magnitude
    matrix_elem = abs(n_matrix[ground_state, final_state])
    
    print(f"0->{final_state}: Energy = {transition_energy:.4f} GHz, "
          f"|<0|n|{final_state}>| = {matrix_elem:.4f}")

# Calculate transition strengths across parameter sweep
ng_vals = np.linspace(-2, 2, 101)
hilbert_space = scq.HilbertSpace([transmon])

sweep = scq.ParameterSweep(
    hilbert_space=hilbert_space,
    paramvals_by_name={'ng': ng_vals},
    evals_count=6
)
sweep.run()

# Extract transition frequencies and matrix elements
transitions_01 = []
matrix_elements_01 = []

for i, ng in enumerate(ng_vals):
    # Get eigensystem at this parameter point
    evals_i, evecs_i = sweep.eigensys(param_index=i)
    
    # Transition frequency
    omega_01 = evals_i[1] - evals_i[0]
    transitions_01.append(omega_01)
    
    # Matrix element
    n_matrix_i = spectrum_utils.get_matrixelement_table(
        transmon, transmon.n_operator, evecs=evecs_i
    )
    matrix_elements_01.append(abs(n_matrix_i[0, 1]))

# Plot results
import matplotlib.pyplot as plt

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

ax1.plot(ng_vals, transitions_01)
ax1.set_ylabel('01 Transition (GHz)')
ax1.set_title('Transition Frequency vs Offset Charge')

ax2.plot(ng_vals, matrix_elements_01)
ax2.set_xlabel('ng')
ax2.set_ylabel('|<0|n|1>|')
ax2.set_title('Transition Matrix Element vs Offset Charge')

plt.tight_layout()
plt.show()

# Matrix elements in composite systems
transmon = scq.Transmon(EJ=25.0, EC=0.2, ng=0.0, ncut=25)
cavity = scq.Oscillator(E_osc=6.0, truncated_dim=4)

system = scq.HilbertSpace([transmon, cavity])
system.add_interaction(
    g_strength=0.1,
    op1=transmon.n_operator,
    op2=cavity.creation_operator + cavity.annihilation_operator
)

# Calculate dressed eigensystem
dressed_evals, dressed_evecs = system.eigensys(evals_count=12)

# Wrap transmon charge operator for composite space
n_wrapped = spectrum_utils.identity_wrap(
    transmon.n_operator, 
    subsystem_index=0,
    subsystem_list=[transmon, cavity]
)

# Calculate dressed matrix elements
dressed_n_matrix = spectrum_utils.get_matrixelement_table(
    system, n_wrapped, evecs=dressed_evecs
)

print("Dressed charge matrix elements (first 6x6 block):")
print(dressed_n_matrix[:6, :6])

# Generate realistic spectrum with lineshapes
from scqubits.utils.spectrum_utils import lorentzian_lineshape

# Define transitions and their strengths
transitions = {
    'frequencies': [4.8, 4.9, 5.0, 5.1],  # GHz
    'strengths': [0.1, 0.8, 1.0, 0.3]     # Relative intensities
}

# Frequency range for spectrum
freq_range = np.linspace(4.5, 5.5, 1000)
spectrum = np.zeros_like(freq_range)

# Add each transition with Lorentzian lineshape
gamma = 0.01  # 10 MHz linewidth

for freq, strength in zip(transitions['frequencies'], transitions['strengths']):
    for i, f in enumerate(freq_range):
        spectrum[i] += strength * lorentzian_lineshape(f, freq, gamma)

# Plot spectrum
plt.figure(figsize=(10, 6))
plt.plot(freq_range, spectrum, 'b-', linewidth=2)
plt.xlabel('Frequency (GHz)')
plt.ylabel('Absorption (arb. units)')
plt.title('Simulated Absorption Spectrum')
plt.grid(True, alpha=0.3)
plt.show()

# Analyze selection rules for different operators
operators = {
    'n': transmon.n_operator,
    'cos_phi': transmon.cos_phi_operator, 
    'sin_phi': transmon.sin_phi_operator
}

# Calculate matrix elements for each operator
for op_name, operator in operators.items():
    matrix_table = spectrum_utils.get_matrixelement_table(
        transmon, operator, evecs=evecs
    )
    
    print(f"\n{op_name} operator matrix elements:")
    print("State transitions with |matrix element| > 0.01:")
    
    for i in range(6):
        for j in range(6):
            if i != j and abs(matrix_table[i, j]) > 0.01:
                print(f"  {i}->{j}: {matrix_table[i, j]:.4f}")