flipcosmo 1.0.0__py3-none-any.whl

flip/covariance/adamsblake17plane/fisher_terms.py

@@ -0,0 +1,300 @@
+import numpy as np
+
+
+def get_partial_derivative_coefficients(
+    model_kind,
+    parameter_values_dict,
+    variant=None,
+    redshift_dict=None,
+):
+    if model_kind == "density":
+        return get_partial_derivative_coefficients_density(
+            parameter_values_dict,
+            variant=variant,
+        )
+    elif model_kind == "velocity":
+        return get_partial_derivative_coefficients_velocity(
+            parameter_values_dict,
+            variant=variant,
+        )
+    elif model_kind == "density_velocity":
+        return get_partial_derivative_coefficients_density_velocity(
+            parameter_values_dict,
+            variant=variant,
+        )
+    elif model_kind == "full":
+        return get_partial_derivative_coefficients_full(
+            parameter_values_dict,
+            variant=variant,
+        )
+
+
+def get_partial_derivative_coefficients_velocity(
+    parameter_values_dict,
+    variant=None,
+):
+    if variant == "growth_index":
+        partial_coefficients_dict = {
+            "Omegam": {
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["gamma"]
+                    * parameter_values_dict["s8"] ** 2
+                    / parameter_values_dict["Omegam"],
+                ],
+            },
+            "gamma": {
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["s8"] ** 2
+                    * np.log(parameter_values_dict["Omegam"]),
+                ],
+            },
+            "s8": {
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["s8"],
+                ],
+            },
+            "bs8": {
+                "vv": [
+                    0,
+                ],
+            },
+        }
+    else:
+        partial_coefficients_dict = {
+            "fs8": {
+                "vv": [
+                    2 * parameter_values_dict["fs8"],
+                ],
+            },
+            "bs8": {
+                "vv": [
+                    0,
+                ],
+            },
+        }
+    return partial_coefficients_dict
+
+
+def get_partial_derivative_coefficients_density(
+    parameter_values_dict,
+    variant=None,
+):
+    if variant == "growth_index":
+        partial_coefficients_dict = {
+            "Omegam": {
+                "gg": [
+                    0,
+                ],
+            },
+            "gamma": {
+                "gg": [
+                    0,
+                ],
+            },
+            "s8": {
+                "gg": [
+                    0,
+                ],
+            },
+            "bs8": {
+                "gg": [
+                    2 * parameter_values_dict["bs8"],
+                ],
+            },
+        }
+    else:
+        partial_coefficients_dict = {
+            "fs8": {
+                "gg": [
+                    0,
+                ],
+            },
+            "bs8": {
+                "gg": [
+                    2 * parameter_values_dict["bs8"],
+                ],
+            },
+        }
+    return partial_coefficients_dict
+
+
+def get_partial_derivative_coefficients_density_velocity(
+    parameter_values_dict,
+    variant=None,
+):
+    if variant == "growth_index":
+        partial_coefficients_dict = {
+            "Omegam": {
+                "gg": [
+                    0,
+                ],
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["gamma"]
+                    * parameter_values_dict["s8"] ** 2
+                    / parameter_values_dict["Omegam"],
+                ],
+            },
+            "gamma": {
+                "gg": [
+                    0,
+                ],
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["s8"] ** 2
+                    * np.log(parameter_values_dict["Omegam"]),
+                ],
+            },
+            "s8": {
+                "gg": [
+                    0,
+                ],
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["s8"],
+                ],
+            },
+            "bs8": {
+                "gg": [
+                    2 * parameter_values_dict["bs8"],
+                ],
+                "vv": [
+                    0,
+                ],
+            },
+        }
+    else:
+        partial_coefficients_dict = {
+            "fs8": {
+                "gg": [
+                    0,
+                ],
+                "vv": [
+                    2 * parameter_values_dict["fs8"],
+                ],
+            },
+            "bs8": {
+                "gg": [
+                    2 * parameter_values_dict["bs8"],
+                ],
+                "vv": [
+                    0,
+                ],
+            },
+        }
+    return partial_coefficients_dict
+
+
+def get_partial_derivative_coefficients_full(
+    parameter_values_dict,
+    variant=None,
+):
+    if variant == "growth_index":
+        partial_coefficients_dict = {
+            "Omegam": {
+                "gg": [
+                    0,
+                ],
+                "gv": [
+                    parameter_values_dict["Omegam"] ** parameter_values_dict["gamma"]
+                    * parameter_values_dict["bs8"]
+                    * parameter_values_dict["gamma"]
+                    * parameter_values_dict["s8"]
+                    / parameter_values_dict["Omegam"],
+                ],
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["gamma"]
+                    * parameter_values_dict["s8"] ** 2
+                    / parameter_values_dict["Omegam"],
+                ],
+            },
+            "gamma": {
+                "gg": [
+                    0,
+                ],
+                "gv": [
+                    parameter_values_dict["Omegam"] ** parameter_values_dict["gamma"]
+                    * parameter_values_dict["bs8"]
+                    * parameter_values_dict["s8"]
+                    * np.log(parameter_values_dict["Omegam"]),
+                ],
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["s8"] ** 2
+                    * np.log(parameter_values_dict["Omegam"]),
+                ],
+            },
+            "s8": {
+                "gg": [
+                    0,
+                ],
+                "gv": [
+                    parameter_values_dict["Omegam"] ** parameter_values_dict["gamma"]
+                    * parameter_values_dict["bs8"],
+                ],
+                "vv": [
+                    2
+                    * parameter_values_dict["Omegam"]
+                    ** (2 * parameter_values_dict["gamma"])
+                    * parameter_values_dict["s8"],
+                ],
+            },
+            "bs8": {
+                "gg": [
+                    2 * parameter_values_dict["bs8"],
+                ],
+                "gv": [
+                    parameter_values_dict["Omegam"] ** parameter_values_dict["gamma"]
+                    * parameter_values_dict["s8"],
+                ],
+                "vv": [
+                    0,
+                ],
+            },
+        }
+    else:
+        partial_coefficients_dict = {
+            "fs8": {
+                "gg": [
+                    0,
+                ],
+                "gv": [
+                    parameter_values_dict["bs8"],
+                ],
+                "vv": [
+                    2 * parameter_values_dict["fs8"],
+                ],
+            },
+            "bs8": {
+                "gg": [
+                    2 * parameter_values_dict["bs8"],
+                ],
+                "gv": [
+                    parameter_values_dict["fs8"],
+                ],
+                "vv": [
+                    0,
+                ],
+            },
+        }
+    return partial_coefficients_dict

