jaxqualin

https://mhycheung.github.io/jaxqualin/
Admin
jaxqualin is a Python package for extracting quasinormal modes from black-hole ringdown simulations.
Tokens:56,127
Snippets:308
Trust Score:9.7
Update:3 weeks ago
Show doc for...
Tokens:
Raw
### Fixed QNMs fit setup

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_summed_waveform

Sets up a QNMFitVaryingStartingTime fitter with fixed QNM frequencies (2,2,0, 3,3,0, 4,4,0 and mirror modes) and a constant mode. Uses a varying starting time array and disables pickle loading.

```python
t0_arr = np.linspace(0, 50, num=101)
    qnm_fixed_list_1 = mode_list(['2.2.0', '2.-2.0', '3.3.0', '3.-3.0', '4.4.0', '4.-4.0', 'constant'], Mf, af)
    run_string_prefix = f'SXS{SXSnum}_lm_2.2_iota_{iota:.7f}_psi_{psi:.7f}'
    N_free = 0
    fitter_1 = QNMFitVaryingStartingTime(h, t0_arr, N_free=N_free, qnm_fixed_list=qnm_fixed_list_1, load_pickle=False, run_string_prefix=run_string_prefix)
    return fitter_1, qnm_fixed_list_1
```

--------------------------------

### Install jaxqualin with pip

Source: https://mhycheung.github.io/jaxqualin

Install the jaxqualin package from PyPI using pip.

```bash
pip install jaxqualin

```

--------------------------------

### QNMFitVaryingStartingTime.get_mirror_ratio_list

Source: https://mhycheung.github.io/jaxqualin/api/fit

Gets the ratios between the prograde and mirror modes from iota and psi.

```APIDOC
## QNMFitVaryingStartingTime.get_mirror_ratio_list

### Description
Gets the ratios between the prograde and mirror modes from iota and psi.

### Method
Method

### Returns
- **List[float]** - list of ratios between prograde and mirror mode amplitudes.
```

--------------------------------

### QNMFitVaryingStartingTime.__init__

Source: https://mhycheung.github.io/jaxqualin/api/fit

Initializes a QNMFitVaryingStartingTime object with parameters for configuring the fit, including waveform, starting times, mode lists, and various fitting options.

```APIDOC
## QNMFitVaryingStartingTime.__init__

### Description
Initializes a QNMFitVaryingStartingTime object with parameters for configuring the fit, including waveform, starting times, mode lists, and various fitting options.

### Method
Constructor

### Parameters
#### Required Parameters
- **h** (waveform) - Required - waveform object to be fitted.
- **t0_arr** (ndarray) - Required - array of starting times for fitting.

#### Optional Parameters
- **N_free** (int) - Optional - number of frequency-free QNMs to include in the model. These modes are completely free, i.e. their mode numbers are not fixed like those in `qnm_free_list`. Default: `0`
- **qnm_fixed_list** (List[mode]) - Optional - list of fixed-frequency QNMs included in the model. Default: `[]`
- **qnm_free_list** (List[mode_free]) - Optional - list of free-frequency QNMs of fixed mode numbers to include in the model, only used for fitting `M` and `a` when `var_M_a = True`. Default: `[]`
- **var_M_a** (bool) - Optional - fit for the mass and spin of the black hole. Warning: Not tested yet. Default: `False`
- **real** (bool) - Optional - whether to fit a real-valued waveform. Default: `False`
- **run_string_prefix** (str) - Optional - prefix of the run name for dumping the `pickle` file. Default: `'Default'`
- **params0** (Optional[ndarray]) - Optional - initial guess for the fit parameters, at least for the earliest `t0` fit. Default: `None`
- **max_nfev** (int) - Optional - maximum number of function evaluations for the fit. Default: `DEFAULT_MAX_NFEV`
- **sequential_guess** (bool) - Optional - whether to use the previous fit as the initial guess for the next fit. Default: `True`
- **load_pickle** (bool) - Optional - whether to load the `pickle` file if it exists. Default: `True`
- **fit_save_prefix** (str) - Optional - prefix of the path to save the `pickle` file. Default: `FIT_SAVE_PATH`
- **nonconvergence_cut** (bool) - Optional - whether to cut the nonconverged fits. Default: `False`
- **A_bound** (float) - Optional - maximum value of the amplitude. Default: `inf`
- **fit_kwargs** (Dict) - Optional - keyword arguments for curve fitting. Default: `{}`
- **initial_num** (int) - Optional - number of initial guesses to use for the first starting time for frequency-free fits. Default: `1`
- **random_initial** (bool) - Optional - whether to generate random initial guesses for the first starting time for frequency-free fits. Default: `False`
- **initial_dict** (Dict) - Optional - key word arguments for `make_initial_guess` method. Default: `{}`
- **A_guess_relative** (bool) - Optional - whether to multiply the initial guess of the amplitude by the peak strain of the waveform. Default: `True`
- **set_seed** (int) - Optional - random seed for generating the initial guesses. Default: `DEFAULT_SEED`
- **weighted** (bool) - Optional - whether to perform a weighted fit. Default: `False`
- **double_skip** (bool) - Optional - whether to skip the next `2^n` `t0` fits when a fit does not converge, where `n` is the number of times the fit did not converge consecutively. Default: `True`
- **include_mirror** (bool) - Optional - whether to include the mirror modes, for fitting waveforms with both waveform polarizations. Default: `False`
- **iota** (Optional[float]) - Optional - inclination angle of the source. Default: `None`
- **psi** (Optional[float]) - Optional - polarization angle of the source. Default: `None`
- **mirror_ignore_phase** (bool) - Optional - whether to ignore the phase difference between the prograde and mirror modes. Default: `True`
- **skip_i_init** (int) - Optional - number of `t0` fits to skip for the first time a nonconvergent fit occured. Default: `1`
- **save_results** (bool) - Optional - whether to save the results. Default: `True`
- **fit_config** (Optional[FitConfig]) - Optional - optional FitConfig dataclass. If provided, overrides the individual max_nfev, weighted, real, include_mirror, iota, and psi parameters. Default: `None`
- **model** (optional) - Optional - optional QNMModel instance for custom parametric models. Default: `None`
- **model_params_guess** (optional) - Optional - dict of initial guesses for model params. Default: `None`
- **model_params_bounds** (optional) - Optional - dict overriding default model param bounds. Default: `None`
```

