diff --git a/src/tools21cm/tau.py b/src/tools21cm/tau.py index 9d286ab..4b4e835 100644 --- a/src/tools21cm/tau.py +++ b/src/tools21cm/tau.py @@ -4,90 +4,106 @@ from tqdm import tqdm def tau(ionfractions, redshifts, num_points = 50): - ''' - Calculate the optical depth to Thomson scattering. - - Parameters: - * ionfractions (numpy array): an array containing the ionized fraction - in various points along the line-of-sight. - * redshifts (numpy array): an array containing the redshifts of - the same points as ionfractions. Must be the same length as - ionfractions - * num_points = 50 (integer): the number of points used for the integration - - Returns: - Tuple containing (tau_0, tau_z) - - tau_0 is the optical depth at each redshift - - tau_z is the corresponding redshift - - Notes: - * The Universe is assumed to be fully ionized below the lowest redshift supplied. - * To get the total optical depth, look at the last value in tau_0 - - ''' - - if len(ionfractions) != len(redshifts): - print('Incorrect length of ionfractions') - raise Exception() - - ionfractions, redshifts = ionfractions[np.argsort(redshifts)], redshifts[np.argsort(redshifts)] - - sigma_T = 6.65e-25 - chi1 = 1.0+const.abu_he - coeff = 2.0*(const.c*1e5)*sigma_T*const.OmegaB/const.Omega0*const.rho_crit_0*\ - chi1/const.mean_molecular/const.m_p/(3.*const.H0cgs) - - tau_z = np.hstack((np.arange(1,num_points+1)/float(num_points)*redshifts[0], redshifts)) - - tau0 = np.zeros(len(redshifts)+num_points) - - #Optical depth for a completely ionized Universe - tau0[0:num_points] = coeff*(np.sqrt(const.Omega0*(1+tau_z[0:num_points])**3+const.lam) - 1) - - for i in range(num_points, len(redshifts)+num_points): - tau0[i] = tau0[i-1]+1.5*coeff*const.Omega0 * \ - (ionfractions[i-1-num_points]*(1+tau_z[i-1])**2/np.sqrt(const.Omega0*(1+tau_z[i-1])**3+const.lam) \ - + ionfractions[i-num_points]*(1+tau_z[i])**2/np.sqrt(const.Omega0*(1+tau_z[i])**3+const.lam) ) * \ - (tau_z[i]-tau_z[i-1])/2 - - return tau0, tau_z + ''' + Calculate the optical depth to Thomson scattering. -def tau_map(ionfractions, redshifts=None, num_points=50, reading_function=None): + Parameters + ---------- + ionfractions : ndarray + An array containing the ionized fraction at various points along the line-of-sight. + + redshifts : ndarray + An array containing the redshifts corresponding to the same points as `ionfractions`. + Must be the same length as `ionfractions`. + + num_points : int, optional + The number of initial points used for the integration to account for high-redshift + contributions where ionization is assumed to be negligible. Default is 50. + + Returns + ------- + tuple + A tuple containing: + + tau_0 : ndarray + Optical depth at each redshift. The length is equal to the length of `redshifts` + `num_points`. + + tau_z : ndarray + The corresponding redshift array. The length is equal to the length of `redshifts` + `num_points`. + + Notes + ----- + - The Universe is assumed to be fully ionized at the lowest redshift supplied. + - To get the total optical depth, refer to the last value in `tau_0`. ''' - Calculate the optical depth to Thomson scattering for lightcones or maps. - Parameters: - * ionfractions (ndarray or dict): Ionized fraction data across various - points along the line-of-sight. + if len(ionfractions) != len(redshifts): + print('Incorrect length of ionfractions') + raise Exception() + + ionfractions, redshifts = ionfractions[np.argsort(redshifts)], redshifts[np.argsort(redshifts)] - * redshifts (ndarray, optional): Array of redshift values corresponding to - the ionfractions data. Must be the same length as ionfractions if `ionfractions` - is an ndarray. If `ionfractions` is a dict and redshifts is None, redshifts are - inferred from the keys of the dictionary. + sigma_T = 6.65e-25 + chi1 = 1.0+const.abu_he + coeff = 2.0*(const.c*1e5)*sigma_T*const.OmegaB/const.Omega0*const.rho_crit_0*\ + chi1/const.mean_molecular/const.m_p/(3.*const.H0cgs) - * num_points (int, optional): Number of initial points used for the integration - to account for high-redshift contributions where ionization is assumed to be negligible. - Default is 50. + tau_z = np.hstack((np.arange(1,num_points+1)/float(num_points)*redshifts[0], redshifts)) - * reading_function (callable, optional): Custom function to read and process ionization - fraction data when ionfractions is provided in a format that requires specialized reading. - The function should take a filename or data identifier as input and return an ndarray of ionization fractions. + tau0 = np.zeros(len(redshifts)+num_points) - Returns: - Tuple containing (tau_0, tau_z) + #Optical depth for a completely ionized Universe + tau0[0:num_points] = coeff*(np.sqrt(const.Omega0*(1+tau_z[0:num_points])**3+const.lam) - 1) - * tau_0 (ndarray): Optical depth values at each spatial position and redshift. - The shape is (N_x, N_y, N_z + num_points), where N_x and N_y are spatial dimensions, - and N_z is the number of redshift slices in output_z. + for i in range(num_points, len(redshifts)+num_points): + tau0[i] = tau0[i-1]+1.5*coeff*const.Omega0 * \ + (ionfractions[i-1-num_points]*(1+tau_z[i-1])**2/np.sqrt(const.Omega0*(1+tau_z[i-1])**3+const.lam) \ + + ionfractions[i-num_points]*(1+tau_z[i])**2/np.sqrt(const.Omega0*(1+tau_z[i])**3+const.lam) ) * \ + (tau_z[i]-tau_z[i-1])/2 - * tau_z (ndarray): Array of redshift values corresponding to each slice in tau_0. - The length is N_z + num_points. + return tau0, tau_z + +def tau_map(ionfractions, redshifts=None, num_points=50, reading_function=None): + ''' + Calculate the optical depth to Thomson scattering for lightcones or maps. - Notes: - * The Universe is assumed to be fully ionized below the lowest redshift supplied. - * The reading_function is useful when working with custom data formats. + Parameters + ---------- + ionfractions : ndarray or dict + Ionized fraction data across various points along the line-of-sight. + + redshifts : ndarray, optional + Array of redshift values corresponding to the ionfractions data. Must be the same length + as ionfractions if `ionfractions` is an ndarray. If `ionfractions` is a dict and redshifts + is None, redshifts are inferred from the keys of the dictionary. + + num_points : int, optional + Number of initial points used for the integration to account for high-redshift contributions + where ionization is assumed to be negligible. Default is 50. + + reading_function : callable, optional + Custom function to read and process ionization fraction data when `ionfractions` is provided + in a format that requires specialized reading. The function should take a filename or data + identifier as input and return an ndarray of ionization fractions. + + Returns + ------- + tuple + A tuple containing: + + tau_0 : ndarray + Optical depth values at each spatial position and redshift. The shape is + (N_x, N_y, N_z + num_points), where N_x and N_y are spatial dimensions, + and N_z is the number of redshift slices in `output_z`. + + tau_z : ndarray + Array of redshift values corresponding to each slice in `tau_0`. + The length is `N_z + num_points`. + + Notes + ----- + - The Universe is assumed to be fully ionized at the lowest redshift supplied. + - The `reading_function` is useful when working with custom data formats. ''' if redshifts is None: assert isinstance(ionfractions, dict), "redshifts must be provided if ionfractions is not a dict."