vectorize HierarchicalHealpixMap evluate

[henrysky]

This PR adds a vectorized version to evaluate HierarchicalHealpixMap.

I have tested the performance of evaluating Combined19 on the whole APOGEE DR17. Simple for loop took 6h56m41s while this PR vectorized function only took 26s. So we can have a speedup of ~1000 times!!

To maintain compatibility, evaluating a single l and b will fall back to the original code while evaluating array input of l and b will use the vectorized code.

I have check the consistency of the result on the whole APOGEE DR17 by np.all(np.isclose(ebv_original, ebv_vectorized, equal_nan=True)) and the code to generate the comparison result is as follow

import numpy as np
from mwdust import Combined19
from astropy.io import fits
from astroNN.apogee import allstar
from galpy.util.coords import radec_to_lb

allstar_f = fits.getdata(allstar(dr=17))

lb = radec_to_lb(allstar_f['RA'], allstar_f['DEC'], degree=True)
ls, bs = lb[:, 0], lb[:, 1]

c19map = Combined19()
ebv = np.zeros_like(ls) * np.nan
for idx, (l, b, d) in enumerate(zip(ls, bs, 1 / allstar_f['GAIAEDR3_PARALLAX'])):
    try:
        ebv[idx] = c19map(l, b, d)
    except:
        pass
np.save("ebv_original.npy", ebv)

ebv_original.zip

import numpy as np
from mwdust import Combined19
from astropy.io import fits
from astroNN.apogee import allstar
from galpy.util.coords import radec_to_lb

allstar_f = fits.getdata(allstar(dr=17))

lb = radec_to_lb(allstar_f['RA'], allstar_f['DEC'], degree=True)
ls, bs = lb[:, 0], lb[:, 1]

c19map = Combined19()
# you don't need to call _evaluate_vectorize directly, but here I will do that to clearly show which code is used
ebv = c19map._evaluate_vectorize(ls, bs, 1 / allstar_f['GAIAEDR3_PARALLAX'])
np.save("ebv_vectorized.npy", ebv)

ebv_vectorized.zip

[jobovy]

I have tested the performance of evaluating Combined19 on the whole APOGEE DR17. Simple for loop took 6h56m41s while this PR vectorized function only took 26s. So we can have a speedup of ~1000 times!!

🤗

This looks great, I'll try to give it a test run later today, but the code looks good to me. One question though: I don't think this code currently allows you to ask for multiple distance moduli at each (l,b), would that be useful to add? I think the code currently allows you to do an array of distance moduli at a single (l,b). Looks like allowing multiple distance moduli at each (l,b) would not be too hard, just need slightly different logic in setting up the results array in the vectorized function I think.

Also, what is the reason not just use the vectorized function for single (l,b) calls? If it works, I think it would be better not to duplicate the code. Or is it much slower to use the vectorized function for a single object?

[jobovy]

Also would be good to add a test.

[henrysky]

I have just added test as well as using the vectorized evaluation even with a single (l, b) as well as single (l,b) with multiple distance.

Also, what is the reason not just use the vectorized function for single (l,b) calls? If it works, I think it would be better not to duplicate the code. Or is it much slower to use the vectorized function for a single object?

I have tested the performance for a single (l, b) with four different distance, it is about 40ms with the original code and about 130ms with the vectorized code. So the difference in time is very small even it is technically a few times slower for a single (l, b).