kNNpy.Auxiliary.TPCF.TracerField2DA

  1####################################################################################################
  2
  3#-------------------  These libraries are required for evaluating the functions  -------------------
  4
  5import numpy as np
  6import healpy as hp
  7from astropy.coordinates import SkyCoord
  8from astropy import units as u
  9from astropy_healpix import HEALPix
 10from astropy.coordinates import ICRS
 11import time
 12import copy
 13import sys
 14import os
 15
 16#Necessary for relative imports (see https://stackoverflow.com/questions/34478398/import-local-function-from-a-module-housed-in-another-directory-with-relative-im)
 17module_path = os.path.abspath(os.path.join('../../'))
 18'''
 19@private
 20'''
 21if module_path not in sys.path:
 22    sys.path.append(module_path)
 23
 24#Importing the required helper function
 25from kNNpy.HelperFunctions_2DA import create_smoothed_field_dict_2DA
 26
 27####################################################################################################
 28
 29#--------------------------------------  Function Definitions  -------------------------------------
 30
 31def CorrelationFunction(BinsRad, MaskedTracerPosRad, FieldSkymap, NR_ND=10, ReturnSmoothedDict=False, Verbose=False):
 32    
 33    r'''
 34    Computes the angular two-point cross-correlation function between the given set of discrete tracers (`MaskedTracerPosRad`) and the given continuous overdensity field (`FieldSkymap`) at the given angular distance scales (`BinsRad`).
 35
 36    Parameters
 37    ----------
 38    BinsRad : list of numpy float array
 39        array of angular distances (in radians) to compute the cross-correlation function at
 40    MaskedTracerPosRad : numpy float array of shape ``(n_tracer, 2)``
 41        array of sky locations for the discrete tracers. For each data point in the array, the first (second) coordinate should be the declination (right ascension) in radians. Please ensure ``-np.pi/2 <= declination <= pi/2`` and ``0 <= right ascension <= 2*np.pi``.
 42    FieldSkymap : numpy float array
 43        the healpy map of the continuous field. The values of the masked pixels, if any, should be set to `hp.UNSEEN`.
 44    NR_ND : int
 45        ratio of number of randoms to number of data points used in the 2PCF calculation to remove biases caused by the presence of an observational mask, by default ``10``. This is similar to the ratio of number of randoms to number data points used in the usual Landy-Szalay estimator of the 2PCF[^1]. See notes for a more detailed explanation of this parameter and how to set it appropriately.
 46    ReturnSmoothedDict : bool, optional
 47        if set to ``True``, the dictionary containing the continuous field masked within the observational footprint, and smoothed at the provided angular distance scales, will be returned along with the nearest-neighbour measurements, by default ``False``.
 48    Verbose : bool, optional
 49        if set to ``True``, the time taken to complete each step of the calculation will be printed, by default ``False``.
 50
 51    Returns
 52    -------
 53    w_theta: numpy float array of shape ``(len(Bins)-1, )``
 54        tracer-field two-point cross-correlation function evaluated at the desired distance bins. Note that the 2PCF can't be estimated at the last bin due to the nature of the algorithm (refer to notes for details).
 55    SmoothedFieldDict : dict
 56        dictionary containing the continuous field masked within the observational footprint and smoothed at the provided angular distance scales, returned only if `ReturnSmoothedDict` is ``True``. For example, ``SmoothedFieldDict['0.215']`` represents the continuous map smoothed at a scale of 0.215 radians.
 57
 58    Raises
 59    ------
 60    ValueError
 61        if declination of any of the tracer points is not in ``[-np.pi/2, np.pi/2]``.
 62    ValueError
 63        if right ascension of any of the tracer points is not in ``[0, 2*np.pi]``.
 64    ValueError
 65        if the given tracer points are not on a two-dimensional grid.
 66    ValueError
 67        if the `NR_ND` is not an integer.
 68
 69    See Also
 70    --------
 71    kNNpy.Auxilliary.TPCF.TracerField2DA.CorrelationFunction_DataVector : computes a data vector of tracer-field two-point cross-correlations in 2D for multiple realisations of the tracer set.
 72    kNNpy.Auxilliary.TPCF.3DTPCF_Tracer-Field.CrossCorr2pt : computes tracer-field two-point cross-correlations in 3D.
 73    kNNpy.kNN_2D_Ang.TracerFieldCross2DA : computes 2D angular tracer-field cross-correlations using the $k$NN formalism.
 74
 75    Notes
 76    -----
 77    Measures the angular two-point cross-correlation function (2pcf) between a set of discrete tracers and a continuous overdensity field using the spherical band-averaging method, as described in Gupta (2024)[^1]. This is a generalisation of the spherical shell-averaging method described in Banerjee & Abel (2023)[^2].
 78    
 79    Data with associated observational footprints are supported, in which case, only tracer positions within the footprint should be provided and the field should be masked appropriately. If the footprints of the tracer set and the field are different, a combined mask representing the intersection of the two footprints should be used. The field at the masked pixels is set to 0 (which is the mean value of an overdensity field) for the purposes of the 2PCF computation to prevent any biases due to the mask.
 80
 81    <enter description for the NR_ND parameter>
 82    <enter description of algorithm and why the 2PCF can't be estimated at the last bin>
 83
 84    References
 85    ----------
 86    [^1]: Kaustubh Rajesh Gupta, Spatial clustering of gravitational wave sources with k-nearest neighbour distributions, MS Thesis, Indian Institute of Science Education and Research Pune [Digital Repository](http://dr.iiserpune.ac.in:8080/xmlui/handle/123456789/8819)
 87
 88    [^2]: Arka Banerjee, Tom Abel, Tracer-field cross-correlations with k-nearest neighbour   distributions, [Monthly Notices of the Royal Astronomical Society](https://doi.org/10.1093/mnras/stac3813), Volume 519, Issue 4, March 2023, Pages 4856–4868
 89    '''
 90    
 91    #-----------------------------------------------------------------------------------------------
 92
 93    if Verbose: total_start_time = time.perf_counter()
 94
 95    #-----------------------------------------------------------------------------------------------
 96        
 97    #Step 0: Check all inputs are consistent with the function requirement
 98
 99    if Verbose: print('Checking inputs ...')
