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

flip/{plot_utils.py → covariance/plot_utils.py}

@@ -14,6 +14,13 @@ def plot_1d_contraction(
     parameter_dict,
     rs_multiplied=True,
 ):
+    """Plot 1D slices of contraction sums for gg/gv/vv blocks.
+
+    Args:
+        contraction (object): Contraction object with `compute_contraction_sum` and `coordinates_dict`.
+        parameter_dict (dict): Parameter values for evaluation.
+        rs_multiplied (bool): Multiply gg by r^2 for visualization.
+    """
     contraction_sum = contraction.compute_contraction_sum(parameter_dict)
     coord = contraction.coordinates_dict
 
@@ -87,6 +94,13 @@ def plot_2d_contraction(
     parameter_dict,
     rs_multiplied=True,
 ):
+    """Plot 2D images of contraction sums for gg/gv/vv blocks.
+
+    Args:
+        contraction (object): Contraction with `compute_contraction_sum` & `coordinates_dict`.
+        parameter_dict (dict): Parameter values for evaluation.
+        rs_multiplied (bool): Multiply gg by r^2 for visualization.
+    """
     contraction_sum = contraction.compute_contraction_sum(parameter_dict)
     coord = contraction.coordinates_dict
 
@@ -155,8 +169,17 @@ def plot_2d_contraction(
 def plot_correlation_from_likelihood(
     likelihood,
     parameter_dict,
+    covariance_prefactor_dict=None,
     **kwargs,
 ):
+    """Plot correlation matrix computed from a likelihood’s covariance.
+
+    Args:
+        likelihood (BaseLikelihood): Likelihood instance providing data/covariance.
+        parameter_dict (dict): Parameter specs; values read to form vector and variance.
+        covariance_prefactor_dict (dict, optional): Prefactors per covariance block.
+        **kwargs: Plot options, e.g., `vmin`, `vmax`.
+    """
     vmin = utils.return_key(kwargs, "vmin", -0.1)
     vmax = utils.return_key(kwargs, "vmax", 0.1)
 
@@ -166,11 +189,13 @@ def plot_correlation_from_likelihood(
     ]
     parameter_values_dict = dict(zip(parameter_names, parameter_values))
 
-    _, vector_variance = likelihood.data(
+    _, vector_variance = likelihood.data.give_data_and_variance(
         parameter_values_dict,
     )
     covariance_sum = likelihood.covariance.compute_covariance_sum(
-        parameter_values_dict, vector_variance
+        parameter_values_dict,
+        vector_variance,
+        covariance_prefactor_dict=covariance_prefactor_dict,
     )
 
     correlation_sum = cov_utils.return_correlation_matrix(covariance_sum)
@@ -189,6 +214,22 @@ def plot_all_fits(
     plot=True,
     **kwargs,
 ):
+    """Scatter all fit results per parameter with errors.
+
+    Args:
+        fit_output (str): Directory containing pickled fit outputs.
+        parameters (list[str]): Parameters to plot.
+        fiducials (list[float]|None): Optional reference lines per parameter.
+        compute_fs8_from_beta (bool): Plot `fs8 = beta_f * bs8` when requested.
+        subset_plot (list[str]|None): Only include fit files whose names contain these substrings.
+        remove_lower (dict|None): Exclude fits where param < threshold.
+        remove_higher (dict|None): Exclude fits where param > threshold.
+        plot (bool): Whether to render the plots.
+        **kwargs: Matplotlib options, e.g., `figsize`.
+
+    Returns:
+        tuple: `(fit_name_to_plot, param_dict, error_dict)`.
+    """
 
     fit_to_plot, fit_name_to_plot = select_valid_fits(
         fit_output,
@@ -258,6 +299,29 @@ def plot_all_mean_fits(
     use_minos=False,
     **kwargs,
 ):
+    """Plot mean parameter values and errors across grouped fits.
+
+    Groups fits by a substring in filenames, computes mean and error metrics, and
+    optionally plots both means and errors.
+
+    Args:
+        fit_output (str): Directory containing pickled fit outputs.
+        parameters (list[str]): Parameters to summarize.
+        fiducials (list[float]|None): Reference lines per parameter.
+        weighted_mean (bool): Weight by inverse Hessian variance when True.
+        compute_fs8_from_beta (bool): Use `fs8 = beta_f * bs8`.
+        plot_std_error (bool): Plot standard deviation instead of mean error.
+        plot_error_bar_of_mean (bool): Plot error-of-the-mean instead of mean error.
+        subset_plot (list[str]|None): Only include fits matching substrings.
+        remove_lower (dict|None): Exclude fits where param < threshold.
+        remove_higher (dict|None): Exclude fits where param > threshold.
+        plot (bool): Whether to render plots.
+        use_minos (bool): Use MINOS errors if available; fall back to Hessian.
+        **kwargs: Matplotlib options, e.g., `figsize`.
+
+    Returns:
+        tuple: `(unique_fit_prop, mean_param_dict, mean_error_dict, error_mean_dict, std_dict, count_dict)`.
+    """
 
     fit_to_plot, fit_name_to_plot = select_valid_fits(
         fit_output,
@@ -330,7 +394,7 @@ def plot_all_mean_fits(
                             error_bs8_high = fits[i][3][index_bs8].upper
                             error_betaf_low = fits[i][3][index_beta_f].lower
                             error_betaf_high = fits[i][3][index_beta_f].upper
-                        except:
+                        except IndexError:
                             print("Minos failed, taking hessian error")
                     errors = [
                         [
@@ -365,7 +429,7 @@ def plot_all_mean_fits(
                             index = np.argwhere(np.array(names) == param_name)[0][0]
                             errors[0].append(abs(fits[i][3][index].lower))
                             errors[1].append(abs(fits[i][3][index].upper))
-                        except:
+                        except IndexError:
                             errors[0].append(fits[i][2][param_name])
                             errors[1].append(fits[i][2][param_name])
                 else:
@@ -447,6 +511,17 @@ def select_valid_fits(
     remove_lower=None,
     remove_higher=None,
 ):
+    """Select valid fit files based on status flags and filters.
+
+    Args:
+        fit_output (str): Directory of pickled fit outputs.
+        subset_plot (list[str]|None): Only include filenames containing these substrings.
+        remove_lower (dict|None): Exclude fits where param < threshold.
+        remove_higher (dict|None): Exclude fits where param > threshold.
+
+    Returns:
+        tuple[list, list]: Fit objects and their filenames.
+    """
     all_fit = glob.glob(os.path.join(fit_output, "*"))
 
     fit_to_plot = []

flip/covariance/symbolic.py

@@ -65,7 +65,7 @@ def simplify_term(
 
 def generate_MN_ab_i_l_function_wide_angle(
     term_B,
-    l,
+    ell,
     l1,
     l2,
 ):
@@ -90,10 +90,10 @@ def generate_MN_ab_i_l_function_wide_angle(
         integral_mu1_M_l * legendre_poly(l2, x=mu2), (mu2, -1, 1)
     )
     term_N_l_l1_l2 = 0
-    for m in range(-l, l + 1):
+    for m in range(-ell, ell + 1):
         for m1 in range(-l1, l1 + 1):
             for m2 in range(-l2, l2 + 1):
-                term_N_l_l1_l2_m_m1_m2 = wigner.gaunt(l, l1, l2, m, m1, m2)
+                term_N_l_l1_l2_m_m1_m2 = wigner.gaunt(ell, l1, l2, m, m1, m2)
                 # The two comments are for the following line, they are the results of an intense
                 # head scratching and are quite important for all the modeling of flip.
 
@@ -107,7 +107,7 @@ def generate_MN_ab_i_l_function_wide_angle(
                 # For the order of term chosen, the sy.pi must be added here.
                 # If not, it will give wrong results for cross-terms (gv).
                 term_N_l_l1_l2_m_m1_m2 *= (
-                    sy.Ynm(l, m, sy.pi - phi, 0)
+                    sy.Ynm(ell, m, sy.pi - phi, 0)
                     * sy.Ynm(l1, m1, theta / 2, 0)
                     * sy.Ynm(l2, m2, theta / 2, sy.pi)
                 )
@@ -126,7 +126,7 @@ def generate_MN_ab_i_l_function_wide_angle(
     return term_M_l_l1_l2, term_N_l_l1_l2
 
 
-def generate_MN_ab_i_l_function_parallel_plane(term_B, l):
+def generate_MN_ab_i_l_function_parallel_plane(term_B, ell):
     """
     The generate_MN_ab_i_l_function_parallel_plane function takes in a term_B and an l value.
     It then generates the M_l and N_l functions for that particular term B, which is used to calculate the parallel plane integral.
@@ -143,12 +143,12 @@ def generate_MN_ab_i_l_function_parallel_plane(term_B, l):
     phi = sy.symbols("phi")
     mu = sy.symbols("mu")
     M_l = sy.Rational(1 / 2) * sy.integrate(
-        term_B * legendre_poly(l, x=mu), (mu, -1, 1)
+        term_B * legendre_poly(ell, x=mu), (mu, -1, 1)
     )
     # The sy.pi term is directly linked to the definition of r chosen in flip.
     # For the order of term chosen, the sy.pi must be added here.
     # If not, it will give wrong results for cross-terms (gv).
-    N_l = sy.Rational(2 * l + 1) * legendre_poly(l, x=sy.cos(sy.pi - phi))
+    N_l = sy.Rational(2 * ell + 1) * legendre_poly(ell, x=sy.cos(sy.pi - phi))
     M_l = simplify_term(
         M_l.expand(func=True),
         simplification_method="simplify_iteration",
@@ -172,7 +172,6 @@ def write_output(
     l1max_list=None,
     l2max_list=None,
     multi_index_model=False,
-    redshift_dependent_model=False,
     regularize_M_terms="None",
 ):
     """
@@ -224,26 +223,26 @@ def write_output(
         dict_terms[f"{type}"] = term_index_list[k]
         dict_lmax[f"{type}"] = lmax_list[k]
         for i, t in enumerate(term_index_list[k]):
-            for l in range(lmax_list[k][i] + 1):
+            for ell in range(lmax_list[k][i] + 1):
                 list_M_ab_i_l = []
                 list_N_ab_i_l = []
                 if wide_angle:
                     for l1 in range(l1max_list[k][i] + 1):
                         for l2 in range(l2max_list[k][i] + 1):
                             M_ab_i_l_l1_l2, N_ab_i_l_l1_l2 = output_pool[
-                                index_pool[f"{type}_{t}_{l}_{l1}_{l2}"]
+                                index_pool[f"{type}_{t}_{ell}_{l1}_{l2}"]
                             ]
                             if (M_ab_i_l_l1_l2 != 0) & (N_ab_i_l_l1_l2 != 0):
                                 list_M_ab_i_l.append(M_ab_i_l_l1_l2)
                                 list_N_ab_i_l.append(N_ab_i_l_l1_l2)
                 else:
                     M_ab_i_l_l1_l2, N_ab_i_l_l1_l2 = output_pool[
-                        index_pool[f"{type}_{t}_{l}"]
+                        index_pool[f"{type}_{t}_{ell}"]
                     ]
                     if (M_ab_i_l_l1_l2 != 0) & (N_ab_i_l_l1_l2 != 0):
                         list_M_ab_i_l.append(M_ab_i_l_l1_l2)
                         list_N_ab_i_l.append(N_ab_i_l_l1_l2)
-                dict_j[f"{type}_{t}_{l}"] = len(list_M_ab_i_l)
+                dict_j[f"{type}_{t}_{ell}"] = len(list_M_ab_i_l)
                 for j in range(len(list_M_ab_i_l)):
                     M_ab_i_l_j = (
                         pycode(list_M_ab_i_l[j])
@@ -261,13 +260,13 @@ def write_output(
                         for add in additional_parameters:
                             additional_str = additional_str + f"{add},"
                     additional_str = additional_str[:-1]
-                    f.write(f"def M_{type}_{t}_{l}_{j}({additional_str}):\n")
-                    f.write(f"    def func(k):\n")
-                    f.write(f"        return({M_ab_i_l_j})\n")
-                    f.write(f"    return(func)\n")
+                    f.write(f"def M_{type}_{t}_{ell}_{j}({additional_str}):\n")
+                    f.write("    def func(k):\n")
+                    f.write(f"        return {M_ab_i_l_j}\n")
+                    f.write("    return func\n")
                     f.write("\n")
 
-                    f.write(f"def N_{type}_{t}_{l}_{j}(theta,phi):\n")
+                    f.write(f"def N_{type}_{t}_{ell}_{j}(theta,phi):\n")
                     f.write(f"    return({N_ab_i_l_j})\n")
                     f.write("\n")
 
@@ -286,8 +285,6 @@ def write_output(
     f.write(f"multi_index_model = {multi_index_model}")
     f.write("\n")
 
-    f.write(f"redshift_dependent_model = {redshift_dependent_model}")
-    f.write("\n")
     f.write(f"regularize_M_terms = {regularize_M_terms}")
     f.write("\n")
     f.close()
@@ -305,7 +302,6 @@ def write_M_N_functions(
     l1max_list=None,
     l2max_list=None,
     multi_index_model=False,
-    redshift_dependent_model=False,
     regularize_M_terms="None",
 ):
     """
@@ -336,17 +332,17 @@ def write_M_N_functions(
     index = 0
     for k, type in enumerate(type_list):
         for i, t in enumerate(term_index_list[k]):
-            for l in range(lmax_list[k][i] + 1):
+            for ell in range(lmax_list[k][i] + 1):
                 B_ab_i = dict_B[f"B_{type}_{t}"]
                 if wide_angle:
                     for l1 in range(l1max_list[k][i] + 1):
                         for l2 in range(l2max_list[k][i] + 1):
-                            params_pool.append([B_ab_i, l, l1, l2])
-                            index_pool[f"{type}_{t}_{l}_{l1}_{l2}"] = index
+                            params_pool.append([B_ab_i, ell, l1, l2])
+                            index_pool[f"{type}_{t}_{ell}_{l1}_{l2}"] = index
                             index = index + 1
                 else:
-                    params_pool.append([B_ab_i, l])
-                    index_pool[f"{type}_{t}_{l}"] = index
+                    params_pool.append([B_ab_i, ell])
+                    index_pool[f"{type}_{t}_{ell}"] = index
                     index = index + 1
 
     if number_worker == 1:
@@ -383,13 +379,58 @@ def write_M_N_functions(
         l1max_list=l1max_list,
         l2max_list=l2max_list,
         multi_index_model=multi_index_model,
-        redshift_dependent_model=redshift_dependent_model,
         regularize_M_terms=regularize_M_terms,
     )
 
 
+def generate_generalized_genericzdep_functions(
+    filename="./analytical/genericzdep/flip_terms.py", number_worker=8
+):
+    """
+    The generate_generalized_genericzdep_functions function generates the flip_terms.py file in the genericzdep directory, which contains functions that calculate M and N terms for a generalized version of Carreres' (2012) model 2 and 3.
+
+    Args:
+        filename: Specify the name of the file that will be generated
+        number_worker: Determine the number of processes to use for multiprocessing
+
+    Returns:
+        A list of functions,
+
+    """
+    mu1, mu2 = sy.symbols("mu1 mu2")
+    k = sy.symbols("k", positive=True, finite=True, real=True)
+    kNL = sy.symbols("kNL", positive=True, finite=True, real=True)
+    additional_parameters = ["kNL"]
+
+    type_list = ["vv"]
+    term_index_list = [["0", "1", "2"]]
+    lmax_list = [[2, 2, 2]]
+    l1max_list = [[1, 1, 1]]
+    l2max_list = [[1, 1, 1]]
+    dict_B = {
+        "B_vv_0": mu1 * mu2 / k**2,
+        "B_vv_1": mu1 * mu2 / kNL**2,
+        "B_vv_2": mu1 * mu2 * k**2 / kNL**4,
+    }
+
+    regularize_M_terms = "None"
+    write_M_N_functions(
+        filename,
+        type_list,
+        term_index_list,
+        lmax_list,
+        dict_B,
+        number_worker=number_worker,
+        wide_angle=True,
+        l1max_list=l1max_list,
+        l2max_list=l2max_list,
+        regularize_M_terms=regularize_M_terms,
+        additional_parameters=additional_parameters,
+    )
+
+
 def generate_generalized_adamsblake17plane_functions(
-    filename="./adamsblake17plane/flip_terms.py", number_worker=8
+    filename="./analytical/adamsblake17plane/flip_terms.py", number_worker=8
 ):
     mu = sy.symbols("mu")
     k = sy.symbols("k", positive=True, finite=True, real=True)
@@ -416,7 +457,7 @@ def generate_generalized_adamsblake17plane_functions(
 
 
 def generate_generalized_adamsblake17_functions(
-    filename="./adamsblake17/flip_terms.py", number_worker=8
+    filename="./analytical/adamsblake17/flip_terms.py", number_worker=8
 ):
 
     mu1, mu2 = sy.symbols("mu1 mu2")
@@ -448,7 +489,7 @@ def generate_generalized_adamsblake17_functions(
 
 
 def generate_generalized_adamsblake20_functions(
-    filename="./adamsblake20/flip_terms.py", number_worker=8
+    filename="./analytical/adamsblake20/flip_terms.py", number_worker=8
 ):
     """
     The generate_generalized_adamsblake20_functions function generates the functions needed to compute the M and N matrices for a generalized version of Adams, Blake & Kitching (2020).
@@ -492,7 +533,7 @@ def generate_generalized_adamsblake20_functions(
 
 
 def generate_generalized_lai22_functions(
-    filename="./lai22/flip_terms.py", number_worker=8
+    filename="./analytical/lai22/flip_terms.py", number_worker=8
 ):
     """
     The generate_generalized_lai22_functions function generates the functions for calculating the M and N terms in
@@ -625,7 +666,7 @@ def generate_generalized_lai22_functions(
 
 
 def generate_generalized_carreres23_functions(
-    filename="./carreres23/flip_terms.py", number_worker=8
+    filename="./analytical/carreres23/flip_terms.py", number_worker=8
 ):
     """
     The generate_generalized_carreres23_functions function generates the flip_terms.py file in the carreres23 directory, which contains functions that calculate M and N terms for a generalized version of Carreres' (2012) model 2 and 3.
@@ -663,7 +704,7 @@ def generate_generalized_carreres23_functions(
 
 
 def generate_generalized_ravouxcarreres_functions(
-    filename="./ravouxcarreres/flip_terms.py", number_worker=8
+    filename="./analytical/ravouxcarreres/flip_terms.py", number_worker=8
 ):
     """
     The generate_generalized_ravouxcarreres_functions function generates the functions needed to compute the generalized Ravoux-Carreres model.
@@ -715,7 +756,7 @@ def generate_generalized_ravouxcarreres_functions(
 
 
 def generate_generalized_rcrk24_functions(
-    filename="./rcrk24/flip_terms.py", number_worker=8
+    filename="./analytical/rcrk24/flip_terms.py", number_worker=8
 ):
     """
     The generate_generalized_rcrk24_functions function generates the flip_terms.py file in the carreres23 directory, which contains functions that calculate M and N terms for a generalized version of Carreres' (2012) model 2 and 3.
@@ -748,7 +789,6 @@ def generate_generalized_rcrk24_functions(
         wide_angle=True,
         l1max_list=l1max_list,
         l2max_list=l2max_list,
-        redshift_dependent_model=True,
         regularize_M_terms=regularize_M_terms,
     )
 
@@ -829,7 +869,7 @@ def write_partial_derivatives(
         )
 
         f.write(
-            "def get_partial_derivative_coefficients(model_kind,parameter_values_dict,variant=None,redshift_dict=None,):\n"
+            "def get_partial_derivative_coefficients(model_kind,parameter_values_dict,variant=None,covariance_prefactor_dict=None,):\n"
         )
         write_one_function(
             f,
@@ -840,7 +880,7 @@ def write_partial_derivatives(
         )
     else:
         f.write(
-            "def get_partial_derivative_coefficients(model_kind,parameter_values_dict,variant=None,):\n"
+            "def get_partial_derivative_coefficients(model_kind,parameter_values_dict,variant=None,covariance_prefactor_dict=None,):\n"
         )
         f.write("    if model_kind == 'density':\n")
         f.write(
@@ -914,7 +954,7 @@ def write_one_function(
     else:
         for i_model, name in enumerate(name_models):
             if name is None:
-                f.write(f"    else:\n")
+                f.write("    else:\n")
             elif i_model > 0:
                 f.write(f"    elif variant == '{name}':\n")
             else:
@@ -939,7 +979,7 @@ def write_one_function(
 
 
 def generate_fisher_coefficients_dictionnary_carreres23(
-    filename="./carreres23/fisher_terms.py",
+    filename="./analytical/carreres23/fisher_terms.py",
 ):
 
     name_models = ["growth_index", None]
@@ -959,7 +999,7 @@ def generate_fisher_coefficients_dictionnary_carreres23(
 
 
 def generate_fisher_coefficients_dictionnary_adamsblake17(
-    filename="./adamsblake17/fisher_terms.py",
+    filename="./analytical/adamsblake17/fisher_terms.py",
 ):
 
     name_models = ["growth_index", None]
@@ -993,7 +1033,7 @@ def generate_fisher_coefficients_dictionnary_adamsblake17(
 
 
 def generate_fisher_coefficients_dictionnary_adamsblake17plane(
-    filename="./adamsblake17plane/fisher_terms.py",
+    filename="./analytical/adamsblake17plane/fisher_terms.py",
 ):
 
     name_models = ["growth_index", None]
@@ -1073,9 +1113,9 @@ def generate_fisher_coefficients_dictionnary_full_nosigmag(
 
 
 def generate_fisher_coefficients_dictionnary_lai22(
-    filename="./lai22/fisher_terms.py",
+    filename="./analytical/lai22/fisher_terms.py",
 ):
-    from flip.covariance.lai22.flip_terms import dictionary_terms
+    from flip.covariance.analytical.lai22.flip_terms import dictionary_terms
 
     name_models = ["growth_index", "growth_index_nobeta", "nobeta", None]
     components = ["gg", "gv", "vv"]
@@ -1262,13 +1302,18 @@ def generate_files():
 
 
 def generate_fisher_files():
+    """Generate Fisher coefficient modules across supported covariance models.
+
+    Writes `fisher_terms.py` files for each model with partial derivative
+    coefficient dictionaries tailored to their parameterizations.
+    """
     generate_fisher_coefficients_dictionnary_carreres23()
     generate_fisher_coefficients_dictionnary_adamsblake17()
     generate_fisher_coefficients_dictionnary_adamsblake17plane()
     generate_fisher_coefficients_dictionnary_full_nosigmag(
-        "./adamsblake20/fisher_terms.py"
+        "./analytical/adamsblake20/fisher_terms.py"
     )
     generate_fisher_coefficients_dictionnary_full_nosigmag(
-        "./ravouxcarreres/fisher_terms.py"
+        "./analytical/ravouxcarreres/fisher_terms.py"
     )
     generate_fisher_coefficients_dictionnary_lai22()

flip/data/__init__.py

@@ -1,3 +1,3 @@
 """Init file of the flip.data package."""
 
-from .data import *
+from . import load_data_test