--------------------------------

### initial_guesses()

Source: https://mhycheung.github.io/jaxqualin/references

Generates initial guesses for the first starting time fit and returns the best guess index, the list of QNMFit objects, and the list of guess parameters.

```APIDOC
## initial_guesses()

### Description
Generate initial guesses for the first `t0` fit.

### Returns
- **best_guess_index** (`int`) - index of the best initial guess.
- **qnm_fit_list** (`List[QNMFit]`) - list of `QNMFit` objects for the initial guesses.
- **guess_list** (`List[ndarray]`) - list of initial guess parameters used.
```

--------------------------------

### Custom waveform example with ModeSearchAllFreeVaryingN

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_custom_waveform

Custom waveform example: sets a specific (l,m) mode, defines a custom potential mode list, and runs mode searches with ModeSearchAllFreeVaryingN. The second cell executes the searches with a spinner, the third plots results, and the final cell runs the app.

```python
h.set_lm(2, 2)
relevant_lm_list = [(2, 2)]
potential_modes_custom = qnm_fixed_list_2 = mode_list(['-2.2.0x2.2.0'], Mf, af)
_run_string_prefix = f'custom_example_lm_2.2'
mode_searcher = ModeSearchAllFreeVaryingN(h, Mf, af, relevant_lm_list=relevant_lm_list, N_list=[3], t0_arr=t0_arr_1, run_string_prefix=_run_string_prefix, BBH_potential_modes=True, potential_modes_custom=potential_modes_custom)
return (mode_searcher,)
```

```python
@app.cell
def _(mo, mode_searcher):
    with mo.status.spinner("Running mode searches..."):
        mode_searcher.do_mode_searches()
    mo.md("**Mode searches complete.**")
    return
```

```python
@app.cell
def _(mode_searcher, plot_mode_searcher_results, plt):
    plot_mode_searcher_results(mode_searcher)
    plt.gcf()
    return
```

```python
if __name__ == "__main__":
    app.run()
```

--------------------------------

### Free-QNM fit with varying starting time

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_custom_waveform

Workflow for fitting a custom waveform with free QNMs (N_free=3) using QNMFitVaryingStartingTime. It creates a fitter with varying starting times, runs the fits, retrieves the full result, and plots the free-mode frequencies against predicted QNM locations for modes 2.2.0, 2.2.1, and 3.2.0.

```python
@app.cell
def _(QNMFitVaryingStartingTime, h, np):
    t0_arr = np.linspace(0, 50, num=101)  # array of starting times to fit for
    qnm_fixed_list = []  # t0 = 0 is the peak of the strain
    _run_string_prefix = f'custom_example_lm_2.2'  # list of QNMs with fixed frequencies in the fit model
    _N_free = 3  # prefix of pickle file for saving the results
    # fitter object
    fitter = QNMFitVaryingStartingTime(h, t0_arr, N_free=_N_free, qnm_fixed_list=qnm_fixed_list, load_pickle=False, run_string_prefix=_run_string_prefix)  # number of free modes to use
    return (fitter,)

```

