example real waveform

import%20marimo%0A%0A__generated_with%20%3D%20%220.19.11%22%0Aapp%20%3D%20marimo.App()%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Example%3A%20Fitting%20a%20real%20(linearly%20polarized)%20waveform%0A%0A%20%20%20%20In%20this%20notebook%20we%20demonstrate%20how%20to%20fit%20a%20real%20(linearly%20polarized)%20waveform%20with%20%60jaxqualin%60.%0A%20%20%20%20This%20is%20useful%20for%20fitting%20head-on%20mergers.%0A%0A%20%20%20%201.%20**Free-frequency%20fit**%20across%20varying%20%24t_0%24%0A%20%20%20%202.%20**Fixed-frequency%20fit**%20across%20varying%20%24t_0%24%0A%20%20%20%203.%20**Variable-mass%20fit**%20across%20varying%20%24t_0%24%20with%20%60real%3DTrue%60%20to%20recover%20%24M%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20from%20jaxqualin.waveforms%20import%20delayed_QNM%2C%20waveform%0A%20%20%20%20from%20jaxqualin.qnmode%20import%20mode_list%2C%20long_str_to_qnms_free%0A%20%20%20%20from%20jaxqualin.fit%20import%20QNMFitVaryingStartingTime%0A%20%20%20%20from%20jaxqualin.plot%20import%20(plot_amplitudes%2C%20plot_phases%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20plot_omega_free%2C%20plot_predicted_qnms)%0A%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20QNMFitVaryingStartingTime%2C%0A%20%20%20%20%20%20%20%20delayed_QNM%2C%0A%20%20%20%20%20%20%20%20long_str_to_qnms_free%2C%0A%20%20%20%20%20%20%20%20mode_list%2C%0A%20%20%20%20%20%20%20%20np%2C%0A%20%20%20%20%20%20%20%20plot_amplitudes%2C%0A%20%20%20%20%20%20%20%20plot_omega_free%2C%0A%20%20%20%20%20%20%20%20plot_phases%2C%0A%20%20%20%20%20%20%20%20plot_predicted_qnms%2C%0A%20%20%20%20%20%20%20%20plt%2C%0A%20%20%20%20%20%20%20%20waveform%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Build%20a%20toy%20real%20waveform%0A%0A%20%20%20%20We%20construct%20a%20real%20waveform%20with%20two%20QNMs%20at%20%24a%20%3D%200%24%3A%20the%20%24(2%2C2%2C0)%24%20fundamental%20mode%20and%20the%20%24(3%2C2%2C0)%24%20mode.%0A%20%20%20%20The%20%60delayed_QNM%60%20function%20adds%20distortions%20near%20%24t%20%3D%200%24%20to%20simulate%20realistic%20merger%20behavior.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(delayed_QNM%2C%20mode_list%2C%20np%2C%20waveform)%3A%0A%20%20%20%20Mf_true%20%3D%201.0%0A%20%20%20%20af_true%20%3D%200.0%0A%20%20%20%20schw_modes%20%3D%20mode_list(%5B'2.2.0'%2C%20'3.2.0'%5D%2C%20Mf_true%2C%20af_true)%0A%0A%20%20%20%20schw_A_phi_dict%20%3D%20%7B%0A%20%20%20%20%20%20%20%20'2.2.0'%3A%20dict(A%3D1.0%2C%20phi%3D0.0)%2C%0A%20%20%20%20%20%20%20%20'3.2.0'%3A%20dict(A%3D0.5%2C%20phi%3Dnp.pi%20%2F%203)%2C%0A%20%20%20%20%7D%0A%0A%20%20%20%20schw_t_arr%20%3D%20np.linspace(0%2C%20120%2C%201000)%0A%20%20%20%20_h_arr%20%3D%20np.zeros(schw_t_arr.shape%2C%20dtype%3Dnp.complex128)%0A%20%20%20%20for%20_i%2C%20_mode%20in%20enumerate(schw_modes)%3A%0A%20%20%20%20%20%20%20%20if%20_i%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_A_delay%2C%20_A_sig%2C%20_phi_sig%20%3D%200%2C%2010%2C%205%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_A_delay%2C%20_A_sig%2C%20_phi_sig%20%3D%205%2C%202%2C%202%0A%20%20%20%20%20%20%20%20_h_arr%20%3D%20_h_arr%20%2B%20delayed_QNM(%0A%20%20%20%20%20%20%20%20%20%20%20%20_mode%2C%20schw_t_arr%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20schw_A_phi_dict%5B_mode.string()%5D%5B'A'%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20schw_A_phi_dict%5B_mode.string()%5D%5B'phi'%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20A_delay%3D_A_delay%2C%20A_sig%3D_A_sig%2C%20phi_sig%3D_phi_sig%2C%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%23%20Take%20the%20real%20part%20for%20a%20real%20waveform%0A%20%20%20%20h_schw%20%3D%20waveform(schw_t_arr%2C%20np.real(_h_arr)%20%2B%200j%20*%20np.real(_h_arr)%2C%20t_peak%3D0)%0A%20%20%20%20return%20Mf_true%2C%20af_true%2C%20h_schw%2C%20schw_modes%0A%0A%0A%40app.cell%0Adef%20_(h_schw%2C%20np%2C%20plt)%3A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots()%0A%20%20%20%20_ax.semilogy(h_schw.time%2C%20np.abs(h_schw.hr))%0A%20%20%20%20_ax.set_xlabel(r'%24t%24')%0A%20%20%20%20_ax.set_ylabel(r'%24%7Ch_r%7C%24')%0A%20%20%20%20_ax.set_title('Toy%20real%20waveform%20(real%20part)')%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Free-frequency%20fit%20(varying%20%24t_0%24)%0A%0A%20%20%20%20First%20we%20fit%20with%20free%20QNMs%20(%24N_f%20%3D%202%24)%20whose%20frequencies%20are%20not%20fixed%2C%20to%20see%20if%20the%20resulting%20frequencies%20converge%20to%20the%20expected%20QNMs%20for%20this%20real%20waveform.%0A%20%20%20%20Different%20colored%20points%20trace%20out%20the%20frequency%20evolution%20of%20each%20free%20QNM%20from%20early%20to%20late%20%24t_0%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(QNMFitVaryingStartingTime%2C%20h_schw%2C%20mo%2C%20np)%3A%0A%20%20%20%20schw_free_t0_arr%20%3D%20np.linspace(0%2C%2050%2C%20num%3D101)%0A%0A%20%20%20%20fitter_schw_free%20%3D%20QNMFitVaryingStartingTime(%0A%20%20%20%20%20%20%20%20h_schw%2C%20schw_free_t0_arr%2C%20N_free%3D2%2C%0A%20%20%20%20%20%20%20%20qnm_fixed_list%3D%5B%5D%2C%0A%20%20%20%20%20%20%20%20real%3DTrue%2C%0A%20%20%20%20%20%20%20%20load_pickle%3DFalse%2C%0A%20%20%20%20%20%20%20%20run_string_prefix%3D'schwarzschild_free_example'%2C%0A%20%20%20%20%20%20%20%20save_results%3DFalse%2C%0A%20%20%20%20)%0A%20%20%20%20with%20mo.status.spinner(%22Running%20free-frequency%20fits...%22)%3A%0A%20%20%20%20%20%20%20%20fitter_schw_free.do_fits()%0A%20%20%20%20result_schw_free%20%3D%20fitter_schw_free.result_full%0A%20%20%20%20return%20(result_schw_free%2C)%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20Mf_true%2C%0A%20%20%20%20af_true%2C%0A%20%20%20%20mode_list%2C%0A%20%20%20%20plot_omega_free%2C%0A%20%20%20%20plot_predicted_qnms%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20result_schw_free%2C%0A)%3A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots()%0A%20%20%20%20_predicted_qnms%20%3D%20mode_list(%5B'2.2.0'%2C%20'3.2.0'%5D%2C%20Mf_true%2C%20af_true)%0A%20%20%20%20plot_omega_free(result_schw_free%2C%20_ax)%0A%20%20%20%20plot_predicted_qnms(_ax%2C%20_predicted_qnms)%0A%20%20%20%20_ax.set_title('Free-frequency%20fit%3A%20QNM%20frequency%20evolution')%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Fixed-frequency%20fit%20(varying%20%24t_0%24)%0A%0A%20%20%20%20Now%20we%20assume%20the%20%24(2%2C2%2C0)%24%20and%20%24(3%2C2%2C0)%24%20modes%20are%20present%20and%20fix%20their%20frequencies%20according%20to%20GR.%0A%20%20%20%20We%20fit%20across%20a%20range%20of%20starting%20times%20%24t_0%24%20and%20plot%20how%20the%20recovered%20amplitudes%20and%20phases%20behave.%0A%20%20%20%20For%20a%20clean%20QNM%20signal%2C%20these%20should%20be%20approximately%20constant%20across%20%24t_0%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(QNMFitVaryingStartingTime%2C%20h_schw%2C%20mo%2C%20np%2C%20schw_modes)%3A%0A%20%20%20%20schw_fixed_t0_arr%20%3D%20np.linspace(0%2C%2050%2C%20num%3D101)%0A%0A%20%20%20%20fitter_schw_fixed_vst%20%3D%20QNMFitVaryingStartingTime(%0A%20%20%20%20%20%20%20%20h_schw%2C%20schw_fixed_t0_arr%2C%20N_free%3D0%2C%0A%20%20%20%20%20%20%20%20qnm_fixed_list%3Dschw_modes%2C%0A%20%20%20%20%20%20%20%20real%3DTrue%2C%0A%20%20%20%20%20%20%20%20load_pickle%3DFalse%2C%0A%20%20%20%20%20%20%20%20run_string_prefix%3D'schwarzschild_fixed_example'%2C%0A%20%20%20%20%20%20%20%20save_results%3DFalse%2C%0A%20%20%20%20)%0A%20%20%20%20with%20mo.status.spinner(%22Running%20fixed-frequency%20fits...%22)%3A%0A%20%20%20%20%20%20%20%20fitter_schw_fixed_vst.do_fits()%0A%20%20%20%20result_schw_fixed_vst%20%3D%20fitter_schw_fixed_vst.result_full%0A%20%20%20%20return%20(result_schw_fixed_vst%2C)%0A%0A%0A%40app.cell%0Adef%20_(plot_amplitudes%2C%20plot_phases%2C%20plt%2C%20result_schw_fixed_vst%2C%20schw_modes)%3A%0A%20%20%20%20_fig%2C%20_axs%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(12%2C%205))%0A%20%20%20%20plot_amplitudes(result_schw_fixed_vst%2C%20fixed_modes%3Dschw_modes%2C%20ax%3D_axs%5B0%5D)%0A%20%20%20%20plot_phases(result_schw_fixed_vst%2C%20fixed_modes%3Dschw_modes%2C%20ax%3D_axs%5B1%5D%2C%20legend%3DFalse)%0A%20%20%20%20_fig.suptitle('Real-waveform%20fixed-frequency%20fit%20across%20%24t_0%24')%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Variable-mass%20fit%20(varying%20%24t_0%24)%0A%0A%20%20%20%20We%20use%20the%20VarMa%20(varying%20mass%20and%20spin)%20fitter%20across%20starting%20times%20with%20%60real%3DTrue%60%20to%20see%20how%20the%20recovered%20mass%20%24M%24%20evolves%20with%20%24t_0%24.%0A%20%20%20%20The%20spin%20will%20not%20be%20recovered%2C%20since%20we%20assumed%20that%20the%20merger%20is%20head-on%2C%20which%20produces%20a%20Schwarzschild%20black%20hole%20remnant.%0A%20%20%20%20At%20later%20starting%20times%20(further%20from%20the%20merger)%2C%20the%20distortions%20are%20smaller%20and%20%24M%24%20should%20converge%20to%20the%20true%20value.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(QNMFitVaryingStartingTime%2C%20h_schw%2C%20long_str_to_qnms_free%2C%20mo%2C%20np)%3A%0A%20%20%20%20schw_varma_t0_arr%20%3D%20np.linspace(0%2C%2050%2C%20num%3D51)%0A%20%20%20%20schw_varma_qnm_free%20%3D%20long_str_to_qnms_free('2.2.0_3.2.0')%0A%0A%20%20%20%20fitter_schw_varma_vst%20%3D%20QNMFitVaryingStartingTime(%0A%20%20%20%20%20%20%20%20h_schw%2C%20schw_varma_t0_arr%2C%20N_free%3D0%2C%0A%20%20%20%20%20%20%20%20qnm_fixed_list%3D%5B%5D%2C%0A%20%20%20%20%20%20%20%20qnm_free_list%3Dschw_varma_qnm_free%2C%0A%20%20%20%20%20%20%20%20var_M_a%3DTrue%2C%0A%20%20%20%20%20%20%20%20real%3DTrue%2C%0A%20%20%20%20%20%20%20%20load_pickle%3DFalse%2C%0A%20%20%20%20%20%20%20%20run_string_prefix%3D'schwarzschild_varma_example'%2C%0A%20%20%20%20%20%20%20%20save_results%3DFalse%2C%0A%20%20%20%20)%0A%20%20%20%20with%20mo.status.spinner(%22Running%20VarMa%20varying%20%24t_0%24%20fits...%22)%3A%0A%20%20%20%20%20%20%20%20fitter_schw_varma_vst.do_fits()%0A%20%20%20%20result_schw_varma_vst%20%3D%20fitter_schw_varma_vst.result_full%0A%20%20%20%20return%20result_schw_varma_vst%2C%20schw_varma_t0_arr%0A%0A%0A%40app.cell%0Adef%20_(Mf_true%2C%20np%2C%20plt%2C%20result_schw_varma_vst%2C%20schw_varma_t0_arr)%3A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(8%2C%204))%0A%20%20%20%20_M_arr%20%3D%20np.array(result_schw_varma_vst.Ma_dict%5B'M'%5D)%0A%20%20%20%20_ax.plot(schw_varma_t0_arr%2C%20_M_arr%2C%20'b-'%2C%20lw%3D2%2C%20label%3D'Recovered%20%24M%24')%0A%20%20%20%20_ax.axhline(Mf_true%2C%20color%3D'r'%2C%20ls%3D'--'%2C%20label%3Df'True%20%24M%20%3D%20%7BMf_true%7D%24')%0A%20%20%20%20_ax.set_xlabel(r'%24t_0%24')%0A%20%20%20%20_ax.set_ylabel(r'%24M%24')%0A%20%20%20%20_ax.set_title('Recovered%20mass%20across%20starting%20times%20(real%20waveform)')%0A%20%20%20%20_ax.legend()%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A