flip/covariance/adamsblake17plane/flip_terms.py

@@ -0,0 +1,75 @@
+import mpmath
+import numpy
+import scipy
+
+
+def set_backend(module):
+    global np, erf
+    if module == "numpy":
+        np = numpy
+        erf = scipy.special.erf
+    elif module == "mpmath":
+        np = mpmath.mp
+        erf = mpmath.erf
+
+
+set_backend("numpy")
+
+
+def M_gg_0_0_0():
+    def func(k):
+        return 1
+
+    return func
+
+
+def N_gg_0_0_0(theta, phi):
+    return 1
+
+
+def M_gv_0_1_0():
+    def func(k):
+        return (100 / 3) / k
+
+    return func
+
+
+def N_gv_0_1_0(theta, phi):
+    return -3 * np.cos(phi)
+
+
+def M_vv_0_0_0():
+    def func(k):
+        return (10000 / 3) / k**2
+
+    return func
+
+
+def N_vv_0_0_0(theta, phi):
+    return 1
+
+
+def M_vv_0_2_0():
+    def func(k):
+        return (4000 / 3) / k**2
+
+    return func
+
+
+def N_vv_0_2_0(theta, phi):
+    return (15 / 4) * np.cos(2 * phi) + 5 / 4
+
+
+dictionary_terms = {"gg": ["0"], "gv": ["0"], "vv": ["0"]}
+dictionary_lmax = {"gg": [0], "gv": [1], "vv": [2]}
+dictionary_subterms = {
+    "gg_0_0": 1,
+    "gv_0_0": 0,
+    "gv_0_1": 1,
+    "vv_0_0": 1,
+    "vv_0_1": 0,
+    "vv_0_2": 1,
+}
+multi_index_model = False
+redshift_dependent_model = False
+regularize_M_terms = None

