Debug power spectra

docs/power_spectra.rst

@@ -0,0 +1,20 @@
+Power Spectra Generator
+=======================
+
+You can use the `flip.power_spectra` module to compute your power spectra using the `flip.power_spectra.compute_power_spectra` function:
+
+```
+k, pdd, pdt, ptt = compute_power_spectra(
+                        power_spectrum_engine,
+                        power_spectrum_settings,
+                        redshift,
+                        minimal_wavenumber,
+                        maximal_wavenumber,
+                        number_points,
+                        logspace=True,
+                        normalize_power_spectrum=True,
+                        power_spectrum_non_linear_model=None,
+                        power_spectrum_model="linearbel",
+                        save_path=None,
+)
+python```

flip/__init__.py

@@ -1,7 +1,7 @@
 """Init file of the flip package."""
 import os
 
-from . import covariance, fitter, gridding, likelihood, power_spectra_generator, utils
+from . import covariance, fitter, gridding, likelihood, utils, power_spectra
 
 __version__ = "1.0.0"
 __flip_dir_path__ = os.path.dirname(__file__)

flip/power_spectra/__init__.py

@@ -1,2 +1,4 @@
 """Init file of the flip.power_spectra package."""
+
 from . import class_engine, models
+from .generator import compute_power_spectra

flip/power_spectra/class_engine.py

