diff --git a/.gitignore b/.gitignore index 93b1f4e..955cfce 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,10 @@ # MacOS .DS_Store +# AI Agent Scripts +sandia-claudecode + + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/examples/example_spin_dependent_squeezing.py b/examples/example_spin_dependent_squeezing.py new file mode 100644 index 0000000..4c8377b --- /dev/null +++ b/examples/example_spin_dependent_squeezing.py @@ -0,0 +1,721 @@ +import ionsim as ism + +import numpy as np +from scipy.sparse import kron as skron +import sys +import matplotlib +import matplotlib.pyplot as plt +import platform +from scipy.special import factorial + +TPI = 2*np.pi + +def four_tone_hamiltonian(basis, qubit_basis, modes, spectator_modes, etas, rabi_rate_x, rabi_rate_y, omega_x_red, omega_x_blue, omega_y_red, omega_y_blue, targetIon: AtomicStructure, + index_mode_i:int, index_mode_j:int, phi:float, sparse=False, mod=None): + ''' Applies 4-tone MS-like gate to a single ion "k" to produce a spin-dependent motional operation. + Consists of 2 sidebands (red and blue) on mode i and another 2 sidebands (red and blue) on mode j: + + H(t) = hbar Omega_x sigma_x a_i exp(-i Delta t) + h.c. + hbar Omega_y sigma_y a_j exp(-i n Delta t + phi) + h.c. + + - currently sets the first MS phase to zero, accepts input for second MS phase as phi + ''' + # Hamiltonian: H = -eta * Omega/2 sum_{k} [-i sigma_+,k exp(iphi) + i sigma_-,k exp(-iphi)] (a e^i delta t + a† e^{-i delta t} ) + # Hamiltonian: H = -eta * Omega J_y (a e^i delta t + a† e^{-i delta t} ) # for phi = 0; Jy = 1/2 (sigma_y,1 + sigma_y,2) + motional_basis = ism.StandardBasis([*modes]) + + DOFs = basis.degrees_of_freedom + + target_ion_index = DOFs.index(targetIon) + + operators = [] + + # build matrices for fundamental raising/lowering spin/Fock operators: + # Raising operator for target ion "k" + raise_spin_k = qubit_basis.enlarge_matrix(ism.Pauli.plus, [targetIon]) + + fock_dimension = len(modes[index_mode_i].energy_levels) + assert fock_dimension == len(modes[index_mode_j].energy_levels) + + # First, include modes that are supposed to be included. (non-spectators) + raise_motion = motional_basis.enlarge_matrix(ism.Fock.raising(fock_dimension), [modes[index_mode_i]]) + lower_motion = motional_basis.enlarge_matrix(ism.Fock.lowering(fock_dimension), [modes[index_mode_j]]) + + # Extract Lamb-Dicke parameter for mode i and j, ion k + eta_ki = etas[target_ion_index, index_mode_i] + eta_kj = etas[target_ion_index, index_mode_j] + + # spin phase = ( phi_R + phi_B )/ 2 + # ------------------- + # x-spin MS-like gate: + # For mode i : + # Can achieve sigma_x by setting spin phase of tones 1 and 2 to zero or each as 2pi + phase_x = np.pi/2. + + x_prefactor = 1j*rabi_rate_x * 0.5 * eta_ki * np.exp(1j*phase_x) + + operator_red_mode_i = x_prefactor * skron(raise_spin_k, lower_motion) + operator_blue_mode_i = x_prefactor * skron(raise_spin_k, raise_motion) + # h.c. is already included by ionsim upon providing these coupling operators to the hamiltonian + + # ------------------- + # y-spin MS-like gate: + # For mode j : + phase_y = 0. + y_prefactor = 1j*rabi_rate_y * 0.5 * eta_kj * np.exp(1j*(phase_y + phi)) + + operator_red_mode_j = y_prefactor * skron(raise_spin_k, lower_motion) + operator_blue_mode_j = y_prefactor * skron(raise_spin_k, raise_motion) + + operators.extend([ + ism.CouplingOperator.from_matrix(basis, operator_red_mode_i, omega_x_red, modulation_function=mod), + ism.CouplingOperator.from_matrix(basis, operator_blue_mode_i, omega_x_blue, modulation_function=mod), + ism.CouplingOperator.from_matrix(basis, operator_red_mode_j, omega_y_red, modulation_function=mod), + ism.CouplingOperator.from_matrix(basis, operator_blue_mode_j, omega_y_blue, modulation_function=mod), + ]) + + for mode in spectator_modes: + mode_index = modes.index(mode) + spectator_raise_motion = motional_basis.enlarge_matrix(ism.Fock.raising(fock_dimension), [mode]) + spectator_lower_motion = motional_basis.enlarge_matrix(ism.Fock.lowering(fock_dimension), [mode]) + + # Extract Lamb-Dicke parameter + eta_ki = etas[target_ion_index, mode_index] + eta_kj = etas[target_ion_index, mode_index] + + x_prefactor = 1j*rabi_rate_x * 0.5 * eta_ki * np.exp(1j*phase_x) + + operator_red_x_spectator = x_prefactor * skron(raise_spin_k, spectator_lower_motion) + operator_blue_x_spectator = x_prefactor * skron(raise_spin_k, spectator_raise_motion) + + y_prefactor = 1j*rabi_rate_y * 0.5 * eta_kj * np.exp(1j*(phase_y + phi)) + + operator_red_y_spectator = y_prefactor * skron(raise_spin_k, spectator_lower_motion) + operator_blue_y_spectator = y_prefactor * skron(raise_spin_k, spectator_raise_motion) + + operators.extend([ + ism.CouplingOperator.from_matrix(basis, operator_red_x_spectator, omega_x_red, modulation_function=mod), + ism.CouplingOperator.from_matrix(basis, operator_blue_x_spectator, omega_x_blue, modulation_function=mod), + ism.CouplingOperator.from_matrix(basis, operator_red_y_spectator, omega_y_red, modulation_function=mod), + ism.CouplingOperator.from_matrix(basis, operator_blue_y_spectator, omega_y_blue, modulation_function=mod), + ]) + + interaction_frame_energies = [-state.energy for state in basis.states] + return ism.Hamiltonian(basis, operators, interaction_frame_energies, sparse=sparse) + + +def thermal_state_populations(nbar: float, fock_dimension): + fock_populations = np.zeros(fock_dimension) + + normalization = 1./(1. + nbar) + for n in range(fock_dimension): + fock_populations[n] = (nbar/(1. + nbar))**(n) + + return fock_populations*normalization + + +def dephasing_dissipator(basis, qubit_basis, modes, dephasing_rates: list[float]): + ''' Dissipator for ion motional mode dephasing ''' + + motional_basis = ism.StandardBasis([*modes]) + qubits = qubit_basis.degrees_of_freedom + + # Build spin basis identity + spin_identities = [] + for spin in qubits: + # Identity matrix for mode m in Fock space, enlarged to fit dimensionality of M modes + spin_identities.append(qubit_basis.enlarge_matrix(ism.Pauli.I, [spin])) + + spin_identity = np.sum(spin_identities, axis=0) + + # Motional mode dephasing is proportional to the phonon number operators + # loop over modes: + lindblad_operators = [] + for mode, dephasing_rate in zip(modes, dephasing_rates): + fock_dim = len(mode.energy_levels) + mode_number_matrix = motional_basis.enlarge_matrix(ism.Fock.number(fock_dim), [mode]) + + # One lindblad operator per mode: L = sqrt(Gamma) * (a†_m a_m) + dephasing_operator = np.sqrt(dephasing_rate) * skron(spin_identity, mode_number_matrix) + + # Lindblad operator is diagonal in the Fock occupancy number basis; use EnergyShiftoperator or GeneralOperator + lindblad_operators.append(ism.EnergyShiftOperator.from_matrix(basis, dephasing_operator)) + + frame_energies = [-state.energy for state in basis.states] + return ism.Dissipator(basis, lindblad_operators, frame_energies) + + + +def heating_dissipator(basis, modes, heating_rates: list[float]): + ''' Dissipator for ion mode heating ''' + + motional_basis = ism.StandardBasis([*modes]) + + # Build spin basis identity + spin_identities = [] + for spin in qubits: + # Identity matrix for mode m in Fock space, enlarged to fit dimensionality of M modes + spin_identities.append(qubit_basis.enlarge_matrix(ism.Pauli.I, [spin])) + + spin_identity = np.sum(spin_identities, axis=0) + + # Phonon raising and lowering operators + # loop over modes: + lindblad_operators = [] + for mode, heating_rate in zip(modes, heating_rates): + fock_dim = len(mode.energy_levels) + mode_lowering_matrix = motional_basis.enlarge_matrix(ism.Fock.lowering(fock_dim), [mode]) + mode_raising_matrix = motional_basis.enlarge_matrix(ism.Fock.raising(fock_dim), [mode]) + + # Two lindblad operators per mode: L1 = sqrt(Gamma) * a_m , L2 = sqrt(Gamma) * a†_m + lowering_operator = np.sqrt(heating_rate) * skron(spin_identity, mode_lowering_matrix) + raising_operator = np.sqrt(heating_rate) * skron(spin_identity, mode_raising_matrix) + #lowering_operator = np.sqrt(decay_rate*(N_thermal+1)) * skron(spin_identity, mode_lowering_matrix) + #raising_operator = np.sqrt(decay_rate*N_thermal) * skron(spin_identity, mode_raising_matrix) + + lindblad_operators.append(ism.CouplingOperator.from_matrix(basis, lowering_operator, 0.)) + lindblad_operators.append(ism.CouplingOperator.from_matrix(basis, raising_operator, 0.)) + + frame_energies = [-state.energy for state in basis.states] + return ism.Dissipator(basis, lindblad_operators, frame_energies) + + +def carrier_hamiltonian(basis, target_ion, rabi_rate: float, phase: float, omega_carrier: float, sparse: bool=False, mod=None): + """ Hamiltonian for carrier spin pulse: (hbar * Omega/2) * sigma_phi e^iphi + h.c. """ + prefactor = rabi_rate * 0.5 * np.exp(-1j*phase) + operator_raising = basis.enlarge_matrix(prefactor * ism.Pauli.plus, [target_ion]) + #raise_spin_k = qubit_basis.enlarge_matrix(ism.Pauli.plus, [target_ion]) + + operators = [ism.CouplingOperator.from_matrix(basis, operator_raising, omega_carrier, modulation_function=mod)] + interaction_frame_energies = [-state.energy for state in basis.states] + return ism.Hamiltonian(basis, operators, interaction_frame_energies, sparse=sparse) + + +def red_sb_pi_pulse_duration(rabi_rate, fock_state, eta, single_ion: bool=True): + # Assume either single or 2-ion sysetm + if single_ion: + t = np.pi / (rabi_rate * np.sqrt(fock_state) * np.abs(eta)) + else: + t = np.pi / (rabi_rate * np.sqrt(fock_state * (fock_state - 1)) * np.abs(eta)) + return t + + + +def compute_N(rho_motional, motional_basis): + + #for state in motional_basis.states: + N_operator = ism.Fock.raising(fock_dimension) @ ism.Fock.lowering(fock_dimension) + return np.trace( N_operator.dot(rho_motional.density_matrix)) + +def estimate_squeezing_magnitude(motional_rho, mode_basis): + ''' Estimate the squeezing parameter from the X and P variances ''' + x, p, x2, p2 = motional_rho.compute_quadratures(True) # includes variance + + avg_X_tilt = x[0].real + avg_X2_tilt = x2[0].real + + X_variance = avg_X2_tilt - avg_X_tilt**2 + #print(f"variance of X for the tilt mode: {X_variance}") + + avg_P_tilt = p[0].real + avg_P2_tilt = p2[0].real + + P_variance = avg_P2_tilt - avg_P_tilt**2 + #print(f"variance of P for the tilt mode: {P_variance}") + if squeezing_phase == 0.: + X_squeezed = False + elif squeezing_phase == np.pi: + X_squeezed = True + + if X_squeezed: + r_X = -0.5 * np.log(2. * X_variance) + r_P = 0.5 * np.log(2. * P_variance) + else: + r_X = 0.5 * np.log(2. * X_variance) + r_P = -0.5 * np.log(2. * P_variance) + + return r_X, r_P + +def squeezed_ground_state_coefficients(r_value, phi, fock_dimension): + wave_function = np.zeros(fock_dimension,dtype=complex) + norm = 1./np.sqrt(np.cosh(r_value)) + for n in range(fock_dimension//2): + wave_function[2*n] = np.exp(1j * n * phi) * np.tanh(r_value)**n + #wave_function[2*n] *= norm * np.sqrt(factorial(2 * n)) + wave_function[2*n] *= norm * np.sqrt(factorial(2 * n))/(2**n) + wave_function[2*n] /= (factorial(n)) + return wave_function + + + +def squeezed_vacuum_populations(fock_dim, r_val): + fock_states = np.arange(fock_dim) + fock_populations = np.zeros(len(fock_states)) + norm = 1./np.cosh(r_val) + for i in range(fock_dim//2): + fock_populations[2*i] = norm * (np.tanh(r_val)**(2*i)) * factorial(2 * i)/(2**(2*i) * factorial(i)**2 ) + return fock_populations + + + +def main(): + import time + from matplotlib import pyplot as plt + + print(f"\n --- Simulation of spin-dependent squeezing by 4-tone sideband protocol --- ") + + include_spectator_mode = False + include_dissipation = False + do_reverse_process = False + thermal_state_IC = False + reverse_phases = False + beam_counter_propagating = True + + num_ions = 2 + + # Create 171Yb+ qubits + qubits = [ + ism.AtomicStructure.from_species(species='171Yb+', term_symbols=['S1/2'], level_names=['S1/2,0,0', 'S1/2,1,0']) + for _ in range(num_ions) + ] + + target_ion = qubits[0] + target_ion_index = qubits.index(target_ion) + + qubit_basis = ism.StandardBasis([*qubits]) + + # Use mode analysis to get motional mode information + omega_x = 2.1 * TPI * 1E6 # MHz -> rad/s + omega_y = 2.6 * TPI * 1E6 # "" + omega_z = 0.50 * TPI * 1E6 # "" + trap_analysis = ism.LinearIonChain.from_atomic_structure_basis(qubit_basis, omega_x, omega_y, omega_z) + trap_analysis.solve_ion_trap_equilibrium() + print() + trap_analysis.print_chain_summary() + + print(trap_analysis.characteristic_parameters) + + + # Truncate the Hilbert space to a single qubit with the two nearby motional modes + # Choose to do oprerations on the radial x tilt mode; option to include the COM as a spectator mode + #print(trap_analysis.get_mode_properties_by_branch('x', 0)) + #print(trap_analysis.get_mode_properties_by_branch('x', 1)) + tilt_mode_index = 0 + COM_mode_index = 1 + + fock_dimension = 10 + branch_dir = 'x' + qubit_basis = ism.StandardBasis([target_ion]) # Now reset the spin basis to just include the target ion + modes = trap_analysis.build_mode_DOFs_from_branch(branch_dir, [tilt_mode_index, COM_mode_index], fock_dimension) + if include_spectator_mode: + basis = ism.StandardBasis([qubits[target_ion_index], *modes]) + motional_basis = ism.StandardBasis(modes) + spectator_modes = [modes[1]] + else: + basis = ism.StandardBasis([qubits[target_ion_index], modes[tilt_mode_index]]) + motional_basis = ism.StandardBasis([modes[tilt_mode_index]]) + modes = [modes[0]] + spectator_modes = [] + + + target_mode_index = 0 # for tilt mode in our basis + + # 4-tone protocol targets 2 modes in general with 2 tones each; for squeezing, we target the same mode with all tones + mode_i_index = target_mode_index + mode_j_index = target_mode_index + + radial_x_tilt_mode_properties = trap_analysis.get_mode_properties_by_branch('x', tilt_mode_index) + mode_i_freq = radial_x_tilt_mode_properties['frequency'] + mode_j_freq = mode_i_freq + + print(f" - Tilt mode frequency: {mode_i_freq/TPI/1E6} [MHz]\n\n") + + # Laser/driving parameters + print(f" --- Laser parameters --- ") + wavelength = 355 # nm + beam_angle = np.pi/4. # relative to axial direction + wavevector = TPI / (355. * 1E-9) * np.array([np.cos(beam_angle), np.sin(beam_angle), 0.]) + if beam_counter_propagating: + # Effective wavevector from + wavevector *= 2. # Raman process: \delta k is 2x basic wavevector + + rabi_rate = 175 * 1E3 * TPI / np.sqrt(2.) #* np.sqrt(0.5) #/ np.sqrt(4.) # kHz -> Hz + x_rabi_rate = rabi_rate + y_rabi_rate = rabi_rate + print(f" - Carrier rabi rate: {rabi_rate / TPI / 1E3} [kHz]\n") + + detuning = 40.0 * 1E3 * TPI # 40 kHz -> rad/s + print(f" - Detuning: {detuning/(TPI * 1E3)} [kHz]\n") + duration = 60 * 1E-6 + print(f" - Gate duration: {duration*1E6:.3f} [µs]\n") + + modulate_amplitude = True + + print(f" - Estimated number of loops: {duration * detuning / TPI}\n") + + detuning_red = detuning + detuning_blue = detuning + + # Blue and red laser frequencies for 1st-sideband tuning: + # For x-spin / operation on mode i, the laser is simply detuned by 1*detuning + offset = 0. * TPI * 1E3 # kHz + omega_qubit = target_ion.energy_levels[1].energy - target_ion.energy_levels[0].energy + omega_b_mode_i = (omega_qubit + (mode_i_freq + detuning_blue)) + offset + omega_r_mode_i = (omega_qubit - (mode_i_freq + detuning_red)) - offset + + # Second MS pulse + # For y-spin / operation on mode j, it is detuned by n*detuning + n = -1 # for squeezing + omega_b_mode_j = (omega_qubit + (mode_j_freq + n*detuning_blue)) - offset + omega_r_mode_j = (omega_qubit - (mode_j_freq + n*detuning_red)) + offset + + ## -- Global parameters for script --- + frequencies = np.array([omega_r_mode_i, omega_r_mode_j, omega_b_mode_i, omega_b_mode_j]) # MHz + frequencies_from_carrier = (frequencies - omega_qubit)/TPI/1E6 # MHz + print(f" - laser frequencies (minus carrier): {frequencies_from_carrier}\n\n") + + + def blackman(t: float, magnitude: float = 1.): + return magnitude*(0.42 - 0.5*np.cos(TPI * t /duration) + 0.08 * np.cos(2.*TPI*t/duration)) + + # Retrieve lamb dicke parameters + lamb_dicke_parameters = trap_analysis.get_radial_lamb_dicke_parameters(wavevector, branch_dir) # matrix of ion, mode + eta = lamb_dicke_parameters[target_ion_index,target_mode_index] + print(f" - |eta| for target ion, target mode: {np.abs(eta)}\n") + print(f" - Sideband rabi rate (eta*Omega): {np.abs(eta) * rabi_rate / 1E3 / TPI} [kHz]\n") + print('\n\n') + print(f" - Lamb dicke parameters: \n{lamb_dicke_parameters}") + mod_function = None + expected_r = duration * ((np.abs(eta) * rabi_rate)**2) / detuning # Square pulse + if modulate_amplitude: + print(f" - Using a Blackman pulse shape\n") + mod_function = blackman + # Area of the blackman pulse shape squared is ~0.3046 + expected_r = 0.3046 * duration * ((np.abs(eta) * rabi_rate)**2) / detuning + else: + print(f" - Using a constant (square) pulse shape\n") + + r = expected_r + print(f" - Expecting squeezing of r = {expected_r}\n") + + + # ============================= + ## Create Hamiltonians: + squeezing_phase = 0. + squeezing_hamiltonian = four_tone_hamiltonian(basis, qubit_basis, modes, spectator_modes, lamb_dicke_parameters, x_rabi_rate, y_rabi_rate, omega_r_mode_i, omega_b_mode_i, omega_r_mode_j, + omega_b_mode_j, target_ion, mode_i_index, mode_j_index, squeezing_phase, mod=mod_function) + + reverse_squeezing_hamiltonian = four_tone_hamiltonian(basis, qubit_basis, modes, spectator_modes, lamb_dicke_parameters, x_rabi_rate, y_rabi_rate, omega_r_mode_i, omega_b_mode_i, omega_r_mode_j, omega_b_mode_j, + target_ion, mode_i_index, mode_j_index, squeezing_phase + np.pi, mod=mod_function) + + pi_pulse_hamiltonian = carrier_hamiltonian(basis, target_ion, rabi_rate, 0., np.abs(omega_qubit), False, None) + mode_dephasing_rates = [5.0, 500.] # quanta/second + if include_dissipation : + # Create dissipator and then full Lindbladian: + #ion_heating_dissipator = heating_dissipator(basis, modes, mode_heating_rates) + _dephasing_dissipator = dephasing_dissipator(basis, qubit_basis, modes, mode_dephasing_rates) + lindbladian_4tone = ism.Lindbladian(squeezing_hamiltonian, _dephasing_dissipator) + #lindbladian_4tone = ism.Lindbladian(squeezing_hamiltonian, ion_heating_dissipator) + # Create reverse processes + lindbladian_4tone_reverse = ism.Lindbladian(reverse_squeezing_hamiltonian, _dephasing_dissipator) + + # Initial fock state, all populated in this state: + starting_fock_state = 0 + excited_qubit1_init_state = False + if excited_qubit1_init_state: + qubit1_index = fock_dimension*num_modes + else: + qubit1_index = 0 + + nbars = [0.1, 0.5] + + IC_fock_pops = [thermal_state_populations(nbar, fock_dimension) for nbar in nbars] # populations + #fock_populations_target_mode_IC = thermal_state(nbar) + if thermal_state_IC: + print(f"\n--- Starting with thermal ground state ---\n") + print(f"Target mode initial populations for nbar = {nbars[0]}, : {IC_fock_pops[0]}") + coefs = np.kron(np.array([1., 0.]), IC_fock_pops[0]) + for m in range(1, len(modes)): + print(f"Spectator mode initial populations for nbar = {nbars[m]}, : {IC_fock_pops[m]}") + coefs = np.kron(coefs, IC_fock_pops[m]) + coefs = np.sqrt(coefs)# populations require square root + else: + print(f"--- Initializing in the ground state of full Hilbert space ---") + IC_fock_pops[0] = np.zeros(fock_dimension) + IC_fock_pops[0][0] = 1. + for m in range(1, len(modes)): + IC_fock_pops[m] = np.zeros(fock_dimension) + IC_fock_pops[m][0] = 1. + coefs = np.zeros(len(basis.states)) + starting_fock_state_mode2 = 0 + coefs[qubit1_index + starting_fock_state + starting_fock_state_mode2] = 1. + + print_state_names = False + if print_state_names: + for indx, state in enumerate(basis.states): + print(indx) + print(state.name) + init_state = ism.State.from_coefficients(basis, list(coefs)) + + times = np.linspace(0, duration, 2000) + + ground_state_level_index = 0 + excited_state_level_index = 1 + ground_state_level_name = qubit_basis.degrees_of_freedom[0].energy_levels[ground_state_level_index].name + excited_state_level_name = qubit_basis.degrees_of_freedom[0].energy_levels[excited_state_level_index].name + print("Ground state level name: " + ground_state_level_name) + print("Excited state level name: " + excited_state_level_name) + + # Now perform cycling to attempt squeezing operation + print(f"\n\n\n ------- Squeezing protocol ------- ") + print(f"Squeezing parameters: r = {r}, phi = {squeezing_phase}") + + # Use rho(t) from previous + # Apply 4 tones with n=-1 detuning relationship to get squeezing + rho_t = init_state + + ### Evolve Schrodinger or master equation dynamics for spin-dependent squeezing via 4-tone protocol + if include_dissipation: + rho_t = rho_t.propagate_using_master_equation(lindbladian_4tone, duration, times) + else: + rho_t = rho_t.propagate_using_schrodinger_equation(squeezing_hamiltonian, duration, times) # no dissipation + + if do_reverse_process: + # Perform pi pulse on qubit: # Omega * t = pi; t = pi/Omega + if reverse_phases: + if include_dissipation: + rho_t = rho_t[-1].propagate_using_master_equation(lindbladian_4tone_reverse, duration, times) + else: + rho_t = rho_t[-1].propagate_using_schrodinger_equation(reverse_squeezing_hamiltonian, duration, times) # no dissipation + + else: + pi_pulse_duration = np.pi/rabi_rate + + # Perform squeezing again, which should return you back to the ground state + if include_dissipation: + pi_pulse_lindbladian = ism.Lindbladian(hamiltonian = pi_pulse_hamiltonian, dissipator = None) + rho_t = rho_t[-1].propagate_using_master_equation(pi_pulse_lindbladian, pi_pulse_duration) + rho_t = rho_t.propagate_using_master_equation(lindbladian_4tone, duration, times) + else: + rho_t = rho_t[-1].propagate_using_schrodinger_equation(pi_pulse_hamiltonian, pi_pulse_duration) # no dissipation + rho_t = rho_t.propagate_using_schrodinger_equation(squeezing_hamiltonian, duration, times) # no dissipation + + # Extract final state + final_state = rho_t[-1] + + # Compute state overlap (fidelity of reverse process) + ## TODO: uncomment after merging branch 57 + #reversibility_fidelity = final_state.motional_state.compute_state_fidelity(init_state.motional_state.density_matrix) + #print(f" - Reversibility fidelity: {reversibility_fidelity}") + + + if len(times) != len(rho_t): + rho_t = rho_t[:len(times)] + print(f"Number of rho(t) timepoints: {len(rho_t)}") + print(f"Number of timepoints: {len(times)}") + + spin_rhos = rho_t.copy() + motional_rhos = rho_t.copy() + # 1. Trace out motional modes and plot spin-state populations + for mode in modes: + spin_rhos = [rho.trace_out_degree_of_freedom(mode) for rho in spin_rhos] + + new_basis = spin_rhos[0].basis # should be 2-qubit basis + + populations = np.array([rho.compute_basis_state_probabilities() for rho in spin_rhos]) + plt.figure(figsize=(6,4)) + for i,state in enumerate(new_basis.states): + plt.plot(times*1E6, populations[:, i], label= state.name) + plt.ylabel('Population', fontsize = 16) + plt.xlabel('Gate Duration (µs)', fontsize = 20) + plt.legend() + plt.show() + + # 2. Trace out spin DOFs and plot Fock state populations + X_modes = np.zeros((len(times),len(modes)),dtype=complex) + P_modes = np.zeros((len(times),len(modes)), dtype=complex) + X2_modes = np.zeros((len(times),len(modes)), dtype=complex) + P2_modes = np.zeros((len(times),len(modes)), dtype=complex) + for spin in [target_ion]: + motional_rhos = [rho.trace_out_degree_of_freedom(spin) for rho in motional_rhos] + # For each time, compute quadrature expectations w. variances + # TODO: Uncomment after merging Wigner/quadrature helper functions branch + # for i, t in enumerate(times): + # x, p, x2, p2 = motional_rhos[i].compute_quadratures(True) # includes variance + # X_modes[i, :] = np.array(x) + # P_modes[i, :] = np.array(p) + # X2_modes[i, :] = np.array(x2) + # P2_modes[i, :] = np.array(p2) + + for mode in spectator_modes: + motional_rhos = [rho.trace_out_degree_of_freedom(mode) for rho in motional_rhos] + + new_basis = motional_rhos[0].basis # should be motional basis + + fock_populations = np.array([rho.compute_basis_state_probabilities() for rho in motional_rhos]) + # print('\nFock populations at end of protocol: \n') + # for state, population in zip(new_basis.states, fock_populations[-1]): + # print(f"\nState: n = {state.name}") + # print(f"Population: {population}") + + fock_states = np.arange(fock_dimension) + squeezed_GS_populations = squeezed_vacuum_populations(fock_dimension, r) + + bar_width = 0.35 + plt.figure(figsize=(6,6)) + plt.bar(fock_states - bar_width/2, fock_populations[-1], bar_width, color = 'b', label = 'Simulation') + if not do_reverse_process: + plt.bar(fock_states + bar_width/2, squeezed_GS_populations, bar_width, color = 'r', label = 'Reference: $r_{target} = ' + str(np.round(r,2)) + '$') + else: + plt.bar(fock_states + bar_width/2, IC_fock_pops[0], bar_width, color = 'r', label = 'Initial distribution') + plt.title('Target mode Fock distribution', fontsize = 14) + plt.ylabel('Population', fontsize = 16) + plt.xlabel('Fock state $n$', fontsize = 20) + plt.legend() + plt.show() + + sys.exit(0) + + + + ############## -------------------------------------------------------------- To include after merging branch 57 -------------------------- ################################# + # Compute wigner distributions: + limit = 4. # phase space limit of x and p + x_grid = np.linspace(-limit, limit, 200) + p_grid = np.linspace(-limit, limit, 200) + + last_index = len(times)-1 + # Compute Wigner distributions over a few selected indices + sampling_indices = [0, last_index//3, last_index//2, int(2*last_index//3), int(3*last_index//4), int(7*last_index//8), last_index] + #sampling_indices = [0, last_index] + + sampled_states = [rho_t[i] for i in sampling_indices] + Wigner_distributions = [rho.compute_wigner_distribution(x_grid, p_grid) for rho in sampled_states] + #print(len(Wigner_distributions[0])) + + if num_modes == 2: + mode_labels = ['Target mode', 'Spectator mode 1'] + else: + mode_labels = ['Target mode'] + + for W, t in zip(Wigner_distributions, times[sampling_indices]): + for i, mode_label in enumerate(mode_labels): + plt.figure(figsize=(6, 5)) + plt.imshow(W[i], extent=[-limit, limit, -limit, limit], origin='lower', cmap='magma') + if do_reverse_process: + plt.title(f"$W(x,p)$: " + mode_label + f", t = {t*2.*1E6:.2f} [µs]", fontsize = 12) + else: + plt.title(f"$W(x,p)$, t = {t*1E6:.2f} [µs]", fontsize = 16) + #plt.title(f"$W(x,p)$: " + mode_label + f", t = {t*1E6:.2f} [µs]", fontsize = 12) + plt.xlabel('Position ($x$)', fontsize = 16) + plt.ylabel('Momentum ($p$)', fontsize = 16) + cbar = plt.colorbar() + cbar.ax.set_ylabel(r'$W(x,p)$', rotation = 0, labelpad = 30, va='center') + if excited_qubit1_init_state: + squeezing_str = 'anti' + else: + squeezing_str = 'fwd' + plt.savefig(squeezing_str + f"_mode_{i}_Wigner_t_{t*1E6:.2f}.pdf", dpi=300) + plt.show() + + avg_X_tilt = X_modes[-1,0].real + avg_X2_tilt = X2_modes[-1,0].real + + print(f" for the tilt mode: {avg_X_tilt}") + print(f" for the tilt mode: {avg_X2_tilt}") + + X_variance = avg_X2_tilt - avg_X_tilt**2 + print(f"variance of X for the tilt mode: {X_variance}") + + avg_P_tilt = P_modes[-1,0].real + avg_P2_tilt = P2_modes[-1,0].real + + print(f"

