kazewong / jim

Gravitational-wave data analysis tools in Jax
MIT License
59 stars 19 forks source link

Reparameterization of extrinsic parameter for better sampling efficiency #131

Closed tsunhopang closed 2 months ago

tsunhopang commented 3 months ago

This pull request makes use of the new priors and transforms system introduced in https://github.com/kazewong/jim/pull/108 for better sampling efficiency and aid the NF learning in flowMC.

See https://arxiv.org/abs/2207.03508 for reference.

This PR now includes

  1. The distance to SNR weighted distance (Eq. 18, but instead of 1 detector, using all of the detectors)
  2. The geocentric arrival time to the detector arrival time (Eq. 13)

more to come

thomasckng commented 3 months ago

@kazewong @xuyuon Actually we may need to add some checking for the additional parameters needed in the transform. For example, sampling spin to Cartesian spin takes masses as input for the transform, but we don't check if masses are in the parameters now. Same problem exist for the transformation implemented in this PR.

tsunhopang commented 3 months ago

@thomasckng this PR is still a work in progress, so you do not need to approve it yet. I will ping u all once everything, e.g. tests etc are ready

tsunhopang commented 3 months ago

Regarding the folding approach for the phasing and angles, I think it would be better to introduce the discrete MCMC rather than the numeric marginalization approach.

tsunhopang commented 3 months ago

I would like to add that it might be better to allow the user to have the freedom to choose the ordering of the forward and inverse transform. That would be very useful for the SNR-weighted distance transform.

thomasckng commented 3 months ago

I would like to add that it might be better to allow the user to have the freedom to choose the ordering of the forward and inverse transform. That would be very useful for the SNR-weighted distance transform.

I think this is related to the discussion in the last development meeting and also #130. Adding a function to switch the direction of transform can be implemented. I can submit a PR for that tomorrow.

tsunhopang commented 3 months ago

This is more of a styling problem that I want to ask for a discussion. @kazewong @thomasckng @xuyuon

Since quite a few transformations for this extrinsic re-parameterization are conditioned on other parameters, e.g., t_det <-> t_c needs also ra, dec, the current implementation for Jim needs to be fixed.

  1. As shown in here, only the parameters in the name_mapping is passed;
  2. However, if I also add the ra, dec into the mapping, it will be removed as in here

Therefore, we would need to pass everything, from start to finish, or an additional condition list. The conditional list might be helpful for the transform ordering, which @thomasckng is working on.

kazewong commented 3 months ago

To clarify, the situation here is we want to define a bijective transform with conditional parameters, in the sense that only the N parameters will be transformed according to the name mapping, but then there are some extra parameters needed, that should not be transformed, right?

I am wondering whether we want to roll out something like the following

class ConditionalBijective(Bijective):

    conditional_params: list[str]

    def __init__(self,
        named_params,
        conditional_params):
      ...

This sounds like it is exactly what we need?

tsunhopang commented 3 months ago

Yes, that's indeed what we need. We can either create that new class or just add the conditional_params to be an attribute for the Transform class.

kazewong commented 3 months ago

I think subclassing bijective transform is a better option here, since normal transform is not necessary conditional. I can get to this today and try merging into this working branch

tsunhopang commented 3 months ago

At line, the jacobian calculation also includes the conditional parameters, but that should be excluded.

kazewong commented 3 months ago

Do you want to just change those lines? I think I am happy with the API, and it doesn't seem to be a complicated change.

tsunhopang commented 3 months ago

@xuyuon, all the transforms are good to go. One note: the transforms for distance and arrival time all have their bound condition on the sky location; therefore, the transform is direct to the unbounded space. Please remember that the sky-frame sky location and the chirp mass must be available at the time of the transformation. This shall affect the order of the transform.

tsunhopang commented 3 months ago

@thomasckng @kazewong good for review now

thomasckng commented 3 months ago

@xuyuon Maybe you can submit a PR to this PR and use the ConditionalBijectiveTransform class for the spin transform.

kazewong commented 3 months ago

I think multiple features or bug fixes better not be merged together too much, since this PR does not depends on the spin transform

thomasckng commented 3 months ago

I think multiple features or bug fixes better not be merged together too much, since this PR does not depends on the spin transform

Let me branch out from this PR branch to implement that then, as I need to use the ConditionalBijectiveTransform for it.

xuyuon commented 2 months ago

I have mentioned the iota problem in the meeting, but I don't think I have presented it well, so let me explain it clearly here.

The issue is related to the two transformations: DistanceToSNRWeightedDistanceTransform and GeocentricArrivalPhaseToDetectorArrivalPhaseTransform. Both of these functions require a conditional parameter iota. However, when I'm working with the Pv2 waveform, I'm sampling the parameter space without iota, so I need to find a way to introduce iota into the problem.

Let's focus on the following parameters: a1, a2, phi12, theta1, theta2, phi_jl, theta_jn, phase_c. These are the parameters on which I have defined priors, and I'll be focusing on the GeocentricArrivalPhaseToDetectorArrivalPhaseTransform.

To get phase_det, we need to input phase_c and iota. We have phase_c, but we're missing iota. It's possible to get iota from the other parameters (a1, a2, phi12, theta1, theta2, phi_jl, theta_jn, phase_c), and I've built a function to do that.

The forward transform path will look like this:

a1, a2, phi12, theta1, theta2, phi_jl, theta_jn, phase_c -> a1, a2, phi12, theta1, theta2, phi_jl, theta_jn, phase_c, iota -> a1, a2, phi12, theta1, theta2, phi_jl, theta_jn, phase_det, iota

This solves the problem for the forward transform. However, for the inverse transform, it gets a bit more complicated. We start with the parameters: a1, a2, phi12, theta1, theta2, phi_jl, theta_jn, phase_det To calculate the transform from phase_det back to phase_c, we need iota as well. But to obtain iota, we need phase_c. This leads to a circular dependency.

The key question therefore is: Is it possible to perform the inverse transform of GeocentricArrivalPhaseToDetectorArrivalPhaseTransform without iota? Or is there a way to get iota from the other available parameters?

Here is a list of the parameters available at that time: M_c, q, a1, a2, phi12, theta1, theta2, phi_jl, theta_jn, phase_det, d_hat_unbounded, t_det_unbounded, zenith, azimuth, psi

xuyuon commented 2 months ago

@tsunhopang I am sorry. I just realised that when I was trying to pull the new features from your branch to mine for testing a few days ago, I accidentally pushed everything in my branch into your branch. My branch is far from ready. I am very sorry that I messed up your branch. Please feel free to revert all the changes I made.

kazewong commented 2 months ago

Closing this PR in favor of #161