```python
@app.cell
def _(fitter, mo):
    with mo.status.spinner("Running free-QNM fits..."):
        fitter.do_fits()
    mo.md("**Free-QNM fits complete.**")
    return

```

```python
@app.cell
def _(fitter):
    # fitter results object
    result = fitter.result_full
    return (result,)

```

```python
@app.cell
def _(Mf, af, mode_list, plot_omega_free, plot_predicted_qnms, plt, result):
    _fig, _ax = plt.subplots()
    predicted_qnms = mode_list(['2.2.0', '2.2.1', '3.2.0'], Mf, af)
    # mode locations to visualize on the plot
    plot_omega_free(result, _ax)
    plot_predicted_qnms(_ax, predicted_qnms)
    _fig
    return

```

--------------------------------

### Build and plot a toy Kerr waveform with jaxqualin

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_varma

Marimo notebook cells that set up the example, import jaxqualin modules, build a toy Kerr waveform from three QNMs, and plot its real and imaginary parts. The cells define the true mass and spin (Mf_true=1.0, af_true=0.7), use delayed_QNM to add distortions, and produce a two-panel semilogy plot.

```python
import marimo

__generated_with = "0.19.11"
app = marimo.App()


@app.cell
def _():
    import marimo as mo

    return (mo,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    # Example: Kerr mass and spin fitting

    In this example, we construct a toy waveform from Kerr quasi-normal modes (QNMs) and fit it using the VarMa (Variable mass and spin) fitter in `jaxqualin`, which simultaneously recovers the black hole mass $M$ and spin $a$ along with mode amplitudes and phases.

    A Kerr black hole has non-zero spin ($a \neq 0$), so the gravitational waveform is complex:
    """)
    return


@app.cell
def _():
    from jaxqualin.waveforms import delayed_QNM, waveform
    from jaxqualin.qnmode import mode_list, long_str_to_qnms_free
    from jaxqualin.fit import QNMFitVarMa, QNMFitVaryingStartingTime
    from jaxqualin.plot import plot_amplitudes, plot_phases

    import numpy as np
    import matplotlib.pyplot as plt

    return (
        QNMFitVarMa,
        QNMFitVaryingStartingTime,
        delayed_QNM,
        long_str_to_qnms_free,
        mode_list,
        np,
        plot_amplitudes,
        plot_phases,
        plt,
        waveform,
    )


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Build a toy Kerr waveform

    We construct a waveform with three Kerr QNMs ($a = 0.7$): the $(2,2,0)$ fundamental, the $(2,2,1)$ first overtone, and the $(3,2,0)$ mode.
    The `delayed_QNM` function adds distortions near $t = 0$ to simulate realistic merger behavior.
    """)
    return


@app.cell
def _(delayed_QNM, mode_list, np, waveform):
    Mf_true = 1.0
    af_true = 0.7
    kerr_modes = mode_list(['2.2.0', '2.2.1', '3.2.0'], Mf_true, af_true)

    kerr_A_phi_dict = {
        '2.2.0': dict(A=1.0, phi=0.0),
        '2.2.1': dict(A=3.0, phi=np.pi / 2),
        '3.2.0': dict(A=1e-2, phi=np.pi),
    }

    kerr_t_arr = np.linspace(0, 120, 1000)
    _h_arr = np.zeros(kerr_t_arr.shape, dtype=np.complex128)
    for _i, _mode in enumerate(kerr_modes):
        if _i == 0:
            _A_delay, _A_sig, _phi_sig = 0, 10, 5
        else:
            _A_delay, _A_sig, _phi_sig = 5, 2, 2
        _h_arr = _h_arr + delayed_QNM(
            _mode, kerr_t_arr,
            kerr_A_phi_dict[_mode.string()]['A'],
            kerr_A_phi_dict[_mode.string()]['phi'],
            A_delay=_A_delay, A_sig=_A_sig, phi_sig=_phi_sig,
        )

    h_kerr = waveform(kerr_t_arr, _h_arr, t_peak=0)
    return Mf_true, af_true, h_kerr


@app.cell
def _(h_kerr, np, plt):
    _fig, _axs = plt.subplots(1, 2, figsize=(12, 4))
    _axs[0].semilogy(h_kerr.time, np.abs(h_kerr.hr))
    _axs[0].set_xlabel(r'$t$')
    _axs[0].set_ylabel(r'$|h_r|$')
    _axs[0].set_title('Real part')
    _axs[1].semilogy(h_kerr.time, np.abs(h_kerr.hi))
    _axs[1].set_xlabel(r'$t$')
    _axs[1].set_ylabel(r'$|h_i|$')
    _axs[1].set_title('Imaginary part')
    _fig.suptitle('Kerr toy waveform')
    _fig.tight_layout()
    _fig
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Case 1: Single starting time VarMa fit

    We use `QNMFitVarMa` with `real=
