Parameter Sweeps

class ParameterSweep:
    def __init__(self, hilbert_space: HilbertSpace, paramvals_by_name: dict, evals_count: int = 6,
                 subsys_update_list: list = None, update_hilbert_space = None, 
                 num_cpus: int = None):
        """
        Parameter sweep for quantum systems.
        
        Parameters:
        - hilbert_space (HilbertSpace): Quantum system to analyze
        - paramvals_by_name (dict): Parameter names and value arrays to sweep
        - evals_count (int): Number of eigenvalues to calculate
        - subsys_update_list (list): Indices of subsystems to update during sweep
        - update_hilbert_space: Function to update system parameters
        - num_cpus (int): Number of CPU cores for parallel processing
        """
        
    def run(self) -> None:
        """Execute the parameter sweep calculation."""
        
    def plot_evals_vs_paramvals(self, which: int = -1, subtract_ground: bool = False,
                               param_name: str = None, **kwargs):
        """
        Plot energy levels vs parameter values.
        
        Parameters:
        - which (int): Which subsystem to plot (-1 for composite system)
        - subtract_ground (bool): Subtract ground state energy
        - param_name (str): Parameter name for x-axis
        """
        
    def transitions(self, initial_state_ind: int, final_state_ind: int) -> np.ndarray:
        """
        Calculate transition frequencies between states.
        
        Parameters:
        - initial_state_ind (int): Index of initial state
        - final_state_ind (int): Index of final state
        
        Returns:
        - np.ndarray: Transition frequencies across parameter sweep
        """
        
    def plot_transitions(self, initial_state_ind: int, final_state_ind: int, 
                        param_name: str = None, **kwargs):
        """Plot transition frequencies vs parameter."""
        
    def eigenvals(self, param_index: int = None) -> np.ndarray:
        """Get eigenvalues at specific parameter point."""
        
    def eigensys(self, param_index: int = None) -> tuple:
        """Get eigensystem at specific parameter point."""
        
    def bare_specdata(self, subsys_index: int) -> SpectrumData:
        """Get bare spectrum data for subsystem."""
        
    def dressed_specdata(self) -> SpectrumData:
        """Get dressed spectrum data for composite system."""

class StoredSweep:
    def __init__(self, spectrumdata: SpectrumData, hilbert_space: HilbertSpace = None,
                 bare_specdata_list: list = None):
        """
        Container for stored parameter sweep data.
        
        Parameters:
        - spectrumdata (SpectrumData): Spectrum data object
        - hilbert_space (HilbertSpace): Associated quantum system
        - bare_specdata_list (list): List of bare spectrum data for subsystems
        """
        
    def plot_evals_vs_paramvals(self, which: int = -1, **kwargs):
        """Plot stored energy levels vs parameters."""
        
    def transitions(self, initial_state_ind: int, final_state_ind: int) -> np.ndarray:
        """Calculate transitions from stored data."""
        
    def new_parametersweep(self, update_hilbert_space, paramvals_by_name: dict) -> ParameterSweep:
        """Create new parameter sweep based on stored data."""

class SpectrumData:
    def __init__(self, energy_table: np.ndarray, system_params: dict = None, 
                 state_table: np.ndarray = None, matrixelem_table: np.ndarray = None):
        """
        Storage for spectrum calculation results.
        
        Parameters:
        - energy_table (np.ndarray): Eigenvalues across parameter sweep
        - system_params (dict): System parameters
        - state_table (np.ndarray): Eigenvectors across parameter sweep  
        - matrixelem_table (np.ndarray): Matrix elements across sweep
        """
        
    def subtract_ground(self) -> SpectrumData:
        """Return new SpectrumData with ground state energy subtracted."""
        
    def eigenvals(self, param_index: int = None) -> np.ndarray:
        """Get eigenvalues at parameter point."""
        
    def eigenvecs(self, param_index: int = None) -> np.ndarray:
        """Get eigenvectors at parameter point."""
        
    def matrixelems(self, operator_name: str, param_index: int = None) -> np.ndarray:
        """Get matrix elements at parameter point."""

import scqubits as scq
import numpy as np