@@ -117,7 +117,7 @@ def compute_power_spectrum(
     if non_linear_model is not None:
         power_spectrum_settings.update(
             {
-                "non_linear": non_linear_model,
+                "non linear": non_linear_model,
             }
         )
     power_spectrum_settings.update({"z_max_pk": redshift})
@@ -138,27 +138,28 @@ def compute_power_spectrum(
         )
         raise error
 
-    power_spectrum_linear, power_spectrum_non_linear = [], []
+    power_spectrum_linear, power_spectrum_non_linear = np.empty((2, number_points))
 
-    if non_linear_model is None:
-        power_spectrum_non_linear = None
-
-    for k in wavenumber:
-        power_spectrum_linear.append(
+    for i, k in enumerate(wavenumber):
+        power_spectrum_linear[i] = (
             model.pk_lin(k * model.h(), redshift) * model.h() ** 3
         )
+
         if non_linear_model is not None:
-            power_spectrum_non_linear.append(
+            power_spectrum_non_linear[i] = (
                 model.pk(k * model.h(), redshift) * model.h() ** 3
             )
 
     fs8_fiducial = get_fiducial_fs8(model, redshift)
     s8_fiducial = get_fiducial_s8(model, redshift)
     fiducial = {"fsigma_8": fs8_fiducial, "sigma_8": s8_fiducial}
 
+    if non_linear_model is None:
+        power_spectrum_non_linear = None
+
     return (
         wavenumber,
-        np.array(power_spectrum_linear),
-        np.array(power_spectrum_non_linear),
+        power_spectrum_linear,
+        power_spectrum_non_linear,
         fiducial,
     )

flip/power_spectra_generator.py → flip/power_spectra/generator.py

@@ -6,6 +6,11 @@
 
 _available_engines = ["class_engine"]
 _available_power_spectrum_model = ["linearbel", "nonlinearbel"]
+_available_power_spectrum_normalizaton = [
+    "no_normalization",
+    "growth_rate",
+    "growth_amplitude",
+]
 
 
 def get_power_spectrum_suffix(
@@ -56,11 +61,35 @@ def compute_power_spectra(
     maximal_wavenumber,
     number_points,
     logspace=True,
-    normalize_power_spectrum=True,
+    normalization_power_spectrum="no_normalization",
     power_spectrum_non_linear_model=None,
     power_spectrum_model="linearbel",
     save_path=None,
 ):
+    """Compute the power spectrum.
+
+    Args:
+        power_spectrum_engine (str): engine to use to compute the power spectrum, see _available_engines.
+        power_spectrum_settings (dic): configuration for the engine.
+        redshift (float): the redshift at which compute the power spectrum.
+        minimal_wavenumber (float): minimum k in h/Mpc.
+        maximal_wavenumber (float): maximum k in h/Mpc.
+        number_points (int): Sampling of the power spectrum.
+        logspace (bool, optional): Sample the power spectrum in logspace or linspace. Defaults to True.
+        normalization_power_spectrum (str, optional): which normalisation to use. Defaults to "no_normalization".
+            Available options are: "no_normalization", "growth_rate" or "growth_amplitude".
+        power_spectrum_non_linear_model (str, optional): Non-linear model to compute. Defaults to None.
+        power_spectrum_model (str, optional): Non-linear model to apply to the computed power spectrum, see _available_power_spectrum_model. Defaults to "linearbel".
+        save_path (str, optional): Path to save the computed power spectrum. Defaults to None.
+
+    Raises:
+        ValueError: power_spectrum_engine is not available
+        ValueError: power_spectrum_model is not available
+        ValueError: _description_
+
+    Returns:
+        _type_: _description_
+    """    
     if power_spectrum_engine not in _available_engines:
         raise ValueError(
             f"The engine {power_spectrum_engine} is not available"
@@ -89,13 +118,6 @@ def compute_power_spectra(
         non_linear_model=power_spectrum_non_linear_model,
     )
 
-    if normalize_power_spectrum:
-        power_spectrum_linear = power_spectrum_linear / fiducial["sigma_8"] ** 2
-        if power_spectrum_non_linear is not None:
-            power_spectrum_non_linear = (
-                power_spectrum_non_linear / fiducial["sigma_8"] ** 2
-            )
-
     power_spectrum_mm, power_spectrum_mt, power_spectrum_tt = eval(
         f"models.get_{power_spectrum_model}_model"
     )(
@@ -105,6 +127,26 @@ def compute_power_spectra(
         **fiducial,
     )
 
+    if normalization_power_spectrum == "growth_rate":
+        power_spectrum_mm = power_spectrum_mm
+        power_spectrum_mt = power_spectrum_mt * (
+            fiducial["fsigma_8"] / fiducial["sigma_8"]
+        )
+        power_spectrum_tt = (
+            power_spectrum_tt * (fiducial["fsigma_8"] / fiducial["sigma_8"]) ** 2
+        )
+    elif normalization_power_spectrum == "growth_amplitude":
+        power_spectrum_mm = power_spectrum_mm / fiducial["sigma_8"] ** 2
+        power_spectrum_mt = power_spectrum_mt / fiducial["sigma_8"] ** 2
+        power_spectrum_tt = power_spectrum_tt / fiducial["sigma_8"] ** 2
+    elif normalization_power_spectrum == "no_normalization":
+        pass
+    else:
+        raise ValueError(
+            f"The normalization {normalization_power_spectrum} of the power spectrum is not available,"
+            f"Please choose in {_available_power_spectrum_normalizaton}."
+        )
+
     if save_path is not None:
         suffix = get_power_spectrum_suffix(
             redshift,

flip/power_spectra/models.py

@@ -22,24 +22,69 @@ def bel_coefficients(sigma_8):
     return a1, a2, a3, invkdelta, b
 
 
-def get_nonlinearbel_model(
+def get_bel_model(
     wavenumber,
     power_spectrum_linear,
     **kwargs,
 ):
+    """
+    The get_bel_model function takes in the linear power spectrum, wavenumber, and sigma_8 value.
+    It then calculates the nonlinear matter-matter power spectrum using a fitting function from Bel et al. (2014).
+    The function also returns the nonlinear matter-matter cross correlation coefficient.
+
+    Args:
+        wavenumber: Calculate the power spectrum
+        power_spectrum_linear: Calculate the nonlinear power spectrum
+        **kwargs: Pass a variable number of keyword arguments to a function
+        : Calculate the nonlinear power spectrum
+
+    Returns:
+        The nonlinear matter power spectrum and the nonlinear galaxy clustering power spectrum
+    """
     if "sigma_8" not in kwargs.keys():
         raise ValueError("Fiducial sigma_8 value is needed for nonlinearbel model")
 
-    if "power_spectrum_non_linear" not in kwargs.keys():
-        raise ValueError("Non linear power spectrum is needed for nonlinearbel model")
-
     a1, a2, a3, invkdelta, b = bel_coefficients(kwargs["sigma_8"])
+
     power_spectrum_tt = power_spectrum_linear * np.exp(
         -wavenumber * (a1 + a2 * wavenumber + a3 * wavenumber**2)
     )
-    power_spectrum_mt = np.sqrt(
-        power_spectrum_linear * kwargs["power_spectrum_non_linear"]
-    ) * np.exp(-invkdelta * wavenumber - b * wavenumber**6)
+
+    power_spectrum_mt = power_spectrum_linear * np.exp(
+        -invkdelta * wavenumber - b * wavenumber**6
+    )
+
+    return power_spectrum_mt, power_spectrum_tt
+
+
+def get_nonlinearbel_model(
+    wavenumber,
+    power_spectrum_linear,
+    **kwargs,
+):
+    """
+    The get_nonlinearbel_model function returns the nonlinear power spectra for a given linear power spectrum.
+
+    Args:
+        wavenumber: Get the wavenumber of the power spectrum
+        power_spectrum_linear: Calculate the power_spectrum_mt and power_spectrum_tt
+        **kwargs: Pass a variable number of keyword arguments to a function
+        : Get the nonlinear power spectrum
+
+    Returns:
+        The nonlinear power spectrum of matter, the cross power spectrum of matter and tracers and the auto-power spectrum of tracers
+
+    """
+    if "power_spectrum_non_linear" not in kwargs.keys():
+        raise ValueError("Non linear power spectrum is needed for nonlinearbel model")
+
+    power_spectrum_mt, power_spectrum_tt = get_bel_model(
+        wavenumber, power_spectrum_linear, **kwargs
+    )
+
+    power_spectrum_mt *= np.sqrt(
+        kwargs["power_spectrum_non_linear"] / power_spectrum_linear
+    )
     power_spectrum_mm = kwargs["power_spectrum_non_linear"]
 
     return power_spectrum_mm, power_spectrum_mt, power_spectrum_tt
@@ -50,15 +95,20 @@ def get_linearbel_model(
     power_spectrum_linear,
     **kwargs,
 ):
-    if "sigma_8" not in kwargs.keys():
-        raise ValueError("Fiducial sigma_8 value is needed for linearbel model")
+    """
+    The get_linearbel_model function returns the linear BEL model for a given power spectrum.
 
-    a1, a2, a3, invkdelta, b = bel_coefficients(kwargs["sigma_8"])
-    power_spectrum_tt = power_spectrum_linear * np.exp(
-        -wavenumber * (a1 + a2 * wavenumber + a3 * wavenumber**2)
-    )
-    power_spectrum_mt = power_spectrum_linear * np.exp(
-        -invkdelta * wavenumber - b * wavenumber**6
+    Args:
+        wavenumber: Calculate the wavenumber in h/mpc
+        power_spectrum_linear: Calculate the power spectrum of the matter field
+        **kwargs: Pass a variable number of keyword arguments to the function
+        : Set the value of the power spectrum
+
+    Returns:
+        The linear power spectrum and the
+    """
+    power_spectrum_mt, power_spectrum_tt = get_bel_model(
+        wavenumber, power_spectrum_linear, **kwargs
     )
     power_spectrum_mm = power_spectrum_linear
 

scripts/flip_compute_power_spectra.py

@@ -1,4 +1,4 @@
-from flip import power_spectra_generator
+from flip.power_spectra import generator
 
 power_spectrum_engine = "class_engine"
 
@@ -29,7 +29,7 @@
     power_spectrum_mm,
     power_spectrum_mt,
     power_spectrum_tt,
-) = power_spectra_generator.compute_power_spectra(
+) = generator.compute_power_spectra(
     power_spectrum_engine,
     power_spectrum_settings,
     redshift,