```

--------------------------------

### QNMFitVaryingStartingTimeResult.fixed_mode_flatness_plot_overlays

Source: https://mhycheung.github.io/jaxqualin/references

Returns the bold dictionary and flat start dictionary for flatness overlays.

```APIDOC
## QNMFitVaryingStartingTimeResult.fixed_mode_flatness_plot_overlays(**kwargs)

### Description
Returns the bold dictionary and flat start dictionary for flatness overlays.

### Method
Function call

### Endpoint
QNMFitVaryingStartingTimeResult.fixed_mode_flatness_plot_overlays

### Parameters
#### Keyword Arguments
- **kwargs** - Optional - Additional keyword arguments passed to the underlying function.

### Response
#### Success Response
- **bold_dict** (dict) - Dictionary mapping mode names to tuples of indices for flatness overlays.
- **t_flat_start_dict** (dict) - Dictionary mapping mode names to optimal starting times.

### Response Example
```
(bold_dict, t_flat_start_dict)
```
```

--------------------------------

### Python CLI to summarize flattest window from fit info

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_fit

A component of a Typer-based CLI that formats and prints a summary line for a flattest window, including amplitude, phase, and earliest flat start time. Handles NaN early start times by printing 'nan (no qualifying window)'. The snippet ends with the standard `if __name__ == "__main__": app.run()` guard.

```python
if np.isnan(_earliest):
    _earliest_txt = "nan (no qualifying window)"
else:
    _earliest_txt = f"{_earliest:.2f}"
_lines.append(
    f"- `{_mode_string_summary}`: flattest window [{info['flattest_start_time']:.2f}, {info['flattest_end_time']:.2f}] M, "
    f"A={info['flattest_amplitude_median']:.4g} (+{info['flattest_amplitude_plus']:.3g}/-{info['flattest_amplitude_minus']:.3g}), "
    f"phi={info['flattest_phase_median']:.4g} (+{info['flattest_phase_plus']:.3g}/-{info['flattest_phase_minus']:.3g}), "
    f"earliest flat start={_earliest_txt} M"
)
mo.md("\n".join(_lines))
return


if __name__ == "__main__":
    app.run()
