vectorize HierarchicalHealpixMap evluate

README.rst

@@ -81,7 +81,7 @@ roughly the following lay-out::
  bayestar2019.h5
  maps/
  SFD_dust_4096_ngp.fits
-  SFD_dust_4096_sgp.fits
+  SFD_dust_4096_sgp.fits
  marshall06/
  ReadMe
  table1.dat
@@ -127,10 +127,15 @@ and they can be plotted as a function of distance at a given (l,b)
 
  combined.plot(55.,0.5) # inputs are (l,b)
 
-(plot not shown). Maps that are derived from the
-``HierarchicalHealpixMap.py`` class (currently all Green-type maps and
-the combined maps) can also be plotted on the sky using a Mollweide
-projection at a given distance using
+(plot not shown). Maps that are derived from the ``HierarchicalHealpixMap.py`` class (currently all Green-type maps and
+the combined maps) can be vectorized to evaluate on array inputs of *l*, *b*, *D*
+
+.. code-block:: python
+
+ combined(numpy.array([30.,40.,50.,60.]),numpy.array([3.,4.,3.,6.]),numpy.array([1.,2.,3.,10.]))
+ array([0.22304147, 0.3780736 , 0.42528571, 0.22258065])
+
+They can also be plotted on the sky using a Mollweide projection at a given distance using
 
 .. code-block:: python
 

mwdust/HierarchicalHealpixMap.py

@@ -32,38 +32,50 @@ def __init__(self,filter=None,sf10=True):
  self._sf10= sf10
  return None
 
- def _evaluate(self,l,b,d):
+
+ def _evaluate(self, ls, bs, ds):
  """
  NAME:
  _evaluate
  PURPOSE:
- evaluate the dust-map
+ evaluate the dust-map for array input
  INPUT:
- l- Galactic longitude (deg)
- b- Galactic latitude (deg)
+ l- Galactic longitude (deg) can be array
+ b- Galactic latitude (deg) can be array
  d- distance (kpc) can be array
  OUTPUT:
  extinction E(B-V)
  HISTORY:
  2015-03-02 - Started - Bovy (IAS)
+ 2023-07-05 - Vectorized - Henry Leung (UofT)
  """
- distmod= 5.*numpy.log10(d)+10.
- if isinstance(l,numpy.ndarray) or isinstance(b,numpy.ndarray):
- raise NotImplementedError("array input for l and b for combined dust map not implemented")
- lbIndx= self._lbIndx(l,b)
- if self._intps[lbIndx] != 0:
- out= self._intps[lbIndx][0](distmod)
- else:
- interpData=\
- interpolate.InterpolatedUnivariateSpline(self._distmods,
- self._best_fit[lbIndx],
- k=self._interpk)
- out= interpData(distmod)
- self._intps[lbIndx]= interpData
- if self._filter is None:
- return out
- else:
- return out*aebv(self._filter,sf10=self._sf10)
+ ls = numpy.atleast_1d(ls)
+ bs = numpy.atleast_1d(bs)
+ ds = numpy.atleast_1d(ds)
+
+ distmod= 5.*numpy.log10(ds)+10.
+ lbIndx= self._lbIndx(ls, bs)
+ if len(ls) == 1 and len(ds) > 1:
+ lbIndx = numpy.tile(lbIndx, len(ds))
+
+ result = numpy.zeros_like(ds)
+ for counter, i, d in zip(numpy.arange(len(result)), lbIndx, distmod):
+ if self._intps[i] != 0:
+ out= self._intps[i](d)
+ else:
+ interpData=\
+ interpolate.InterpolatedUnivariateSpline(self._distmods,
+ self._best_fit[i],
+ k=self._interpk)
+ out= interpData(d)
+ self._intps[i]= interpData
+ result[counter] = out
+ if self._filter is not None:
+ result = result * aebv(self._filter,sf10=self._sf10)
+ # set nan for invalid indices
+ result[lbIndx==-1] = numpy.nan
+ return result
+
 
  def dust_vals_disk(self,lcen,bcen,dist,radius):
  """
@@ -122,19 +134,27 @@ def dust_vals_disk(self,lcen,bcen,dist,radius):
  extinction= extinction*aebv(self._filter,sf10=self._sf10) 
  return (pixarea,extinction)
 
- def _lbIndx(self,l,b):
- """Return the index in the _combineddata array corresponding to this (l,b)"""
+ def _lbIndx(self, ls, bs):
+ """Return the indices in the _combineddata array corresponding to arrays of (l, b)"""
+ stop_mask = numpy.zeros(len(ls), dtype=bool) # mask to be accumulated when looping through nside
+ indx_result = numpy.ones(len(ls), dtype=int) * -1 # -1 for bad star, int array has no nan
  for nside in self._nsides:
- # Search for the pixel in this Nside level
- tpix= ang2pix(nside,(90.-b)*_DEGTORAD,
- l*_DEGTORAD,nest=True)
- indx= (self._pix_info['healpix_index'] == tpix)\
- *(self._pix_info['nside'] == nside)
- if numpy.sum(indx) == 1:
- return self._indexArray[indx]
- elif numpy.sum(indx) > 1:
- raise IndexError("Given (l,b) pair has multiple matches!")
- raise IndexError("Given (l,b) pair not within the region covered by the extinction map")
+ tpix = ang2pix(nside, (90.-bs)*_DEGTORAD, ls*_DEGTORAD, nest=True)
+ nside_idx = numpy.where(self._pix_info['nside'] == nside)[0]
+ healpix_index_nside = self._pix_info['healpix_index'][nside_idx]
+ sorted_order = numpy.argsort(healpix_index_nside)
+ # use searchsorted to find the index of tpix in healpix_index_nside efficiently
+ result = numpy.searchsorted(healpix_index_nside, tpix, sorter=sorted_order)
+ # need to deal with indices where they are larger than the largest healpix_index_nside
+ known_bad_idx = (result == len(nside_idx))
+ result[known_bad_idx] = 0 # wrap around
+ result = sorted_order[result] # reverse the sorting before indexing
+ result = nside_idx[result]
+ good_result_idx = ((self._pix_info['healpix_index'][result] == tpix) & (self._pix_info['nside'][result] == nside) & (~known_bad_idx))
+ indx_result = numpy.where(~stop_mask & good_result_idx, result, indx_result)
+ indx_result = numpy.where(known_bad_idx & ~stop_mask, -1, indx_result) # set bad star to -1
+ stop_mask = stop_mask | good_result_idx # update mask for the next nside
+ return indx_result
 
  def plot_mollweide(self,d,**kwargs):
  """

