Add redshift dependency

flip/covariance/adamsblake17plane/coefficients.py

@@ -2,6 +2,8 @@ def get_coefficients(
     model_type,
     parameter_values_dict,
     variant=None,
+    redshift_dict=None,
+    power_spectrum_amplitude_function=None,
 ):
     coefficients_dict = {}
     if model_type in ["density", "full", "density_velocity"]:

flip/covariance/adamsblake17plane/fisher_terms.py

@@ -5,6 +5,8 @@ def get_partial_derivative_coefficients(
     model_type,
     parameter_values_dict,
     variant=None,
+    redshift_dict=None,
+    power_spectrum_amplitude_function=None,
 ):
     if model_type == "density":
         return get_partial_derivative_coefficients_density(

flip/covariance/adamsblake17plane/flip_terms.py

@@ -57,3 +57,4 @@ def N_vv_0_2_0(theta, phi):
     "vv_0_2": 1,
 }
 multi_index_model = False
+redshift_dependent_model = False

flip/covariance/adamsblake17plane/generator.py

@@ -253,5 +253,11 @@ def generate_covariance(
         )
 
     los_definition = "bisector"
-
-    return covariance_dict, number_densities, number_velocities, los_definition
+    redshift_dict = None
+    return (
+        covariance_dict,
+        number_densities,
+        number_velocities,
+        los_definition,
+        redshift_dict,
+    )

flip/covariance/adamsblake20/coefficients.py

@@ -2,6 +2,8 @@ def get_coefficients(
     model_type,
     parameter_values_dict,
     variant=None,
+    redshift_dict=None,
+    power_spectrum_amplitude_function=None,
 ):
     coefficients_dict = {}
     if model_type in ["density", "full", "density_velocity"]:

flip/covariance/adamsblake20/fisher_terms.py

