format change again

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."