mwdust/util/healpix.py

@@ -38,7 +38,7 @@ def check_nside(nside, nest=False):
  # nside can only be a power of 2 for nest, but generally less than 2**29
  nside_arr = np.array(nside).astype(np.int64)
  is_ok = True
- if nside == nside_arr and 0 < nside and nside <= 2**29:
+ if np.all(np.logical_and(np.logical_and(nside == nside_arr, np.all(np.less(0, nside))), nside_arr <= 2**29)):
  if nest:
  is_ok = (nside_arr & (nside_arr - 1)) == 0
  else:

tests/test_green19.py

@@ -4,7 +4,7 @@
 rng= default_rng()
 
 
-def test_distance_dependent():
+def test_distance_dependent_forloop():
  # Test that extinction from the Green19 map monotonicaly increases 
  # with distance for a given line of sight.
  from mwdust import Green19
@@ -16,12 +16,27 @@ def test_distance_dependent():
  for ii in range(len(glons))]
  for jj in range(len(dists))]).T 
  assert numpy.all(numpy.fabs(ebvs-ebvs[0])>=0), \
- 'Green19 extinction does not monotonically increase with distance for at lease one given line of sight'
+ 'For-loop Green19 extinction does not monotonically increase with distance for at lease one given line of sight'
  return None
 
 
-def test_against_known_values():
- # Test that the SFD extinction agrees with a table of known values
+def test_distance_dependent_vectorized():
+ # Test that extinction from the Green19 map monotonicaly increases 
+ # with distance for a given line of sight.
+ from mwdust import Green19
+ green19= Green19()
+ glons= rng.uniform(50.,220.,size=20)
+ glats= rng.uniform(-90.,90.,size=20)
+ dists= rng.uniform(0.01,1.,size=5)
+ ebvs= numpy.array([[green19(glons[ii],glats[ii],dists) \
+ for ii in range(len(glons))]]).T 
+ assert numpy.all(numpy.fabs(ebvs-ebvs[0])>=0), \
+ 'Vectorized Green19 extinction does not monotonically increase with distance for at lease one given line of sight'
+ return None
+
+
+def test_against_known_values_forloop():
+ # Test that the Green19 extinction agrees with a table of known values
  # These were computed using mwdust on Linux using mwdust 1.2
  # before custom healpix functions were implemented (thus computed using healpy)
  from mwdust import Green19
@@ -33,5 +48,22 @@ def test_against_known_values():
  ebvs= known.T[3]
  assert numpy.all(ebvs-[green19(glon,glat,dist)[0]
  for glon,glat,dist in zip(glons,glats,dists)] < 10.**-8.), \
- 'Green19 extinction does not agree with known values'
+ 'For-loop Green19 extinction does not agree with known values'
+ return None
+
+
+def test_against_known_values_vectorized():
+ # Test that the Green19 extinction agrees with a table of known values
+ # These were computed using mwdust on Linux using mwdust 1.2
+ # before custom healpix functions were implemented (thus computed using healpy)
+ # and before the vectorized version of the Green19 map was implemented
+ from mwdust import Green19
+ green19= Green19()
+ known= numpy.loadtxt(pathlib.Path(__file__).parent / 'green19_benchmark.dat',delimiter=',')
+ glons= known.T[0]
+ glats= known.T[1]
+ dists= known.T[2]
+ ebvs= known.T[3]
+ assert numpy.all(ebvs-green19(glons,glats,dists) < 10.**-8.), \
+ 'Vectorized Green19 extinction does not agree with known values'
  return None