# Create transmon
transmon = scq.Transmon(EJ=25.0, EC=0.2, ng=0.0, ncut=30)
hilbert_space = scq.HilbertSpace([transmon])

# Define parameter sweep  
ng_vals = np.linspace(-2.0, 2.0, 101)
sweep = scq.ParameterSweep(
    hilbert_space=hilbert_space,
    paramvals_by_name={'ng': ng_vals},
    evals_count=6
)

# Run calculation
sweep.run()

# Plot results
sweep.plot_evals_vs_paramvals(
    subtract_ground=True,
    param_name='ng'
)

# Calculate transition frequencies
transition_01 = sweep.transitions(0, 1)
transition_12 = sweep.transitions(1, 2)

print("01 transition range:", transition_01.min(), "to", transition_01.max())
print("12 transition range:", transition_12.min(), "to", transition_12.max())

# Two-parameter sweep
EJ_vals = np.linspace(20, 30, 26)
ng_vals = np.linspace(-1, 1, 51)

sweep_2d = scq.ParameterSweep(
    hilbert_space=hilbert_space,
    paramvals_by_name={
        'EJ': EJ_vals,
        'ng': ng_vals
    },
    evals_count=4
)

sweep_2d.run()

# Plot 2D parameter space
import matplotlib.pyplot as plt

# Extract data for plotting
energies = sweep_2d.dressed_specdata().energy_table
EJ_grid, ng_grid = np.meshgrid(EJ_vals, ng_vals, indexing='ij')

# Plot ground state energy
plt.figure(figsize=(10, 8))
plt.contourf(EJ_grid, ng_grid, energies[:, :, 0], levels=50)
plt.xlabel('EJ (GHz)')
plt.ylabel('ng')
plt.title('Ground State Energy')
plt.colorbar()
plt.show()

# Sweep coupled qubit-cavity system
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
)

# Sweep transmon frequency
EJ_vals = np.linspace(20, 30, 51)
sweep = scq.ParameterSweep(
    hilbert_space=system,
    paramvals_by_name={'EJ': EJ_vals},
    evals_count=12,
    subsys_update_list=[0]  # Only update transmon
)

sweep.run()

# Plot dressed spectrum
sweep.plot_evals_vs_paramvals(
    which=-1,  # Composite system
    subtract_ground=True
)

# Plot bare transmon levels for comparison
sweep.plot_evals_vs_paramvals(
    which=0,   # First subsystem (transmon)
    subtract_ground=True
)

# Enable parallel processing for large sweeps
import scqubits as scq

# Set global multiprocessing settings
scq.settings.NUM_CPUS = 4
scq.settings.MULTIPROC = 'multiprocessing'

# Large parameter sweep
flux_vals = np.linspace(0, 1, 201)
fluxonium = scq.Fluxonium(EJ=2.0, EC=0.5, EL=0.8, flux=0.0)

sweep = scq.ParameterSweep(
    hilbert_space=scq.HilbertSpace([fluxonium]),
    paramvals_by_name={'flux': flux_vals},
    evals_count=10,
    num_cpus=4  # Parallel processing
)

sweep.run()

# Save sweep results
scq.io_utils.write(sweep, 'transmon_sweep.h5')

# Load saved sweep
loaded_sweep = scq.io_utils.read('transmon_sweep.h5')

# Create StoredSweep for analysis
stored_sweep = scq.StoredSweep(
    spectrumdata=loaded_sweep.dressed_specdata(),
    hilbert_space=loaded_sweep.hilbert_space
)

# Analyze stored data
stored_sweep.plot_evals_vs_paramvals()
transitions = stored_sweep.transitions(0, 1)

def update_system(system, param_val):
    """Custom function to update system parameters."""
    # Update both EJ and EC simultaneously
    EJ_new = param_val
    EC_new = param_val * 0.01  # Scale EC with EJ
    
    system.subsystem_list[0].EJ = EJ_new
    system.subsystem_list[0].EC = EC_new

# Use custom update function
EJ_vals = np.linspace(15, 35, 41)
sweep = scq.ParameterSweep(
    hilbert_space=system,
    paramvals_by_name={'EJ': EJ_vals},
    update_hilbert_space=update_system,
    evals_count=6
)

sweep.run()