```

--------------------------------

### Set up and run fixed-QNM fit with QNMFitVaryingStartingTime

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_fit

Sets up and runs a fixed-frequency QNM fit for modes 2,2,0, 2,2,1, and 3,2,0 using QNMFitVaryingStartingTime with N_free=0, then extracts the full result.

```python
@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Case 2: Fixed QNMs (fixed frequencies)

    Here we assume that the $2{,}2{,}0$, $2{,}2{,}1$ and $3{,}2{,}0$ modes are present, so we use a fit model including these three modes and fix their frequencies according to GR with the help of the `qnm` package.
    We will use $N_f = 0$, meaning that we do not include additional free QNMs on top of the three fixed modes.
    """)
    return

```

```python
@app.cell
def _(Mf, QNMFitVaryingStartingTime, SXSnum, af, h, l, m, mode_list, np):
    _t0_arr = np.linspace(0, 50, num=501)
    qnm_fixed_list_1 = mode_list(['2.2.0', '2.2.1', '3.2.0'], Mf, af)
    _run_string_prefix = f'SXS{SXSnum}_lm_{l}.{m}'
    _N_free = 0
    fitter_1 = QNMFitVaryingStartingTime(h, _t0_arr, N_free=_N_free, qnm_fixed_list=qnm_fixed_list_1, load_pickle=False, run_string_prefix=_run_string_prefix)
    return fitter_1, qnm_fixed_list_1

```

```python
@app.cell
def _(fitter_1, mo):
    with mo.status.spinner("Running fixed-QNM fits..."):
        fitter_1.do_fits()
    mo.md("**Fixed-QNM fits complete.**")
    return

```

```python
@app.cell
def _(fitter_1):
    result_1 = fitter_1.result_full
    return (result_1,)

```

--------------------------------

### Plot recovered mass vs starting time

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_real_waveform

Plots the recovered mass M as a function of starting time t0 for the real waveform fit. Includes a horizontal line for the true mass and labels axes and title. Returns the figure.

```python
_fig, _ax = plt.subplots(figsize=(8, 4))
_M_arr = np.array(result_schw_varma_vst.Ma_dict['M'])
_ax.plot(schw_varma_t0_arr, _M_arr, 'b-', lw=2, label='Recovered $M$')
_ax.axhline(Mf_true, color='r', ls='--', label=f'True $M = {Mf_true}$')
_ax.set_xlabel(r'$t_0$')
_ax.set_ylabel(r'$M$')
_ax.set_title('Recovered mass across starting times (real waveform)')
_ax.legend()
_fig.tight_layout()
_fig
return
```

--------------------------------

### Free-frequency fit across starting times

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_real_waveform

Runs free-frequency QNM fits across a range of starting times t0 for a Schwarzschild waveform, using QNMFitVaryingStartingTime with N_free=2 and no fixed modes. Plots the frequency evolution of the free QNMs.

```python
@app.cell(hide_code=True)
def _(QNMFitVaryingStartingTime, h_schw, mo, np):
    schw_free_t0_arr = np.linspace(0, 50, num=101)

    fitter_schw_free = QNMFitVaryingStartingTime(
        h_schw, schw_free_t0_arr, N_free=2,
        qnm_fixed_list=[],
        real=True,
        load_pickle=False,
        run_string_prefix='schwarzschild_free_example',
        save_results=False,
    )
    with mo.status.spinner("Running free-frequency fits..."):
        fitter_schw_free.do_fits()
    result_schw_free = fitter_schw_free.result_full
    return (result_schw_free,)


@app.cell
def _(
    Mf_true,
    af_true,
    mode_list,
    plot_omega_free,
    plot_predicted_qnms,
    plt,
    result_schw_free,
):
    _fig, _ax = plt.subplots()
    _predicted_qnms = mode_list(['2.2.0', '3.2.0'], Mf_true, af_true)
    plot_omega_free(result_schw_free, _ax)
    plot_predicted_qnms(_ax, _predicted_qnms)
    _ax.set_title('Free-frequency fit: QNM frequency evolution')
    _fig
    return


@app.cell
def _(mo):
    mo.md(r"""
    ## Fixed-frequency fit (varying $t_0$)

    Now we assume the $(2,2,0)$ and $(3,2,0)$ modes are present and fix their frequencies according to GR.
    We fit across a range of starting times $t_0$ and plot how the recovered amplitudes and phases behave.
    For a clean QNM signal, these should be approximately constant across $t_0$.
    """)
    return


@app.cell
def _(QNMFitVaryingStartingTime, h_schw, mo, np, schw_modes):
    # ... (truncated in source, but the visible part is the fixed-frequency fit)
    # The source shows a fitter with qnm_fixed_list=schw_modes, real=True, load_pickle=False, run_string_prefix='schwarzschild_fixed_example', save_results=False, and a spinner.
    # Since the code is truncated, we only include the visible portion.
    # The visible code is:
    # fitter_schw_fixed_vst = QNMFitVaryingStartingTime(
    #     h_schw, schw_fixed_t0_arr, N_free=0,
    #     qnm_fixed_list=schw_modes,
    #     real=True,
    #     load_pickle=False,
    #     run_string_prefix='schwarzschild_fixed_example',
    #     save_results=False,
    # )
    # with mo.status.spinner("Running fixed-frequency fits..."):
    #     fitter_schw_fixed_vst.do_fits()
    # result_schw_fixed_vst = fitter_schw_fixed_vst.result_full
    # return (result_schw_fixed_vst,)
    pass


@app.cell
def _(mo):
    mo.md(r"""
    ## Variable-mass fit (varying $t_0$)

    We use the VarMa (varying mass and spin) fitter across starting times with `real=True` [...]
    """)
    return


@app.cell
def _(QNMFitVaryingStartingTime, h_schw, mo, np, long_str_to_qnms_free):
    # ... (truncated in source, but the visible part is the variable-mass fit)
    # The source shows a fitter with qnm_fixed_list=long_str_to_qnms_free('2.2.0_3.2.0'), N_free=0, real=True, load_pickle=False, run_string_prefix='schwarzschild_varma_vst', save_results=False, and a spinner.
    # Since the code is truncated, we only include the visible portion.
    # The visible code is:
    # fitter_schw_varma_vst = QNMFitVaryingStartingTime(
    #     h_schw, schw_varma_t0_arr, N_free=0,
    #     qnm_fixed_list=long_str_to_qnms_free('2.2.0_3.2.0'),
    #     real=True,
    #     load_pickle=False,
    #     run_string_prefix='schwarzschild_varma_vst',
    #     save_results=False,
    # )
    # with mo.status.spinner("Running variable-mass fits..."):
    #     fitter_schw_varma_vst.do_fits()
    # result_schw_varma_vst = fitter_schw_varma_vst.result_full
    # return (result_schw_varma_vst,)
    pass

```

--------------------------------

### Run Schwarzschild Varma fits on real waveform

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_real_waveform

Configures and runs the Schwarzschild Varma fitter on real waveform data, varying the starting time t0. Uses the provided parameters and displays a spinner during execution. Returns the fit result and the array of starting times.

```python
qnm_fixed_list=[],
qnm_free_list=schw_varma_qnm_free,
var_M_a=True,
real=True,
load_pickle=False,
run_string_prefix='schwarzschild_varma_example',
save_results=False,
)
with mo.status.spinner("Running VarMa varying $t_0$ fits..."):
    fitter_schw_varma_vst.do_fits()
result_schw_varma_vst = fitter_schw_varma_vst.result_full
return result_schw_varma_vst, schw_varma_t0_arr
```

--------------------------------

### Import marimo and create app

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_custom_waveform

Initializes a marimo notebook with version 0.19.11 and creates the app object.

```python
import marimo

__generated_with = "0.19.11"
app = marimo.App()
```

--------------------------------

### QNMFitVaryingStartingTime

Source: https://mhycheung.github.io/jaxqualin/api/fit

Documentation for the QNMFitVaryingStartingTime class, including its methods and attributes.

```APIDOC
## QNMFitVaryingStartingTime

### Description
A class for fitting the postmerger waveform with a varying starting time. It provides methods to generate initial guesses, perform fits, and handle nonconvergence.

### Methods

#### get_mirror_ratio_list
- **Description**: Get `mirror_ratio_list` from `iota` and `psi`.

#### initial_guesses
- **Description**: Generate initial guesses for the first `t0` fit.

#### make_nan_result
- **Description**: Generate a `QNMFitVaryingStartingTimeResult` object with `nan` values.

#### do_fits
- **Description**: Perform the fits.

### Attributes
- `t0_arr` (ndarray) - Array of starting times for fitting.
- `h` (waveform) - Waveform object to be fitted.
- `var_M_a` (bool) - Fit for the mass and spin of the black hole.
- `Warning` (bool) - Not tested yet.
- `real` (bool) - Whether to fit a real-valued waveform.
- `N_free` (int) - Number of frequency-free QNMs to include in the model.
- `qnm_fixed_list` (List[mode]) - List of fixed-frequency QNMs included in the model.
- `qnm_free_list` (List[mode_free]) - List of free-frequency QNMs of fixed mode numbers to include in the model.
- `run_string_prefix` (str) - Prefix of the run name for dumping the `pickle` file.
- `nonconvergence_cut` (bool) - Whether to cut the nonconverged fits.
- `nonconvergence_indx` (List[int]) - Indices of the nonconverged fits.
- `initial_num` (int) - Number of initial guesses to use for the first starting time for frequency-free fits.
- `include_mirror` (bool) - Whether to include the mirror modes.
- `mirror_ratio_list` (List[float]) - List of ratios between prograde and mirror mode amplitudes.
- `iota` (float) - Inclination angle of the source.
- `psi` (float) - Polarization angle of the source.
- `save_results` (bool) - Whether to save the results.
- `params0` (ndarray) - Initial guess for the fit parameters.
- `max_nfev` (int) - Maximum number of function evaluations for the fit.
- `sequential_guess` (bool) - Whether to use the previous fit as the initial guess for the next fit.
- `load_pickle` (bool) - Whether to load the `pickle` file if it exists.
- `fit_save_prefix` (str) - Prefix of the path to save the `pickle` file.
- `A_bound` (float) - Maximum value of the amplitude.
- `fit_kwargs` (Dict[str, Any]) - Keyword arguments for curve fitting.
- `initial_dict` (Dict[str, Any]) - Keyword arguments for `make_initial_guess` method.
- `A_guess_relative` (bool) - Whether to multiply the initial guess of the amplitude by the peak strain of the waveform.
- `set_seed` (int) - Random seed for generating the initial guesses.
- `weighted` (bool) - Whether to perform a weighted fit.
- `double_skip` (bool) - Whether to skip the next `2^n` `t0` fits when a fit does not converge.
- `skip_i_init` (int) - Number of `t0` fits to skip for the first time a nonconvergent fit occurred.
- `result_full` (QNMFitVaryingStartingTimeResult) - Object for storing the fit results.
```

--------------------------------

### plot_phases

Source: https://mhycheung.github.io/jaxqualin/references

Plots the phases of the QNM modes as a function of starting time of the fit.

```APIDOC
## `plot_phases(results_full, fixed_modes=None, ax=None, alpha=1.0, ls='-', use_label=True, shift_phase=True, legend=True, color_dict={}, lw=2.0, bold_dict={}, lw_bold=4.0, alpha_bold=1.0, t_flat_start_dict={}, flat_start_s=20, flat_start_marker='o', plot_mirror_pred=False, iota=None, psi=0.0, af=None)`

### Description
Plots the phases of the QNM modes as a function of starting time of the fit.

### Parameters
- **results_full** (`Union[QNMFitVaryingStartingTimeResult, QNMFitVaryingStartingTimeResultModel, QNMFitVaryingStartingTimeResultVarMa]`) - Required - The results object of the fit.
- **fixed_modes** (`Optional[Union[QNMFitVaryingStartingTimeResult, QNMFitVaryingStartingTimeResultModel, QNMFitVaryingStartingTimeResultVarMa]]`) - Optional - The list of fixed modes to plot. Default: `None`.
- **ax** (`Axes`) - Optional - The matplotlib axes to plot on. If None, a new figure and axes will be created. Default: `None`.
- **alpha** (`float`) - Optional - The alpha value of the lines. Default: `1.0`.
- **ls** (`str`) - Optional - The line style of the lines. Default: `'-'`.
- **use_label** (`bool`) - Optional - Whether to use the mode name as the legend label. Default: `True`.
- **shift_phase** (`bool`) - Optional - Whether to shift the phase by `pi` if the amplitude is negative. Default: `True`.
- **legend** (`bool`) - Optional - Whether to include the legend. Default: `True`.
- **color_dict** (`Dict[str, str]`) - Optional - A dictionary of the colors of the mode names. Default: `{}`.
- **lw** (`float`) - Optional - The line width of the lines. Default: `2.0`.
- **bold_dict** (`Dict[str, Tuple[int, int]]`) - Optional - A dictionary of mode names to tuples of indices indicating the starting and ending index of the region where the amplitude is flat. The line between the indices will be plotted with a larger line width. Default: `{}`.
- **lw_bold** (`float`) - Optional - The line width of the lines corresponding to the indices in `bold_dict`. Default: `4.0`.
- **alpha_bold** (`float`) - Optional - The alpha value of the lines corresponding to the indices in `bold_dict`. Default: `1.0`.
- **t_flat_start_dict** (`Dict[str, float]`) - Optional - A dictionary of mode names to the optimal starting times. A scatter point will be plotted at the starting time of the fit for the corresponding mode. Default: `{}`.
- **flat_start_s** (`int`) - Optional - The size of the scatter points corresponding to the indices in `t_flat_start_dict`. Default: `20`.
- **flat_start_marker** (`str`) - Optional - The marker of the scatter points corresponding to the indices in `t_flat_start_dict`. Default: `'o'`.
- **plot_mirror_pred** (`bool`) - Optional - Whether to plot the predicted phase of the mirror modes. Default: `False`.
- **iota** (`Optional[float]`) - Optional - The inclination angle of the source. Default: `None`.
- **psi** (`float`) - Optional - polarization angle of the source. Default: `0.0`.
- **af** (`Optional[float]`) - Optional - The remnant spin of the black hole. Default: `None`.

### Returns
None (plots the data on the provided axes or a new figure).
```

--------------------------------

### plot_omega_free

Source: https://mhycheung.github.io/jaxqualin/references

Plots the complex QNM frequencies as a function of starting time of the fit.

```APIDOC
## `plot_omega_free(results_full, ax=None, plot_indxs=[], t0_min=None, t0_max=None, indicate_start=False, color=None, line_alpha=0.3, scatter_alpha=0.5, scatter_size=1.0, color_indicate=False, color_indicate_list=[])`

### Description
Plots the complex QNM frequencies as a function of starting time of the fit.

### Parameters
- **results_full** (`Union[QNMFitVaryingStartingTimeResult, QNMFitVaryingStartingTimeResultModel, QNMFitVaryingStartingTimeResultVarMa]`) - Required - The results object of the fit.
- **ax** (`Axes`) - Optional - The matplotlib axes to plot on. If None, a new figure and axes will be created. Default: `None`.
- **plot_indxs** (`List[int]`) - Optional - The indices of the modes to plot. If empty, all modes will be plotted. Default: `[]`.
- **t0_min** (`Optional[float]`) - Optional - The minimum starting time to plot. Default: `None`.
- **t0_max** (`Optional[float]`) - Optional - The maximum starting time to plot. Default: `None`.
- **indicate_start** (`bool`) - Optional - Whether to indicate the starting point of the fit with a marker. Default: `False`.
- **color** (`Optional[str]`) - Optional - The color of the scatter points and line. Default: `None`.
- **line_alpha** (`float`) - Optional - The alpha value of the line. Default: `0.3`.
- **scatter_alpha** (`float`) - Optional - The alpha value of the scatter points. Default: `0.5`.
- **scatter_size** (`float`) - Optional - The size of the scatter points. Default: `1.0`.
- **color_indicate** (`bool`) - Optional - Whether to color the scatter points and line according to `color_indicate_list`. Default: `False`.
- **color_indicate_list** (`List[int]`) - Optional - Which default color to use for each mode, if `color_indicate` is `True`. Default: `[]`.

### Returns
None (plots the data on the provided axes or a new figure).
```

--------------------------------

### Fit SXS waveform with free QNMs using jaxqualin

Source: https://mhycheung.github.io/jaxqualin/notebooks/example_fit

Marimo notebook cells that import jaxqualin modules, download an SXS waveform, set up a free-QNM fit with QNMFitVaryingStartingTime, run the fits, and plot the resulting frequencies against predicted QNMs.

```python
import marimo

__generated_with = "0.19.11"
app = marimo.App()


@app.cell
def _():
    import marimo as mo

    return (mo,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    # Example: fitting a waveform in the SXS catalog
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Import modules
    """)
    return