@@ -5,6 +5,8 @@ def get_partial_derivative_coefficients(
     model_type,
     parameter_values_dict,
     variant=None,
+    redshift_dict=None,
+    power_spectrum_amplitude_function=None,
 ):
     if model_type == "density":
         return get_partial_derivative_coefficients_density(

flip/covariance/adamsblake20/flip_terms.py

@@ -72,11 +72,7 @@ def M_gg_1_2_0(sig_g):
     def func(k):
         return (
             -5 * np.exp(-(k**2) * sig_g**2) / (k**2 * sig_g**2)
-            - 5
-            / 4
-            * np.sqrt(np.pi)
-            * scipy.special.erf(k * sig_g)
-            / (k**3 * sig_g**3)
+            - 5 / 4 * np.sqrt(np.pi) * scipy.special.erf(k * sig_g) / (k**3 * sig_g**3)
             - 45 / 4 * np.exp(-(k**2) * sig_g**2) / (k**4 * sig_g**4)
             + (45 / 8)
             * np.sqrt(np.pi)
@@ -359,3 +355,4 @@ def N_vv_0_2_0(theta, phi):
     "vv_0_2": 1,
 }
 multi_index_model = False
+redshift_dependent_model = False

flip/covariance/carreres23/coefficients.py

@@ -5,6 +5,8 @@ def get_coefficients(
     model_type,
     parameter_values_dict,
     variant=None,
+    redshift_dict=None,
+    power_spectrum_amplitude_function=None,
 ):
     coefficients_dict = {}
     coefficients_dict["vv"] = [parameter_values_dict["fs8"] ** 2]

flip/covariance/carreres23/fisher_terms.py

@@ -5,6 +5,8 @@ def get_partial_derivative_coefficients(
     model_type,
     parameter_values_dict,
     variant=None,
+    redshift_dict=None,
+    power_spectrum_amplitude_function=None,
 ):
     if variant == "growth_index":
         partial_coefficients_dict = {

flip/covariance/carreres23/flip_terms.py

@@ -28,3 +28,4 @@ def N_vv_0_2_0(theta, phi):
 dictionary_lmax = {"vv": [2]}
 dictionary_subterms = {"vv_0_0": 1, "vv_0_1": 0, "vv_0_2": 1}
 multi_index_model = False
+redshift_dependent_model = False

flip/covariance/carreres23/generator.py

@@ -121,10 +121,12 @@ def generate_covariance(
     )
 
     los_definition = "bisector"
+    redshift_dict = None
 
     return (
         {"vv": np.array([cov_vv])},
         number_densities,
         number_velocities,
         los_definition,
+        redshift_dict,
     )

flip/covariance/contraction.py

@@ -17,6 +17,8 @@ def __init__(
         coordinates_dict=None,
         basis_definition=None,
         endpoint_los_definition=None,
+        redshift_dict=None,
+        power_spectrum_amplitude_function=None,
         variant=None,
     ):
         self.model_name = model_name
@@ -25,6 +27,8 @@ def __init__(
         self.coordinates_dict = coordinates_dict
         self.basis_definition = basis_definition
         self.endpoint_los_definition = endpoint_los_definition
+        self.redshift_dict = redshift_dict
+        self.power_spectrum_amplitude_function = power_spectrum_amplitude_function
         self.variant = variant
 
     @classmethod
@@ -41,12 +45,15 @@ def init_from_flip(
         additional_parameters_values=None,
         basis_definition="bisector",
         endpoint_los_definition="bisector",
+        redshift=None,
+        power_spectrum_amplitude_function=None,
         variant=None,
         **kwargs,
     ):
         (
             contraction_dict,
             coordinates_dict,
+            redshift_dict,
         ) = contract_covariance(
             model_name,
             model_type,
@@ -59,6 +66,7 @@ def init_from_flip(
             additional_parameters_values=additional_parameters_values,
             basis_definition=basis_definition,
             endpoint_los_definition=endpoint_los_definition,
+            redshift=redshift,
             **kwargs,
         )
 
@@ -69,6 +77,8 @@ def init_from_flip(
             coordinates_dict=coordinates_dict,
             basis_definition=basis_definition,
             endpoint_los_definition=endpoint_los_definition,
+            redshift_dict=redshift_dict,
+            power_spectrum_amplitude_function=power_spectrum_amplitude_function,
             variant=variant,
         )
 
@@ -125,7 +135,9 @@ def compute_contraction_sum(
         coefficients_dict = coefficients.get_coefficients(
             self.model_type,
             parameter_values_dict,
-            self.variant,
+            variant=self.variant,
+            redshift_dict=self.redshift_dict,
+            power_spectrum_amplitude_function=self.power_spectrum_amplitude_function,
         )
         contraction_covariance_sum_dict = {}
         if self.model_type == "density":
@@ -293,6 +305,7 @@ def contract_covariance(
     additional_parameters_values=None,
     basis_definition="bisector",
     endpoint_los_definition="bisector",
+    redshift=None,
     number_worker=8,
     hankel=True,
 ):
@@ -339,5 +352,10 @@ def contract_covariance(
             number_worker=number_worker,
             hankel=hankel,
         )[:, 1:].reshape(-1, len(coord_1), len(coord_2))
-
-    return contraction_dict, coordinates_dict
+    redshift_dict = generator_flip.generate_redshift_dict(
+        model_name,
+        model_type,
+        redshift_velocity=redshift,
+        redshift_density=redshift,
+    )
+    return contraction_dict, coordinates_dict, redshift_dict

flip/covariance/covariance.py

@@ -20,6 +20,8 @@ def __init__(
         full_matrix=False,
         number_densities=None,
         number_velocities=None,
+        redshift_dict=None,
+        power_spectrum_amplitude_function=None,
         variant=None,
     ):
         """
@@ -49,6 +51,8 @@ def __init__(
         self.full_matrix = full_matrix
         self.number_densities = number_densities
         self.number_velocities = number_velocities
+        self.redshift_dict = redshift_dict
+        self.power_spectrum_amplitude_function = power_spectrum_amplitude_function
         self.variant = variant
 
     @classmethod
@@ -61,6 +65,7 @@ def init_from_flip(
         coordinates_velocity=None,
         additional_parameters_values=None,
         los_definition="bisector",
+        power_spectrum_amplitude_function=None,
         variant=None,
         **kwargs,
     ):
@@ -93,6 +98,7 @@ def init_from_flip(
             covariance_dict,
             number_densities,
             number_velocities,
+            redshift_dict,
         ) = generator_flip.generate_covariance(
             model_name,
             model_type,
@@ -115,6 +121,8 @@ def init_from_flip(
             full_matrix=False,
             number_densities=number_densities,
             number_velocities=number_velocities,
+            redshift_dict=redshift_dict,
+            power_spectrum_amplitude_function=power_spectrum_amplitude_function,
             variant=variant,
         )
 
@@ -127,6 +135,7 @@ def init_from_generator(
         coordinates_velocity=None,
         coordinates_density=None,
         additional_parameters_values=None,
+        power_spectrum_amplitude_function=None,
         variant=None,
         **kwargs,
     ):
@@ -162,6 +171,7 @@ def init_from_generator(
             number_densities,
             number_velocities,
             los_definition,
+            redshift_dict,
         ) = generator.generate_covariance(
             model_type,
             power_spectrum_dict,
@@ -181,6 +191,8 @@ def init_from_generator(
             full_matrix=False,
             number_densities=number_densities,
             number_velocities=number_velocities,
+            redshift_dict=redshift_dict,
+            power_spectrum_amplitude_function=power_spectrum_amplitude_function,
             variant=variant,
         )
 
@@ -299,6 +311,8 @@ def compute_covariance_sum(
             self.model_type,
             parameter_values_dict,
             variant=self.variant,
+            redshift_dict=self.redshift_dict,
+            power_spectrum_amplitude_function=self.power_spectrum_amplitude_function,
         )
         coefficients_dict_diagonal = coefficients.get_diagonal_coefficients(
             self.model_type,

flip/covariance/generator.py

@@ -12,6 +12,7 @@
 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.rcrk24 import flip_terms as flip_terms_rcrk24
 from flip.covariance.ravouxcarreres import flip_terms as flip_terms_ravouxcarreres
 from flip.utils import create_log
 
@@ -22,6 +23,7 @@
     "lai22",
     "carreres23",
     "ravouxcarreres",
+    "rcrk24",
 ]
 
 
@@ -45,10 +47,7 @@ def correlation_integration(l, r, k, integrand):
     """
     kr = np.outer(k, r)
     integrand = (
-        (-1) ** (l // 2)
-        * (k**2 / (2 * np.pi**2))
-        * integrand
-        * spherical_jn(l, kr).T
+        (-1) ** (l // 2) * (k**2 / (2 * np.pi**2)) * integrand * spherical_jn(l, kr).T
     )
     return (-1) ** (l % 2) * integrate.simps(integrand, x=k)
 
@@ -409,6 +408,49 @@ def compute_cov(
     return covariance
 
 
+def generate_redshift_dict(
+    model_name,
+    model_type,
+    redshift_velocity=None,
+    redshift_density=None,
+    coordinates_velocity=None,
+    coordinates_density=None,
+):
+    redshift_dependent_model = eval(f"flip_terms_{model_name}.redshift_dependent_model")
+    if redshift_dependent_model:
+        redshift_dict = {}
+    else:
+        return None
+
+    if model_type in ["density", "full", "density_velocity"]:
+        if redshift_dependent_model:
+            if redshift_density is not None:
+                redshift_dict["g"] = redshift_density
+            else:
+                if len(coordinates_density) < 4:
+                    raise ValueError(
+                        "You are using a model which is redshift dependent."
+                        "Please provide redshifts as the fourth field"
+                        "of the coordinates_density value"
+                    )
+                else:
+                    redshift_dict["g"] = coordinates_density[3]
+    if model_type in ["velocity", "full", "density_velocity"]:
+        if redshift_dependent_model:
+            if redshift_velocity is not None:
+                redshift_dict["v"] = redshift_velocity
+            else:
+                if len(coordinates_velocity) < 4:
+                    raise ValueError(
+                        "You are using a model which is redshift dependent."
+                        "Please provide redshifts as the fourth field"
+                        "of the coordinates_velocity value"
+                    )
+                else:
+                    redshift_dict["v"] = coordinates_velocity[3]
+    return redshift_dict
+
+
 def generate_covariance(
     model_name,
     model_type,
@@ -446,6 +488,7 @@ def generate_covariance(
         coordinates_velocity,
     )
     covariance_dict = {}
+
     if model_type in ["density", "full", "density_velocity"]:
         covariance_dict["gg"] = compute_cov(
             model_name,
@@ -493,4 +536,11 @@ def generate_covariance(
             hankel=hankel,
             los_definition=los_definition,
         )
-    return covariance_dict, number_densities, number_velocities
+
+    redshift_dict = generate_redshift_dict(
+        model_name,
+        model_type,
+        coordinates_velocity=coordinates_velocity,
+        coordinates_density=coordinates_density,
+    )
+    return covariance_dict, number_densities, number_velocities, redshift_dict

flip/covariance/lai22/coefficients.py

@@ -7,6 +7,8 @@ def get_coefficients(
     model_type,
     parameter_values_dict,
     variant=None,
+    redshift_dict=None,
+    power_spectrum_amplitude_function=None,
 ):
     coefficients_dict = {}
     if model_type in ["density", "full", "density_velocity"]:

flip/covariance/lai22/fisher_terms.py

@@ -5,6 +5,8 @@ def get_partial_derivative_coefficients(
     model_type,
     parameter_values_dict,
     variant=None,
+    redshift_dict=None,
+    power_spectrum_amplitude_function=None,
 ):
     if model_type == "density":
         return get_partial_derivative_coefficients_density(

flip/covariance/lai22/flip_terms.py

@@ -15984,3 +15984,4 @@ def N_vv_0_0_2_0(theta, phi):
     "vv_0_0_2": 1,
 }
 multi_index_model = True
+redshift_dependent_model = False