flip/covariance/adamsblake17plane/generator.py

@@ -0,0 +1,263 @@
+import multiprocessing as mp
+from functools import partial
+
+import numpy as np
+from scipy.special import spherical_jn
+
+from flip.covariance import cov_utils
+
+
+def angle_between(ra_0, ra_1, dec_0, dec_1):
+    """Compute cos of the angle between r0 and r1."""
+    cos_alpha = np.cos(ra_1 - ra_0) * np.cos(dec_0) * np.cos(dec_1) + np.sin(
+        dec_0
+    ) * np.sin(dec_1)
+    return cos_alpha
+
+
+def separation(r_0, r_1, cos_alpha):
+    """Compute separation between r_0 and r_1."""
+    return np.sqrt(r_0**2 + r_1**2 - 2 * r_0 * r_1 * cos_alpha)
+
+
+def window_vv(r_0, r_1, cos_alpha, sep, j0kr, j2kr):
+    """Note: here, the bisector angle definition is used to compute"""
+    win = 1 / 3 * (j0kr + j2kr)
+    alpha = np.arccos(np.clip(cos_alpha, -1.0, 1.0))
+    phi = cov_utils.compute_phi_bisector_theorem(sep, alpha, r_0, r_1)
+    win += -j2kr * np.cos(phi) ** 2
+    return win
+
+
+def window_vg(r_0, r_1, cos_alpha, sep, j1kr):
+    """Note: here, the bisector angle definition is used to compute"""
+    alpha = np.arccos(np.clip(cos_alpha, -1.0, 1.0))
+    phi = cov_utils.compute_phi_bisector_theorem(sep, alpha, r_0, r_1)
+    win = j1kr * np.cos(phi)
+    return win
+
+
+def intp(win, k, pk):
+    pint = win.T * pk
+    return np.trapz(pint, x=k)
+
+
+def compute_coef_gg(k, pk, coord):
+    cos = angle_between(coord[0], coord[1], coord[2], coord[3])
+    sep = separation(coord[4], coord[5], cos)
+    ksep = np.outer(k, sep)
+    j0 = spherical_jn(0, ksep)
+    res = intp(j0, k, pk * k**2)
+    return res
+
+
+def compute_coef_gv(k, pk, coord):
+    cos = angle_between(coord[0], coord[1], coord[2], coord[3])
+    sep = separation(coord[4], coord[5], cos)
+    ksep = np.outer(k, sep)
+    j1 = spherical_jn(1, ksep)
+    res = window_vg(coord[4], coord[5], cos, sep, j1)
+    res = intp(res, k, pk * k)
+    return res
+
+
+def compute_coef_vv(k, pk, coord):
+    cos = angle_between(coord[0], coord[1], coord[2], coord[3])
+    sep = separation(coord[4], coord[5], cos)
+    ksep = np.outer(k, sep)
+    j0 = spherical_jn(0, ksep)
+    j2 = spherical_jn(2, ksep)
+    res = window_vv(coord[4], coord[5], cos, sep, j0, j2)
+    res = intp(res, k, pk)
+    return res
+
+
+def covariance_vv(
+    ra_v,
+    dec_v,
+    rcomov_v,
+    wavenumber,
+    power_spectrum,
+    size_batch=100_000,
+    number_worker=8,
+):
+    N = len(ra_v)
+    n_task = int((N * (N + 1)) / 2) - N
+    batches = []
+    for n in range(0, n_task, size_batch):
+        brange = np.arange(n, np.min((n + size_batch, n_task)))
+        i_list, j_list = cov_utils.compute_i_j(N, brange)
+        r_comovi, rai, deci = rcomov_v[i_list], ra_v[i_list], dec_v[i_list]
+        r_comovj, raj, decj = rcomov_v[j_list], ra_v[j_list], dec_v[j_list]
+        batches.append([rai, raj, deci, decj, r_comovi, r_comovj])
+
+    with mp.Pool(number_worker) as pool:
+        func = partial(compute_coef_vv, wavenumber, power_spectrum)
+        pool_results = pool.map(func, batches)
+    values = np.concatenate(pool_results)
+
+    var_val = np.trapz(power_spectrum / 3, x=wavenumber)
+    cov = np.insert(values, 0, var_val)
+    cov = 100**2 / (2 * np.pi**2) * cov
+    return cov
+
+
+def covariance_gv(
+    ra_g,
+    dec_g,
+    rcomov_g,
+    ra_v,
+    dec_v,
+    rcomov_v,
+    wavenumber,
+    power_spectrum,
+    size_batch=100_000,
+    number_worker=8,
+):
+    number_objects_g = len(ra_g)
+    number_objects_v = len(ra_v)
+
+    n_task = int(number_objects_g * number_objects_v)
+    batches = []
+    for n in range(0, n_task, size_batch):
+        brange = np.arange(n, np.min((n + size_batch, n_task)))
+        i_list, j_list = cov_utils.compute_i_j_cross_matrix(number_objects_v, brange)
+        r_comovi, rai, deci = rcomov_g[i_list], ra_g[i_list], dec_g[i_list]
+        r_comovj, raj, decj = rcomov_v[j_list], ra_v[j_list], dec_v[j_list]
+        batches.append([rai, raj, deci, decj, r_comovi, r_comovj])
+
+    with mp.Pool(number_worker) as pool:
+        func = partial(compute_coef_gv, wavenumber, power_spectrum)
+        pool_results = pool.map(func, batches)
+    values = np.concatenate(pool_results)
+    var_val = compute_coef_gv(
+        wavenumber, power_spectrum, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+    )
+    cov = np.insert(values, 0, var_val)
+    cov = 100 / (2 * np.pi**2) * cov
+    return cov
+
+
+def covariance_gg(
+    ra_g,
+    dec_g,
+    rcomov_g,
+    wavenumber,
+    power_spectrum,
+    size_batch=100_000,
+    number_worker=8,
+):
+    N = len(ra_g)
+    n_task = int((N * (N + 1)) / 2) - N
+    batches = []
+    for n in range(0, n_task, size_batch):
+        brange = np.arange(n, np.min((n + size_batch, n_task)))
+        i_list, j_list = cov_utils.compute_i_j(N, brange)
+        r_comovi, rai, deci = rcomov_g[i_list], ra_g[i_list], dec_g[i_list]
+        r_comovj, raj, decj = rcomov_g[j_list], ra_g[j_list], dec_g[j_list]
+        batches.append([rai, raj, deci, decj, r_comovi, r_comovj])
+
+    with mp.Pool(number_worker) as pool:
+        func = partial(compute_coef_gg, wavenumber, power_spectrum)
+        pool_results = pool.map(func, batches)
+    values = np.concatenate(pool_results)
+    var_val = compute_coef_gg(
+        wavenumber, power_spectrum, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+    )
+    cov = np.insert(values, 0, var_val)
+    cov = 1 / (2 * np.pi**2) * cov
+    return cov
+
+
+def generate_covariance(
+    model_kind,
+    power_spectrum_dict,
+    coordinates_velocity=None,
+    coordinates_density=None,
+    **kwargs,
+):
+    """
+    The generate_covariance function generates the covariance matrix for a given model type, power spectrum, and coordinates.
+
+    Args:
+        model_kind: Determine which covariance matrices are generated, and the coordinates_density and coordinates_velocity parameters are used to generate the covariance matrices
+        power_spectrum_dict: Pass the power spectrum to the function
+        coordinates_velocity: Define the coordinates of the velocity field
+        coordinates_density: Define the coordinates of the density field
+        **kwargs: Pass keyword arguments to the function
+        : Generate the covariance matrix for a given model
+        The wide angle definition is bisector.
+
+    Returns:
+        A dictionary of covariance matrices
+
+    Doc Author:
+        Trelent
+    """
+    cov_utils.check_generator_need(
+        model_kind,
+        coordinates_density,
+        coordinates_velocity,
+    )
+    covariance_dict = {}
+
+    if model_kind in ["density", "full", "density_velocity"]:
+        covariance_dict["gg"] = np.array(
+            [
+                covariance_gg(
+                    coordinates_density[0],
+                    coordinates_density[1],
+                    coordinates_density[2],
+                    power_spectrum_dict["gg"][0][0],
+                    power_spectrum_dict["gg"][0][1],
+                    **kwargs,
+                )
+            ]
+        )
+        number_densities = len(coordinates_density[0])
+    else:
+        number_densities = None
+
+    if model_kind in ["velocity", "full", "density_velocity"]:
+        covariance_dict["vv"] = np.array(
+            [
+                covariance_vv(
+                    coordinates_velocity[0],
+                    coordinates_velocity[1],
+                    coordinates_velocity[2],
+                    power_spectrum_dict["vv"][0][0],
+                    power_spectrum_dict["vv"][0][1],
+                    **kwargs,
+                )
+            ]
+        )
+        number_velocities = len(coordinates_velocity[0])
+    else:
+        number_velocities = None
+
+    if model_kind == "full":
+        covariance_dict["gv"] = np.array(
+            [
+                covariance_gv(
+                    coordinates_density[0],
+                    coordinates_density[1],
+                    coordinates_density[2],
+                    coordinates_velocity[0],
+                    coordinates_velocity[1],
+                    coordinates_velocity[2],
+                    power_spectrum_dict["gv"][0][0],
+                    power_spectrum_dict["gv"][0][1],
+                    **kwargs,
+                )
+            ]
+        )
+
+    los_definition = "bisector"
+    redshift_dict = None
+    return (
+        covariance_dict,
+        number_densities,
+        number_velocities,
+        los_definition,
+        redshift_dict,
+    )