for the tilt mode: {avg_P_tilt}") + print(f" for the tilt mode: {avg_P2_tilt}") + + P_variance = avg_P2_tilt - avg_P_tilt**2 + print(f"variance of P for the tilt mode: {P_variance}") + + print(f"\nVerifying r from variances:") + # If X is squeezed: + if squeezing_phase == 0.: + X_squeezed = False + else: + X_squeezed = True + + if X_squeezed: + r_X = -0.5 * np.log(2. * X_variance) + r_P = 0.5 * np.log(2. * P_variance) + else: + r_X = 0.5 * np.log(2. * X_variance) + r_P = -0.5 * np.log(2. * P_variance) + + print(f"r estimate from X variance: {r_X}") + print(f"r estimate from P variance: {r_P}") + + target_motional_state = ism.State.from_coefficients(motional_rhos[0].basis, list(squeezed_ground_state_coefficients(r, squeezing_phase, fock_dimension))) + squeezed_state_fidelities = np.array([rho_mot.compute_state_fidelity(target_motional_state.density_matrix) for rho_mot in motional_rhos]) + + # Check purity: + print(f"Purity: {np.trace(motional_rhos[-1].density_matrix @ motional_rhos[-1].density_matrix)}") + + print('\n') + print(f'State fidelity with squeezed ground state: {squeezed_state_fidelities[-1]}') + + r_X = np.zeros_like(times) + r_P = np.zeros_like(times) + n_t = np.zeros_like(times) + for i, t in enumerate(times): + r_X[i], r_P[i] = estimate_squeezing_magnitude(motional_rhos[i], motional_rhos[i].basis) + n_t[i] = compute_N(motional_rhos[i], motional_rhos[i].basis) + + plt.figure(figsize=(5,5)) + plt.plot(times*1E6, r_X, marker = 's', linestyle = 'solid', color = 'k', markersize = 3, linewidth = 1.5, label = r'$r_{X}$') + plt.plot(times*1E6, r_P, marker = 'o', linestyle = 'solid', color = 'r', markersize = 3, linewidth = 1.5, label = r'$r_{P}$') + plt.axhline(y = r, color = 'k', linestyle='dashed', linewidth = 1.5, label = r'$r_{target} = ' + str(np.round(r, 4)) + '$') + plt.xlabel('Time [µs] ', fontsize = 20) + plt.ylabel(r'$r$', fontsize = 24, rotation = 0., labelpad=15) + plt.legend() + plt.show() + + plt.figure(figsize=(5,5)) + plt.plot(times*1E6, squeezed_state_fidelities, marker = 's', linestyle = 'solid', color = 'k', markersize = 3, linewidth = 1.5, label = r'State fidelity') + plt.xlabel('Time [µs] ', fontsize = 20) + plt.ylabel(r'$\mathcal{F}$', fontsize = 24, rotation = 0., labelpad=15) + plt.legend() + plt.show() + + plt.figure(figsize=(5,5)) + plt.plot(times*1E6, n_t, marker = 's', linestyle = 'solid', color = 'k', markersize = 3, linewidth = 1.5, label = r'Simulation: $N(t)$') + plt.axhline(y = np.sinh(r)**2, color = 'k', linestyle='dashed', linewidth = 1.5, label = r'$\langle N \rangle = \sinh(r)^2 $') + plt.xlabel('Time [µs] ', fontsize = 20) + plt.ylabel(r'$N$', fontsize = 24, rotation = 0., labelpad=15) + plt.legend() + plt.show() + + + fock_states = np.arange(fock_dimension) + squeezed_GS_populations = squeezed_vacuum_populations(fock_dimension, r) + + bar_width = 0.35 + plt.figure(figsize=(6,6)) + plt.bar(fock_states - bar_width/2, fock_populations[-1], bar_width, color = 'b', label = 'Simulation') + if not do_reverse_process: + plt.bar(fock_states + bar_width/2, squeezed_GS_populations, bar_width, color = 'r', label = 'Reference: $r_{target} = ' + str(np.round(r,2)) + '$') + else: + plt.bar(fock_states + bar_width/2, IC_fock_pops[0], bar_width, color = 'r', label = 'Initial distribution') + plt.ylabel('Population', fontsize = 16) + plt.xlabel('Fock state $n$', fontsize = 20) + plt.legend() + plt.show() + + if modulate_amplitude: + plt.figure(figsize=(5,5)) + plt.plot(times*1E6, mod_function(times), marker = 's', linestyle = 'solid', color = 'k', markersize = 2, linewidth = 1.5) + plt.xlabel('Time [µs] ', fontsize = 20) + plt.ylabel(r'$\Omega$', fontsize = 24, rotation = 0., labelpad=15) + plt.title('Amplitude Modulation')#: $\sigma = ' + str(variance)) + #plt.legend() + plt.show() + +if __name__ == '__main__': + main() diff --git a/src/ionsim/__init__.py b/src/ionsim/__init__.py index 94caf0e..b31db74 100644 --- a/src/ionsim/__init__.py +++ b/src/ionsim/__init__.py @@ -30,3 +30,4 @@ def patched_kron(a, b, *args, **kwargs): from .zeeman_solver import ZeemanHyperfineSolver from .composite_operator import CompositeOperator from .dissipator import Dissipator, DissipatorSpontaneousEmission, Lindbladian +from .trapped_ion_mode_analysis import TrappedIonCrystal, LinearIonChain diff --git a/src/ionsim/degree_of_freedom.py b/src/ionsim/degree_of_freedom.py index beaec95..31a7c7a 100644 --- a/src/ionsim/degree_of_freedom.py +++ b/src/ionsim/degree_of_freedom.py @@ -28,12 +28,16 @@ class DegreeOfFreedom(ABC): """A degree of freedom in a basis of states.""" energy_levels: Sequence[EnergyLevel] - name: str | None = None # TODO: will we use these names? + #name: str | None = None # TODO: will we use these names? @dataclass(frozen=True, eq=False) class AtomicStructure(DegreeOfFreedom): """An atomic structure object, containing atomic internal energy levels corresponding to angular momentum eigenstates.""" energy_levels: list[AtomicInternalEnergyLevel] + # Store mass and atomic number + atomic_mass: float + atomic_charge: int + name: str | None = None # TODO: will we use these names? @classmethod def from_species(cls, species: str, term_symbols: list[str] | None = None, level_names: list[str] | None = None, @@ -42,8 +46,9 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level config_data = cls.get_config_data(species) nuclear_spin = config_data['nuclear_spin'] levels_data = config_data['levels'] - mass = config_data['mass'] # Daltons - z = config_data['Z'] # Atomic number, number of protons + atomic_mass = config_data['mass'] # Daltons + atomic_charge = config_data['charge'] # Daltons + atomic_number = config_data['Z'] # Atomic number, number of protons magnetic_moment = config_data['magnetic_moment'] # units of \mu_{N} structure = 'fine' if nuclear_spin == 0 else 'hyperfine' @@ -79,11 +84,11 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level gj = 2. * (gj1 - 1.) * (k*(k+1) + j1*(j1+1) - l2*(l2 + 1))/((2*j + 1)*(2*k + 1)) gj += (3*j*(j+1) - k*(k+1) + s2*(s2+1))/(2.*j*(j+1)) fine_data['gj'] = gj - Zeeman_solver = ZeemanHyperfineSolver(nuclear_spin, j, None, s2, fine_data['hyperfine_A']*2.*np.pi, mass, magnetic_moment, z, gj = gj) + Zeeman_solver = ZeemanHyperfineSolver(nuclear_spin, j, None, s2, fine_data['hyperfine_A']*2.*np.pi, atomic_mass, magnetic_moment, atomic_number, gj = gj) else: s = fine_data['s'] l = fine_data['l'] - Zeeman_solver = ZeemanHyperfineSolver(nuclear_spin, j, l, s, fine_data['hyperfine_A']*2.*np.pi, mass, magnetic_moment, z) + Zeeman_solver = ZeemanHyperfineSolver(nuclear_spin, j, l, s, fine_data['hyperfine_A']*2.*np.pi, atomic_mass, magnetic_moment, atomic_number) zeeman_energy_shifts, zeeman_eigenvecs = Zeeman_solver.solve_at_field(magnetic_field) zeeman_energy_shifts *= np.pi*2. # convert to rad/s @@ -111,7 +116,7 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level level = HyperfineLevel(**fine_data, i=nuclear_spin, f=f, mf=mf, external_energy_shift = zeeman_shift_energy) if level_names is None or level.name in level_names: levels.append(level) - return cls(levels, name) + return cls(levels, atomic_mass, atomic_charge, name) @classmethod def get_level_factory(cls, coupling_scheme: str): @@ -194,6 +199,7 @@ def check_uniqueness_of_term_symbol(term_symbol: str, levels_data: list[dict]): class MotionalMode(DegreeOfFreedom): """An normal mode of motion for a linear chain of ions.""" energy_levels: list[CollectiveMotionalEnergyLevel] + name: str | None = None # TODO: will we use these names? @classmethod def from_frequency(cls, frequency: float, fock_dimension: int, name: str | None = None): diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py new file mode 100644 index 0000000..161d6f5 --- /dev/null +++ b/src/ionsim/trapped_ion_mode_analysis.py @@ -0,0 +1,1091 @@ +import numpy as np +from numpy.typing import NDArray +import scipy.constants as const +import warnings +import scipy.optimize as opt +from ionsim.custom_types import Matrix, Vector +from ionsim.degree_of_freedom import AtomicStructure, MotionalMode +from ionsim.basis import StandardBasis + +## References: +# https://arxiv.org/abs/2007.12725 +# https://arxiv.org/abs/quant-ph/9702053 + +def characteristic_length(q: float, mass: float, omega: float) -> float: + """ Computes characteristic length in trapped ion system """ + k_e = 1 / (4 * np.pi * const.epsilon_0) # Coulomb constant + l0 = ((k_e * q ** 2) / (.5 * mass * omega ** 2)) ** (1 / 3) + return l0 + + +class TrappedIonCrystal: + def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector | float, atomic_charges: Vector | int, + mode_organization: str = 'frequency_only', reindexing_strategy: str = 'distance'): + """ Class for determining properties of trapped-ion phonon modes, requiring the following system parameters: + + - num_ions: the number of ions + - omega_x: harmonic trap frequency in the x direction in units of rad/s. + - omega_y: harmonic trap frequency in the y direction in units of rad/s. + - omega_z: harmonic trap frequency in the z direction. This is the "axial" direction used for 1D chains. + - atomic masses: array of atomic masses or a single number => same mass for each ion. Units are amu (atomic mass units) + - atomic charges: array of (q) atomic charges (# of protons of an element) or a single number => all ions are the same. + - mode_organization: strategy for organizing modes ('frequency_only' or 'branch_sorted') + - reindexing_strategy: strategy for reindexing ions ('distance' or 'z_axis') + + """ + self.num_ions = num_ions + + # Store configuration strategies + self.mode_organization = mode_organization + self.reindexing_strategy = reindexing_strategy + + # Validate configuration + if mode_organization not in ['frequency_only', 'branch_sorted']: + raise ValueError(f"Invalid mode_organization: {mode_organization}. Must be 'frequency_only' or 'branch_sorted'") + if reindexing_strategy not in ['distance', 'z_axis']: + raise ValueError(f"Invalid reindexing_strategy: {reindexing_strategy}. Must be 'distance' or 'z_axis'") + + self.atomic_masses = self.convert_to_array(atomic_masses) * const.u # kg from amu + self.atomic_charges = self.convert_to_array(atomic_charges) + + # Safety checks: + assert len(self.atomic_masses) == num_ions + assert len(self.atomic_charges) == num_ions + + # Charge of all the protons -> total nuclear charge: + self.nuclear_charges = np.ones_like(self.atomic_masses) * const.e # Z * e, const.e ==> elementary charge in Coulombs + #self.nuclear_charges = self.atomic_charges * const.e # Z * e, const.e ==> elementary charge in Coulombs + + # Trapping frequencies for each ion + self.omega_x = self.calculate_species_trap_frequencies(omega_x) + self.omega_y = self.calculate_species_trap_frequencies(omega_y) + self.omega_z = self.calculate_species_trap_frequencies(omega_z) + + # Check that the trap is stable: + if not self.trap_is_stable(): + raise IonSimError(f"Trap is unstable from negative trap frequencies.") + + def calculate_species_trap_frequencies(self, omega: float) -> Vector: + """ Computes relative trap frequencies for each species: + + w_i = sqrt(q_{i} m_0 / m_{i} q_0) omega + + from omega_(secular, i) = sqrt(q_i V_0 / (2 m_i d^2) ) for the i'th ion. + + - An atom will experience a potentially different pseudopotential based on its mass and charge. + - This is a valid approach when micromotion is small compared to secular motion. + """ + # assume that the trapping frequency given corresponds to the first ion species + q0 = self.nuclear_charges[0] # charge of first ion + m0 = self.atomic_masses[0] # charge of first ion + species_trap_frequencies = np.sqrt(self.nuclear_charges * m0 /(q0 * self.atomic_masses)) * omega + return species_trap_frequencies + + def trap_is_stable(self): + # Checks that all trap frequencies are positive + return np.all(self.omega_z > 0) and np.all(self.omega_y > 0) and np.all(self.omega_x > 0) + + def convert_to_array(self, x: float | int | list | NDArray) -> NDArray: + """ Converts a scalar to a vector or returns the vector """ + if not hasattr(x, "__len__"): + x = np.ones(self.num_ions) * x + return x + return np.array(x) + + def set_up_dimensionless_parameters(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.): + """ Compute dimensionless parameters, using first ion's properties and axial (z) trapping frequency. + + - characteristic length + - characteristic time + - characteristic velocity + - characteristic energy + + """ + # Check positivity: + if trap_freq_scale <= 0: + raise IonSimError(f"Trap frequency scale should be positive. Received {trap_freq_scale}") + if mass_scale <= 0: + raise IonSimError(f"Mass scale should be positive. Received {mass_scale}") + + # Store the scales so they are retrievable by the user + self.charge_scale = charge_scale + self.mass_scale = mass_scale + self.trap_freq_scale = trap_freq_scale + + # Normalize to the first ion species and axial (z) trap frequency + # ion properties + self.m = self.atomic_masses / self.mass_scale + self.q = self.nuclear_charges / self.charge_scale + + # trap frequencies + self.wz = self.omega_z / self.trap_freq_scale + self.wy = self.omega_y / self.trap_freq_scale + self.wx = self.omega_x / self.trap_freq_scale + self.trap_frequencies_ND = [self.wx, self.wy, self.wz] + + # Compute characteristic length, time, velocity, and energy scales and store in a dictionary + self.characteristic_parameters = {} + self.characteristic_parameters['length'] = characteristic_length(self.charge_scale, self.mass_scale, self.trap_freq_scale) + self.characteristic_parameters['time'] = 1 / self.trap_freq_scale # characteristic time + self.characteristic_parameters['velocity'] = self.characteristic_parameters['length'] * self.trap_freq_scale # characteristic velocity + self.characteristic_parameters['energy'] = 0.5 * self.mass_scale * self.characteristic_parameters['velocity'] ** 2 # characteristic energy + + def get_eigenvector_norm(self, eigenvector: Vector, H: Matrix) -> float: + """ Computes the norm of the eigenvector en w.r.t. the Hamiltonian H. """ + norm = np.sqrt(eigenvector.T.conj() @ H @ eigenvector) + return norm + + + def get_canonical_transformation(self, H, eigenvectors, eigenvalues: Vector | None=None): + """ Given the eigensolve of the dynamical matrix, get the transform matrix to the canonical coordinates. + X = S X', where X' = (Q,P)^T and X = (q,p)^T + """ + # the sign due to the convention of writing the time evolution as exp(-iwt) + sign = -1 + num_coords, num_eigenvalues = np.shape(eigenvectors) + assert num_coords //2 == num_eigenvalues + T = np.zeros((num_coords,num_coords),dtype=complex) + eigenvectors = self.normalize_eigenvectors(eigenvectors, H, eigenvalues) + T = np.sqrt(2)*np.concatenate((np.real(eigenvectors), sign*np.imag(eigenvectors)), axis=1) + return T + + def check_for_zero_modes(self): + assert np.all(self.eigvals > 0), "All eigenvalues must be positive" + + def check_outer_relation(self, H: Matrix): + D = self.build_symplectic_matrix() @ H + Eval, Evec = np.linalg.eig(D) + _, en = self.sort_modes(Eval,Evec) + en = self.normalize_eigenvectors(en,H) + Outers = np.zeros((6*self.num_ions,6*self.num_ions),dtype=complex) + for i in range(6*self.num_ions): + norm = self.get_eigenvector_norm(en[:,i],H) + Outers = Outers + np.outer(en[:,i],en[:,i].conj())/norm + I_right = H @ Outers + I_left = Outers @ H + eye = np.eye(6*self.num_ions,dtype=complex) + try: + assert np.allclose(I_left,eye) + assert np.allclose(I_right,eye) + except AssertionError: + warnings.warn("Outer relation check failed") + + def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool: + M = np.linalg.inv(T) @ S + H_diag = M.T @ E @ M + H_diag_check = np.diag(np.tile(self.eigvals,2)) + try: + assert np.allclose(H_diag,H_diag_check) + except AssertionError: + has_duplicate_eigvals = np.any(np.triu(np.isclose(self.eigvals[:, None], self.eigvals[None, :], atol = 1E-6), k=1)) + warnings.warn("Diagnolization check failed") + print("has duplicate eigenvalues: ", has_duplicate_eigvals) + + def solve_ion_trap_equilibrium(self): + """ Solve for equilibrium positions and analyze normal modes for a linear chain by minizing the Coulomb + + harmonic trap energy functional for a system of ions + + This method: + 1. Sets up dimensionless parameters + 2. Solves for equilibrium positions + 3. Reindexes ions by their z-position (closest to center first) + 4. Computes normal modes with branch sorting + 5. Performs validation checks + + """ + # Convert to dimensionless units using axial trap frequency and first ion's mass and charge + self.set_up_dimensionless_parameters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0]) + self.equilibrium_positions = self.solve_for_equilibrium_positions() + + # Use configurable reindexing strategy + if self.reindexing_strategy == 'z_axis': + self.reindex_ions_by_z() + else: + self.reindex_ions() + + # Compute helper matrices for computing normal mode properties + mass_matrix = self.build_mass_matrix(self.m) + E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix) + T_matrix = self.build_momentum_transform_matrix(mass_matrix) + H_matrix = self.compute_H_matrix(T_matrix, E_matrix) + + self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix) + self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs) + self.check_for_zero_modes() + S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) + + # Perform checks + self.check_outer_relation(H_matrix) + self.check_diagonalization(T_matrix, S_matrix, E_matrix) + + @property + def normal_mode_frequencies(self): + return self.eigvals * self.trap_freq_scale + + + def normalize_eigenvectors(self, eigvecs, H: Matrix, eigenvalues: Vector | None=None): + """ Rescale the eigenvectors ens w.r.t. the Hamiltonian H. """ + eigvecs_rescaled = np.zeros_like(eigvecs,dtype=complex) + num_coords, num_eigenvalues = np.shape(eigvecs) + if eigenvalues is None: + eigenvalues = np.ones(num_eigenvalues) + for i in range(num_eigenvalues): + en = eigvecs[:,i].reshape(num_coords,1) + norm = self.get_eigenvector_norm(en,H) + eigvecs_rescaled[:,i] = en[:,0]/norm + eigvecs_rescaled[:,i] *= np.sqrt(eigenvalues[i]) + return eigvecs_rescaled + + + def get_eigenvectors_xv_coords(self,T,ens): + ens_vel = np.zeros_like(ens,dtype=complex) + ens_vel[:,:] = np.linalg.inv(T) @ ens + return ens_vel + + + def calculate_normal_modes(self, H_matrix: Matrix) -> (Vector, list[Vector]): + """ Compute normal mode frequencies from eigenvalues of the matrix D. Returns eigenvalues and eigenvectors. + + Solves det( D + i*omega*Identity) == 0 for omega. + + The matrix "D" is defind as D = JH, where J = (0 , I ; I, 0) (eq. 1.67 of thesis) + + """ + J = self.build_symplectic_matrix() + D_matrix = J @ H_matrix + # u_n(t) = exp(-i w_n t) u_n(0), minus by convention + eigvals, eigvecs = np.linalg.eig(-D_matrix) + eigvals, eigvecs = self.organize_modes(eigvals, eigvecs) + eigvecs = self.normalize_eigenvectors(eigvecs, H_matrix) + + return eigvals, eigvecs + + + def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) -> tuple[Vector, Vector, Vector]: + """ Returns arrays corresponding to an ion's x, y, and z coordinates, respectively, from a flattened vector u: + x = u[0:N], y = u[N:2N], z[2N:] + + e.g. x[1], y[1], z[1] is the x, y, z coordinate values for ion 2. + """ + x = flattened_coordinate_vector[0:self.num_ions] + y = flattened_coordinate_vector[self.num_ions:2*self.num_ions] + z = flattened_coordinate_vector[2*self.num_ions:] + return x, y, z + + + def reindex_ions(self): + """Re-indexes the ions to order based on the distance from the center of the trap, where smallest index is closest to the center.""" + x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions) + r = np.sqrt(x**2 + y**2 + z**2) + idx = np.argsort(r) + + reindexed_positions = np.hstack((x[idx], y[idx], z[idx])) + self.equilibrium_positions = reindexed_positions + self.m = self.m[idx] + self.q = self.q[idx] + self.atomic_masses = self.atomic_masses[idx] + self.nuclear_charges = self.nuclear_charges[idx] + self.atomic_charges = self.atomic_charges[idx] + + def reindex_ions_by_z(self): + """Re-indexes ions based on their position along the z-axis, with the ion closest to the center having index 0.""" + x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions) + idx = np.argsort(z) + + self.equilibrium_positions = np.hstack((x[idx], y[idx], z[idx])) + + # Reindex all other arrays accordingly + self.m = self.m[idx] + self.q = self.q[idx] + self.atomic_masses = self.atomic_masses[idx] + self.nuclear_charges = self.nuclear_charges[idx] + self.atomic_charges = self.atomic_charges[idx] + + # trapping frequencies + self.omega_x = self.omega_x[idx] + self.omega_y = self.omega_y[idx] + self.omega_z = self.omega_z[idx] + # uses first ion's properties as the characteristic scalings + self.set_up_dimensionless_parameters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0]) + + + def solve_for_equilibrium_positions(self, positions_guess: Vector | None=None): + """ Solves for equilibrium position vector: u, which represents a flattened spatial grid. """ + # Set the initial guess for the solver + if positions_guess == None: + # Initialize a random guess + u0 = np.zeros(3*self.num_ions) + u0[:] = (np.random.rand(3*self.num_ions) * 2 - 1) * self.num_ions + else: + u0 = positions_guess + + # Solve for the equilibrium positions by minimizing the potential energy (trap + Coulomb) + bfgs_tolerance = 1e-34 + solver_output = opt.minimize(self.potential_energy, u0, method='BFGS', jac=self.force, + options={'gtol': bfgs_tolerance, 'disp': False}) + equilibrium_positions = solver_output.x + + # Get potential energy at equilibrium positions + self.p0 = self.potential_energy(equilibrium_positions) + return equilibrium_positions + + + def potential_trap(self, positions: Vector) -> Vector: + """ Computes trapping potential, takes in a flattened (1D) array of all position coordinates """ + r = self.ion_coordinates_from_flattened(positions) + V_trap = 0. + for coord, omega in zip(r, self.trap_frequencies_ND): + V_trap += 0.5 * np.sum((self.m * omega**2) * coord**2) + return V_trap + + def potential_coulomb(self, positions: Vector) -> Vector: + """ Computes Coulomb potential, takes in a flattened (1D) array of all position coordinates """ + r = self.ion_coordinates_from_flattened(positions) + dr = [] + for coord in r: + dr.append(coord[:, np.newaxis] - coord) + + rsep = np.sqrt(dr[0]**2 + dr[1]**2 + dr[2]**2).astype(np.float64) + qq = (self.q * self.q[:, np.newaxis]).astype(np.float64) + with np.errstate(divide='ignore'): + V_Coulomb = np.sum( np.where(rsep != 0., qq / rsep, 0) ) / 2 # divide by 2 to avoid double counting + V_Coulomb *= .5 + return V_Coulomb + + def potential_energy(self, positions: Vector) -> Vector: + """ Computes the potential energy of each ion in its coordinate basis """ + return self.potential_trap(positions) + self.potential_coulomb(positions) + + def force_trap(self, positions: Vector): + r = self.ion_coordinates_from_flattened(positions) + + F_r = [] + for coord, omega in zip(r, self.trap_frequencies_ND): + F_r.append(self.m * omega**2 * coord) + + force_trap = np.hstack(F_r) + return force_trap + + def force_coulomb(self, positions: Vector): + """ Computes the Coulomb force by derivative w.r.t. ion coordinates """ + r = self.ion_coordinates_from_flattened(positions) + dr = [] + + for coord in r: + dr.append(coord[:, np.newaxis] - coord) + rsep = np.sqrt(dr[0]**2 + dr[1]**2 + dr[2]**2).astype(np.float64) + qq = (self.q * self.q[:, np.newaxis]).astype(np.float64) + + with np.errstate(divide='ignore', invalid='ignore'): + rsep3 = np.where(rsep != 0., rsep ** (-3), 0) + + F_r = [] # Force in each direction + for dx in dr: + f_r = dx * rsep3 * qq + F_r.append(-np.sum(np.array(f_r), axis=1)) + + force_coulomb = np.hstack(F_r) + force_coulomb *= 0.5 + return force_coulomb + + def force(self, positions: Vector) -> Vector: + """ Computes the force vector, e.g. derivative of the potential energy w.r.t. ion position: dU/dx, dU/dy, dU/dz. + + - returns a one-dimensional length-3N array characterizing forces w.r.t. ion coordinates for N ions. + + """ + Force = self.force_coulomb(positions) + self.force_trap(positions) + return Force + + def hessian_trap(self, positions: Vector) -> Matrix: + """ Computes the Hessian of the trap """ + H_rr = [] + for i, omega in enumerate(self.trap_frequencies_ND): + H_rr.append(np.diag(self.m * omega**2 * np.ones(self.num_ions))) + zeros = np.zeros((self.num_ions, self.num_ions)) + H = np.block([[H_rr[0], zeros, zeros], [zeros, H_rr[1], zeros], [zeros, zeros, H_rr[2]]]) + return H + + def hessian_coulomb(self, positions: Vector) -> Matrix: + """ Computes the Hessian of the Coulomb interaction """ + r = self.ion_coordinates_from_flattened(positions) + dr = [] + + for coord in r: + dr.append(coord[:, np.newaxis] - coord) + rsep = np.sqrt(dr[0]**2 + dr[1]**2 + dr[2]**2).astype(np.float64) + qq = (self.q * self.q[:, np.newaxis]).astype(np.float64) + + with np.errstate(divide='ignore'): + rsep5 = np.where(rsep != 0., rsep ** (-5), 0) + + dr_sq = np.array(dr)**2 + rsep2 = rsep ** 2 + + # X derivatives, Y derivatives for alpha != beta + H_rr = (rsep2 - 3 * dr_sq) * rsep5 # form: [Hxx, Hyy, Hzz] + + # Above, for alpha == beta + # Compute diagonals: + for i in range(3): + H_rr[i][np.diag_indices(self.num_ions)] = -np.sum(H_rr[i], axis=0) + dx = dr[0] + dy = dr[1] + dz = dr[2] + + # Off-diagonal elements: + Hxy = -3 * dx * dy * rsep5 + Hxy[np.diag_indices(self.num_ions)] = 3 * np.sum(dx * dy * rsep5, axis=0) + Hxz = -3 * dx * dz * rsep5 + Hxz[np.diag_indices(self.num_ions)] = 3 * np.sum(dx * dz * rsep5, axis=0) + Hyz = -3 * dy * dz * rsep5 + Hyz[np.diag_indices(self.num_ions)] = 3 * np.sum(dy * dz * rsep5, axis=0) + + H_rr *= qq + Hxy *= qq + Hxz *= qq + Hyz *= qq + + H_coulomb = np.block([[H_rr[0], Hxy, Hxz], [Hxy, H_rr[1], Hyz], [Hxz, Hyz, H_rr[2]]]) + H_coulomb /= 2 + return H_coulomb + + + def hessian(self, positions: Vector) -> Matrix: + """ Computes total Hessian (trap + Coulomb interaction) """ + H = self.hessian_coulomb(positions) + self.hessian_trap(positions) + return H + + def build_mass_matrix(self, masses: Vector) -> Matrix: + """ Builds a diagonal mass matrix representing the ions """ + return np.diag(np.tile(masses, 3)) + + def build_E_matrix(self, positions: Vector, mass_matrix: Matrix) -> Matrix: + """ 6N x 6N energy matrix for N ions """ + PE_matrix = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128) + KE_matrix = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128) + E_matrix = np.zeros((6*self.num_ions, 6*self.num_ions), dtype=np.complex128) + + PE_matrix = self.hessian(positions) + KE_matrix = mass_matrix + zeros = np.zeros((3*self.num_ions, 3*self.num_ions)) + E_matrix = np.block([[PE_matrix, zeros], [zeros, KE_matrix]]) + return E_matrix + + def compute_H_matrix(self, T_matrix, E_matrix): + """ Computes 6N x 6N Hermitian Hamiltonian matrix in (q, p) canonical coordinates """ + T_matrix_inv = np.linalg.inv(T_matrix) + H_matrix = T_matrix_inv.T @ E_matrix @ T_matrix_inv + return H_matrix + + def build_momentum_transform_matrix(self, mass_matrix: Matrix) -> Matrix: + # assuming no magnetic field + eye = np.eye(3*self.num_ions) + zeros = np.zeros((3*self.num_ions, 3*self.num_ions)) + T = np.block([[eye, zeros], [zeros, mass_matrix]]) + return T + + def build_symplectic_matrix(self) -> Matrix: + """ Computes J a 6N x 6N symplectic matrix defined as J = (0 , I ; I, 0) (eq. 1.67 of thesis) """ + zeros = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128) + I = np.eye(3*self.num_ions, dtype=np.complex128) + J = np.block([[zeros, I], [-I, zeros]]) + return J + + def sort_modes(self, eigvals, eigvecs): + eigvals = np.imag(eigvals) + sort_dex = np.argsort(eigvals) + eigvals = eigvals[sort_dex] + eigvecs = eigvecs[:,sort_dex] + return eigvals, eigvecs + + def split_modes(self, eigvals, eigvecs): + half = len(eigvals) // 2 + eigvals = eigvals[half:] + eigvecs = eigvecs[:,half:] + return eigvals, eigvecs + + + def organize_modes(self, eigvals, eigvecs): + eigvals, eigvecs = self.sort_modes(eigvals, eigvecs) + eigvals, eigvecs = self.split_modes(eigvals, eigvecs) + + if self.mode_organization == 'branch_sorted': + eigvals, eigvecs = self.sort_by_branch(eigvals, eigvecs) + + return eigvals, eigvecs + + def _xyz_classify_modes(self, evecs): + """Classify modes by their dominant spatial direction (x, y, z).""" + N_coords, N_modes = np.shape(evecs) + N_ions = N_coords // 6 + classifier = np.zeros(N_modes, dtype=int) + for mode_index in range(N_modes): + x_score = np.sum(np.abs(evecs[0:N_ions, mode_index])) + np.sum(np.abs(evecs[3*N_ions:4*N_ions, mode_index])) + y_score = np.sum(np.abs(evecs[N_ions:2*N_ions, mode_index])) + np.sum(np.abs(evecs[4*N_ions:5*N_ions, mode_index])) + z_score = np.sum(np.abs(evecs[2*N_ions:3*N_ions, mode_index])) + np.sum(np.abs(evecs[5*N_ions:6*N_ions, mode_index])) + scores = [x_score, y_score, z_score] + max_index = np.argmax(scores) + classifier[mode_index] = max_index + for direction in range(3): + count = np.sum(classifier == direction) + assert count == N_modes // 3, f"Classification error: direction {direction} has {count} modes, expected {N_modes // 3}" + return classifier + + def sort_by_branch(self, evals, evecs): + """ Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)). + Only makes sense for linear chain configurations where Hessian is block diagonal. """ + classifier = self._xyz_classify_modes(evecs) + # within each branch, sort by frequency + N_ions = len(evals) // 3 + sorted_by_branch_evals = np.zeros_like(evals) + sorted_by_branch_evecs = np.zeros_like(evecs) + for direction in range(3): + direction_indices = np.where(classifier == direction)[0] + # they are already sorted by frequency, so take them in order + sorted_by_branch_evals[direction*N_ions:(direction+1)*N_ions] = evals[direction_indices] + sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices] + return sorted_by_branch_evals, sorted_by_branch_evecs + + + def calculate_mode_participation_factors(self) -> Matrix: + """ Computes ion-mode participation factors, related to the Lamb-Dicke parameters. + + Mode participation factors (eta) take the following matrix form: + + shape: (dimension, ion, mode), for N ions this is (3, N, 3N). This is the most general case. + e.g. eta[1, 2, 3] is eta in the "y" direction, ion 1, and the 2nd mode. + + For N ions, this form organizes the 3N modes into N modes per direction "d", where d = x, y, z. + + """ + eigvecs = self.eigvecs + num_coords, num_modes = np.shape(eigvecs) + num_ions = num_modes // 3 + mode_participation_factors = np.zeros((3, num_ions, num_modes), dtype = np.complex128) + for mode_index in range(num_modes): + for pos_coord in range(num_coords//2): + direction_index = pos_coord // num_ions + ion_index = pos_coord % num_ions + prefactor = np.sqrt(2 * self.m[ion_index] * self.eigvals[mode_index] ) + zpm_dimensionful = np.sqrt(const.hbar / (2 * self.mass_scale * self.trap_freq_scale)) + mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * prefactor * eigvecs[pos_coord, mode_index] + return mode_participation_factors + + + # Derived properties + def compute_reference_single_ion_lamb_dicke_factors(self, wavenumber: float) -> (float, float, float): + """ Computes analytical Lamb-Dicke parameter by eta = k * sqrt(hbar / m omega) for an ion in a light-field with wavevector |k| + - wavenumber: wavevector magnitude |k| in units of 1 / m + """ + # Convention to use first value of trap frequency arrays (representing one of the ions) + trap_frequencies = np.array([self.omega_x[0], self.omega_y[0], self.omega_z[0]]) + mass = self.atomic_masses[0] + eta_x, eta_y, eta_z = wavenumber * np.sqrt(const.hbar / (2 * mass * trap_frequencies)) + return eta_x, eta_y, eta_z + + def return_equilibrium_positions(self, dimensionless: bool) -> Vector: + """Return equilibrium positions in either dimensionless or SI units (inverse meters).""" + if dimensionless: + return self.equilibrium_positions + else: + return self.equilibrium_positions * self.characteristic_parameters['length'] + + @classmethod + def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y: float, omega_z: float): + """Build the mode analysis class from a species name, corresponding to a system of N ions under harmonic trapping.""" + # Import necessary data for the species + species_data = AtomicStructure.get_config_data(species_name) + atomic_mass = species_data['mass'] + atomic_charge = species_data['charge'] + # Construct the class + return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_charge) + + @classmethod + def from_atomic_structure_basis(cls, atomic_structure_basis: StandardBasis, omega_x: float, omega_y: float, omega_z: float): + """ Build the mode analysis class from a basis of AtomicStructure degrees of freedom under harmonic trapping. """ + # Extract number of ions and the mass and atomic number from the DOF in the basis + atomic_struct_DOFs = atomic_structure_basis.degrees_of_freedom + num_ions = len(atomic_struct_DOFs) + atomic_masses = [] + atomic_charges = [] + + for DOF in atomic_struct_DOFs: + if not isinstance(DOF, AtomicStructure): + raise IonSimError("Atomic structure basis should only contain AtomicStructure objects. No motional modes should be included.") + atomic_masses.append(DOF.atomic_mass) + atomic_charges.append(DOF.atomic_charge) + return cls(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_charges) + + def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]: + """ Builds and returns a list of IonSim Motional Degrees of Freedom. + + - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes + """ + modes = [] + fock_dimensions = self.convert_to_array(fock_dimensions).astype(int) + for idx, fock_dim in zip(mode_indices, fock_dimensions): + mode_index = idx # or some function of this index + modes.append(MotionalMode.from_frequency(self.eigvals[mode_index], fock_dim)) + + return modes + + def _ensure_mode_analysis_run(self): + """Ensure that mode analysis has ran. If not, run it automatically.""" + if not hasattr(self, 'eigvals') or not hasattr(self, 'eigvecs'): + self.solve_ion_trap_equilibrium() + + def get_mode_participation_factors_for_mode(self, mode_index: int) -> Matrix: + """ Returns mode participation factors for a specific mode by its index. + + Outpout: Array of shape (3, num_ions) containing participation factors + where result[dimension, ion] gives the participation factor + for that dimension and ion in the specified mode. + """ + self._ensure_mode_analysis_run() + + # Get full participation factor matrix + full_participation = self.calculate_mode_participation_factors() + + # Extract participation factors for the specific mode + mode_participation = full_participation[:, :, mode_index] + + return mode_participation + + def get_lamb_dicke_parameters_for_mode(self, mode_index: int, wavevector: Vector) -> Matrix: + """ Return Lamb-Dicke parameters for a specific mode by its index. + + - mode_index: Index of the mode of interest + - wavevector: Wavevector k as a 3-element array (kx, ky, kz) or scalar wavenumber + + Returns an array of shape (num_ions,) containing Lamb-Dicke parameters + where result[ion] gives the LD parameter for that ion in the specified mode. + + Note: Lamb Dicke parameters have an overall global phase freedom. + """ + self._ensure_mode_analysis_run() + + # Get participation factors for this mode + mode_participation = self.get_mode_participation_factors_for_mode(mode_index) + + # Calculate Lamb-Dicke parameters + if wavevector.shape != (3,): + raise ValueError("Wavevector must be a 3-element array (kx, ky, kz)") + + # Compute dot product: (3, num_ions) · (3,) -> (num_ions,) + lamb_dicke_params = np.zeros(self.num_ions, dtype=complex) + for i in range(3): + lamb_dicke_params += wavevector[i] * mode_participation[i, :] + + return lamb_dicke_params + + def get_mode_properties(self, mode_index: int) -> dict: + """Get comprehensive properties for a specific mode. Returns a dictionary containing frequency, + participation factors, and other properties for the specified mode. + + - mode_index: Index of the mode of interest + + Returns: + Dictionary containing: + - 'frequency': Mode frequency in rad/s + - 'participation_factors': Array of shape (3, num_ions) + - 'frequency_MHz': Mode frequency in MHz (convenience) + """ + self._ensure_mode_analysis_run() + + # Get basic mode properties + frequency = self.eigvals[mode_index] * self.trap_freq_scale + participation_factors = self.get_mode_participation_factors_for_mode(mode_index) + + return { + 'frequency': frequency, # rad/s + 'frequency_MHz': frequency / (2 * np.pi * 1e6), # MHz + 'participation_factors': participation_factors, + 'mode_index': mode_index + } + + +class LinearIonChain(TrappedIonCrystal): + """ Subclass for setting up and analyzing linear ion chains. + + In a linear ion chain, ions are typically aligned along the z-axis (axial direction), + with minimal displacement in the x and y directions (radial directions). + + This class accomplishes: + 1. Enforce linear chain configuration + 2. Provide specialized methods for linear chain analysis + 3. Offer convenience methods for common linear chain scenarios + """ + + def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, + atomic_masses: np.ndarray | float, atomic_charges: np.ndarray | int): + """ Initialize a linear ion chain analysis, defaulting to branch sorted modes. """ + mode_organization = 'branch_sorted' + reindexing_strategy = 'z_axis' + super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_charges, mode_organization, reindexing_strategy) + + def solve_ion_trap_equilibrium(self): + """ Solve for equilibrium positions and analyze normal modes for a linear chain. It reindexes ions by their axial position. """ + # Set up dimensionless parameters using first ion's properties and axial trap frequency + self.set_up_dimensionless_parameters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0]) + + + # Solve for equilibrium positions + self.equilibrium_positions = self.solve_for_equilibrium_positions() + + # Reindex ions by their z-position (important for linear chains) + self.reindex_ions_by_z() + + # Build necessary matrices + mass_matrix = self.build_mass_matrix(self.m) + E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix) + T_matrix = self.build_momentum_transform_matrix(mass_matrix) + H_matrix = self.compute_H_matrix(T_matrix, E_matrix) + + # Calculate normal modes with branch sorting (x, y, z branches) + self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix) + self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs) + + # Validate results + self.check_for_zero_modes() + self.check_outer_relation(H_matrix) + + # Compute canonical transformation matrix + S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) + self.check_diagonalization(T_matrix, S_matrix, E_matrix) + + def get_axial_modes(self) -> tuple: + """ Returns eigenvalues, eigenvectors for the axial normal modes """ + # Axial modes are the last third of the modes (after branch sorting) + n_modes_per_branch = self.num_ions + axial_start = 2 * n_modes_per_branch + axial_end = 3 * n_modes_per_branch + + axial_eigvals = self.eigvals[axial_start:axial_end] + axial_eigvecs = self.eigvecs[:, axial_start:axial_end] + + return axial_eigvals, axial_eigvecs + + def get_radial_modes(self, direction='x') -> tuple: + """ Returns eigenvalues, eigenvectors for the radial normal modes in the specified direction. + - direction : str; 'x' or 'y' for radial direction """ + if direction not in ['x', 'y']: + raise ValueError("Direction must be 'x' or 'y'") + + n_modes_per_branch = self.num_ions + if direction == 'x': + start_idx = 0 + else: # direction == 'y' + start_idx = n_modes_per_branch + + end_idx = start_idx + n_modes_per_branch + + radial_eigvals = self.eigvals[start_idx:end_idx] + radial_eigvecs = self.eigvecs[:, start_idx:end_idx] + + return radial_eigvals, radial_eigvecs + + def get_axial_ion_positions(self): + """ Returns an array of ion axial positions along the z-axis. """ + # Get z positions (dimensionful) + x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions) + z_dimensionful = z * self.characteristic_parameters['length'] + sorted_z = np.sort(z_dimensionful) + return sorted_z + + + def get_ion_spacing(self): + """ Returns an array of ion spacings between consecutive ions along the z-axis. """ + # Get z positions (dimensionful) + sorted_z_coords = self.get_axial_ion_positions() + ion_spacing = np.diff(sorted_z_coords) + + return ion_spacing + + def get_center_of_mass_position(self) -> tuple: + """ Get the center of mass position of the ion chain. """ + # Get positions (dimensionful) + x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions) + x_dim = x * self.characteristic_parameters['length'] + y_dim = y * self.characteristic_parameters['length'] + z_dim = z * self.characteristic_parameters['length'] + + # Calculate center of mass + total_mass = np.sum(self.atomic_masses) + com_x = np.sum(x_dim * self.atomic_masses) / total_mass + com_y = np.sum(y_dim * self.atomic_masses) / total_mass + com_z = np.sum(z_dim * self.atomic_masses) / total_mass + + return com_x, com_y, com_z + + def get_axial_mode_frequencies(self): + """ Returns the frequencies [rad/s] of the axial modes. """ + axial_eigvals, _ = self.get_axial_modes() + return axial_eigvals * self.trap_freq_scale + + def get_radial_mode_frequencies(self, direction='x'): + """ Returns the frequencies of the radial modes in rad/s for the specified direction ('x' or 'y'). """ + radial_eigvals, _ = self.get_radial_modes(direction) + return radial_eigvals * self.trap_freq_scale + + def get_mode_participation_factors_by_branch(self): + """ Returns a dictionary of mode participation factors organized by branch (x, y, z). + + Output dictionary with keys 'x', 'y', 'z' containing mode participation factors + for each branch. Each value is a 2D array of shape (num_ions, num_modes_per_branch) + """ + # Get full mode participation factors + mode_pf = self.calculate_mode_participation_factors() + + # Organize by branch + n_modes_per_branch = self.num_ions + + result = { + 'x': mode_pf[0, :, :n_modes_per_branch], + 'y': mode_pf[1, :, n_modes_per_branch:2*n_modes_per_branch], + 'z': mode_pf[2, :, 2*n_modes_per_branch:3*n_modes_per_branch] + } + + return result + + def get_mode_properties_by_branch(self, branch: str, mode_index_in_branch: int) -> dict: + """ Returns a dictionary containing properties (freq, participations) for a specific mode within a branch (x, y, or z). + + For linear chains with branch-sorted modes, this provides an intuitive way + to access mode properties by branch and mode number within that branch. + + - branch: 'x', 'y', or 'z' for the mode branch + - mode_index_in_branch: Index of the mode within the specified branch (0 to num_ions-1) + """ + self._ensure_mode_analysis_run() + + if branch not in ['x', 'y', 'z']: + raise ValueError("Branch must be 'x', 'y', or 'z'") + + if mode_index_in_branch < 0 or mode_index_in_branch >= self.num_ions: + raise ValueError(f"Mode index in branch must be between 0 and {self.num_ions-1}") + + # Calculate the global mode index + n_modes_per_branch = self.num_ions + if branch == 'x': + global_mode_index = mode_index_in_branch + elif branch == 'y': + global_mode_index = n_modes_per_branch + mode_index_in_branch + else: # branch == 'z' + global_mode_index = 2 * n_modes_per_branch + mode_index_in_branch + + # Get properties using the parent class method + properties = super().get_mode_properties(global_mode_index) + + # Add branch-specific information + properties['branch'] = branch + properties['mode_index_in_branch'] = mode_index_in_branch + + return properties + + def get_lamb_dicke_parameters_by_branch(self, branch: str, mode_index_in_branch: int, + wavevector: Vector | float) -> Vector: + """Get Lamb-Dicke parameters for a specific mode within a branch. + + Parameters: + branch: 'x', 'y', or 'z' for the mode branch + mode_index_in_branch: Index of the mode within the specified branch + wavevector: Wavevector k as a 3-element array or scalar wavenumber + + Returns: + Array of Lamb-Dicke parameters for each ion in the specified mode + """ + self._ensure_mode_analysis_run() + + # Get the global mode index + n_modes_per_branch = self.num_ions + if branch == 'x': + global_mode_index = mode_index_in_branch + elif branch == 'y': + global_mode_index = n_modes_per_branch + mode_index_in_branch + else: # branch == 'z' + global_mode_index = 2 * n_modes_per_branch + mode_index_in_branch + + # Use parent class method to get LD parameters + return super().get_lamb_dicke_parameters_for_mode(global_mode_index, wavevector) + + def calculate_lamb_dicke_parameters_by_branch(self, wavevector: Vector) -> dict: + """ Calculate Lamb-Dicke parameters for all ions and modes, organized by branch. + + The Lamb-Dicke parameter η = k · Δr_0 represents the ratio of the spatial + extent of the ion's zero-point motion to the wavelength of the laser. + + Parameters: + wavevector : k = (kx, ky, kz) as a 3-element array in units of 1/m + + Returns: + Dictionary with keys 'x', 'y', 'z' containing Lamb-Dicke parameters + for each branch. Each value is a 2D array of shape (num_ions, num_modes_per_branch) + representing η_{direction}[ion_index, mode_index] + """ + # Get mode participation factors (which are proportional to zero-point motion) + mode_pf_by_branch = self.get_mode_participation_factors_by_branch() + + if wavevector.shape != (3,): + raise ValueError("Wavevector must be a 3-element array (kx, ky, kz)") + + # Get full mode participation factors + full_mode_pf = self.calculate_mode_participation_factors() + + # Compute Lamb-Dicke parameters as dot product: η = k · Δr_0 + # This gives us shape (num_ions, num_modes) for the total LD parameter + num_modes = 3 * self. num_ions + + # Reshape for dot product: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes) + total_ld_params = np.zeros((self.num_ions, num_modes), dtype=complex) + for i in range(3): + total_ld_params += wavevector[i] * full_mode_pf[i, :, :] + + # Organize by branch for consistency with scalar case + n_modes_per_branch = self.num_ions + lamb_dicke_parameters = { + 'x': total_ld_params[:, :n_modes_per_branch], + 'y': total_ld_params[:, n_modes_per_branch:2*n_modes_per_branch], + 'z': total_ld_params[:, 2*n_modes_per_branch:3*n_modes_per_branch] + } + + return lamb_dicke_parameters + + def build_mode_DOFs_from_branch(self, branch: str, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]: + """ Builds and returns a list of IonSim Motional Degrees of Freedom. + + - mode_indices corresponds to a mode's index within a branch (instead of the entire set of 3N modes for N ions) + - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes + - ex] branch = 'x', mode_indices = [0,1] (tilt and COM), fock_dimensions = 6 + """ + modes = [] + + # Convert list of Fock dimensions to an array if it's not already an array + fock_dimensions = self.convert_to_array(fock_dimensions).astype(int) + for mode_idx, fock_dim in zip(mode_indices, fock_dimensions): + mode_props = self.get_mode_properties_by_branch(branch, mode_idx) + frequency = mode_props['frequency'] # rad/s + modes.append(MotionalMode.from_frequency(frequency, fock_dim)) + + return modes + + def calculate_lamb_dicke_parameters_full(self, wavevector: Vector) -> Matrix: + """ Calculate full Lamb-Dicke parameter matrix from a laser wavevector k = (kx, ky, kz) + + Returns: + - Total LD parameter matrix of shape (num_ions, num_modes) where eta[ion, mode] + gives the total LD parameter k · Δr_0. + + Note: Lamb Dicke parameters have an overall global phase freedom. + """ + # Get full mode participation factors + mode_pf = self.calculate_mode_participation_factors() + + if wavevector.shape != (3,): + raise ValueError(f"Wavevector must be a 3-element array (kx, ky, kz). Received {wavevector}") + + # Compute total Lamb-Dicke parameter as dot product + # Shape: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes) + num_modes = 3 * self.num_ions + total_ld_params = np.zeros((self.num_ions, num_modes), dtype=complex) + + for i in range(3): + total_ld_params += wavevector[i] * mode_pf[i, :, :] + + return total_ld_params + + def get_axial_lamb_dicke_parameters(self, wavevector: Vector) -> Vector: + """ Get Lamb-Dicke parameters for axial (z) modes only. + + Returns: Lamb-Dicke parameters for axial modes, NDArray of shape (num_ions, num_axial_modes) + """ + lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters_by_branch(wavevector) + return lamb_dicke_by_branch['z'] + + def get_radial_lamb_dicke_parameters(self, wavevector: Vector | float, direction: str = 'x') -> Vector: + """ Get Lamb-Dicke parameters for radial modes modes in a specified direction. + + Returns: Lamb-Dicke parameters for radial modes, array of shape (num_ions, num_radial_modes) + + Note: Lamb Dicke parameters have an overall global phase freedom. + """ + if direction not in ['x', 'y']: + raise ValueError("Direction must be 'x' or 'y'") + + lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters_by_branch(wavevector) + return lamb_dicke_by_branch[direction] + + def get_two_central_ion_separation(self) -> float: + """ Calculate the separation between the two central ions in the chain. Returns: Distance between the two central ions in meters """ + # Get z positions (dimensionful) + x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions) + z_dimensionful = z * self.characteristic_parameters['length'] + + # Find the two central ions (closest to the center) + central_ions = np.argsort(np.abs(z_dimensionful))[:2] + + # Calculate separation + dl = np.abs(z_dimensionful[central_ions[0]] - z_dimensionful[central_ions[1]]) + return dl + + def find_axial_frequency_for_desired_central_ion_separation(self, target_separation: float, + bounds: tuple = (0.1e6, 0.5e6)) -> float: + """ Computes the axial frequency (wz) [rad/s] required to achieve a desired separation between central ions. + + Parameters: + - target_separation : float; Desired separation between central ions in meters + - bounds : tuple; (min_freq, max_freq) in Hz for the root finding algorithm + + """ + from scipy.optimize import root_scalar + + def separation_error(wz_Hz): + # Create a temporary analysis object with the current wz + temp_analysis = LinearIonChain( + num_ions=self.num_ions, + omega_x=self.omega_x[0], + omega_y=self.omega_y[0], + omega_z=2*np.pi*wz_Hz, + atomic_masses=self.atomic_masses[0], + atomic_charges=self.atomic_charges[0] + ) + temp_analysis.solve_ion_trap_equilibrium() + current_separation = temp_analysis.get_two_central_ion_separation() + return current_separation - target_separation + + # Convert bounds from Hz to rad/s for the root finding + bounds_rad_s = (2*np.pi*bounds[0], 2*np.pi*bounds[1]) + + result = root_scalar(separation_error, bracket=bounds_rad_s, method='bisect') + if not result.converged: + raise ValueError("Root finding did not converge. Try adjusting the bounds or check the function behavior.") + + return result.root + + def print_chain_summary(self): + """ Print a summary of the linear ion chain configuration. """ + print() + print(f" ----- Linear Ion Chain Analysis Summary -----") + print(f"Number of ions: {self.num_ions}") + print(f"Trap frequencies: x={self.omega_x[0]/(2*np.pi)/1e6:.2f} MHz, " + f"y={self.omega_y[0]/(2*np.pi)/1e6:.2f} MHz, " + f"z={self.omega_z[0]/(2*np.pi)/1e6:.2f} MHz") + + # Get ion spacing + spacing = self.get_ion_spacing() + z_coordinates = self.get_axial_ion_positions() + print(f"Ion axial coordinates (z-axis): {z_coordinates*1e6} µm") + print(f"Ion spacing (z-axis): {spacing*1e6} µm") + + # Get COM position + com_x, com_y, com_z = self.get_center_of_mass_position() + print(f"Center of mass position: ({com_x*1e6:.2f} µm, {com_y*1e6:.2f} µm, {com_z*1e6:.2f} µm)") + + # Get mode frequencies + axial_freqs = self.get_axial_mode_frequencies() + radial_x_freqs = self.get_radial_mode_frequencies('x') + radial_y_freqs = self.get_radial_mode_frequencies('y') + + print(f"\nAxial mode frequencies: {axial_freqs/(2*np.pi)/1e6} MHz") + print(f"Radial X mode frequencies: {radial_x_freqs/(2*np.pi)/1e6} MHz") + print(f"Radial Y mode frequencies: {radial_y_freqs/(2*np.pi)/1e6} MHz") diff --git a/tests/__pycache__/__init__.cpython-312.pyc b/tests/__pycache__/__init__.cpython-312.pyc index fe3ce6e..265d7e6 100644 Binary files a/tests/__pycache__/__init__.cpython-312.pyc and b/tests/__pycache__/__init__.cpython-312.pyc differ diff --git a/tests/test_dissipators.py b/tests/test_dissipators.py index b2bffbc..284eedb 100644 --- a/tests/test_dissipators.py +++ b/tests/test_dissipators.py @@ -14,7 +14,6 @@ from scipy.sparse import kron as skron from ionsim.zeeman_solver import ZeemanHyperfineSolver -from ionsim.testing import assert_array_close from ionsim.basis import StandardBasis from ionsim.operator import CouplingOperator, EnergyShiftOperator from ionsim.named_operators import Pauli, Fock diff --git a/tests/test_mode_analysis.py b/tests/test_mode_analysis.py new file mode 100644 index 0000000..704a651 --- /dev/null +++ b/tests/test_mode_analysis.py @@ -0,0 +1,197 @@ +import unittest +import numpy as np +from scipy.linalg import expm +from ionsim.trapped_ion_mode_analysis import TrappedIonCrystal, LinearIonChain +from ionsim.testing import assert_array_close + +class TestModeAnalysis(unittest.TestCase): + + def setUp(self): + TPI = 2*np.pi + """Set up the necessary objects for testing.""" + # Set up tests with a list of dictionaries: + self.test_cases = [ + { + 'test name' : 'single ion 171Yb', + 'mass' : 170.936, # amu + 'omega x' : 10. * TPI * 1E6, # rad/s + 'omega y' : 1.9 * TPI * 1E6, + 'omega z' : 0.5 * TPI * 1E6, + 'N' : 1, # number of ions + 'wavenumber' : TPI / (355.0*1E-9), # 1/m + 'charge' : 1 + }, + { + 'test name' : 'two ion 171Yb', + 'mass' : 170.936, # amu + 'omega x' : 2. * TPI * 1E6, # rad/s + 'omega y' : 1.5 * TPI * 1E6, + 'omega z' : 0.5 * TPI * 1E6, + 'N' : 2, # number of ions + 'wavenumber' : TPI / (355.0*1E-9), # 1/m + 'charge' : 1 + } + ] + self.mode_analyzers = { + case['test name'] : TrappedIonCrystal(case['N'], case['omega x'], case['omega y'], case['omega z'], case['mass'], case['charge']) + for case in self.test_cases + } + + # Compute Lamb-Dicke parameter analytical reference for each case + self.references = {} + + # 1. Single ion case: + k = self.test_cases[0]['wavenumber'] + eta_x, eta_y, eta_z = self.mode_analyzers['single ion 171Yb'].compute_reference_single_ion_lamb_dicke_factors(k) + etas_analytical = np.zeros((3, 1, 3), dtype = np.complex128) + etas_analytical[0, 0, 2] = eta_x + etas_analytical[1, 0, 1] = eta_y + etas_analytical[2, 0, 0] = eta_z + self.references['single ion 171Yb'] = etas_analytical + + # 2. Two-ion case: + k = self.test_cases[1]['wavenumber'] + eta_x, eta_y, eta_z = self.mode_analyzers['two ion 171Yb'].compute_reference_single_ion_lamb_dicke_factors(k) + wx = self.test_cases[1]['omega x'] + wy = self.test_cases[1]['omega y'] + wz = self.test_cases[1]['omega z'] + + wy_tilt = np.sqrt(wy**2 - wz**2) + wx_tilt = np.sqrt(wx**2 - wz**2) + eta_COM_x = eta_x / np.sqrt(2) + eta_COM_y = eta_y / np.sqrt(2) + eta_COM_z = eta_z / np.sqrt(2) + eta_stretch_z = eta_z /np.sqrt(2) /3**(1/4) + # I don't know the analytical expressions for the tilt modes off the top of my head... let's assume its the square root of the normalized mode frequency + eta_tilt_x = eta_x / np.sqrt(2) / np.sqrt(wx_tilt / wx) + eta_tilt_y = eta_y / np.sqrt(2) / np.sqrt(wy_tilt / wy) + etas_analytical = np.zeros((3, 2, 6), dtype = np.complex128) + etas_analytical[0, 0, 5] = eta_COM_x + etas_analytical[0, 1, 5] = eta_COM_x + etas_analytical[0, 0, 4] = eta_tilt_x + etas_analytical[0, 1, 4] = -eta_tilt_x + etas_analytical[1, 0, 3] = eta_COM_y + etas_analytical[1, 1, 3] = eta_COM_y + etas_analytical[1, 0, 2] = eta_tilt_y + etas_analytical[1, 1, 2] = -eta_tilt_y + etas_analytical[2, 0, 0] = eta_COM_z + etas_analytical[2, 1, 0] = eta_COM_z + etas_analytical[2, 0, 1] = eta_stretch_z + etas_analytical[2, 1, 1] = -eta_stretch_z + self.references['two ion 171Yb'] = etas_analytical + + + def test_mode_analysis_solvers(self): + # Test functionality of mode analysis for computing Lamb-Dicke parameters as compared to a reference + # Compute Lamb-Dicke parameters for all test cases + for case, mode_analyzer, reference_values in zip(self.test_cases, self.mode_analyzers.values(), self.references.values()): + # Solve the ion-trap equilibrium problem + mode_analyzer.solve_ion_trap_equilibrium() + # Compute and store Lamb-Dicke parameters: + computed_Lamb_Dicke_parameters = case['wavenumber'] * mode_analyzer.calculate_mode_participation_factors() + assert_array_close(np.abs(computed_Lamb_Dicke_parameters), np.abs(reference_values), atol = 1E-10, rtol=None) + + def test_linear_ion_chain_analysis(self): + """Test the LinearIonChain class with a 5-ion Yb+ chain.""" + # Create a 5-ion linear chain of Yb+ ions + num_ions = 5 + TPI = 2*np.pi + omega_x = TPI * 9e6 # 10 MHz + omega_y = TPI * 10e6 # 10 MHz + omega_z = TPI * 1e6 # 1 MHz (axial frequency typically lower) + atomic_mass = 170.936 # Yb+ mass in amu + atomic_charge = 1 + + # Create and analyze the linear chain + chain = LinearIonChain(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_charge) + chain.solve_ion_trap_equilibrium() + #chain.print_chain_summary() + + # Test that we can get axial and radial modes + axial_eigvals, axial_eigvecs = chain.get_axial_modes() + radial_x_eigvals, radial_x_eigvecs = chain.get_radial_modes('x') + radial_y_eigvals, radial_y_eigvecs = chain.get_radial_modes('y') + + # Verify shapes + self.assertEqual(axial_eigvals.shape, (num_ions,)) + self.assertEqual(axial_eigvecs.shape, (6*num_ions, num_ions)) + self.assertEqual(radial_x_eigvals.shape, (num_ions,)) + self.assertEqual(radial_x_eigvecs.shape, (6*num_ions, num_ions)) + + # Test that we can get ion spacing + spacing = chain.get_ion_spacing() + self.assertEqual(spacing.shape, (num_ions-1,)) + self.assertTrue(np.all(spacing > 0)) # All spacings should be positive + + # Test that we can get center of mass + com_x, com_y, com_z = chain.get_center_of_mass_position() + self.assertIsInstance(com_x, (float, np.floating)) + self.assertIsInstance(com_y, (float, np.floating)) + self.assertIsInstance(com_z, (float, np.floating)) + + # Test that we can get mode frequencies + axial_freqs = chain.get_axial_mode_frequencies() + radial_x_freqs = chain.get_radial_mode_frequencies('x') + radial_y_freqs = chain.get_radial_mode_frequencies('y') + + self.assertEqual(axial_freqs.shape, (num_ions,)) + self.assertEqual(radial_x_freqs.shape, (num_ions,)) + self.assertEqual(radial_y_freqs.shape, (num_ions,)) + + # Test mode participation factors by branch + mode_pf_by_branch = chain.get_mode_participation_factors_by_branch() + self.assertIn('x', mode_pf_by_branch) + self.assertIn('y', mode_pf_by_branch) + self.assertIn('z', mode_pf_by_branch) + + # Test that frequencies are positive + self.assertTrue(np.all(axial_freqs > 0)) + self.assertTrue(np.all(radial_x_freqs > 0)) + self.assertTrue(np.all(radial_y_freqs > 0)) + + # Test Lamb-Dicke parameter calculations + wavenumber = 2 * np.pi / (355.0 * 1e-9) # example laser wavelength used with Yb+ + + # Test full Lamb-Dicke parameter matrix + wavevector = wavenumber * np.array([np.cos(np.pi/4.), np.sin(np.pi/4.), 0.]) # perpendicular to the axial chain + full_ld_params = chain.calculate_lamb_dicke_parameters_full(wavevector) + num_modes = 3*num_ions + self.assertEqual(full_ld_params.shape, (num_ions, num_modes)) + + # Test branch-organized Lamb-Dicke parameters + ld_by_branch = chain.calculate_lamb_dicke_parameters_by_branch(wavevector) + self.assertIn('x', ld_by_branch) + self.assertIn('y', ld_by_branch) + self.assertIn('z', ld_by_branch) + self.assertEqual(ld_by_branch['x'].shape, (num_ions, num_ions)) + self.assertEqual(ld_by_branch['y'].shape, (num_ions, num_ions)) + self.assertEqual(ld_by_branch['z'].shape, (num_ions, num_ions)) + + # Test axial Lamb-Dicke parameters + axial_ld = chain.get_axial_lamb_dicke_parameters(wavevector) + self.assertEqual(axial_ld.shape, (num_ions, num_ions)) + + # Test radial Lamb-Dicke parameters + radial_x_ld = chain.get_radial_lamb_dicke_parameters(wavevector, 'x') + radial_y_ld = chain.get_radial_lamb_dicke_parameters(wavevector, 'y') + self.assertEqual(radial_x_ld.shape, (num_ions, num_ions)) + self.assertEqual(radial_y_ld.shape, (num_ions, num_ions)) + + # Verify that branch-organized LD params match the corresponding parts of full matrix + n_modes_per_branch = num_ions + np.testing.assert_array_almost_equal(ld_by_branch['x'], full_ld_params[:, :n_modes_per_branch]) + np.testing.assert_array_almost_equal(ld_by_branch['y'], full_ld_params[:, n_modes_per_branch:2*n_modes_per_branch]) + np.testing.assert_array_almost_equal(ld_by_branch['z'], full_ld_params[:, 2*n_modes_per_branch:3*n_modes_per_branch]) + + # Test that Lamb-Dicke parameters are related to mode participation factors + mode_pf = chain.get_mode_participation_factors_by_branch() + #print(mode_pf['x']) + #print(ld_by_branch['z']) + #mode_pf = chain.get_mode_participation_factors_by_branch() + #expected_full_ld = wavenumber * mode_pf + #np.testing.assert_array_almost_equal(full_ld_params, expected_full_ld) + + +if __name__ == '__main__': + # Run the unit tests + unittest.main() diff --git a/tests/test_zeeman_solver.py b/tests/test_zeeman_solver.py index d77dfbe..1b71d38 100644 --- a/tests/test_zeeman_solver.py +++ b/tests/test_zeeman_solver.py @@ -11,8 +11,12 @@ import numpy as np from scipy.linalg import expm from ionsim.zeeman_solver import ZeemanHyperfineSolver +#<<<<<<< HEAD +#from ionsim.degree_of_freedom import AtomicSpin +#======= from ionsim.testing import assert_array_close from ionsim.degree_of_freedom import AtomicStructure +#>>>>>>> main class TestZeemanSolver(unittest.TestCase): @@ -57,7 +61,7 @@ def setUp(self): } # Test 3: Use AtomicStructure from_species() - self.spin = AtomicStructure.from_species(species='171Yb', term_symbols=['S0'], level_names=['S0,1/2,1/2', 'S0,1/2,-1/2'], magnetic_field = 400.) + self.Yb = AtomicStructure.from_species(species='171Yb', term_symbols=['S0'], level_names=['S0,1/2,1/2', 'S0,1/2,-1/2'], magnetic_field = 400.) def test_explicit_zeeman_solvers(self): # Test functionality when using explicit construction of Zeeman Solver objects for each test case. @@ -107,7 +111,7 @@ def test_explicit_zeeman_solvers(self): def test_AtomicStructure_ZeemanShift(self): """Test the Zeemaen shift functionality within AtomicStructure class.""" expected_frequency = 149.98214416459987*2. # kHz - qubit_frequency = self.spin.energy_levels[1].energy - self.spin.energy_levels[0].energy + qubit_frequency = self.Yb.energy_levels[1].energy - self.Yb.energy_levels[0].energy qubit_frequency /= (2.* np.pi) # convert from rad/s to Hz self.assertAlmostEqual(expected_frequency, np.abs(qubit_frequency)*1E-3, places=6)