Time and phase marginalization

[tsunhopang]

This commit is in response to the issue #2

You may find the comparison corner plots as follows:

[tsunhopang]

@kazewong this git conflict might need your intervene

[kazewong]

I fixed the conflicts and moved the internal likelihood function into likelihood.py. I also add typing information in the function.

Before merging this, I think it will be beneficial to also migrate the body of the heterodyne likelihood to the same API.

[tsunhopang]

Yes, that was the plan to have it also included there in another pull request. I can have it also add it there within this one.

For relative binning, only phase marginalization will be added; as the time marginalization will then use a non-uniform grid fft, for which I am not sure if it is beneficial (e.g. see Appendix A in https://arxiv.org/abs/2312.06009)

[kazewong]

If it is not a lot of changes, perhaps we can just have it in this PR.

On relative binning, I agree we can first add phase marginalization.

For time marginalization, I think at some point we can consider using a non-uniform fft package jax-finufft, which to my understand is quite performant. The main issue there currently is the installation will get a bit more complicated than simple pip install, and asaik the API and functionality is not necessarily stable yet. So this is more likely to go into future PR

[tsunhopang]

For the relative binning likelihood; the optimizer finds the following parameters;

Injected parameters = {
'M_c': 28.095555795460438,
'eta': 0.24710059171597634,
's1_z': -0.016246134632590996,
's2_z': -0.04423844417344978,
'd_L': 445.4687493344804,
't_c': 0.0,
'phase_c': 4.100547786091415,
'iota': 2.920749029724262,
'psi': 0.787166765060988,
'ra': 0.15039344639787625,
'dec': -0.5355555460283846
}
ref_parameters_original = {
'M_c': 28.125974539678566,
'eta': 0.25,
's1_z': -0.99,
's2_z': 0.9191400868807302,
'd_L': 223.59820615113975,
't_c': 0.009947852095418,
'phase_c': 6.283173520274804,
'iota': 1.3616503037593668,
'psi': 3.141298463313976,
'ra': 2.2766191364329975,
'dec': 1.1116867914351234
}
ref_parameters_margphi = {
'M_c': 28.13080784160065,
'eta': 0.24999999999599654,
's1_z': -0.99,
's2_z': 0.9200876980166237,
'd_L': 223.01512154400882,
't_c': 0.010463400533415948,
'iota': 1.3685230276718776,
'psi': 1.5009806917550423,
'ra': 2.202961964586829,
'dec': 1.1170486830648356
}

I was expecting the marginalization over phase would help the optimizer to get closer to the true parameters, yet, the impact is minimal.

For PE, I have two sets of runs with both the reference parameters fixed at the injected parameters, so that we have an apple-to-apple comparison; the corner plots are as follows

[kazewong]

I would prefer to put the likelihood function in the likelihood.py file instead of utils.py. You can leave the logi0 in.