Add utility to compute the from Yi Tan et al. 2023

Author: aymgal

coolest/api/analysis.py

@@ -87,7 +87,7 @@ def effective_einstein_radius(self, center=None, initial_guess=1, initial_delta_
         else:
             center_x, center_y = center
 
-        def loop_until_overshoot(r_Ein, delta, direction, runningtotal, total_pix):
+        def _loop_until_overshoot(r_Ein, delta, direction, runningtotal, total_pix):
             """
             this subfunction iteratively adjusts the mask radius by delta either inward or outward until the sign flips on mean_kappa-area
             """
@@ -145,7 +145,7 @@ def only_between_r1r2(r1, r2, r_grid):
 
         for n in range(n_iter):
             #overshoots, turn around and backtrack at higher precision
-            r_Ein, delta, direction, runningtotal, total_pix = loop_until_overshoot(r_Ein, delta, direction, runningtotal, total_pix)
+            r_Ein, delta, direction, runningtotal, total_pix = _loop_until_overshoot(r_Ein, delta, direction, runningtotal, total_pix)
             direction=direction*-1
             delta=delta/2
         accuracy=grid_res/2 #after testing, accuracy is about grid_res/2
@@ -216,10 +216,9 @@ def effective_radial_slope(self, r_eval=None, center=None, r_vec=np.linspace(0,
         if r_eval==None:
             return slope 
         else:
-            closest_r = self.find_nearest(r_vec,r_eval) #just takes closest r. Could rebuild it to interpolate.
+            closest_r = self._find_nearest(r_vec,r_eval) #just takes closest r. Could rebuild it to interpolate.
             return slope[r_vec==closest_r]
 
-        
     def effective_radius_light(self, outer_radius=10, center=None, coordinates=None,
                                initial_guess=1, initial_delta_pix=10, 
                                n_iter=10, return_model=False, return_accuracy=False,
@@ -307,14 +306,6 @@ def effective_radius_light(self, outer_radius=10, center=None, coordinates=None,
             return r_eff, accuracy
         return r_eff
 
-    
-    def find_nearest(self, array, value):
-        """subfunction to find nearest closest element in array to value"""
-        array = np.asarray(array)
-        idx = (np.abs(array - value)).argmin()
-        return array[idx]
-    
-
     def two_point_correlation(self, Nbins=100, rmax=None, normalize=False, 
                               use_profile_coordinates=True, coordinates=None, 
                               min_flux=None, min_flux_frac=None, 
@@ -508,7 +499,6 @@ def total_magnitude(self, outer_radius=10, center=None, coordinates=None,
         mag_tot = -2.5*np.log10(flux_tot) + mag_zero_point
 
         return mag_tot
-    
 
     def ellipticity_from_moments(self, center=None, coordinates=None, **kwargs_selection):
         """Estimates the axis ratio and position angle of the model map 
@@ -571,6 +561,27 @@ def ellipticity_from_moments(self, center=None, coordinates=None, **kwargs_selec
 
         return q, phi
 
+    def lensing_information(self, a=16, b=0, 
+                            noise_map=None, arc_mask=None, theta_E=None,
+                            entity_idx_theta_E=0, profile_idx_theta_E=0):
+        """
+        Computes the 'lensing information' defined in Yi Tan et al. 2023, Equations (8) and (9).
+        https://ui.adsabs.harvard.edu/abs/2023arXiv231109307T/abstract
+        """
+        data = self.coolest.observation.pixels.get_pixels()
+        # TODO: subtract the lens light
+        print("WARNING: no lens light subtracted; assuming the data contains only the arcs.")
+        if noise_map is None:
+            raise NotImplementedError("Getting the noise map from the COOLEST instance is yet implemented.")
+        if theta_E is None:
+            mass_profile = self.coolest.lensing_entities[entity_idx_theta_E].mass_model[profile_idx_theta_E]
+            theta_E = mass_profile.parameters['theta_E'].point_estimate.value
+        x, y = self.coordinates.pixel_coordinates
+        I, theta_E, phi_ref, mask = util.lensing_information(
+            data, x, y, theta_E, noise_map, a=a, b=b, arc_mask=arc_mask
+        )
+        return I, theta_E, phi_ref, mask
+
 
     @staticmethod
     def effective_radius(light_map, x, y, outer_radius=10, initial_guess=1, initial_delta_pix=10, n_iter=10):
@@ -633,3 +644,10 @@ def effective_radius(light_map, x, y, outer_radius=10, initial_guess=1, initial_
 
         return r_eff, grid_res
 
+    @staticmethod
+    def _find_nearest(array, value):
+        """subfunction to find nearest closest element in array to value"""
+        array = np.asarray(array)
+        idx = (np.abs(array - value)).argmin()
+        return array[idx]
+    

coolest/api/util.py

@@ -135,6 +135,66 @@ def downsampling(image, factor=1):
         raise ValueError(f"Downscaling factor {factor} is not possible with shape ({nx}, {ny})")
 
 
+def lensing_information(data_lens_sub, x, y, theta_E, noise_map, center_x_lens=0, center_y_lens=0,
+                        a=16, b=0, arc_mask=None):
+    """
+    Computes the 'lensing information' defined in Yi Tan et al. 2023, Equations (8) and (9).
+    https://ui.adsabs.harvard.edu/abs/2023arXiv231109307T/abstract
+
+    Parameters
+    ----------
+    data_lens_sub : np.ndarray
+        Imaging data as a 2D array. It is assumed to contain no lens light.
+    x : np.ndarray
+        2D array of x coordinates, in arcsec.
+    y : np.ndarray
+        2D array of y coordinates, in arcsec.
+    theta_E : float
+        Einstein radius in arcsec, by default None
+    noise_map : np.ndarray
+        2D array with 1-sigma noise level per pixel (same units as `data_lens_sub`), by default None
+    center_x_lens : int, optional
+        x coordinates of the center of the lens, by default 0
+    center_y_lens : int, optional
+        y coordinates of the center of the lens, by default 0
+    a : int, optional
+        Exponent in Eq. (9) from Yi Tan et al. 2023, by default 16
+    b : int, optional
+        Exponent in Eq. (9) from Yi Tan et al. 2023, by default 0
+    arc_mask : np.ndarray, optional
+        Binary 2D array with 1s where there is are lensed arcs, by default None
+
+    Returns
+    -------
+    4-tuple
+        Lensing information I, Einstein radius, reference azimuthal angle, total mask used for computing I
+    """
+    if arc_mask is None:
+        arc_mask = np.ones_like(data_lens_sub)
+    # estimate background noise from one corner of the noise map
+    sigma_bkg = np.mean(noise_map[:10, :10])
+    # build a mask to only consider pixels at least 3 times the background noise level
+    snr_mask = np.where(data_lens_sub > 3.*sigma_bkg, 1., 0.)
+    # combine user mask and SNR mask
+    arc_mask_tot = snr_mask * arc_mask
+    # shift coordinates so that lens is at (0, 0)
+    theta_x, theta_y = x - center_x_lens, y - center_y_lens
+    # compute polar coordinates centered on the lens
+    theta_r = np.hypot(theta_x, theta_y)
+    phi = np.arctan2(theta_y, theta_x)
+    # find index of the brightest pixel (within the arc mask)
+    max_idx = np.where(data_lens_sub == (data_lens_sub*arc_mask_tot).max())
+    # get azimuthal angle corresponding to the brightest pixel
+    phi_ref = float(np.arctan2(theta_y[max_idx], theta_x[max_idx]))
+    # compute the weights following Eq. (9) from Yi Tan et al. 2023
+    weights = ( 1. + np.abs(theta_r - theta_E) / theta_E * (1 + np.abs(phi - phi_ref) / phi_ref)**b )**a
+    # compute the weighted sum
+    numerator = np.sum(arc_mask_tot*weights*data_lens_sub)
+    denominator = np.sqrt(np.sum(arc_mask_tot*noise_map**2))
+    lens_I = numerator / denominator
+    return lens_I, theta_E, phi_ref, arc_mask_tot
+
+
 def split_lens_source_params(coolest_list, name_list, lens_light=False):
     """
     Read several json files already containing a model with the results of this model fitting
@@ -147,7 +207,7 @@ def split_lens_source_params(coolest_list, name_list, lens_light=False):
 
     OUTPUT
     ------
-     param_all_lens, param_all_source: organized dictionnaries readable by plotting function
+     param_all_lens, param_all_source: organized dictionaries readable by plotting function
     """
 
     param_all_lens = {}
@@ -417,6 +477,7 @@ def read_sersic(light, param={}, prefix='Sersic_0_'):
 
     return param
 
+
 def find_critical_lines(coordinates, mag_map):
     # invert and find contours corresponding to infinite magnification (i.e., changing sign)
     inv_mag = 1. / np.array(mag_map)
@@ -428,6 +489,7 @@ def find_critical_lines(coordinates, mag_map):
         lines.append((np.array(curve_x), np.array(curve_y)))
     return lines
 
+
 def find_caustics(crit_lines, composable_lens):
     """`composable_lens` can be an instance of `ComposableLens` or `ComposableMass`"""
     lines = []
@@ -436,6 +498,7 @@ def find_caustics(crit_lines, composable_lens):
         lines.append((np.array(cl_src_x), np.array(cl_src_y)))
     return lines
 
+
 def find_all_lens_lines(coordinates, composable_lens):
     """`composable_lens` can be an instance of `ComposableLens` or `ComposableMass`"""
     from coolest.api.composable_models import ComposableLensModel, ComposableMassModel  # avoiding circular imports