@app.cell
def _():
    from jaxqualin.qnmode import mode, mode_list
    from jaxqualin.waveforms import get_waveform_SXS
    from jaxqualin.fit import QNMFitVaryingStartingTime
    from jaxqualin.plot import plot_amplitudes, plot_phases, plot_omega_free, plot_predicted_qnms

    import numpy as np
    import matplotlib.pyplot as plt

    return (
        QNMFitVaryingStartingTime,
        get_waveform_SXS,
        mode_list,
        np,
        plot_amplitudes,
        plot_omega_free,
        plot_phases,
        plot_predicted_qnms,
        plt,
    )


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Getting the SXS waveform

    The waveform will be downloaded with the `sxs` package, if not already.
    """)
    return


@app.cell
def _(get_waveform_SXS):
    SXSnum = "0305"
    l = 2
    m = 2

    h, Mf, af, Lev = get_waveform_SXS(SXSnum, l, m)
    return Mf, SXSnum, af, h, l, m


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Case 1: Fit model with free QNMs (unfixed frequencies)

    Here we will fit with three free QNMs ($N_{f} = 3$) and see if the resulting frequencies approach any of the QNMs we expect to find.
    By default, the fit results are saved with `pickle` into `./.jaxqualin_cache/fits/`.
    If `load_pickle` is `True`, the fitter will load the pickled result that matches `run_string_prefix` and the list of modes used.
    """)
    return


