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%20Kerr%20mass%20and%20spin%20fitting%0A%0A%20%20%20%20In%20this%20example%2C%20we%20construct%20a%20toy%20waveform%20from%20Kerr%20quasi-normal%20modes%20(QNMs)%20and%20fit%20it%20using%20the%20VarMa%20(Variable%20mass%20and%20spin)%20fitter%20in%20%60jaxqualin%60%2C%20which%20simultaneously%20recovers%20the%20black%20hole%20mass%20%24M%24%20and%20spin%20%24a%24%20along%20with%20mode%20amplitudes%20and%20phases.%0A%0A%20%20%20%20A%20Kerr%20black%20hole%20has%20non-zero%20spin%20(%24a%20%5Cneq%200%24)%2C%20so%20the%20gravitational%20waveform%20is%20complex%3A%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%20QNMFitVarMa%2C%20QNMFitVaryingStartingTime%0A%20%20%20%20from%20jaxqualin.plot%20import%20plot_amplitudes%2C%20plot_phases%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%20QNMFitVarMa%2C%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_phases%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%20Kerr%20waveform%0A%0A%20%20%20%20We%20construct%20a%20waveform%20with%20three%20Kerr%20QNMs%20(%24a%20%3D%200.7%24)%3A%20the%20%24(2%2C2%2C0)%24%20fundamental%2C%20the%20%24(2%2C2%2C1)%24%20first%20overtone%2C%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.7%0A%20%20%20%20kerr_modes%20%3D%20mode_list(%5B'2.2.0'%2C%20'2.2.1'%2C%20'3.2.0'%5D%2C%20Mf_true%2C%20af_true)%0A%0A%20%20%20%20kerr_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'2.2.1'%3A%20dict(A%3D3.0%2C%20phi%3Dnp.pi%20%2F%202)%2C%0A%20%20%20%20%20%20%20%20'3.2.0'%3A%20dict(A%3D1e-2%2C%20phi%3Dnp.pi)%2C%0A%20%20%20%20%7D%0A%0A%20%20%20%20kerr_t_arr%20%3D%20np.linspace(0%2C%20120%2C%201000)%0A%20%20%20%20_h_arr%20%3D%20np.zeros(kerr_t_arr.shape%2C%20dtype%3Dnp.complex128)%0A%20%20%20%20for%20_i%2C%20_mode%20in%20enumerate(kerr_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%20kerr_t_arr%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20kerr_A_phi_dict%5B_mode.string()%5D%5B'A'%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20kerr_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%20h_kerr%20%3D%20waveform(kerr_t_arr%2C%20_h_arr%2C%20t_peak%3D0)%0A%20%20%20%20return%20Mf_true%2C%20af_true%2C%20h_kerr%0A%0A%0A%40app.cell%0Adef%20_(h_kerr%2C%20np%2C%20plt)%3A%0A%20%20%20%20_fig%2C%20_axs%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(12%2C%204))%0A%20%20%20%20_axs%5B0%5D.semilogy(h_kerr.time%2C%20np.abs(h_kerr.hr))%0A%20%20%20%20_axs%5B0%5D.set_xlabel(r'%24t%24')%0A%20%20%20%20_axs%5B0%5D.set_ylabel(r'%24%7Ch_r%7C%24')%0A%20%20%20%20_axs%5B0%5D.set_title('Real%20part')%0A%20%20%20%20_axs%5B1%5D.semilogy(h_kerr.time%2C%20np.abs(h_kerr.hi))%0A%20%20%20%20_axs%5B1%5D.set_xlabel(r'%24t%24')%0A%20%20%20%20_axs%5B1%5D.set_ylabel(r'%24%7Ch_i%7C%24')%0A%20%20%20%20_axs%5B1%5D.set_title('Imaginary%20part')%0A%20%20%20%20_fig.suptitle('Kerr%20toy%20waveform')%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%20Case%201%3A%20Single%20starting%20time%20VarMa%20fit%0A%0A%20%20%20%20We%20use%20%60QNMFitVarMa%60%20with%20%60real%3DFalse%60%20(the%20default)%20to%20fit%20the%20complex%20waveform%20at%20a%20single%20starting%20time%20%24t_0%20%3D%205%24.%0A%20%20%20%20The%20fitter%20recovers%20the%20mode%20amplitudes%2C%20phases%2C%20the%20black%20hole%20mass%20%24M%24%2C%20and%20spin%20%24a%24.%0A%0A%20%20%20%20The%20parameter%20vector%20%60popt%60%20has%20the%20structure%3A%20%24%5BA_0%2C%20%5Cphi_0%2C%20A_1%2C%20%5Cphi_1%2C%20%5Cldots%2C%20M%2C%20a%5D%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Mf_true%2C%20QNMFitVarMa%2C%20af_true%2C%20h_kerr%2C%20long_str_to_qnms_free%2C%20mo%2C%20np)%3A%0A%20%20%20%20kerr_qnm_free%20%3D%20long_str_to_qnms_free('2.2.0_2.2.1_3.2.0')%0A%0A%20%20%20%20fitter_kerr_varma%20%3D%20QNMFitVarMa(%0A%20%20%20%20%20%20%20%20h_kerr%2C%20t0%3D5.0%2C%20qnm_free_list%3Dkerr_qnm_free%2C%0A%20%20%20%20%20%20%20%20real%3DFalse%2C%0A%20%20%20%20%20%20%20%20guess_free%3D%5B1.0%2C%200.5%5D%2C%0A%20%20%20%20%20%20%20%20guess_M_a%3D%5BMf_true%20*%200.9%2C%20af_true%20*%200.9%5D%2C%0A%20%20%20%20)%0A%20%20%20%20fitter_kerr_varma.do_fit()%0A%0A%20%20%20%20_popt_kerr_varma%20%3D%20np.array(fitter_kerr_varma.popt)%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20**VarMa%20fit%20results%20(single%20%24t_0%20%3D%205%24)%3A**%0A%0A%20%20%20%20-%20Mismatch%3A%20%60%7Bfitter_kerr_varma.mismatch%3A.2e%7D%60%0A%20%20%20%20-%20Recovered%20%24M%24%3A%20%60%7B_popt_kerr_varma%5B-2%5D%3A.6f%7D%60%20(true%3A%20%60%7BMf_true%7D%60)%0A%20%20%20%20-%20Recovered%20%24a%24%3A%20%60%7B_popt_kerr_varma%5B-1%5D%3A.6f%7D%60%20(true%3A%20%60%7Baf_true%7D%60)%0A%20%20%20%20%22%22%22)%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%20Case%202%3A%20Varying%20starting%20time%20VarMa%20fit%0A%0A%20%20%20%20We%20now%20fit%20across%20a%20range%20of%20starting%20times%20%24t_0%24%20to%20see%20how%20the%20recovered%20mass%20%24M%24%20and%20spin%20%24a%24%20evolve.%0A%20%20%20%20At%20later%20starting%20times%20(further%20from%20the%20merger)%2C%20the%20distortions%20are%20smaller%20and%20the%20recovered%20parameters%20should%20converge%20to%20the%20true%20values.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(QNMFitVaryingStartingTime%2C%20h_kerr%2C%20long_str_to_qnms_free%2C%20mo%2C%20np)%3A%0A%20%20%20%20kerr_varma_t0_arr%20%3D%20np.linspace(0%2C%2050%2C%20num%3D51)%0A%20%20%20%20kerr_varma_qnm_free%20%3D%20long_str_to_qnms_free('2.2.0_2.2.1_3.2.0')%0A%0A%20%20%20%20fitter_kerr_varma_vst%20%3D%20QNMFitVaryingStartingTime(%0A%20%20%20%20%20%20%20%20h_kerr%2C%20kerr_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%3Dkerr_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%3DFalse%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'kerr_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_kerr_varma_vst.do_fits()%0A%20%20%20%20result_kerr_varma_vst%20%3D%20fitter_kerr_varma_vst.result_full%0A%20%20%20%20return%20kerr_varma_t0_arr%2C%20result_kerr_varma_vst%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%23%20Recovered%20mass%20and%20spin%20across%20starting%20times%0A%0A%20%20%20%20The%20dashed%20red%20lines%20indicate%20the%20true%20values%20%24M%20%3D%201.0%24%20and%20%24a%20%3D%200.7%24.%0A%20%20%20%20At%20early%20starting%20times%2C%20the%20distortions%20from%20%60delayed_QNM%60%20cause%20deviations%20from%20the%20true%20values.%0A%20%20%20%20At%20later%20starting%20times%2C%20the%20fit%20converges%20as%20the%20waveform%20approaches%20a%20pure%20QNM%20sum.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Mf_true%2C%20af_true%2C%20kerr_varma_t0_arr%2C%20np%2C%20plt%2C%20result_kerr_varma_vst)%3A%0A%20%20%20%20_fig%2C%20_axs%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(12%2C%204))%0A%0A%20%20%20%20_M_arr%20%3D%20np.array(result_kerr_varma_vst.Ma_dict%5B'M'%5D)%0A%20%20%20%20_axs%5B0%5D.plot(kerr_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_axs%5B0%5D.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_axs%5B0%5D.set_xlabel(r'%24t_0%24')%0A%20%20%20%20_axs%5B0%5D.set_ylabel(r'%24M%24')%0A%20%20%20%20_axs%5B0%5D.legend()%0A%0A%20%20%20%20_a_arr%20%3D%20np.array(result_kerr_varma_vst.Ma_dict%5B'a'%5D)%0A%20%20%20%20_axs%5B1%5D.plot(kerr_varma_t0_arr%2C%20_a_arr%2C%20'b-'%2C%20lw%3D2%2C%20label%3D'Recovered%20%24a%24')%0A%20%20%20%20_axs%5B1%5D.axhline(af_true%2C%20color%3D'r'%2C%20ls%3D'--'%2C%20label%3Df'True%20%24a%20%3D%20%7Baf_true%7D%24')%0A%20%20%20%20_axs%5B1%5D.set_xlabel(r'%24t_0%24')%0A%20%20%20%20_axs%5B1%5D.set_ylabel(r'%24a%24')%0A%20%20%20%20_axs%5B1%5D.legend()%0A%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%23%20Recovered%20amplitudes%20and%20phases%0A%0A%20%20%20%20We%20can%20also%20look%20at%20how%20the%20mode%20amplitudes%20and%20phases%20evolve%20across%20starting%20times.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(plot_amplitudes%2C%20plot_phases%2C%20plt%2C%20result_kerr_varma_vst)%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_kerr_varma_vst%2C%20ax%3D_axs%5B0%5D)%0A%20%20%20%20plot_phases(result_kerr_varma_vst%2C%20ax%3D_axs%5B1%5D%2C%20legend%3DFalse)%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
557facbf136b02ea169f5a4a7034c864