example summed waveform

import%20marimo%0A%0A__generated_with%20%3D%20%220.19.11%22%0Aapp%20%3D%20marimo.App(width%3D%22full%22)%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%20waveform%20with%20summed%20multipoles%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%20get_SXS_waveform_summed%0A%20%20%20%20from%20jaxqualin.qnmode%20import%20mode_list%0A%20%20%20%20from%20jaxqualin.fit%20import%20QNMFitVaryingStartingTime%0A%20%20%20%20import%20jaxqualin.fit%20as%20fit_module%0A%20%20%20%20from%20jaxqualin.plot%20import%20(plot_amplitudes%2C%20plot_phases%2C%20%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%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20QNMFitVaryingStartingTime%2C%0A%20%20%20%20%20%20%20%20get_SXS_waveform_summed%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)%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%20Make%20waveform%0A%0A%20%20%20%20We%20will%20use%20a%20waveform%20that%20is%20a%20superposition%20of%20multipoles%20(up%20to%20%60l_max%20%3D%204%60)%20of%20the%20%60SXS%3ABBH%3A0305%60%20simulation.%0A%20%20%20%20The%20observer%20is%20at%20the%20angular%20coordinate%20%24(%5Ciota%2C%20%5Cpsi)%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(get_SXS_waveform_summed%2C%20np)%3A%0A%20%20%20%20SXSnum%20%3D%20'0305'%0A%20%20%20%20iota%20%3D%20np.pi%2F3%20%23%20Cotesta's%20angle%0A%20%20%20%20psi%20%3D%20np.pi%2F2%0A%20%20%20%20h%2C%20Mf%2C%20af%20%3D%20get_SXS_waveform_summed(SXSnum%2C%20iota%2C%20psi%2C%20l_max%3D4%2C%20res%3D0%2C%20N_ext%3D2)%0A%20%20%20%20return%20Mf%2C%20SXSnum%2C%20af%2C%20h%2C%20iota%2C%20psi%0A%0A%0A%40app.cell%0Adef%20_(h%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.time%2C%20np.abs(h.hr))%0A%20%20%20%20_ax.set_xlabel('%24t%24')%0A%20%20%20%20_ax.set_ylabel('%24%7Ch_r%7C%24')%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%20Free%20QNMs%20(unfixed%20frequencies)%0A%0A%20%20%20%20This%20takes%20a%20bit%20of%20time%20because%20we%20need%20to%20fit%20with%20twice%20the%20number%20of%20modes%20to%20include%20both%20%24%5Cpm%20m%24%20modes.%20In%20future%20updates%20it%20might%20be%20possible%20to%20fit%20with%20a%20single%20mode%20with%20elliptical%20polarization.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(QNMFitVaryingStartingTime%2C%20SXSnum%2C%20h%2C%20iota%2C%20np%2C%20psi)%3A%0A%20%20%20%20_t0_arr%20%3D%20np.linspace(0%2C%2050%2C%20num%3D101)%20%20%23%20array%20of%20starting%20times%20to%20fit%20for%0A%20%20%20%20qnm_fixed_list%20%3D%20%5B%5D%20%20%23%20t0%20%3D%200%20is%20the%20peak%20of%20the%20strain%0A%20%20%20%20_run_string_prefix%20%3D%20f'SXS%7BSXSnum%7D_lm_2.2_iota_%7Biota%3A.7f%7D_psi_%7Bpsi%3A.7f%7D'%20%20%23%20list%20of%20QNMs%20with%20fixed%20frequencies%20in%20the%20fit%20model%0A%20%20%20%20_N_free%20%3D%206%20%20%23%20prefix%20of%20pickle%20file%20for%20saving%20the%20results%0A%20%20%20%20%23%20fitter%20object%0A%20%20%20%20fitter%20%3D%20QNMFitVaryingStartingTime(h%2C%20_t0_arr%2C%20N_free%3D_N_free%2C%20qnm_fixed_list%3Dqnm_fixed_list%2C%20load_pickle%3DFalse%2C%20run_string_prefix%3D_run_string_prefix)%20%20%23%20number%20of%20free%20modes%20to%20use%0A%20%20%20%20return%20(fitter%2C)%0A%0A%0A%40app.cell%0Adef%20_(fitter%2C%20mo)%3A%0A%20%20%20%20with%20mo.status.spinner(%22Running%20free-QNM%20fits...%22)%3A%0A%20%20%20%20%20%20%20%20fitter.do_fits()%0A%20%20%20%20mo.md(%22**Free-QNM%20fits%20complete.**%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(fitter)%3A%0A%20%20%20%20%23%20fitter%20results%20object%0A%20%20%20%20result%20%3D%20fitter.result_full%0A%20%20%20%20return%20(result%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%23%23%23%20Plotting%20the%20results%0A%20%20%20%20We%20will%20see%20some%20of%20the%20mirror%20modes%20(modes%20with%20negative%20%24m%24)%20because%20we%20summed%20all%20of%20the%20multipoles%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Mf%2C%20af%2C%20mode_list%2C%20plot_omega_free%2C%20plot_predicted_qnms%2C%20plt%2C%20result)%3A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots()%0A%20%20%20%20predicted_qnms%20%3D%20mode_list(%5B'2.2.0'%2C%20'2.2.1'%2C%20'3.2.0'%2C%20'3.3.0'%2C%20'2.-2.0'%2C%20'3.-3.0'%2C%20'4.-4.0'%2C%20'4.4.0'%5D%2C%20Mf%2C%20af)%0A%20%20%20%20%23%20mode%20locations%20to%20visualize%20on%20the%20plot%0A%20%20%20%20plot_omega_free(result%2C%20_ax)%0A%20%20%20%20plot_predicted_qnms(_ax%2C%20predicted_qnms)%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%20Fixed%20QNMs%20(fixed%20frequencies)%0A%20%20%20%20Let%20us%20assume%20that%20the%20%242%7B%2C%7D2%7B%2C%7D0%2C%203%7B%2C%7D3%7B%2C%7D0%2C%204%7B%2C%7D4%7B%2C%7D0%24%20modes%20and%20their%20mirror%20counter-parts%20exist%20in%20the%20waveform%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Mf%2C%20QNMFitVaryingStartingTime%2C%20SXSnum%2C%20af%2C%20h%2C%20iota%2C%20mode_list%2C%20np%2C%20psi)%3A%0A%20%20%20%20_t0_arr%20%3D%20np.linspace(0%2C%2050%2C%20num%3D101)%0A%20%20%20%20qnm_fixed_list_1%20%3D%20mode_list(%5B'2.2.0'%2C%20'2.-2.0'%2C%20'3.3.0'%2C%20'3.-3.0'%2C%20'4.4.0'%2C%20'4.-4.0'%2C%20'constant'%5D%2C%20Mf%2C%20af)%0A%20%20%20%20_run_string_prefix%20%3D%20f'SXS%7BSXSnum%7D_lm_2.2_iota_%7Biota%3A.7f%7D_psi_%7Bpsi%3A.7f%7D'%0A%20%20%20%20_N_free%20%3D%200%0A%20%20%20%20fitter_1%20%3D%20QNMFitVaryingStartingTime(h%2C%20_t0_arr%2C%20N_free%3D_N_free%2C%20qnm_fixed_list%3Dqnm_fixed_list_1%2C%20load_pickle%3DFalse%2C%20run_string_prefix%3D_run_string_prefix)%0A%20%20%20%20return%20fitter_1%2C%20qnm_fixed_list_1%0A%0A%0A%40app.cell%0Adef%20_(fitter_1%2C%20mo)%3A%0A%20%20%20%20with%20mo.status.spinner(%22Running%20fixed-QNM%20fits...%22)%3A%0A%20%20%20%20%20%20%20%20fitter_1.do_fits()%0A%20%20%20%20result_1%20%3D%20fitter_1.result_full%0A%20%20%20%20mo.md(%22**Fixed-QNM%20fits%20complete.**%22)%0A%20%20%20%20return%20(result_1%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%23%23%23%20Plotting%20the%20results%0A%20%20%20%20Because%20our%20waveform%20is%20non-precessing%2C%20by%20PT%20symmetry%2C%20the%20amplitudes%20of%20the%20mirror%20modes%20can%20be%20predicted%20from%20those%20of%20the%20prograde%20modes%2C%20i.e.%0A%0A%20%20%20%20%24%5Ctilde%7BA%7D_%7B%5Cell%7B%2C%7D-m%7B%2C%7Dn%7D%20%3D%20%5Cdfrac%7BS_%7B%5Cell%7B%2C%7D-m%7B%2C%7Dn%7D(%5Ciota%2C%20%5Cpsi)%7D%7BS_%7B%5Cell%7B%2C%7Dm%7B%2C%7Dn%7D%5E*(%5Ciota%2C%20%5Cpsi)%7D%5Ctilde%7BA%7D%5E*_%7B%5Cell%7B%2C%7Dm%7B%2C%7Dn%7D%20%2C%24%0A%0A%20%20%20%20Where%20%24%5Ctilde%7BA%7D%24%20denotes%20the%20complex%20amplitude%2C%20and%20%24S_%7B%5Cell%7B%2C%7Dm%7B%2C%7Dn%7D%24%20are%20the%20spin%20weighted%20(%24s%20%3D%20-2%24%20in%20this%20case)%20spheroidal%20harmonics.%0A%20%20%20%20We%20compute%20the%20predicted%20mirror%20mode%20amplitudes%20and%20plot%20them%20as%20dashed%20lines%20with%20colors%20corresponding%20to%20their%20prograde%20mode%20counter-parts.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20af%2C%0A%20%20%20%20iota%2C%0A%20%20%20%20plot_amplitudes%2C%0A%20%20%20%20plot_phases%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20psi%2C%0A%20%20%20%20qnm_fixed_list_1%2C%0A%20%20%20%20result_1%2C%0A)%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_1%2C%20fixed_modes%3Dqnm_fixed_list_1%2C%20ax%3D_axs%5B0%5D%2C%20plot_mirror_pred%3DTrue%2C%20iota%3Diota%2C%20psi%3Dpsi%2C%20af%3Daf)%0A%20%20%20%20plot_phases(result_1%2C%20fixed_modes%3Dqnm_fixed_list_1%2C%20ax%3D_axs%5B1%5D%2C%20legend%3DFalse%2C%20plot_mirror_pred%3DTrue%2C%20iota%3Diota%2C%20psi%3Dpsi%2C%20af%3Daf)%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%20Including%20a%20pair%20of%20prograde%20and%20mirror%20modes%20as%20the%20same%20mode%0A%20%20%20%20Given%20the%20relationship%20between%20the%20complex%20amplitudes%20of%20the%20prograde%20and%20mirror%20modes%20for%20non-precessing%20mergers%2C%20we%20can%20include%20both%20modes%20as%20one%20mode%20into%20our%20fit%20model.%0A%20%20%20%20This%20can%20be%20done%20with%20%60include_mirror%20%3D%20True%60.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Mf%2C%20QNMFitVaryingStartingTime%2C%20SXSnum%2C%20af%2C%20h%2C%20iota%2C%20mode_list%2C%20np%2C%20psi)%3A%0A%20%20%20%20_t0_arr%20%3D%20np.linspace(0%2C%2050%2C%20num%3D101)%0A%20%20%20%20qnm_fixed_list_2%20%3D%20mode_list(%5B'2.2.0'%2C%20'3.3.0'%2C%20'4.4.0'%2C%20'constant'%5D%2C%20Mf%2C%20af)%0A%20%20%20%20_run_string_prefix%20%3D%20f'SXS%7BSXSnum%7D_lm_2.2_iota_%7Biota%3A.7f%7D_psi_%7Bpsi%3A.7f%7D_incl_mirror'%0A%20%20%20%20_N_free%20%3D%200%0A%20%20%20%20fitter_2%20%3D%20QNMFitVaryingStartingTime(h%2C%20_t0_arr%2C%20N_free%3D_N_free%2C%20qnm_fixed_list%3Dqnm_fixed_list_2%2C%20load_pickle%3DFalse%2C%20run_string_prefix%3D_run_string_prefix%2C%20include_mirror%3DTrue%2C%20iota%3Diota%2C%20psi%3Dpsi)%0A%20%20%20%20return%20fitter_2%2C%20qnm_fixed_list_2%0A%0A%0A%40app.cell%0Adef%20_(fitter_2%2C%20mo)%3A%0A%20%20%20%20with%20mo.status.spinner(%22Running%20mirror-mode%20fits...%22)%3A%0A%20%20%20%20%20%20%20%20fitter_2.do_fits()%0A%20%20%20%20result_2%20%3D%20fitter_2.result_full%0A%20%20%20%20mo.md(%22**Mirror-mode%20fits%20complete.**%22)%0A%20%20%20%20return%20(result_2%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%23%23%23%20Plotting%20the%20results%0A%20%20%20%20Only%20the%20amplitude%20of%20the%20prograde%20modes%20are%20shown%2C%20but%20the%20mirror%20modes%20have%20been%20included%20in%20the%20fit%2C%20with%20amplitudes%20fixed%20by%20PT%20symmetry%20as%20explained%20above.%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%20qnm_fixed_list_2%2C%20result_2)%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_2%2C%20fixed_modes%3Dqnm_fixed_list_2%2C%20ax%3D_axs%5B0%5D%2C%20plot_mirror_pred%3DFalse)%0A%20%20%20%20plot_phases(result_2%2C%20fixed_modes%3Dqnm_fixed_list_2%2C%20ax%3D_axs%5B1%5D%2C%20legend%3DFalse%2C%20plot_mirror_pred%3DFalse)%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