Implementation of time marginalisation for importance sampling.
To turn on time marginalisation, add time_marginalization: {} to the is_settings.yaml file. This will automatically drop the geocent_time axis from the dingo samples (if present) and train the NDE only on the remaining 14 parameters.
The time-marginalised (tm) likelihood assumes a uniform prior over geocent_time. This is directly loaded from the samples, which in turn inherit this from the model. An error is raised if the prior is not uniform. Note: For GNPE with exact time-translational equivariance the prior is set in the train_settings, but has no effect in practice. We could possibly set it in the is settings instead, but I think loading it from the model is more consistent.
The tm-likelihood is computed via FFT, which has a time resolution of 1 / domain.f_max. This is often not sufficient. I added the option to perform n_fft shifted FFTs to increase the resolution by a factor of n_fft. To modify this, set time_marginalization: {n_fft: 5} in is_settings.yaml. It seems that n_fft=5 is an appropriate value for f_max=1024, leading to a time resolution of ~0.2ms.
The time marginalization requires a few initialisation steps in order to keep likelihood evaluations as fast as possible. This includes discretising the time prior, and importantly creating n_fft time-shifted versions of the strain data (as opposed to time shifting the signals n_fft times, which can only be done at likelihood evaluation time). This is implemented with the method initialize_time_marginalization.
The log_prob method automatically selects time_marginalization if initialize_time_marginalization was called before.
At some point we should add unit tests for the likelihood methods.
Implementation of time marginalisation for importance sampling.
time_marginalization: {}
to theis_settings.yaml
file. This will automatically drop thegeocent_time
axis from the dingo samples (if present) and train the NDE only on the remaining 14 parameters.1 / domain.f_max
. This is often not sufficient. I added the option to performn_fft
shifted FFTs to increase the resolution by a factor ofn_fft
. To modify this, settime_marginalization: {n_fft: 5}
inis_settings.yaml
. It seems thatn_fft=5
is an appropriate value forf_max=1024
, leading to a time resolution of~0.2ms
.n_fft
time-shifted versions of the strain data (as opposed to time shifting the signalsn_fft
times, which can only be done at likelihood evaluation time). This is implemented with the methodinitialize_time_marginalization
.log_prob
method automatically selects time_marginalization ifinitialize_time_marginalization
was called before.At some point we should add unit tests for the likelihood methods.