Speed up Issue with HAL

[omodei]

Feature Request

During the HAWC collaboration meeting in Puerto Vallarta, it was noted that 3ML analysis was much slower with respect gammapy analysis. Quentin looked into possible bottleneck of the HAL code. Here is a summary (verbatim from him):

• The PSF for extended sources is taken at the center of the ROI and not at the position of each source, could be a problem for large ROIs:
  https://github.com/search?q=repo%3AthreeML/hawc_hal%20point_source_image&type=code
  

  hawc_hal/hawc_hal/convolved_source/convolved_point_source.py

  Line 59 in ce74038

  def get_source_map(self, response_bin_id, tag=None, integrate=False, psf_integration_method='fast'):

• The PSF convolutor for extended source uses the default psf_integration_method=‘exact’ while for point source it is set to fast
  

  hawc_hal/hawc_hal/HAL.py

  Line 209 in ce74038

  central_response_bins = self._response.get_response_dec_bin(self._roi.ra_dec_center[1])

  
  

  hawc_hal/hawc_hal/HAL.py

  Lines 215 to 217 in ce74038

  	self._psf_convolutors[bin_id] = PSFConvolutor(
  	central_response_bins[bin_id].psf, self._flat_sky_projection
  	)

  
  

  hawc_hal/hawc_hal/psf_fast/psf_convolutor.py

  Lines 24 to 25 in ce74038

  	interpolator = PSFInterpolator(psf_wrapper, flat_sky_proj)
  	psf_stamp = interpolator.point_source_image(flat_sky_proj.ra_center, flat_sky_proj.dec_center)

• For extended sources the psf convolution is not cached if the source parameters do not change
  

  hawc_hal/hawc_hal/HAL.py

  Lines 999 to 1030 in ce74038

  	for ext_id in range(n_ext_sources):
  	this_conv_src = self._convolved_ext_sources[ext_id]
  	expectation_per_transit = this_conv_src.get_source_map(energy_bin_id)
  	if this_ext_model_map is None:
  	# First addition
  	this_ext_model_map = expectation_per_transit
  	else:
  	this_ext_model_map += expectation_per_transit
  	# Now convolve with the PSF
  	if this_model_map is None:
  	# Only extended sources
  	this_model_map = (
  	self._psf_convolutors[energy_bin_id].extended_source_image(this_ext_model_map)
  	* data_analysis_bin.n_transits
  	)
  	else:
  	this_model_map += (
  	self._psf_convolutors[energy_bin_id].extended_source_image(this_ext_model_map)
  	* data_analysis_bin.n_transits
  	)

• The following should be moved to a function that return npred for each source so it can be cached if its parameters do not change, it seems that only the psf interpolation is cached
  

  hawc_hal/hawc_hal/HAL.py

  Lines 974 to 992 in ce74038

  	this_conv_src = self._convolved_point_sources[pts_id]
  	expectation_per_transit = this_conv_src.get_source_map(
  	energy_bin_id,
  	tag=None,
  	psf_integration_method=self._psf_integration_method,
  	)
  	expectation_from_this_source = expectation_per_transit * data_analysis_bin.n_transits
  	if this_model_map is None:
  	# First addition
  	this_model_map = expectation_from_this_source
  	else:
  	this_model_map += expectation_from_this_source

• many functions use this to compute npred but there is no caching
  

  hawc_hal/hawc_hal/HAL.py

  Line 966 in ce74038

  def _get_expectation(self, data_analysis_bin, energy_bin_id, n_point_sources, n_ext_sources):

• the likelihood could be evaluated on the flat geometry without need to reproject to healpix, could help if reprojection is slow
  

  hawc_hal/hawc_hal/HAL.py

  Lines 1032 to 1045 in ce74038

  	# Now transform from the flat sky projection to HEALPiX
  	if this_model_map is not None:
  	# First divide for the pixel area because we need to interpolate brightness
  	# this_model_map = old_div(this_model_map, self._flat_sky_projection.project_plane_pixel_area)
  	this_model_map = this_model_map / self._flat_sky_projection.project_plane_pixel_area
  	this_model_map_hpx = self._flat_sky_to_healpix_transform[energy_bin_id](
  	this_model_map, fill_value=0.0
  	)
  	# Now multiply by the pixel area of the new map to go back to flux
  	this_model_map_hpx *= hp.nside2pixarea(data_analysis_bin.nside, degrees=True)

• this loop could be parallelized
  

  hawc_hal/hawc_hal/HAL.py

  Line 844 in ce74038

  for bin_id in self._active_planes:

• this loop could be parallelized
  

  hawc_hal/hawc_hal/HAL.py

  Line 972 in ce74038

  for pts_id in range(n_point_sources):

• not sure to understand how the parameter change property is used but it could differentiate spatial and spectral parameters
  because if only spectral parameters change it is not necessary to repeat the PSF convolution, one can rescale the cached value with a different weight in each energy bin. (edited)
  

  hawc_hal/hawc_hal/convolved_source/convolved_extended_source.py

  Line 133 in ce74038

  def _parameter_change_callback(self, this_parameter):

[omodei]

Before starting changing the code, we should keep track of our progresses. I suggest to set up a "catalog-style" example of one ROI where the convergence/speed is problematic, and work our way from there.