@app.cell
def _(QNMFitVaryingStartingTime, SXSnum, h, l, m, np):
    _t0_arr = np.linspace(0, 50, num=501)  # array of starting times to fit for
    qnm_fixed_list = []  # t0 = 0 is the peak of the strain
    _run_string_prefix = f'SXS{SXSnum}_lm_{l}.{m}'  # list of QNMs with fixed frequencies in the fit model
    _N_free = 3  # prefix of pickle file for saving the results
    # fitter object
    fitter = QNMFitVaryingStartingTime(h, _t0_arr, N_free=_N_free, qnm_fixed_list=qnm_fixed_list, load_pickle=False, run_string_prefix=_run_string_prefix)  # number of free modes to use
    return (fitter,)


@app.cell
def _(fitter, mo):
    with mo.status.spinner("Running free-QNM fits..."):
        fitter.do_fits()
    mo.md("**Free-QNM fits complete.**")
    return


@app.cell
def _(fitter):
    # fitter results object
    result = fitter.result_full
    return (result,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Plotting the results

    Different colored points trace out the frequency evolution of a free QNM from early to late $t_0$.
    """)
    return


@app.cell
def _(Mf, af, mode_list, plot_omega_free, plot_predicted_qnms, plt, result):
    _fig, _ax = plt.subplots()
    predicted_qnms = mode_list(['2.2.0', '2.2.1', '3.2.0'], Mf, af)
    # mode locations to visualize on the plot
    plot_omega_free(result, _ax)
    plot_predicted_qnms(_ax, predicted_qnms)
    return
```
Tokens:
Raw