flip/covariance/adamsblake20/__init__.py

@@ -0,0 +1,10 @@
+_variant = [None, "baseline", "nobeta"]
+
+_free_par = {
+    "fs8": ["velocity@all", "density@nobeta"],
+    "bs8": "density@all",
+    "sigv": "velocity@all",
+    "beta_f": "density@baseline",
+}
+
+_coordinate_keys = ["ra", "dec", "rcom_zobs"]

flip/covariance/adamsblake20/coefficients.py

@@ -0,0 +1,46 @@
+def get_coefficients(
+    parameter_values_dict,
+    model_kind,
+    variant=None,
+    redshift_dict=None,
+):
+    coefficients_dict = {}
+    if model_kind in ["density", "full", "density_velocity"]:
+        if variant == "nobeta":
+            coefficients_dict["gg"] = [
+                parameter_values_dict["bs8"] ** 2,
+                parameter_values_dict["bs8"] * parameter_values_dict["fs8"],
+                parameter_values_dict["fs8"] ** 2,
+            ]
+        else:
+            coefficients_dict["gg"] = [
+                parameter_values_dict["bs8"] ** 2,
+                parameter_values_dict["bs8"] ** 2 * parameter_values_dict["beta_f"],
+                parameter_values_dict["bs8"] ** 2
+                * parameter_values_dict["beta_f"] ** 2,
+            ]
+    if model_kind in ["full"]:
+        if variant == "nobeta":
+            coefficients_dict["gv"] = [
+                parameter_values_dict["bs8"] * parameter_values_dict["fs8"],
+                parameter_values_dict["fs8"] ** 2,
+            ]
+        else:
+            coefficients_dict["gv"] = [
+                parameter_values_dict["bs8"] * parameter_values_dict["fs8"],
+                parameter_values_dict["bs8"]
+                * parameter_values_dict["fs8"]
+                * parameter_values_dict["beta_f"],
+            ]
+    if model_kind in ["velocity", "full", "density_velocity"]:
+        coefficients_dict["vv"] = [parameter_values_dict["fs8"] ** 2]
+    return coefficients_dict
+
+
+def get_diagonal_coefficients(parameter_values_dict, model_kind):
+    coefficients_dict = {}
+    if model_kind in ["density", "full", "density_velocity"]:
+        coefficients_dict["gg"] = 0.0
+    if model_kind in ["velocity", "full", "density_velocity"]:
+        coefficients_dict["vv"] = parameter_values_dict["sigv"] ** 2
+    return coefficients_dict