correlation function

src/tools21cm/corr_function.py

@@ -2,15 +2,40 @@
 from scipy.interpolate import splrep,splev
 from scipy.spatial import cKDTree
 from tqdm import tqdm
-from .power_spectrum import power_spectrum_1d
+from .power_spectrum import power_spectrum_1d, cross_power_spectrum_1d
 from . import conv
 
+def ps_to_corr(ps, ks, rbins=10, box_dims=None, n_grid=None, binning='log'):
+	'''
+	Function that converts the spherically averaged power spectrum into 
+	the correlation function using the following equation:
+	.. math:: \\xi(r) = \\int 4\\pi k^2 P(k) \\mathrm{sinc}(kr)dk
+	'''
+	if isinstance(rbins, (int)): 
+		rbins = 10**np.linspace(np.log10(2*box_dims/n_grid), np.log10(box_dims/2), rbins) if binning=='log' \
+					else np.linspace(2*box_dims/n_grid, box_dims/2, rbins)
+
+	# if binning=='log':
+	# 	ps_tck = splrep(np.log10(ks[np.isfinite(ps)]), ps[np.isfinite(ps)], k=1)
+	# 	ps_fun = lambda k: splev(np.log10(k),ps_tck)
+	# else:
+	# 	ps_tck = splrep(ks[np.isfinite(ps)], ps[np.isfinite(ps)], k=1)
+	# 	ps_fun = lambda k: splev(k,ps_tck)
+	# intergral_bins = kbins*10
+	# knew = np.linspace(ks[0],ks[-1],intergral_bins)
+	# pnew = ps_fun(knew)
+	knew, pnew = ks[np.isfinite(ps)], ps[np.isfinite(ps)]
+
+	integ = knew**2*pnew*np.sinc(knew[None,:]*rbins[:,None])
+	corr  = np.trapz(integ, knew, axis=1) 
+
+	return corr*4*np.pi, rbins
 
-def correlation_function(input_array_nd, rbins=10, kbins=10, box_dims=None, binning='log'):
+def correlation_function(input_array_nd, rbins=10, kbins=10, box_dims=None, binning='log', k_binning='log'):
 	'''
 	Function estimates the spherically averaged power spectrum and 
 	then calulates the correlation function using the following equation:
-	.. math:: \\xi(r) = (2\\pi^2)^{-1} \\int k^2 P(k) \\mathrm{sinc}(kr)dk
+	.. math:: \\xi(r) = \\int 4\\pi k^2 P(k) \\mathrm{sinc}(kr)dk
 
 	Parameters: 
 				input_array_nd (numpy array): the data array
@@ -25,29 +50,58 @@ def correlation_function(input_array_nd, rbins=10, kbins=10, box_dims=None, binn
 				    dimensions. If it is a float, this is taken as the box length
 				    along all dimensions. If it is an array-like, the elements are
 				    taken as the box length along each axis.
-				binning = 'log' : It defines the type of binning in k-space. The other option is 
+				binning = 'log' (str): It defines the type of binning in k-space. The other option is 
 				                'linear' or 'mixed'.
+				intergral_bins = 100 (int): The number bins to use for numerically solving the intergral.
 	Returns: 
 				A tuple with (Pk, bins), where Pk is an array with the 
 				power spectrum and bins is an array with the k bin centers.
 	'''
 
 	ps, ks, n_modes = power_spectrum_1d(input_array_nd, 
-											kbins=rbins if kbins is None else kbins, 
-											box_dims=None, 
+											kbins=rbins*10 if kbins is None else kbins, 
+											box_dims=box_dims, 
 											return_n_modes=True, 
-											binning='log', 
+											binning=k_binning, 
 											)
-	if type(rbins)==int: 
-		rbins = 10**np.linspace(np.log10(2*box_dims/max(input_array_nd.shape)), np.log10(box_dims/2), rbins) if binning=='log' \
-					else np.linspace(2*box_dims/max(input_array_nd.shape), box_dims/2, rbins)
+	return ps_to_corr(ps, ks, rbins=rbins, box_dims=box_dims, n_grid=input_array_nd.shape[0], binning=binning)
 
-	ps_tck = splrep(ks[np.isfinite(ps)], ps[np.isfinite(ps)])
-	knew  = np.linspace(ks[0],ks[-1],100)
-	integ = knew**2*splev(knew,ps_tck)*np.sinc(knew[None,:]*rbins[:,None])
-	corr  = np.trapz(integ, knew, axis=1) 
+def cross_correlation_function(input_array1_nd, input_array2_nd, rbins=10, kbins=100, box_dims=None, binning='log', k_binning='log'):
+	'''
+	Function estimates the spherically averaged cross power spectrum and 
+	then calulates the cross correlation function using the following equation:
+	.. math:: \\xi(r) = \\int 4\\pi k^2 P(k) \\mathrm{sinc}(kr)dk
+
+	Parameters: 
+				input_array1_nd (numpy array): the data array 1
+				input_array2_nd (numpy array): the data array 2
+				rbins = 10 (integer or array-like): The number of bins,
+				    or a list containing the bin edges. If an integer is given, the bins
+				    are logarithmically spaced.
+				kbins = 100 (integer): The number of bins,
+				    used to estimate the power spectrum. If an integer is given, the bins
+				    are logarithmically spaced.
+				box_dims = None (float or array-like): the dimensions of the 
+				    box in Mpc. If this is None, the current box volume is used along all
+				    dimensions. If it is a float, this is taken as the box length
+				    along all dimensions. If it is an array-like, the elements are
+				    taken as the box length along each axis.
+				binning = 'log' (str): It defines the type of binning in k-space. The other option is 
+				                'linear' or 'mixed'.
+				intergral_bins = 100 (int): The number bins to use for numerically solving the intergral.
+	Returns: 
+				A tuple with (Pk, bins), where Pk is an array with the 
+				power spectrum and bins is an array with the k bin centers.
+	'''
+
+	ps, ks, n_modes = cross_power_spectrum_1d(input_array1_nd, input_array2_nd,
+											kbins=rbins*10 if kbins is None else kbins, 
+											box_dims=box_dims, 
+											return_n_modes=True, 
+											binning=k_binning, 
+											)
 
-	return corr/2/np.pi**2, rbins
+	return ps_to_corr(ps, ks, rbins=rbins, box_dims=box_dims, n_grid=input_array1_nd.shape[0], binning=binning)
 
 def landy_szalay_estimator(data, randoms=None, rbins=10, box_dims=None, binning='log', **kwargs):
 	'''