100
101    if np.any((MaskedTracerPosRad[:, 0]<-np.pi/2) | (MaskedTracerPosRad[:, 0]>np.pi/2)):
102        raise ValueError('Invalid tracer point position(s): please ensure -pi/2 <= declination <= pi/2.')
103
104    if np.any((MaskedTracerPosRad[:, 1]<0) | (MaskedTracerPosRad[:, 0]>2*np.pi)):
105        raise ValueError('Invalid tracer point position(s): please ensure 0 <= right ascension <= 2*pi.')
106
107    if MaskedTracerPosRad.shape[1]!=2: 
108        raise ValueError('Incorrect spatial dimension for tracers: array containing the tracer positions must be of shape (n_tracer, 2), where n_tracer is the number of tracers.')
109    
110    if not isinstance(NR_ND, int): 
111        raise ValueError("Please input an integer value for 'NR_ND'.")
112
113    if Verbose: print('\tdone.')
114
115    #-----------------------------------------------------------------------------------------------
116
117    #Getting the mask, sky coverage and HEALPix NSIDE
118
119    mask = np.ones_like(FieldSkymap).astype(int)
120    mask_ind = np.where(FieldSkymap==hp.UNSEEN)
121    mask[mask_ind] = 0
122    sky_frac = len(np.where(mask==1)[0])/len(mask)
123    NSIDE = hp.get_nside(FieldSkymap)
124
125    #-----------------------------------------------------------------------------------------------
126
127    #Calculate the smoothed field dictionary
128
129    if Verbose: 
130        smooth_start_time = time.perf_counter()
131        print('\nSmoothing the continuous field at the given angular distance scales...')
132
133    #Note: QueryMask is set to 2 everywhere to avoid masking out any additional pixels. The pixels outside the footprint are automatically masked during the smoothing process (refer to the HelperFunctions module documentation for more details).
134    SmoothedFieldDict = create_smoothed_field_dict_2DA(FieldSkymap, [BinsRad], query_mask=2*np.ones_like(FieldSkymap).astype(int), Verbose=False)
135
136    if Verbose: print('\tdone; time taken: {:.2e} s.'.format(time.perf_counter()-smooth_start_time))
137    
138    #-----------------------------------------------------------------------------------------------
139
140    #Computing the 2pcf
141
142    if Verbose: 
143        tpcf_start_time = time.perf_counter()
144        print('\nComputing the 2PCF...')
145
146    w_theta = np.zeros(len(BinsRad)-1)
147    
148    for i, ss in enumerate(BinsRad[:-1]):
149
150        sdgm_low = copy.deepcopy(SmoothedFieldDict[str(ss)])
151        sdgm_low[sdgm_low==hp.UNSEEN] = 0
152        sdgm_high = copy.deepcopy(SmoothedFieldDict[str(BinsRad[i+1])])
153        sdgm_high[sdgm_high==hp.UNSEEN] = 0
154
155        hp_object = HEALPix(nside=NSIDE, order='ring', frame=ICRS())
156
157        #Discrete tracers
158        ra = MaskedTracerPosRad[:, 1]*u.radian
159        dec = MaskedTracerPosRad[:, 0]*u.radian
160        coords = SkyCoord(ra, dec, frame='icrs')
161        ds2 = hp_object.interpolate_bilinear_skycoord(coords, sdgm_high)
162        ds1 = hp_object.interpolate_bilinear_skycoord(coords, sdgm_low)
163
164        #Randoms
165        n_rand = MaskedTracerPosRad.shape[0]*NR_ND
166        ra_rand = np.random.uniform(low=0, high=2*np.pi, size=2*int(n_rand/sky_frac))
167        dec_rand = np.arcsin(np.random.uniform(low=-1, high=1, size=2*int(n_rand/sky_frac)))
168        ipix = hp.ang2pix(NSIDE, np.rad2deg(ra_rand), np.rad2deg(dec_rand), nest=False, lonlat=True)
169        mask_val = mask[ipix]
170        ind_masked = np.where(mask_val==1)[0]
171        ra_rand_masked = ra_rand[ind_masked]
172        dec_rand_masked = dec_rand[ind_masked]
173        ind_ds = np.random.choice(len(ra_rand_masked), n_rand, replace=False)
174        ra_rand_masked_ds = ra_rand_masked[ind_ds]*u.radian
175        dec_rand_masked_ds = dec_rand_masked[ind_ds]*u.radian
176        coords_rand = SkyCoord(ra_rand_masked_ds, dec_rand_masked_ds, frame='icrs')
177        ds2_rand = hp_object.interpolate_bilinear_skycoord(coords_rand, sdgm_high)
178        ds1_rand = hp_object.interpolate_bilinear_skycoord(coords_rand, sdgm_low)
179        
180        A2 = 2*np.pi*(1-np.cos(BinsRad[i+1]))
181        A1 = 2*np.pi*(1-np.cos(BinsRad[i]))
182
183        w_theta[i] = np.mean((A2*ds2-A1*ds1)/(A2-A1)) - np.mean((A2*ds2_rand-A1*ds1_rand)/(A2-A1))
184
185    if Verbose: print('\tdone; time taken: {:.2e} s.'.format(time.perf_counter()-tpcf_start_time))
186
187    #-----------------------------------------------------------------------------------------------
188
189    if Verbose: print('\ntotal time taken: {:.2e} s.'.format(time.perf_counter()-total_start_time))
190    
191    if ReturnSmoothedDict: 
192        return w_theta, SmoothedFieldDict
193    else: 
194        return w_theta
195
196####################################################################################################
197    
198def CorrelationFunction_DataVector(BinsRad, MaskedTracerPosVectorRad, FieldSkymap, NR_ND=10, ReturnSmoothedDict=False, Verbose=False):
199    
200    r'''
201    Returns a 'data vector' of the angular two-point cross-correlation function between multiple realisations of the given set of discrete tracers (`MaskedTracerPosVectorRad`) and a single realisation of the given continuous overdensity field (`FieldSkymap`) at the given angular distance scales (`BinsRad`). Please refer to notes to understand why this might be useful.
202
203    Parameters
204    ----------
205    BinsRad : list of numpy float array
206        array of angular distances (in radians) to compute the cross-correlation function at
207    MaskedTracerPosVectorRad : numpy float array of shape ``(n_realisations, n_tracer, 2)``
208        array of sky locations for the first set of discrete tracers. For each data point in the array, the first (second) coordinate should be the declination (right ascension) in radians. Please ensure ``-np.pi/2 <= declination <= pi/2`` and ``0 <= right ascension <= 2*np.pi``.
209    FieldSkymap : numpy float array
210        the healpy map of the continuous field. The values of the masked pixels, if any, should be set to `hp.UNSEEN`.
211    NR_ND : int
212        ratio of number of randoms to number of data points used in the 2PCF calculation to remove biases caused by the presence of an observational mask, by default ``10``. This is similar to the ratio of number of randoms to number data points used in the usual Landy-Szalay estimator of the 2PCF[^1]. See notes for a more detailed explanation of this parameter and how to set it appropriately.
213    ReturnSmoothedDict : bool, optional
214        if set to ``True``, the dictionary containing the continuous field masked within the observational footprint, and smoothed at the provided angular distance scales, will be returned along with the nearest-neighbour measurements, by default ``False``.
215    Verbose : bool, optional
216        if set to ``True``, the time taken to complete each step of the calculation will be printed, by default ``False``.
217
218    Returns
219    -------
220    w_theta_vector: numpy float array of shape ``(n_realisations, len(BinsRad)-1)``
221        data vector containing the tracer-field two-point cross-correlation function for multiple realisations of the discrete tracer set evaluated at the desired distance bins. Note that the 2PCF can't be estimated at the last bin due to the nature of the algorithm (refer to notes for details).
222    SmoothedFieldDict : dict
223        dictionary containing the continuous field masked within the observational footprint and smoothed at the provided angular distance scales, returned only if `ReturnSmoothedDict` is ``True``. For example, ``SmoothedFieldDict['0.215']`` represents the continuous map smoothed at a scale of 0.215 radians.
224
225    Raises
226    ------
227    ValueError
228        if declination of any of the tracer points is not in ``[-np.pi/2, np.pi/2]``.
229    ValueError
230        if right ascension of any of the tracer points is not in ``[0, 2*np.pi]``.
231    ValueError
232        if the given tracer points are not on a two-dimensional grid.
233    ValueError
234        if the `NR_ND` is not an integer.
235
236    See Also
237    --------
238    kNNpy.Auxilliary.TPCF.TracerField2DA.CorrelationFunction : computes tracer-field two-point cross-correlations in 2D for a single realisation of the tracer set.
239    kNNpy.Auxilliary.TPCF.3DTPCF_Tracer-Field.CrossCorr2pt : computes tracer-field two-point cross-correlations in 3D.
240    kNNpy.kNN_2D_Ang.TracerFieldCross2DA_DataVector : computes a data vector of 2D angular tracer-field cross-correlations for multiple tracer realisations using the $k$NN formalism.
241
242    Notes
243    -----
244    Please refer to the documentation of kNNpy.Auxilliary.TPCF.TracerFieldCross2DA.CorrelationFunction for important usage notes that also apply to this function and references. <Explain why cross-correlating multiple realisations of tracer with single realisation of field might be useful>
245
246    References
247    ----------
248    [^1]: Kaustubh Rajesh Gupta, Spatial clustering of gravitational wave sources with k-nearest neighbour distributions, MS Thesis, Indian Institute of Science Education and Research Pune [Digital Repository](http://dr.iiserpune.ac.in:8080/xmlui/handle/123456789/8819)
249    '''
250    
251    #-----------------------------------------------------------------------------------------------
252
253    if Verbose: total_start_time = time.perf_counter()
254
255    #-----------------------------------------------------------------------------------------------
256        
257    #Step 0: Check all inputs are consistent with the function requirement
258
259    if Verbose: print('Checking inputs ...')
260
261    if np.any((MaskedTracerPosVectorRad[:, :, 0]<-np.pi/2) | (MaskedTracerPosVectorRad[:, :, 0]>np.pi/2)):
262        raise ValueError('Invalid tracer point position(s): please ensure -pi/2 <= declination <= pi/2.')
263
264    if np.any((MaskedTracerPosVectorRad[:, :, 1]<0) | (MaskedTracerPosVectorRad[:, :, 0]>2*np.pi)):
265        raise ValueError('Invalid tracer point position(s): please ensure 0 <= right ascension <= 2*pi.')
266
267    if MaskedTracerPosVectorRad.shape[2]!=2: 
268        raise ValueError('Incorrect spatial dimension for tracers: array containing the tracer positions must be of shape (n_tracer, 2), where n_tracer is the number of tracers.')
269    
270    if not isinstance(NR_ND, int): 
271        raise ValueError("Please input an integer value for 'NR_ND'.")
272
273    if Verbose: print('\tdone.')
274
275    #-----------------------------------------------------------------------------------------------
276
277    #Getting the mask, sky coverage and HEALPix NSIDE
278
279    mask = np.ones_like(FieldSkymap).astype(int)
280    mask_ind = np.where(FieldSkymap==hp.UNSEEN)
281    mask[mask_ind] = 0
282    sky_frac = len(np.where(mask==1)[0])/len(mask)
283    NSIDE = hp.get_nside(FieldSkymap)
284
285    #-----------------------------------------------------------------------------------------------
286
287    #Calculate the smoothed field dictionary
288
289    if Verbose: 
290        smooth_start_time = time.perf_counter()
291        print('\nSmoothing the continuous field at the given angular distance scales...')
292
293    #Note: QueryMask is set to 2 everywhere to avoid masking out any additional pixels. The pixels outside the footprint are automatically masked during the smoothing process (refer to the HelperFunctions module documentation for more details).
294    SmoothedFieldDict = create_smoothed_field_dict_2DA(FieldSkymap, [BinsRad], query_mask=2*np.ones_like(FieldSkymap).astype(int), Verbose=False)
295
296    if Verbose: print('\tdone; time taken: {:.2e} s.'.format(time.perf_counter()-smooth_start_time))
297    
298    #-----------------------------------------------------------------------------------------------
299
300    #Looping the 2PCF calculation over the multiple realisations
301
302    n_reals = MaskedTracerPosVectorRad.shape[0]
303    w_theta_vector = np.zeros((n_reals, len(BinsRad)-1))
304
305    for realisation, MaskedTracerPosRad in enumerate(MaskedTracerPosVectorRad):
306
307        if Verbose:
308            start_time_real = time.perf_counter()
309            print(f'\n\n--------------  Realisation {realisation+1}/{n_reals}  --------------\n')
310
311        #-------------------------------------------------------------------------------------------
312
313        #Computing the 2pcf
314
315        if Verbose: 
316            print('Computing the 2PCF...')
317
318        w_theta_vector[realisation] = np.zeros(len(BinsRad)-1)
319        
320        for i, ss in enumerate(BinsRad[:-1]):
321
322            sdgm_low = copy.deepcopy(SmoothedFieldDict[str(ss)])
323            sdgm_low[sdgm_low==hp.UNSEEN] = 0
324            sdgm_high = copy.deepcopy(SmoothedFieldDict[str(BinsRad[i+1])])
325            sdgm_high[sdgm_high==hp.UNSEEN] = 0
326
327            hp_object = HEALPix(nside=NSIDE, order='ring', frame=ICRS())
328
329            #Discrete tracers
330            ra = MaskedTracerPosRad[:, 1]*u.radian
331            dec = MaskedTracerPosRad[:, 0]*u.radian
332            coords = SkyCoord(ra, dec, frame='icrs')
333            ds2 = hp_object.interpolate_bilinear_skycoord(coords, sdgm_high)
334            ds1 = hp_object.interpolate_bilinear_skycoord(coords, sdgm_low)
335
336            #Randoms
337            n_rand = MaskedTracerPosRad.shape[0]*NR_ND
338            ra_rand = np.random.uniform(low=0, high=2*np.pi, size=2*int(n_rand/sky_frac))
339            dec_rand = np.arcsin(np.random.uniform(low=-1, high=1, size=2*int(n_rand/sky_frac)))
340            ipix = hp.ang2pix(NSIDE, np.rad2deg(ra_rand), np.rad2deg(dec_rand), nest=False, lonlat=True)
341            mask_val = mask[ipix]
342            ind_masked = np.where(mask_val==1)[0]
343            ra_rand_masked = ra_rand[ind_masked]
344            dec_rand_masked = dec_rand[ind_masked]
345            ind_ds = np.random.choice(len(ra_rand_masked), n_rand, replace=False)
346            ra_rand_masked_ds = ra_rand_masked[ind_ds]*u.radian
347            dec_rand_masked_ds = dec_rand_masked[ind_ds]*u.radian
348            coords_rand = SkyCoord(ra_rand_masked_ds, dec_rand_masked_ds, frame='icrs')
349            ds2_rand = hp_object.interpolate_bilinear_skycoord(coords_rand, sdgm_high)
350            ds1_rand = hp_object.interpolate_bilinear_skycoord(coords_rand, sdgm_low)
351            
352            A2 = 2*np.pi*(1-np.cos(BinsRad[i+1]))
353            A1 = 2*np.pi*(1-np.cos(BinsRad[i]))
354
355            w_theta_vector[realisation, i] = np.mean((A2*ds2-A1*ds1)/(A2-A1)) - np.mean((A2*ds2_rand-A1*ds1_rand)/(A2-A1))
356
357        if Verbose: print('\tdone. time taken for realisation {}: {:.2e} s.'.format(realisation+1, time.perf_counter()-start_time_real))
358
359    #-----------------------------------------------------------------------------------------------
360
361    if Verbose: print('\ntotal time taken: {:.2e} s.'.format(time.perf_counter()-total_start_time))
362    
363    if ReturnSmoothedDict: 
364        return w_theta_vector, SmoothedFieldDict
365    else: 
366        return w_theta_vector
367
368####################################################################################################
369
370#----------------------------------------  END OF PROGRAM!  ----------------------------------------
371
372####################################################################################################