diff --git a/scripts/simulations/002_calibration_demonstration.py b/scripts/simulations/002_calibration_demonstration.py new file mode 100644 index 0000000..81989fe --- /dev/null +++ b/scripts/simulations/002_calibration_demonstration.py @@ -0,0 +1,426 @@ +"""A demonstration of the ringdown spectroscopy protocol.""" + +from rabifun.system import * +from rabifun.plots import * +from rabifun.utilities import * +from ringfit.utils import WelfordAggregator +from rabifun.analysis import * +import multiprocessing +import copy + + +def solve_shot( + params: Params, t: np.ndarray, t_before: np.ndarray, t_after: np.ndarray +): + """A worker function to solve for the time evolution in separate processes. + + :param params: The parameters of the system. + :param t: The time axis. + :param t_before: The time axis before the EOM is switched off. + :param t_after: The time axis after the EOM is switched off. + """ + + solution = solve(t, params) + amps = solution.y[::, len(t_before) - 1 :] + + return t_after, amps + + +def make_shots( + params: Params, + total_lifetimes: float, + eom_range: tuple[float, float], + eom_steps: int, + num_freq: int = 1, + randomize_off_time: bool = True, +): + """Generate a series of shots with varying EOM frequencies. + + The implementation here slightly varies the off time of the laser + so as to introduce some random relative phases of the modes. + + :param params: The parameters of the system. + :param total_lifetimes: The total time of the experiment in + lifetimes. + :param eom_range: The range of EOM frequencies in units of + :any:`params.Ω`. + :param eom_steps: The number of steps in the EOM frequency range. + :param num_freq: The number of frequencies to drive. If a number + greater than 1 is given, the EOM will be driven at multiple + frequencies where the highest frequency is the base frequency + plus an consecutive integer multiples of :any:`params.Ω`. + :param randomize_off_time: Whether to randomize the off time of the + laser. This is useful to introduce some random relative phases + for the modes. + """ + + solutions = [] + shot_params = [] + + rng = np.random.default_rng(seed=0) + off_time = params.laser_off_time or 0 + analyze_time = params.lifetimes(total_lifetimes) - off_time + t_after = time_axis(params, total_time=analyze_time, resolution=0.01) + + for step in range(eom_steps): + base = params.Ω * ( + eom_range[0] + (eom_range[1] - eom_range[0]) * step / eom_steps + ) + + current_params = copy.deepcopy(params) + current_params.drive_override = ( + base + params.Ω * np.arange(num_freq), + np.ones(num_freq), + ) + + current_params.drive_phases = rng.uniform(0, 2 * np.pi, size=num_freq) + off_time += rng.normal(0, 0.1 * off_time) if randomize_off_time else 0 + + current_params.laser_off_time = None # off_time + current_params.drive_off_time = off_time + + t_before = time_axis(params, total_time=off_time, resolution=0.01) + t = np.concatenate([t_before[:-1], t_after + t_before[-1]]) + shot_params.append((current_params, t, t_before, t_after)) + + pool = multiprocessing.Pool() + solutions = pool.starmap(solve_shot, shot_params) + return solutions + + +def process_shots( + params: Params, + solutions: list[tuple[np.ndarray, np.ndarray]], + noise_amplitude: float, + num_freq: int, +): + """ + Calculates the normalized average Fourier power spectrum of a + series of experimental (simulated) shots. + + :param params: The parameters of the system. + :param solutions: A list of solutions to process returned by + :any:`solve_shot`. + :param noise_amplitude: The amplitude of the noise to add to the + signal. + + The amplitude is normalized by 2/η which is roughly the steady + state signal amplitude if a bath mode is excited resonantly by + a unit-strength input. + + :param num_freq: The number of frequencies to drive. See + :any:`make_shots` for details. + """ + + rng = np.random.default_rng(seed=0) + + noise_amplitude /= params.η * np.pi + + aggregate = None + for t, amps in solutions: + signal = output_signal(t, amps, params) + signal += rng.normal(scale=noise_amplitude, size=len(signal)) + window = (0, t[-1]) + + freq, fft = fourier_transform( + t, + signal, + low_cutoff=0.3 * params.Ω, + high_cutoff=params.Ω * (1 + num_freq), + window=window, + ) + + power = np.abs(fft) ** 2 + + # ugly hack because shape is hard to predict + if aggregate is None: + aggregate = WelfordAggregator(power) + else: + aggregate.update(power) + + assert aggregate is not None # appease pyright + + max_power = np.max(aggregate.mean) + return (freq, aggregate.mean / max_power, aggregate.ensemble_std / max_power) + + +def process_and_plot_results( + params: Params, + ax: plt.Axes, + freq: np.ndarray, + average_power_spectrum: np.ndarray, + σ_power_spectrum: np.ndarray, + annotate: bool = True, +): + """ + Fits the ringdown spectrum and plots the results. + + :param params: The parameters of the system. + :param ax: The axis to plot on. + :param freq: The frequency array. + :param average_power_spectrum: The average power spectrum obtained from :any:`process_shots`. + :param σ_power_spectrum: The standard deviation of the power + spectrum. + :param annotate: Whether to annotate the plot with peak and mode positions. + """ + + ringdown_params = RingdownParams( + fω_shift=params.measurement_detuning, + mode_window=(params.N, params.N), + fΩ_guess=params.Ω, + fδ_guess=params.Ω * params.δ, + η_guess=0.5, + absolute_low_cutoff=0.3 * params.Ω, + ) + + peak_info = find_peaks( + freq, + average_power_spectrum, + ringdown_params, + prominence=0.05, + height=0.1, + σ_power=σ_power_spectrum, + ) + + peak_info = refine_peaks(peak_info, ringdown_params, height_cutoff=0.05) + + plot_spectrum_and_peak_info(ax, peak_info, annotate=annotate) + + if peak_info.lm_result is not None: + fine_freq = np.linspace(freq.min(), freq.max(), 5000) + fine_fit = peak_info.lm_result.eval(ω=fine_freq) + ax.plot( + fine_freq, + fine_fit - peak_info.noise_floor, + color="C3", + zorder=-100, + label="Fit", + ) + ax.set_ylim(-0.1, max(1, fine_fit.max() * 1.1)) + + ax.set_xlabel("Frequency (MHz)") + ax.legend() + + if annotate: + annotate_ringodown_mode_positions(params, ax) + + +def generate_data( + Ω=13, + η=0.2, + g_0=0.5, + η_factor=5, + noise_amplitude=0.3, + laser_detuning=0, + laser_on_mode=0, + N=10, + eom_ranges=(0.5, 2.5), + eom_steps=20, + excitation_lifetimes=2, + measurement_lifetimes=4, + num_freq=3, + extra_title="", + save: str | None = None, + randomize_off_time: bool = True, +): + """Simulate and plot the ringdown spectroscopy protocol. + + The idea is to have the laser on ``laser_on_mode`` and to sweep + the EOM frequency over a range of values given in ``eom_ranges`` + in ``eom_steps`` steps. For each step, the laser and EOM are + inputting into the system for a time given by + ``excitation_lifetimes``. Then, the ringdown signal is collected + for a time given by ``measurement_lifetimes``. (Lifetime units + are given by ``η``.) The resulting power spectra are averaged and + then fitted. + + :param Ω: The FSR of the system. + :param η: The decay rate of the system. + :param g_0: The coupling strength of the system in units of + :any:`Ω`. Note that the effective coupling strength between + the ``A`` site and the bath modes is reduced by a factor of + :math:`\sqrt{2}`. + + :param η_factor: The factor by which the decay rate of the A site + is greater. + :param noise_amplitude: The amplitude of the noise to add to the + signal. See :any:`process_shots` for details. + :param laser_detuning: The detuning of the laser from the the mode + it is exciting. + :param laser_on_mode: The mode that the laser is exciting. + :param N: The number of bath modes. + :param eom_ranges: The range of EOM frequencies in units of + :any:`Ω`. + :param eom_steps: The number of steps in the EOM frequency range. + :param excitation_lifetimes: The time the EOM is driving the + system. + :param measurement_lifetimes: The time the system is left to ring + down. + + Note that the laser is not turned off during the ringdown. + + :param num_freq: The number of frequencies to drive. See + :any:`make_shots` for details. + :param extra_title: A string to add to the title of the plot. + :param save: The filename to save the figure to. If None, + the figure is not saved. + :param randomize_off_time: Whether to randomize the off time of the + laser. See :any:`make_shots` for details. + + :returns: The figure containing the plot. + """ + + final_laser_detuning = laser_detuning + ( + 0 if laser_on_mode == 0 else (laser_on_mode - 1 / 4) * Ω + ) + + params = Params( + η=η, + η_hybrid=η_factor * η, + Ω=Ω, + δ=1 / 4, + ω_c=0.1, + g_0=g_0 * num_freq, # as it would be normalized otherwise + laser_detuning=final_laser_detuning, + N=N, + N_couplings=N, + measurement_detuning=0, + α=0, + rwa=False, + flat_energies=False, + correct_lamb_shift=0, + laser_off_time=None, + small_loop_detuning=0, + drive_override=(np.array([]), np.array([])), + ) + + params.laser_off_time = params.lifetimes(excitation_lifetimes) + params.drive_off_time = params.lifetimes(excitation_lifetimes) + + solutions = make_shots( + params, + excitation_lifetimes + measurement_lifetimes, + eom_ranges, + eom_steps, + num_freq, + randomize_off_time=randomize_off_time, + ) + + freq, average_power_spectrum, σ_power_spectrum = process_shots( + params, + solutions, + noise_amplitude, + num_freq, + ) + + fig = make_figure(extra_title, figsize=(12, 6)) + fig.clear() + ax = fig.subplots() + + process_and_plot_results(params, ax, freq, average_power_spectrum, σ_power_spectrum) + ax.text( + 0.01, + 0.95, + f"""$Ω/2π = {params.Ω}$MHz +$η/2π = {params.η}MHz$ +$g_0 = {params.g_0}Ω$ +$N = {params.N}$ +noise = ${noise_amplitude * 2}$ +$η_A = {η_factor}η$ +EOM range = {eom_ranges[0]:.2f}Ω to {eom_ranges[1]:.2f}Ω +EOM steps = {eom_steps} +excitation time = {excitation_lifetimes} lifetimes +measurement time = {measurement_lifetimes} lifetimes +on mode = {laser_on_mode} +laser detuning = {laser_detuning} +num freq = {num_freq} +total time = {(excitation_lifetimes + measurement_lifetimes) * eom_steps / (params.η * 1e6)}s +randomize off time = {'yes' if randomize_off_time else 'no'} """, + transform=ax.transAxes, + ha="left", + va="top", + size=10, + bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.9), + ) + ax.set_title(extra_title) + + fig.tight_layout() + + save_figure(fig, f"002_{save}") + quick_save_pickle( + locals(), + f"002_{save}", + ) + + return fig + + +# %% save +if __name__ == "__main__": + generate_data( + g_0=0.5, + η_factor=5, + noise_amplitude=0.3, + N=5, + eom_ranges=(0.7, 0.9), + eom_steps=100, + laser_detuning=0, + laser_on_mode=0, + excitation_lifetimes=2, + measurement_lifetimes=4, + num_freq=4, + randomize_off_time=True, + extra_title="Laser on A site", + save="001_laser_on_A", + ) + + fig = generate_data( + g_0=0.5, + η_factor=5, + noise_amplitude=0.3, + N=5, + eom_ranges=(1.2, 1.3), + eom_steps=200, + laser_detuning=0, + laser_on_mode=-1, + excitation_lifetimes=2, + measurement_lifetimes=4, + num_freq=1, + extra_title="Laser on Bath Mode", + randomize_off_time=True, + save="002_laser_on_bath", + ) + + fig = generate_data( + g_0=0.3, + η_factor=5, + noise_amplitude=0.0, + N=5, + eom_ranges=(1.2, 1.3), + eom_steps=100, + laser_detuning=0, + laser_on_mode=-1, + excitation_lifetimes=2, + measurement_lifetimes=4, + num_freq=1, + randomize_off_time=False, + extra_title="Laser on Bath Mode, No randomized off-time.", + save="003_laser_on_bath_no_random", + ) + + fig = generate_data( + g_0=0.3, + η_factor=5, + noise_amplitude=0.0, + N=5, + eom_ranges=(1.2, 1.3), + eom_steps=100, + laser_detuning=0, + laser_on_mode=-1, + excitation_lifetimes=2, + measurement_lifetimes=4, + num_freq=1, + randomize_off_time=True, + extra_title="Laser on Bath Mode, Randomized off-time.", + save="004_laser_on_bath_random", + ) diff --git a/scripts/simulations/figs/002_001_laser_on_A.pdf b/scripts/simulations/figs/002_001_laser_on_A.pdf new file mode 100644 index 0000000..242082a Binary files /dev/null and b/scripts/simulations/figs/002_001_laser_on_A.pdf differ diff --git a/scripts/simulations/figs/002_001_laser_on_A.pdf.meta.yaml b/scripts/simulations/figs/002_001_laser_on_A.pdf.meta.yaml new file mode 100644 index 0000000..c5a8a3d --- /dev/null +++ b/scripts/simulations/figs/002_001_laser_on_A.pdf.meta.yaml @@ -0,0 +1,27 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: bc45229f9e8466056fd4d5143f87e739dd3785ae +description: 'RESULT: calibration simulation' +extra_meta: null +function: generate_data +function_args: + N: 5 + eom_ranges: !!python/tuple + - 0.7 + - 0.9 + eom_steps: 100 + excitation_lifetimes: 2 + extra_title: Laser on A site + g_0: 0.5 + laser_detuning: 0 + laser_on_mode: 0 + measurement_lifetimes: 4 + noise_amplitude: 0.3 + num_freq: 4 + randomize_off_time: true + save: 001_laser_on_A + Ω: 13 + η: 0.2 + η_factor: 5 +name: 002_001_laser_on_A +refers_to: ./figs/002_001_laser_on_A.pdf +source: scripts/simulations/002_calibration_demonstration.py diff --git a/scripts/simulations/figs/002_001_laser_on_A.png b/scripts/simulations/figs/002_001_laser_on_A.png new file mode 100644 index 0000000..e09d5e6 Binary files /dev/null and b/scripts/simulations/figs/002_001_laser_on_A.png differ diff --git a/scripts/simulations/figs/002_002_laser_on_bath.pdf b/scripts/simulations/figs/002_002_laser_on_bath.pdf new file mode 100644 index 0000000..790b200 Binary files /dev/null and b/scripts/simulations/figs/002_002_laser_on_bath.pdf differ diff --git a/scripts/simulations/figs/002_002_laser_on_bath.pdf.meta.yaml b/scripts/simulations/figs/002_002_laser_on_bath.pdf.meta.yaml new file mode 100644 index 0000000..4d22b76 --- /dev/null +++ b/scripts/simulations/figs/002_002_laser_on_bath.pdf.meta.yaml @@ -0,0 +1,27 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: 79069d5f978fcd3e48fa401aad6dbcfcf2d8f917 +description: 'RESULT: calibration simulation' +extra_meta: null +function: generate_data +function_args: + N: 5 + eom_ranges: !!python/tuple + - 1.2 + - 1.3 + eom_steps: 200 + excitation_lifetimes: 2 + extra_title: Laser on Bath Mode + g_0: 0.5 + laser_detuning: 0 + laser_on_mode: -1 + measurement_lifetimes: 4 + noise_amplitude: 0.3 + num_freq: 1 + randomize_off_time: true + save: 002_laser_on_bath + Ω: 13 + η: 0.2 + η_factor: 5 +name: 002_002_laser_on_bath +refers_to: ./figs/002_002_laser_on_bath.pdf +source: scripts/simulations/002_calibration_demonstration.py diff --git a/scripts/simulations/figs/002_002_laser_on_bath.png b/scripts/simulations/figs/002_002_laser_on_bath.png new file mode 100644 index 0000000..8c09133 Binary files /dev/null and b/scripts/simulations/figs/002_002_laser_on_bath.png differ diff --git a/scripts/simulations/figs/002_003_laser_on_bath_no_random.pdf b/scripts/simulations/figs/002_003_laser_on_bath_no_random.pdf new file mode 100644 index 0000000..39cc10a Binary files /dev/null and b/scripts/simulations/figs/002_003_laser_on_bath_no_random.pdf differ diff --git a/scripts/simulations/figs/002_003_laser_on_bath_no_random.pdf.meta.yaml b/scripts/simulations/figs/002_003_laser_on_bath_no_random.pdf.meta.yaml new file mode 100644 index 0000000..506c658 --- /dev/null +++ b/scripts/simulations/figs/002_003_laser_on_bath_no_random.pdf.meta.yaml @@ -0,0 +1,27 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: ec06ea8db6bf543984fa0cf3b39843c05044e654 +description: 'RESULT: calibration simulation' +extra_meta: null +function: generate_data +function_args: + N: 5 + eom_ranges: !!python/tuple + - 1.2 + - 1.3 + eom_steps: 100 + excitation_lifetimes: 2 + extra_title: Laser on Bath Mode, No randomized off-time. + g_0: 0.3 + laser_detuning: 0 + laser_on_mode: -1 + measurement_lifetimes: 4 + noise_amplitude: 0.0 + num_freq: 1 + randomize_off_time: false + save: 003_laser_on_bath_no_random + Ω: 13 + η: 0.2 + η_factor: 5 +name: 002_003_laser_on_bath_no_random +refers_to: ./figs/002_003_laser_on_bath_no_random.pdf +source: scripts/simulations/002_calibration_demonstration.py diff --git a/scripts/simulations/figs/002_003_laser_on_bath_no_random.png b/scripts/simulations/figs/002_003_laser_on_bath_no_random.png new file mode 100644 index 0000000..03c61cf Binary files /dev/null and b/scripts/simulations/figs/002_003_laser_on_bath_no_random.png differ diff --git a/scripts/simulations/figs/002_004_laser_on_bath_random.pdf b/scripts/simulations/figs/002_004_laser_on_bath_random.pdf new file mode 100644 index 0000000..3d230b5 Binary files /dev/null and b/scripts/simulations/figs/002_004_laser_on_bath_random.pdf differ diff --git a/scripts/simulations/figs/002_004_laser_on_bath_random.pdf.meta.yaml b/scripts/simulations/figs/002_004_laser_on_bath_random.pdf.meta.yaml new file mode 100644 index 0000000..d05d50a --- /dev/null +++ b/scripts/simulations/figs/002_004_laser_on_bath_random.pdf.meta.yaml @@ -0,0 +1,27 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: 7d0e9be06057bfe7332a83bf4997b15565a4f832 +description: 'RESULT: calibration simulation' +extra_meta: null +function: generate_data +function_args: + N: 5 + eom_ranges: !!python/tuple + - 1.2 + - 1.3 + eom_steps: 100 + excitation_lifetimes: 2 + extra_title: Laser on Bath Mode, Randomized off-time. + g_0: 0.3 + laser_detuning: 0 + laser_on_mode: -1 + measurement_lifetimes: 4 + noise_amplitude: 0.0 + num_freq: 1 + randomize_off_time: true + save: 004_laser_on_bath_random + Ω: 13 + η: 0.2 + η_factor: 5 +name: 002_004_laser_on_bath_random +refers_to: ./figs/002_004_laser_on_bath_random.pdf +source: scripts/simulations/002_calibration_demonstration.py diff --git a/scripts/simulations/figs/002_004_laser_on_bath_random.png b/scripts/simulations/figs/002_004_laser_on_bath_random.png new file mode 100644 index 0000000..80dbbda Binary files /dev/null and b/scripts/simulations/figs/002_004_laser_on_bath_random.png differ diff --git a/scripts/simulations/outputs/002_001_laser_on_A.pkl.meta.yaml b/scripts/simulations/outputs/002_001_laser_on_A.pkl.meta.yaml new file mode 100644 index 0000000..9c292d3 --- /dev/null +++ b/scripts/simulations/outputs/002_001_laser_on_A.pkl.meta.yaml @@ -0,0 +1,7 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: f9eaf01d1e1c5b0925d718c69a7f84e0e14ce875 +description: 'RESULT: calibration simulation' +function: generate_data +function_args: {} +refers_to: outputs/002_001_laser_on_A.pkl +source: scripts/simulations/002_calibration_demonstration.py diff --git a/scripts/simulations/outputs/002_002_laser_on_bath.pkl.meta.yaml b/scripts/simulations/outputs/002_002_laser_on_bath.pkl.meta.yaml new file mode 100644 index 0000000..7d89d92 --- /dev/null +++ b/scripts/simulations/outputs/002_002_laser_on_bath.pkl.meta.yaml @@ -0,0 +1,7 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: 471ae089658da27f23b939a65b24607786ebd880 +description: 'RESULT: calibration simulation' +function: generate_data +function_args: {} +refers_to: outputs/002_002_laser_on_bath.pkl +source: scripts/simulations/002_calibration_demonstration.py diff --git a/scripts/simulations/outputs/002_003_laser_on_bath_no_random.pkl.meta.yaml b/scripts/simulations/outputs/002_003_laser_on_bath_no_random.pkl.meta.yaml new file mode 100644 index 0000000..77a07b3 --- /dev/null +++ b/scripts/simulations/outputs/002_003_laser_on_bath_no_random.pkl.meta.yaml @@ -0,0 +1,7 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: 97126fb016e5811d4ba2682ca4371ae561cb3e5b +description: 'RESULT: calibration simulation' +function: generate_data +function_args: {} +refers_to: outputs/002_003_laser_on_bath_no_random.pkl +source: scripts/simulations/002_calibration_demonstration.py diff --git a/scripts/simulations/outputs/002_003_laser_on_bath_random.pkl.meta.yaml b/scripts/simulations/outputs/002_003_laser_on_bath_random.pkl.meta.yaml new file mode 100644 index 0000000..bb3474d --- /dev/null +++ b/scripts/simulations/outputs/002_003_laser_on_bath_random.pkl.meta.yaml @@ -0,0 +1,7 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: df99bc5460c3241dc76a628d40dde8f5c9242ca2 +description: 'RESULT: calibration simulation' +function: generate_data +function_args: {} +refers_to: outputs/002_003_laser_on_bath_random.pkl +source: scripts/simulations/002_calibration_demonstration.py diff --git a/scripts/simulations/outputs/002_004_laser_on_bath_random.pkl.meta.yaml b/scripts/simulations/outputs/002_004_laser_on_bath_random.pkl.meta.yaml new file mode 100644 index 0000000..f3a90f0 --- /dev/null +++ b/scripts/simulations/outputs/002_004_laser_on_bath_random.pkl.meta.yaml @@ -0,0 +1,7 @@ +change_id: nnyqrroulxpmtpxksrmpyoozouxtluur +commit_id: 7a4a4d061cb43345a0984d34b902b485fa956704 +description: 'RESULT: calibration simulation' +function: generate_data +function_args: {} +refers_to: outputs/002_004_laser_on_bath_random.pkl +source: scripts/simulations/002_calibration_demonstration.py diff --git a/src/rabifun/plots.py b/src/rabifun/plots.py index 2fc696f..29f9ca4 100644 --- a/src/rabifun/plots.py +++ b/src/rabifun/plots.py @@ -219,6 +219,8 @@ def plot_spectrum_and_peak_info(ax, peaks: RingdownPeakData, annotate=False): :param ax: The axis to plot on. :param peaks: The peak data. :param annotate: Whether to annotate the peaks. + + If ``true`` each peak will be annotated with it's estimated frequency. """ ax.clear() @@ -243,7 +245,8 @@ def plot_spectrum_and_peak_info(ax, peaks: RingdownPeakData, annotate=False): freq, lorentz[0] * 1.1, rf"{i} (${roundfreq}\pm {rounderr}$)", - fontsize=20, + fontsize=10, + clip_on=True, ) t.set_bbox(dict(facecolor="white", alpha=0.8, edgecolor="white"))