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%20Custom%20QNM%20modes%20and%20parametric%20models%0A%0A%20%20%20%20This%20notebook%20demonstrates%20two%20features%20in%20%60jaxqualin%60%3A%0A%0A%20%20%20%201.%20**Custom%20fixed-frequency%20modes**%20--%20fit%20a%20waveform%20with%20user-specified%20complex%20frequencies.%0A%20%20%20%202.%20**Fully%20custom%20model**%20--%20define%20a%20QNM%20model%20as%20a%20function%20of%20arbitrary%20parameters.%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%20(%0A%20%20%20%20%20%20%20%20mode_list%2C%20custom_mode%2C%20custom_mode_list%2C%0A%20%20%20%20%20%20%20%20KerrModel%2C%20QNMModel%2C%20model_mode_free%2C%20model_mode%2C%0A%20%20%20%20)%0A%20%20%20%20from%20jaxqualin.fit%20import%20(%0A%20%20%20%20%20%20%20%20QNMFit%2C%20QNMFitModel%2C%20QNMFitVaryingStartingTime%2C%0A%20%20%20%20)%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%20QNMFitVaryingStartingTime%2C%0A%20%20%20%20%20%20%20%20QNMModel%2C%0A%20%20%20%20%20%20%20%20custom_mode%2C%0A%20%20%20%20%20%20%20%20delayed_QNM%2C%0A%20%20%20%20%20%20%20%20model_mode_free%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%20Part%201%3A%20Custom%20fixed-frequency%20QNMs%0A%0A%20%20%20%20If%20you%20want%20to%20use%20non-Kerr%20frequencies%2C%20you%20can%20create%20a%20%60custom_mode%60%0A%20%20%20%20objects%20with%20arbitrary%20%24%5Comega%24%20and%20optional%20labels%2C%20then%20use%20them%20directly%20in%20%60QNMFit%60%20or%0A%20%20%20%20%60QNMFitVaryingStartingTime%60.%0A%0A%20%20%20%20Here%20we%20use%20arbitrary%20frequencies%20that%20have%20nothing%20to%20do%20with%20Kerr%20BHs%3A%0A%20%20%20%20-%20%60%22fundamental%22%60%3A%20%20%24%5Comega_a%20%3D%200.50%20-%200.08i%24%0A%20%20%20%20-%20%60%22overtone%22%60%3A%20%20%24%5Comega_b%20%3D%200.30%20-%200.12i%24%0A%0A%20%20%20%20We%20then%20use%20%60delayed_QNM%60%20to%20build%20a%20toy%20waveform%20and%20%60QNMFitVaryingStartingTime%60%20to%20fit%20at%20multiple%20starting%20times%20with%20these%20custom%20modes.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(custom_mode%2C%20delayed_QNM%2C%20np%2C%20waveform)%3A%0A%20%20%20%20omega_1a%20%3D%200.50%20-%200.08j%0A%20%20%20%20omega_1b%20%3D%200.30%20-%200.12j%0A%0A%20%20%20%20custom_modes_1%20%3D%20%5B%0A%20%20%20%20%20%20%20%20custom_mode(omega_1a%2C%20label%3D%22fundamental%22)%2C%0A%20%20%20%20%20%20%20%20custom_mode(omega_1b%2C%20label%3D%22overtone%22)%2C%0A%20%20%20%20%5D%0A%0A%20%20%20%20A_phi_1%20%3D%20%5B(1.0%2C%200.0)%2C%20(3.0%2C%20np.pi%20%2F%202)%5D%0A%0A%20%20%20%20t_arr_1%20%3D%20np.linspace(0%2C%20120%2C%201000)%0A%20%20%20%20_h_arr%20%3D%20np.zeros(t_arr_1.shape%2C%20dtype%3Dnp.complex128)%0A%20%20%20%20for%20_i%2C%20(_mode%2C%20(_A%2C%20_phi))%20in%20enumerate(zip(custom_modes_1%2C%20A_phi_1))%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%20t_arr_1%2C%20_A%2C%20_phi%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%20%20%20%20h_1%20%3D%20waveform(t_arr_1%2C%20_h_arr%2C%20t_peak%3D0)%0A%20%20%20%20return%20custom_modes_1%2C%20h_1%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%20Fit%20with%20custom%20fixed%20modes%0A%0A%20%20%20%20We%20use%20%60QNMFitVaryingStartingTime%60%20with%20frequencies%20fixed%20to%20the%20custom%20modes%2C%20i.e.%20%60N_free%3D0%60%0A%20%20%20%20and%20%60qnm_fixed_list%60%20containing%20the%20custom%20modes.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20QNMFitVaryingStartingTime%2C%0A%20%20%20%20custom_modes_1%2C%0A%20%20%20%20h_1%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plot_amplitudes%2C%0A%20%20%20%20plot_phases%2C%0A%20%20%20%20plt%2C%0A)%3A%0A%20%20%20%20t0_arr_1%20%3D%20np.linspace(0%2C%2050%2C%20num%3D51)%0A%0A%20%20%20%20fitter_1%20%3D%20QNMFitVaryingStartingTime(%0A%20%20%20%20%20%20%20%20h_1%2C%20t0_arr_1%2C%20N_free%3D0%2C%0A%20%20%20%20%20%20%20%20qnm_fixed_list%3Dcustom_modes_1%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'custom_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(%22Fitting%20with%20custom%20fixed%20modes...%22)%3A%0A%20%20%20%20%20%20%20%20fitter_1.do_fits()%0A%0A%20%20%20%20result_1%20%3D%20fitter_1.result_full%0A%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%3Dcustom_modes_1%2C%20ax%3D_axs%5B0%5D)%0A%20%20%20%20plot_phases(result_1%2C%20fixed_modes%3Dcustom_modes_1%2C%20ax%3D_axs%5B1%5D%2C%20legend%3DFalse)%0A%20%20%20%20_fig.suptitle('Part%201%3A%20Custom%20fixed-omega%20modes')%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%0A%20%20%20%20return%20(result_1%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20np%2C%20result_1)%3A%0A%20%20%20%20_keys%20%3D%20list(result_1.A_dict.keys())%0A%20%20%20%20_vals%20%3D%20%5Bf%22%60%7Bk%7D%60%3A%20A%20%3D%20%7Bnp.array(v)%5B-1%5D%3A.4f%7D%22%20for%20k%2C%20v%20in%20result_1.A_dict.items()%5D%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20**Result%20dict%20keys%3A**%0A%0A%20%20%20%20%7Bchr(10).join('-%20'%20%2B%20v%20for%20v%20in%20_vals)%7D%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%20Part%202%3A%20Fitting%20a%20custom%20model%0A%0A%20%20%20%20You%20can%20also%20define%20a%20QNM%20model%20from%20scratch%2C%20where%20%24%5Comega%24%20is%20a%20function%20of%0A%20%20%20%20completely%20arbitrary%20parameters%20(not%20necessarily%20%24M%24%20and%20%24a%24).%0A%0A%20%20%20%20As%20a%20toy%20example%2C%20consider%20a%20model%20where%20each%20mode%20has%20the%20frequency%3A%0A%0A%20%20%20%20%24%24%5Comega%20%3D%20%5Calpha%20%2B%20i%5C%2C%5Cbeta%24%24%0A%0A%20%20%20%20where%20%24%5Calpha%24%20and%20%24%5Cbeta%24%20are%20the%20free%20parameters%0A%20%20%20%20to%20be%20fitted.%20Again%2C%20we%20use%20%60delayed_QNM%60%20to%20build%20a%20toy%20waveform%20and%0A%20%20%20%20%60QNMFitVaryingStartingTime%60%20to%20fit%20at%20multiple%20starting%20times.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(QNMModel)%3A%0A%20%20%20%20class%20SimpleFreqModel(QNMModel)%3A%0A%20%20%20%20%20%20%20%20param_names%20%3D%20%5B%22alpha%22%2C%20%22beta%22%5D%0A%0A%20%20%20%20%20%20%20%20def%20compute_omega(self%2C%20lmnx%2C%20alpha%2C%20beta%2C%20**kwargs)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20alpha%20%2B%201j%20*%20beta%0A%0A%20%20%20%20%20%20%20%20def%20param_bounds(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22alpha%22%3A%20(0%2C%2010)%2C%20%22beta%22%3A%20(-5%2C%200)%7D%0A%0A%20%20%20%20return%20(SimpleFreqModel%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(SimpleFreqModel%2C%20custom_mode%2C%20delayed_QNM%2C%20np%2C%20waveform)%3A%0A%20%20%20%20alpha_true%20%3D%200.5%0A%20%20%20%20beta_true%20%3D%20-0.08%0A%0A%20%20%20%20sfm_cm%20%3D%20custom_mode(alpha_true%20%2B%201j%20*%20beta_true%2C%20label%3D%22toy_mode%22)%0A%0A%20%20%20%20_A_true%2C%20_phi_true%20%3D%202.0%2C%200.3%0A%20%20%20%20t_arr_3%20%3D%20np.linspace(0%2C%20120%2C%201000)%0A%20%20%20%20_h_arr%20%3D%20delayed_QNM(%0A%20%20%20%20%20%20%20%20sfm_cm%2C%20t_arr_3%2C%20_A_true%2C%20_phi_true%2C%0A%20%20%20%20%20%20%20%20A_delay%3D3%2C%20A_sig%3D5%2C%20phi_sig%3D3%2C%0A%20%20%20%20)%0A%20%20%20%20h_3%20%3D%20waveform(t_arr_3%2C%20_h_arr%2C%20t_peak%3D0)%0A%0A%20%20%20%20sfm%20%3D%20SimpleFreqModel()%0A%20%20%20%20return%20alpha_true%2C%20beta_true%2C%20h_3%2C%20sfm%0A%0A%0A%40app.cell%0Adef%20_(QNMFitVaryingStartingTime%2C%20h_3%2C%20mo%2C%20model_mode_free%2C%20np%2C%20sfm)%3A%0A%20%20%20%20sfm_modes%20%3D%20%5Bmodel_mode_free(%5B%5B2%2C%202%2C%200%5D%5D%2C%20model%3Dsfm)%5D%0A%0A%20%20%20%20t0_arr_3%20%3D%20np.linspace(0%2C%2050%2C%20num%3D51)%0A%0A%20%20%20%20fitter_3%20%3D%20QNMFitVaryingStartingTime(%0A%20%20%20%20%20%20%20%20h_3%2C%20t0_arr_3%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%3Dsfm_modes%2C%0A%20%20%20%20%20%20%20%20var_M_a%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'custom_simple_freq'%2C%0A%20%20%20%20%20%20%20%20save_results%3DFalse%2C%0A%20%20%20%20%20%20%20%20model%3Dsfm%2C%0A%20%20%20%20%20%20%20%20model_params_guess%3D%7B%22alpha%22%3A%200.4%2C%20%22beta%22%3A%20-0.1%7D%2C%0A%20%20%20%20)%0A%20%20%20%20with%20mo.status.spinner(%22Fitting%20with%20fully%20custom%20model...%22)%3A%0A%20%20%20%20%20%20%20%20fitter_3.do_fits()%0A%0A%20%20%20%20result_3%20%3D%20fitter_3.result_full%0A%20%20%20%20return%20result_3%2C%20t0_arr_3%0A%0A%0A%40app.cell%0Adef%20_(alpha_true%2C%20beta_true%2C%20np%2C%20plt%2C%20result_3%2C%20t0_arr_3)%3A%0A%20%20%20%20_fig%2C%20_axs%20%3D%20plt.subplots(1%2C%203%2C%20figsize%3D(15%2C%204))%0A%0A%20%20%20%20_alpha_arr%20%3D%20np.array(result_3.model_params_dict%5B'alpha'%5D)%0A%20%20%20%20_axs%5B0%5D.plot(t0_arr_3%2C%20_alpha_arr%2C%20'b-'%2C%20lw%3D2)%0A%20%20%20%20_axs%5B0%5D.axhline(alpha_true%2C%20color%3D'r'%2C%20ls%3D'--'%2C%20label%3Drf'True%20%24%5Calpha%20%3D%20%7Balpha_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'%24%5Calpha%24')%0A%20%20%20%20_axs%5B0%5D.legend()%0A%0A%20%20%20%20_beta_arr%20%3D%20np.array(result_3.model_params_dict%5B'beta'%5D)%0A%20%20%20%20_axs%5B1%5D.plot(t0_arr_3%2C%20_beta_arr%2C%20'b-'%2C%20lw%3D2)%0A%20%20%20%20_axs%5B1%5D.axhline(beta_true%2C%20color%3D'r'%2C%20ls%3D'--'%2C%20label%3Drf'True%20%24%5Cbeta%20%3D%20%7Bbeta_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'%24%5Cbeta%24')%0A%20%20%20%20_axs%5B1%5D.legend()%0A%0A%20%20%20%20_mismatch%20%3D%20np.array(result_3.mismatch_arr)%0A%20%20%20%20_axs%5B2%5D.semilogy(t0_arr_3%2C%20_mismatch%2C%20'k-'%2C%20lw%3D2)%0A%20%20%20%20_axs%5B2%5D.set_xlabel(r'%24t_0%24')%0A%20%20%20%20_axs%5B2%5D.set_ylabel('Mismatch')%0A%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
072a5545bb05190aa43a3607f6d7aaa1