flipcosmo 1.0.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

flip/covariance/generator.py

@@ -1,3 +1,4 @@
+import importlib
 import multiprocessing as mp
 from functools import partial
 
@@ -9,14 +10,6 @@ from scipy.signal import savgol_filter
 from scipy.special import spherical_jn
 
 from flip.covariance import cov_utils
-from flip.covariance.adamsblake17 import flip_terms as flip_terms_adamsblake17
-from flip.covariance.adamsblake17plane import flip_terms as flip_terms_adamsblake17plane
-from flip.covariance.adamsblake20 import flip_terms as flip_terms_adamsblake20
-from flip.covariance.carreres23 import flip_terms as flip_terms_carreres23
-from flip.covariance.lai22 import flip_terms as flip_terms_lai22
-from flip.covariance.ravouxcarreres import flip_terms as flip_terms_ravouxcarreres
-from flip.covariance.ravouxnoanchor25 import flip_terms as flip_terms_ravouxnoanchor25
-from flip.covariance.rcrk24 import flip_terms as flip_terms_rcrk24
 from flip.utils import create_log
 
 log = create_log()
@@ -38,7 +31,7 @@ _avail_regularization_option = [
 ]
 
 
-def correlation_integration(l, r, k, integrand):
+def correlation_integration(ell, r, k, integrand):
     """
     The correlation_integration function is used to calculate the correlation function for a given multipole.
     It does this by integrating over k, which is the magnitude of the wavevector. The integration is performed using
@@ -53,17 +46,20 @@ def correlation_integration(l, r, k, integrand):
         integrand: Calculate the integrand of the correlation function
 
     Returns:
-        The integral of the integrand
+        ndarray: Correlation function $xi_\ell(r)$ via Simpson integration.
 
     """
     kr = np.outer(k, r)
     integrand = (
-        (-1) ** (l // 2) * (k**2 / (2 * np.pi**2)) * integrand * spherical_jn(l, kr).T
+        (-1) ** (ell // 2)
+        * (k**2 / (2 * np.pi**2))
+        * integrand
+        * spherical_jn(ell, kr).T
     )
-    return (-1) ** (l % 2) * integrate.simpson(integrand, x=k)
+    return (-1) ** (ell % 2) * integrate.simpson(integrand, x=k)
 
 
-def correlation_hankel(l, r, k, integrand, hankel_overhead_coefficient=2, kmin=None):
+def correlation_hankel(ell, r, k, integrand, hankel_overhead_coefficient=2, kmin=None):
     """
     The correlation_hankel function is a wrapper for the cosmoprimo.fftlog.PowerToCorrelation function,
     which computes the correlation function from power spectrum using FFTLog (Hamilton 2000).
@@ -79,12 +75,12 @@ def correlation_hankel(l, r, k, integrand, hankel_overhead_coefficient=2, kmin=N
         hankel_overhead_coefficient: Determine the minimum r value for which to use the hankel transform
 
     Returns:
-        The correlation function
+        ndarray: Correlation function $xi_\ell(r)$ combining Hankel and direct integration.
 
     Note:
         If l is odd, count a 1j term in the integrand, without the need for adding it
     """
-    Hankel = cosmoprimo.fftlog.PowerToCorrelation(k, ell=l, q=0, complex=False)
+    Hankel = cosmoprimo.fftlog.PowerToCorrelation(k, ell=ell, q=0, complex=False)
     Hankel.set_fft_engine("numpy")
     r_hankel, xi_hankel = Hankel(integrand)
     mask = r < np.min(r_hankel) * hankel_overhead_coefficient
@@ -93,13 +89,13 @@ def correlation_hankel(l, r, k, integrand, hankel_overhead_coefficient=2, kmin=N
             "Min pw spectrum k is too high, please take a lower one. Use kmin parameter to lower bound integration."
         )
     output = np.empty_like(r)
-    output[mask] = correlation_integration(l, r[mask], k, integrand)
-    output[~mask] = (-1) ** (l % 2) * np.interp(r[~mask], r_hankel, xi_hankel)
+    output[mask] = correlation_integration(ell, r[mask], k, integrand)
+    output[~mask] = (-1) ** (ell % 2) * np.interp(r[~mask], r_hankel, xi_hankel)
 
     # Regularization
     if kmin is not None:
         kreg = np.geomspace(np.min(k), kmin, int(len(k) / 20))
-        output -= correlation_integration(l, r, kreg, np.interp(kreg, k, integrand))
+        output -= correlation_integration(ell, r, kreg, np.interp(kreg, k, integrand))
     return output
 
 
@@ -134,23 +130,24 @@ def coefficient_hankel(
         : Define the model name
 
     Returns:
-        The covariance of the a-th and b-th terms
+        ndarray: Covariance contribution for term `term_index`.
 
     """
+    flip_terms = importlib.import_module(
+        f"flip.covariance.analytical.{model_name}.flip_terms", package=__package__
+    )
+
     cov_ab_i = 0
-    flip_terms = eval(f"flip_terms_{model_name}")
     flip_terms.set_backend("numpy")
     dictionary_subterms = flip_terms.dictionary_subterms
     regularize_M_terms = flip_terms.regularize_M_terms
-    Z_ab_i = 1
-    if eval(f"flip_terms.redshift_dependent_model"):
-        Z_ab_i = eval(f"flip_terms.Z_{covariance_type}_{term_index}")(
-            wavenumber, coord[3], coord[4], *additional_parameters_values
-        )  # Shape issue: Z_ab_i is an outer product. How to include that in the Hankel transform?
-    for l in range(lmax + 1):
-        number_terms = dictionary_subterms[f"{covariance_type}_{term_index}_{l}"]
+
+    for ell in range(lmax + 1):
+        number_terms = dictionary_subterms[f"{covariance_type}_{term_index}_{ell}"]
         for j in range(number_terms):
-            M_ab_i_l_j = eval(f"flip_terms.M_{covariance_type}_{term_index}_{l}_{j}")
+            M_ab_i_l_j = getattr(
+                flip_terms, f"M_{covariance_type}_{term_index}_{ell}_{j}"
+            )
             M_ab_i_l_j_evaluated = regularize_M(
                 M_ab_i_l_j,
                 wavenumber,
@@ -159,18 +156,18 @@ def coefficient_hankel(
                 flip_terms,
                 additional_parameters_values,
             )
-            M_ab_i_l_j_evaluated = M_ab_i_l_j_evaluated * Z_ab_i
-            N_ab_i_l_j = eval(f"flip_terms.N_{covariance_type}_{term_index}_{l}_{j}")(
-                coord[1], coord[2]
-            )
+
+            N_ab_i_l_j = getattr(
+                flip_terms, f"N_{covariance_type}_{term_index}_{ell}_{j}"
+            )(coord[1], coord[2])
             hankel_ab_i_l_j = correlation_hankel(
-                l,
+                ell,
                 coord[0],
                 wavenumber,
                 M_ab_i_l_j_evaluated * power_spectrum,
                 **kwargs,
             )
-            cov_ab_i = cov_ab_i + N_ab_i_l_j * hankel_ab_i_l_j
+            cov_ab_i += N_ab_i_l_j * hankel_ab_i_l_j
     return cov_ab_i
 
 
@@ -200,23 +197,23 @@ def coefficient_trapz(
         : Determine the model to be used
 
     Returns:
-        A matrix
+        ndarray: Covariance contribution for term `term_index`.
 
     """
     cov_ab_i = 0
-    flip_terms = eval(f"flip_terms_{model_name}")
+    flip_terms = importlib.import_module(
+        f"flip.covariance.analytical.{model_name}.flip_terms", package=__package__
+    )
     flip_terms.set_backend("numpy")
     dictionary_subterms = flip_terms.dictionary_subterms
     regularize_M_terms = flip_terms.regularize_M_terms
-    Z_ab_i = 1
-    if eval(f"flip_terms.redshift_dependent_model"):
-        Z_ab_i = eval(f"flip_terms.Z_{covariance_type}_{term_index}")(
-            wavenumber, coord[3], coord[4], *additional_parameters_values
-        )
-    for l in range(lmax + 1):
-        number_terms = dictionary_subterms[f"{covariance_type}_{term_index}_{l}"]
+
+    for ell in range(lmax + 1):
+        number_terms = dictionary_subterms[f"{covariance_type}_{term_index}_{ell}"]
         for j in range(number_terms):
-            M_ab_i_l_j = eval(f"flip_terms.M_{covariance_type}_{term_index}_{l}_{j}")
+            M_ab_i_l_j = getattr(
+                flip_terms, f"M_{covariance_type}_{term_index}_{ell}_{j}"
+            )
             M_ab_i_l_j_evaluated = regularize_M(
                 M_ab_i_l_j,
                 wavenumber,
@@ -225,20 +222,19 @@ def coefficient_trapz(
                 flip_terms,
                 additional_parameters_values,
             )
-            M_ab_i_l_j_evaluated = M_ab_i_l_j_evaluated * Z_ab_i
-            N_ab_i_l_j = eval(f"flip_terms.N_{covariance_type}_{term_index}_{l}_{j}")(
-                coord[1], coord[2]
-            )
+            N_ab_i_l_j = getattr(
+                flip_terms, f"N_{covariance_type}_{term_index}_{ell}_{j}"
+            )(coord[1], coord[2])
             kr = np.outer(wavenumber, coord[0])
             integrand = (
-                (-1) ** (l // 2)
+                (-1) ** (ell // 2)
                 * (wavenumber**2 / (2 * np.pi**2))
                 * M_ab_i_l_j_evaluated
                 * power_spectrum
-                * spherical_jn(l, kr).T
+                * spherical_jn(ell, kr).T
             )
-            hankel_ab_i_l_j = (-1) ** (l % 2) * np.trapz(integrand, x=wavenumber)
-            cov_ab_i = cov_ab_i + N_ab_i_l_j * hankel_ab_i_l_j
+            hankel_ab_i_l_j = (-1) ** (ell % 2) * np.trapezoid(integrand, x=wavenumber)
+            cov_ab_i += N_ab_i_l_j * hankel_ab_i_l_j
     return cov_ab_i
 
 
@@ -255,7 +251,28 @@ def regularize_M(
     lowk_unstable_mean_filtering=10,
     mpmmath_decimal_precision=50,
 ):
+    """Evaluate and optionally regularize M(k) term.
+
+    Applies one of the supported regularizations to stabilize numerical
+    behavior (mpmath high-precision, Savitzky–Golay smoothing, or low-k
+    asymptote detection).
 
+    Args:
+        M_function (callable): Function returning M(k) given additional parameters.
+        wavenumber (ndarray): k samples.
+        regularize_M_terms (dict|None): Per type regularization option.
+        covariance_type (str): `"gg"`, `"gv"`, or `"vv"`.
+        flip_terms (module): Terms module to switch backend when needed.
+        additional_parameters_values (dict|tuple): Extra parameters for M.
+        savgol_window (int): Window size for Savitzky–Golay.
+        savgol_polynomial (int): Polynomial order for Savitzky–Golay.
+        lowk_unstable_threshold (float): Threshold for low-k asymptote detection.
+        lowk_unstable_mean_filtering (int): Window for mean filtering indices.
+        mpmmath_decimal_precision (int): Decimal precision for mpmath.
+
+    Returns:
+        ndarray: Evaluated M(k) after regularization.
+    """
     if regularize_M_terms is None:
         return M_function(*additional_parameters_values)(wavenumber)
     else:
@@ -331,7 +348,6 @@ def compute_coordinates(
     covariance_type,
     coordinates_density=None,
     coordinates_velocity=None,
-    redshift_dict=None,
     size_batch=10_000,
     los_definition="bisector",
 ):
@@ -346,21 +362,17 @@ def compute_coordinates(
         : Compute the angle separation between two objects
 
     Returns:
-        A list of parameters
+        list: Per-batch arrays `[r, theta, phi]`.
     """
     if covariance_type == "gg":
         ra = coordinates_density[0]
         dec = coordinates_density[1]
         comoving_distance = coordinates_density[2]
-        if redshift_dict is not None:
-            redshift = redshift_dict["g"]
         number_objects = len(ra)
     elif covariance_type == "vv":
         ra = coordinates_velocity[0]
         dec = coordinates_velocity[1]
         comoving_distance = coordinates_velocity[2]
-        if redshift_dict is not None:
-            redshift = redshift_dict["v"]
         number_objects = len(ra)
     elif covariance_type == "gv":
         ra_g = coordinates_density[0]
@@ -368,9 +380,6 @@ def compute_coordinates(
         comoving_distance_g = coordinates_density[2]
         ra_v = coordinates_velocity[0]
         dec_v = coordinates_velocity[1]
-        if redshift_dict is not None:
-            redshift_g = redshift_dict["g"]
-            redshift_v = redshift_dict["v"]
         comoving_distance_v = coordinates_velocity[2]
         number_objects_g = len(ra_g)
         number_objects_v = len(ra_v)
@@ -399,9 +408,6 @@ def compute_coordinates(
                 dec_v[j_list],
                 comoving_distance_v[j_list],
             )
-            if redshift_dict is not None:
-                z_i = redshift_g[i_list]
-                z_j = redshift_v[j_list]
         else:
             i_list, j_list = cov_utils.compute_i_j(number_objects, batches)
             ra_i, dec_i, r_i = (
@@ -414,16 +420,10 @@ def compute_coordinates(
                 dec[j_list],
                 comoving_distance[j_list],
             )
-            if redshift_dict is not None:
-                z_i = redshift[i_list]
-                z_j = redshift[j_list]
         r, theta, phi = cov_utils.angle_separation(
             ra_i, ra_j, dec_i, dec_j, r_i, r_j, los_definition=los_definition
         )
-        if redshift_dict is not None:
-            parameters.append([r, theta, phi, z_i, z_j])
-        else:
-            parameters.append([r, theta, phi])
+        parameters.append([r, theta, phi])
     return parameters
 
 
@@ -451,18 +451,21 @@ def compute_coeficient(
         : Define the number of threads used to compute the covariance matrix
 
     Returns:
-        A list of arrays
+        ndarray: Stacked covariance arrays per term index.
     """
     if additional_parameters_values is None:
-        additional_parameters_values = ()
+        additional_parameters_values = {}
     if hankel:
         coefficient = coefficient_hankel
     else:
         coefficient = coefficient_trapz
 
-    term_index_list = eval(f"flip_terms_{model_name}.dictionary_terms")[covariance_type]
-    lmax_list = eval(f"flip_terms_{model_name}.dictionary_lmax")[covariance_type]
-    multi_index_model = eval(f"flip_terms_{model_name}.multi_index_model")
+    flip_terms = importlib.import_module(
+        f"flip.covariance.analytical.{model_name}.flip_terms", package=__package__
+    )
+    term_index_list = getattr(flip_terms, "dictionary_terms")[covariance_type]
+    lmax_list = getattr(flip_terms, "dictionary_lmax")[covariance_type]
+    multi_index_model = getattr(flip_terms, "multi_index_model")
 
     function_covariance_dict = {}
     for i, index in enumerate(term_index_list):
@@ -540,7 +543,6 @@ def compute_cov(
     power_spectrum_list,
     coordinates_density=None,
     coordinates_velocity=None,
-    redshift_dict=None,
     additional_parameters_values=None,
     size_batch=10_000,
     number_worker=8,
@@ -564,7 +566,7 @@ def compute_cov(
         : Compute the covariance matrix for a given model
 
     Returns:
-        The covariance matrix for a given model and set of parameters
+        ndarray: Covariance arrays per term index.
 
     """
     if model_name not in _avail_models:
@@ -577,7 +579,6 @@ def compute_cov(
         covariance_type,
         coordinates_density=coordinates_density,
         coordinates_velocity=coordinates_velocity,
-        redshift_dict=redshift_dict,
         size_batch=size_batch,
         los_definition=los_definition,
     )
@@ -595,17 +596,6 @@ def compute_cov(
     return covariance
 
 
-def get_redshift_dependent_model_flag(model_name):
-    """
-    The get_redshift_dependent_model_flag function returns a boolean indicating whether the model is redshift dependent or not.
-
-    Args:
-        model_name: Determine which model to use
-    """
-    redshift_dependent_model = eval(f"flip_terms_{model_name}.redshift_dependent_model")
-    return redshift_dependent_model
-
-
 def generate_covariance(
     model_name,
     model_kind,
@@ -635,7 +625,7 @@ def generate_covariance(
         : Define the number of workers to use for the computation
 
     Returns:
-        A dictionary with the covariance matrices and their dimensions
+        tuple: `(covariance_dict, number_densities, number_velocities)`.
 
     """
     cov_utils.check_generator_need(
@@ -645,17 +635,6 @@ def generate_covariance(
     )
     covariance_dict = {}
 
-    redshift_dependent_model = get_redshift_dependent_model_flag(model_name)
-    if redshift_dependent_model:
-        redshift_dict = cov_utils.generate_redshift_dict(
-            redshift_dependent_model,
-            model_kind,
-            coordinates_velocity=coordinates_velocity,
-            coordinates_density=coordinates_density,
-        )
-    else:
-        redshift_dict = None
-
     if model_kind in ["density", "full", "density_velocity"]:
         covariance_dict["gg"] = compute_cov(
             model_name,
@@ -663,7 +642,6 @@ def generate_covariance(
             power_spectrum_dict["gg"],
             coordinates_density=coordinates_density,
             coordinates_velocity=coordinates_velocity,
-            redshift_dict=redshift_dict,
             additional_parameters_values=additional_parameters_values,
             size_batch=size_batch,
             number_worker=number_worker,
@@ -682,7 +660,6 @@ def generate_covariance(
             power_spectrum_dict["vv"],
             coordinates_density=coordinates_density,
             coordinates_velocity=coordinates_velocity,
-            redshift_dict=redshift_dict,
             additional_parameters_values=additional_parameters_values,
             size_batch=size_batch,
             number_worker=number_worker,
@@ -701,7 +678,6 @@ def generate_covariance(
             power_spectrum_dict["gv"],
             coordinates_density=coordinates_density,
             coordinates_velocity=coordinates_velocity,
-            redshift_dict=redshift_dict,
             additional_parameters_values=additional_parameters_values,
             size_batch=size_batch,
             number_worker=number_worker,
@@ -710,4 +686,4 @@ def generate_covariance(
             kmin=kmin,
         )
 
-    return covariance_dict, number_densities, number_velocities, redshift_dict
+    return covariance_dict, number_densities, number_velocities