[Feature Req] Vectorised redshift calculations in hmf+halomod #discussion

[andrejdvornik]

Is your feature request related to a problem? Please describe.
This feature request is not related to a problem, but would significantly speed up calculations, with relatively little development (I think).

Describe the solution you'd like
It would be quite practical to implement an option to pass a redshift array and evaluate all the quantities at all redshifts in the array. While one can easily use the update() method to loop over redshifts, some initial testing shows that passing an array is around factor of 10 faster.

Additional context
I have tried to rough hack the implementation in hmf, and it shows significant speed up with minimal development, as majority of functions seem to work out of the box on (all numpy array broadcasting is mostly there, CAMB seems to deal with that without an issue either, there is only a couple of interpolation functions that need some further changes).
The same would need to be passed on to halomod, where it seems that the implementation is a bit more involved, but not impossible.

I am willing to have a proper go at this myself in the following weeks/months and can open PRs for implementation. Any thoughts welcome, or reasons why this might not be a good idea...

[steven-murray]

Hi @andrejdvornik, this is an interesting idea. Is your proposed plan to have e.g.

mf = MassFunction(
    z = np.linspace(0, 10, 1),
    ...
)

print(mf.growth_factor)  # prints 10 numbers
print(mf.dndlog10m)  # prints a 10xlen(m) array

?

I think the biggest drawback is the increased mental load for the user in remembering that the output arrays like .sigma etc are now 2D and the first axis corresponds to redshift.

[andrejdvornik]

Indeed, or generally any 1d array:

mf = MassFunction(
    z = np.array([0.0, 2.0]),
    ...
)

Furthermore, the proposed output for spectra would be also then (nz, nk), and for instance the c(M) relations from halomod shape (nz, nm) and profiles (nz, nk, nm).

The increased mental load would indeed be a drawback to consider, but I think the implementation can be completely backwards compatible: if float is passed, then the existing format of output arrays will also be returned. And one could still use the update() if needed. What would also potentially increase the complexity is the fact that some other parameters would need to match the shape of input redshift, if desired. For instance, delta_c and overdensity comes to mind, and potentially more.

[steven-murray]

I see a couple of drawbacks here:

1. Arrays like profiles (nz, nk, nm) could quickly become quite large memory-wise, since both nk and nm can be reasonably large.
2. If we have an interface where sometimes the return array is 2D and other times its 1D, it puts a lot more complexity into the code -- we have to check each time which axis belongs to m etc. It's better if the returned array always has the same number of dimensions so it can be consistently indexed. It's not impossible, but I have found many times that having different return types for different inputs can lead to all kinds of confusion.
3. This is perhaps less important, but we could ask "why redshift?" -- why not vectorize over e.g. some cosmological param, or delta_c, or anything else? The answer could be that redshift is more likely to be looped over, and it could be that there are use-cases where you need to integrate over redshift (e.g. for integrated source counts for a survey) which is not true for other parameters. But we should be clear about why we think this in particular is important.

In saying that, increased speed is obviously not a bad thing. I guess the question is whether the increased speed will enable new use-cases, given that the run-time of computing a massfunction is not that large already.

One thing I have been considering (not at great length since I don't spend too much time on this library these days) is that it might be better to always set up all the quantities in hmf as interpolators (over mass and/or k), so instead of simply doing mf.dndm you do mf.dndm(m=[masses,...]) and it already has a pre-computed result sitting in the background, and it just interpolates for you at those masses. This way, we could ensure we have large mass coverage, and optimize the mass coverage for particular integral cases etc. If we do this, perhaps it would be useful to introduce the vectorization over redshift here (modulo some of the above arguments).

What do you think? It would be a bit of work, and necessitate a new major version of hmf...

[andrejdvornik]

Indeed, the profile arrays can become quite large (speaking from experience), but for such cases I could imagine the need might be also connected with the access to a compute cluster / server with enough memory. Not always the case, but likelihood of really precise calculations goes hand in hand with resources needed to perform such calculations. Would argue that this aspect worries me the least. For the second point, I would argue the same, better to have in such case an uniform return interface, that in case of a single redshift would still be consistent in shape and number of dimensions as the case for an vector input. This should not cause any trouble necessary with backward compatibility, as if for instance the return is now in the shape of (1, 200) for mass functions instead of (200,), as numpy does know how to properly broadcast this.

As for the third point, I think the main reason would simply be this equation (and similar for any other cross terms and combinations):

I think this is quite a common situation in weak lensing where gaining speed in creation of power spectra as a function of redshift would be beneficial, especially where other parameters are supposed to be the same.

(And as a side note, this vectorisation would also allow me personally interface with hmf and halomod in a way that I could use them as a parent class in the much more involved halo model code built on top of them. Currently I do interfacing using the get_halomodel() function as it allows for efficient updates of params, but I need to pass a lot of properties further down the line and I need to build such constructs:

def dndlnm(self):
        """
        Return the differential mass function.
        """
        return np.array([x.dndlnm for x in self.hmf]) # pretty much adding in the z-dependence.

Instead, the vectorisation would allow for simple inheritance by initialisation of the parent class, as it functions in the hmf and halomod.)

Now, the quick and dirty hack I have played with somewhat works fine with hmf, and with halomod I only really touched the concentrations and profiles, as I do not rely on the rest. The speed gain comes really from the fact that CAMB is already vectorised for redshifts, from the reduced calls to the concentration calculation, and growth and transfer methods. The gain with that is really from 2 seconds to 0.2 seconds, and that is only for 15 redshifts! When running MCMC this is of significant importance.

Interpolators would be in my opinion a good step further, as they would allow much more flexible interaction with code built on top of hmf and halomod. Speaking of time to work on that, I think I could contribute the vectorisation for the current code setup, I am afraind I do not have much time to do significant development on interpolation scheme, at least not in the next few months.