example custom 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%20your%20own%20waveform%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%2C%20mode_list%0A%20%20%20%20from%20jaxqualin.fit%20import%20QNMFitVaryingStartingTime%0A%20%20%20%20from%20jaxqualin.selection%20import%20ModeSearchAllFreeVaryingN%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%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_mode_searcher_results)%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%20ModeSearchAllFreeVaryingN%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%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_mode_searcher_results%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%23%20Make%20waveform%0A%0A%20%20%20%20As%20example%2C%20we%20will%20make%20a%20waveform%20that%20contains%20three%20QNMs%20that%20are%20distorted%20close%20to%20the%20merger%20peak.%0A%20%20%20%20You%20can%20replace%20this%20with%20whatever%20time%20domain%20waveform%20you%20would%20like%20to%20fit.%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)%3A%0A%20%20%20%20Mf%20%3D%201%0A%20%20%20%20af%20%3D%200.7%0A%20%20%20%20modes%20%3D%20mode_list(%5B'2.2.0'%2C%20'2.2.1'%2C%20'3.2.0'%5D%2C%20Mf%2C%20af)%0A%20%20%20%20A_phi_dict%20%3D%20%7B'2.2.0'%3A%20dict(A%3D1.0%2C%20phi%3D0.0)%2C%20'2.2.1'%3A%20dict(A%3D3.0%2C%20phi%3Dnp.pi%20%2F%202)%2C%20'3.2.0'%3A%20dict(A%3D0.01%2C%20phi%3Dnp.pi)%7D%0A%20%20%20%20t_arr%20%3D%20np.linspace(0%2C%20120%2C%201000)%0A%20%20%20%20h_arr%20%3D%20np.zeros(t_arr.shape%2C%20dtype%3Dnp.complex128)%0A%20%20%20%20for%20i%2C%20_mode%20in%20enumerate(modes)%3A%0A%20%20%20%20%20%20%20%20if%20i%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20A_delay%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20A_sig%20%3D%2010%0A%20%20%20%20%20%20%20%20%20%20%20%20phi_sig%20%3D%205%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20A_delay%20%3D%205%0A%20%20%20%20%20%20%20%20%20%20%20%20A_sig%20%3D%202%0A%20%20%20%20%20%20%20%20%20%20%20%20phi_sig%20%3D%202%0A%20%20%20%20%20%20%20%20h_arr%20%3D%20h_arr%20%2B%20delayed_QNM(_mode%2C%20t_arr%2C%20A_phi_dict%5B_mode.string()%5D%5B'A'%5D%2C%20A_phi_dict%5B_mode.string()%5D%5B'phi'%5D%2C%20A_delay%3DA_delay%2C%20A_sig%3DA_sig%2C%20phi_sig%3Dphi_sig)%0A%20%20%20%20return%20Mf%2C%20af%2C%20h_arr%2C%20t_arr%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%20object%0A%0A%20%20%20%20You%20can%20pass%20in%20whatever%20time%20domain%20waveform.%0A%20%20%20%20If%20you%20know%20that%20the%20peak%20strain%20is%20located%20at%20t%20%3D%200%2C%20then%20pass%20in%20%60t_peak%20%3D%200%60%2C%20or%20else%20the%20waveform%20object%20detects%20the%20peak%20on%20its%20own.%0A%20%20%20%20If%20%60t_peak%60%20is%20not%20passed%2C%20the%20first%20%60remove_num%20%3D%20500%60%20data%20points%20are%20removed%20by%20default%20because%20BBH%20merger%20simulations%20often%20contain%20junk%20radiation%20in%20the%20beginning.%0A%20%20%20%20You%20can%20set%20%60remove_num%20%3D%200%60%20if%20you%20do%20not%20want%20to%20remove%20any%20data%20points.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h_arr%2C%20t_arr%2C%20waveform)%3A%0A%20%20%20%20h%20%3D%20waveform(t_arr%2C%20h_arr%2C%20t_peak%20%3D%200)%0A%20%20%20%20return%20(h%2C)%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%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(QNMFitVaryingStartingTime%2C%20h%2C%20np)%3A%0A%20%20%20%20t0_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'custom_example_lm_2.2'%20%20%23%20list%20of%20QNMs%20with%20fixed%20frequencies%20in%20the%20fit%20model%0A%20%20%20%20_N_free%20%3D%203%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%20t0_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%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'%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%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Mf%2C%20QNMFitVaryingStartingTime%2C%20af%2C%20h%2C%20mode_list%2C%20np)%3A%0A%20%20%20%20t0_arr_1%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.1'%2C%20'3.2.0'%5D%2C%20Mf%2C%20af)%0A%20%20%20%20_run_string_prefix%20%3D%20f'custom_example_lm_2.2'%0A%20%20%20%20_N_free%20%3D%200%0A%20%20%20%20fitter_1%20%3D%20QNMFitVaryingStartingTime(h%2C%20t0_arr_1%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%2C%20t0_arr_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%0Adef%20_(plot_amplitudes%2C%20plot_phases%2C%20plt%2C%20qnm_fixed_list_1%2C%20result_1)%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)%0A%20%20%20%20plot_phases(result_1%2C%20fixed_modes%3Dqnm_fixed_list_1%2C%20ax%3D_axs%5B1%5D%2C%20legend%3DFalse)%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%20Mode%20search%0A%0A%20%20%20%20With%20%60BBH_potential_modes%20%3D%20True%60%20(default)%2C%20the%20mode%20searcher%20uses%20a%20list%20of%20modes%20that%20are%20expected%20to%20be%20seen%20in%20the%20waveform%20in%20a%20binary-black-hole%20merger%2C%20which%20depends%20on%20the%20waveform%20harmonic%20in%20question.%20For%20a%20custom%20waveform%2C%20the%20harmonic%20can%20be%20specified%20by%20%60h.set_lm(l%2C%20m)%60.%20The%20modesearcher%20will%20include%20all%20the%20overtones%20in%20that%20harmonic%2C%20spheroidal%20mixing%20modes%20with%20the%20same%20%60m%60%20but%20another%20%60l%60%2C%20and%20quadratic%20modes%20with%20%60m_1%20%2B%20m_2%20%3D%20m%60.%20mode%20mixing%20from%20other%20harmonics%20(e.g.%20due%20to%20not%20working%20in%20the%20superrest%20BMS%20frame)%20can%20be%20included%20by%20specifying%20%60relevant_lm_list%60.%20Consult%20the%20source%20code%20of%20the%20%60potential_modes%60%20function%20in%20%60jaxqualin.qnmode%60%20for%20more%20details.%0A%0A%20%20%20%20Additional%20custom%20modes%20can%20be%20included%20with%20the%20%60potential_modes_custom%60%20keyword%20argument.%20If%20%60BBH_potential_modes%20%3D%20False%60%20then%20these%20will%20be%20the%20only%20potential%20modes%20that%20the%20mode%20searcher%20will%20try%20to%20find.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Mf%2C%20ModeSearchAllFreeVaryingN%2C%20af%2C%20h%2C%20mode_list%2C%20t0_arr_1)%3A%0A%20%20%20%20h.set_lm(2%2C%202)%0A%20%20%20%20relevant_lm_list%20%3D%20%5B(2%2C%202)%5D%0A%20%20%20%20potential_modes_custom%20%3D%20qnm_fixed_list_2%20%3D%20mode_list(%5B'-2.2.0x2.2.0'%5D%2C%20Mf%2C%20af)%0A%20%20%20%20_run_string_prefix%20%3D%20f'custom_example_lm_2.2'%0A%20%20%20%20mode_searcher%20%3D%20ModeSearchAllFreeVaryingN(h%2C%20Mf%2C%20af%2C%20relevant_lm_list%3Drelevant_lm_list%2C%20N_list%3D%5B3%5D%2C%20t0_arr%3Dt0_arr_1%2C%20run_string_prefix%3D_run_string_prefix%2C%20BBH_potential_modes%3DTrue%2C%20potential_modes_custom%3Dpotential_modes_custom)%0A%20%20%20%20return%20(mode_searcher%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20mode_searcher)%3A%0A%20%20%20%20with%20mo.status.spinner(%22Running%20mode%20searches...%22)%3A%0A%20%20%20%20%20%20%20%20mode_searcher.do_mode_searches()%0A%20%20%20%20mo.md(%22**Mode%20searches%20complete.**%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mode_searcher%2C%20plot_mode_searcher_results%2C%20plt)%3A%0A%20%20%20%20plot_mode_searcher_results(mode_searcher)%0A%20%20%20%20plt.gcf()%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