Vectorizing GP bases/design matrices for chromatic modeling (interpolation and ecorr)

[blarsen10]

Hopefully I'm not too off base with this, but a couple of the functions to create GP bases in enterprise seem like they are unnecessarily slow because they are not vectorized (i.e. they contain for loops). My understanding is this is usually okay because the bases are typically created once during the first likelihood call and then cached. However, its now becoming more common in noise modeling to allow the radio frequency scaling index of a chromatic process to vary, which I think means recreating the GP basis for every new idx sample. For example, calling enterprise.signals.createfourierdesignmatrix_chromatic and allowing the idx parameter to vary in MCMC should be recreating the design matrix several times. Thankfully the function to create the Fourier design matrix, createfourierdesignmatrix_red is already vectorized, therefore this is relatively efficient. However, not all the functions for creating bases are vectorized yet, and if they are used for chromatic noise modeling this is an extra inefficiency that can be fixed. The only 2 commonly used bases I found that are not vectorized are the interpoaltion and ECORR bases:

• enterprise.signals.utils.linear_interp_basis. This seems like a problem if using linear_interp_basis_chromatic in enterprise_extensions with varied idx during MCMC. Vectorizing this seems doable.
• enterprise.signals.utils.create_quantization_matrix. This is usually used for ECORR. As far I know this is not used yet for varied chromatic index modeling, but it might be useful for that purpose in the near future! The way this is written looks like it might be tricky to vectorize though?

Maybe there are other bases that should be included in this list? Jeremy has also been working on better caching of the chromatic basis which will also help a lot with this, though I think that would belong better in a separate issue

[vhaasteren]

Hi @blarsen10,

I honestly did not even know what linear_interp_basis was doing. In our FFT-Interp PR we use a similar routine. Looking at them, it's actually the same. Look at: create_fft_time_basis in the FFT branch. We utilize scipy.interpolate, and I assume that's a pretty efficient implementation. I guess, to avoid overlap, I should look into merging those before the PR is merged.

As for ECORR, I would always recommend ECORR kernel (if we fix the IPTA issues that is), which should much much faster. The function create_quantization_matrix is already quite fast for what it is doing. How would you modify it? My guess is that any efficient modification you will come up with would probably split it up into a factor (first for loop) and a frequency-dependent (second for loop) function that we don't both want to run every MCMC step.

[blarsen10]

@vhaasteren thanks for the response! Yeah I agree it makes sense to compare and eventually combine those implementations of the interpolation bases (and probably linear_interp_basis should not be in utils?). Ideally they should do the same thing, though one thing I noticed is that to cut down on basis size, linear_interp_basis removes any empty columns that show up in the design matrix if there are no TOAs represented in that column (i.e. due to a large data gap), whereas it I'm not sure that your FFT version does that. In terms of efficiency to create the basis, without having actually checked it, I agree your FFT version is probably more efficient thanks to scipy. We are currently still using linear_interp_basis for the 15yr custom chromatic modeling.

For the quantization matrix, I still need to think about how it could be done more efficiently, but I assume there might be a way that doesn't use for loops...? I'm not sure yet though, because the method of setting up binned simultaneous observations in the quantization matrix is different from the linearly-spaced interpolation basis, since the bin locations are dependent on the previous TOA positions. I also agree it's better to use the ECORR kernel (especially since it's fast now!), although this quantization matrix method also shows up in a weird place besides just ECORR, which is the model of radio-frequency-dependent DM correlations in enterprise_extensions (though thankfully we are saving ourselves the headache of using that model with varied chromatic index).

[vhaasteren]

Hi @blarsen10 ,

The FFT version does not remove any empty columns that show up in the design matrix: it does not need to. Are you familiar with the combine=True option? Enterprise tries to be clever when you do that, and it automatically combines all design matrix columns that are identical. That is useful for signals that have a common basis, such as a GWB and the red noise. They have the same Fourier modes, so combining such columns reduces the size of the Phi matrix. This is done in _combine_basis_columns in signal_base.py. You would end up with a single column of zeros in the Phi/Sigma matrices, and that is ok. Even though it seems illogical to combine columns of different signals (which can happen here), if you go through the math it all works out and is ok.

I am not sure whether those mechanics works the same for parameterized bases, but I would think so. We should test that. My guess is we should indeed just merge the two functions. The proper place would indeed be gp_bases.gp, where I would have had a larger probability of noticing it.

[blarsen10]

Hi @vhaasteren ,

That makes sense, we wouldn't want to remove empty columns for red noise because combine = True is more efficient. However, I think we still will want to have the option to remove the empty columns for DM and chromatic models, as we have found some pulsars prefer a large number of tightly-spaced interpolation knots. In this case, the runtime hugely benefits from removing the empty columns, and the basis would not combine with the red noise basis anyways. So I propose instead that there should be a remove_empty_cols keyword argument to cover both use cases.

Returning to the original optimization issue, I agree and am convinced now that create_quantization_matrix can't run much faster than it does already because of the greedy binning, so I don't think anything should be done there. As for linear_interp_basis, I strung together a version that constructs the basis just using array indexing, and it runs about twice as fast as the one currently in utils. I also did a quick comparison of my version of linear_interp_basis with your create_fft_time_basis in the FFT branch. I confirmed both functions produce the same output using J1909's TOAs from the 15yr dataset. I also checked the runtime. On my machine it looks like linear_interp_basis actually runs a bit faster, so it must be that scipy.interpolate.interp1d is introducing some extra overhead there. That said, they are both fast for small nknots, and your version is more useful because it generalizes to other orders. I've attached a snippet of this test below (I can also share the notebook if you'd like to take a look).

So, let me know if you'd like me to update linear_interp_basis and move it to the appropriate spot in gp_bases, or if you would prefer to merge them together on your branch (or if you have some other idea in mind).