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

flip/covariance/analytical/carreres23/generator.py

@@ -0,0 +1,198 @@
+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 two directions.
+
+    Args:
+        ra_0: Right ascension of first object (radians).
+        ra_1: Right ascension of second object (radians).
+        dec_0: Declination of first object (radians).
+        dec_1: Declination of second object (radians).
+
+    Returns:
+        Cosine of the angle between the two directions.
+    """
+    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 comoving separation given distances and angular cosine.
+
+    Args:
+        r_0: Comoving distance of first object.
+        r_1: Comoving distance of second object.
+        cos_alpha: Cosine of the angular separation.
+
+    Returns:
+        Scalar separation ``|r_0 - r_1|`` generalized to non-colinear case.
+    """
+    return np.sqrt(r_0**2 + r_1**2 - 2 * r_0 * r_1 * cos_alpha)
+
+
+def window(r_0, r_1, cos_alpha, sep, j0kr, j2kr):
+    """Wide-angle window function using bisector definition.
+
+    Args:
+        r_0: Comoving distance of first object.
+        r_1: Comoving distance of second object.
+        cos_alpha: Cosine of angle between directions.
+        sep: Comoving separation.
+        j0kr: Spherical Bessel ``j_0(k r)`` evaluated on grid.
+        j2kr: Spherical Bessel ``j_2(k r)`` evaluated on grid.
+
+    Returns:
+        Window values per k contributing to the vv covariance integral.
+    """
+    win = 1 / 3 * (j0kr - 2 * j2kr) * cos_alpha
+    win += j2kr * r_0 * r_1 / sep**2 * (1 - cos_alpha**2)
+    return win
+
+
+def intp(win, k, pk):
+    """Integrate window times power spectrum over k.
+
+    Args:
+        win: Window array per k (shape compatible with k).
+        k: Wavenumber grid.
+        pk: Power spectrum values at k.
+
+    Returns:
+        Scalar integral value via trapezoidal rule.
+    """
+    pint = win.T * pk
+    return np.trapz(pint, x=k)
+
+
+def covariance_vv(
+    ra_in,
+    dec_in,
+    rcomov_in,
+    k_in,
+    pk_in,
+    grid_window_in=None,
+    size_batch=100_000,
+    number_worker=8,
+):
+    """Compute Carreres23 vv covariance using bisector wide-angle.
+
+    Applies optional grid window to the velocity power spectrum and integrates
+    per pair in batches using multiprocessing.
+
+    Args:
+        ra_in: Right ascensions (radians).
+        dec_in: Declinations (radians).
+        rcomov_in: Comoving distances.
+        k_in: Wavenumber grid.
+        pk_in: Velocity power spectrum values at ``k_in``.
+        grid_window_in: Optional window to apply to ``pk_in``.
+        size_batch: Number of pairs per batch for processing.
+        number_worker: Number of parallel workers.
+
+    Returns:
+        Flattened covariance vector: variance at index 0, then upper triangle.
+    """
+    N = len(ra_in)
+
+    if grid_window_in is not None:
+        pk = pk_in * grid_window_in**2
+    else:
+        pk = pk_in
+
+    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_in[i_list], ra_in[i_list], dec_in[i_list]
+        r_comovj, raj, decj = rcomov_in[j_list], ra_in[j_list], dec_in[j_list]
+        batches.append([rai, raj, deci, decj, r_comovi, r_comovj])
+
+    with mp.Pool(number_worker) as pool:
+        func = partial(compute_coef, k_in, pk)
+        pool_results = pool.map(func, batches)
+    values = np.concatenate(pool_results)
+
+    var_val = np.trapz(pk / 3, x=k_in)
+    cov = np.insert(values, 0, var_val)
+    cov = 100**2 / (2 * np.pi**2) * cov
+    return cov
+
+
+def compute_coef(k, pk, coord):
+    """Compute covariance coefficient for a single pair.
+
+    Args:
+        k: Wavenumber grid.
+        pk: Power spectrum values at k.
+        coord: Tuple/list ``(ra_i, ra_j, dec_i, dec_j, r_i, r_j)``.
+
+    Returns:
+        Scalar integral value contributing to the covariance.
+    """
+    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(coord[4], coord[5], cos, sep, j0, j2)
+    res = intp(res, k, pk)
+    return res
+
+
+def generate_covariance(
+    model_kind,
+    power_spectrum_dict,
+    coordinates_density=False,
+    coordinates_velocity=None,
+    **kwargs,
+):
+    """Generate Carreres23 covariance for velocity-only model.
+
+    Wide-angle definition uses the bisector. Only ``vv`` block is produced.
+
+    Args:
+        model_kind: Must be ``"velocity"``.
+        power_spectrum_dict: Dict with ``{"vv": [(k,), (pk,)]}`` entries.
+        coordinates_density: Ignored; kept for interface compatibility.
+        coordinates_velocity: Tuple ``(ra, dec, rcom)`` for velocity tracers.
+        **kwargs: Forwarded to ``covariance_vv`` (e.g., batching/workers/window).
+
+    Returns:
+        Tuple ``(covariance_dict, number_densities, number_velocities, los_definition)``.
+    """
+    assert model_kind == "velocity"
+    cov_utils.check_generator_need(
+        model_kind,
+        coordinates_density,
+        coordinates_velocity,
+    )
+    number_densities = None
+    number_velocities = len(coordinates_velocity[0])
+    cov_vv = covariance_vv(
+        coordinates_velocity[0],
+        coordinates_velocity[1],
+        coordinates_velocity[2],
+        power_spectrum_dict["vv"][0][0],
+        power_spectrum_dict["vv"][0][1],
+        **kwargs,
+    )
+
+    los_definition = "bisector"
+
+    return (
+        {"vv": np.array([cov_vv])},
+        number_densities,
+        number_velocities,
+        los_definition,
+    )

flip/covariance/analytical/genericzdep/__init__.py

@@ -0,0 +1,5 @@
+_variant = [None]
+
+_free_par = {"sigv": "velocity@all"}
+
+_coordinate_keys = ["ra", "dec", "rcom_zobs", "zobs"]

flip/covariance/analytical/genericzdep/coefficients.py

@@ -0,0 +1,53 @@
+from ..._config import __use_jax__
+
+if __use_jax__:
+    try:
+        import jax.numpy as jnp
+
+        jax_installed = True
+
+    except ImportError:
+        import numpy as jnp
+
+        jax_installed = False
+else:
+    import numpy as jnp
+
+    jax_installed = False
+
+
+def get_coefficients(
+    parameter_values_dict,
+    model_kind,
+    variant=None,
+    covariance_prefactor_dict=None,
+):
+    coefficients_dict = {}
+
+    redshift_velocity = covariance_prefactor_dict["redshift_velocity"]
+    f_z = covariance_prefactor_dict["f_z"]
+    E_z = covariance_prefactor_dict["E_z"]
+    D_growth_z = covariance_prefactor_dict["D_growth_z"]
+
+    coefficients_dict["vv"] = get_cov_matrix_prefactor(
+        redshift_velocity, f_z, E_z, D_growth_z
+    )
+    return coefficients_dict
+
+
+def get_cov_matrix_prefactor(redshift, f_z, E_z, D_growth_z):
+    Ai_z = f_z * D_growth_z * E_z / (1 + redshift)
+    Aij_z = jnp.outer(Ai_z, Ai_z)
+    Delta_D_growth_squared = jnp.subtract.outer(D_growth_z, D_growth_z) ** 2
+
+    return [
+        Aij_z * jnp.ones_like(Delta_D_growth_squared),
+        Aij_z * Delta_D_growth_squared,
+        Aij_z * 1 / 2 * Delta_D_growth_squared**2,
+    ]
+
+
+def get_diagonal_coefficients(parameter_values_dict, model_kind):
+    coefficients_dict = {}
+    coefficients_dict["vv"] = parameter_values_dict["sigv"] ** 2
+    return coefficients_dict

flip/covariance/analytical/genericzdep/flip_terms.py

@@ -0,0 +1,99 @@
+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_vv_0_0_0(kNL):
+    def func(k):
+        return (10000 / 9) / k**2
+
+    return func
+
+
+def N_vv_0_0_0(theta, phi):
+    return 3 * np.cos(theta)
+
+
+def M_vv_0_2_0(kNL):
+    def func(k):
+        return (10000 / 9) / k**2
+
+    return func
+
+
+def N_vv_0_2_0(theta, phi):
+    return (9 / 2) * np.cos(2 * phi) + (3 / 2) * np.cos(theta)
+
+
+def M_vv_1_0_0(kNL):
+    def func(k):
+        return (10000 / 9) / kNL**2
+
+    return func
+
+
+def N_vv_1_0_0(theta, phi):
+    return 3 * np.cos(theta)
+
+
+def M_vv_1_2_0(kNL):
+    def func(k):
+        return (10000 / 9) / kNL**2
+
+    return func
+
+
+def N_vv_1_2_0(theta, phi):
+    return (9 / 2) * np.cos(2 * phi) + (3 / 2) * np.cos(theta)
+
+
+def M_vv_2_0_0(kNL):
+    def func(k):
+        return (10000 / 9) * k**2 / kNL**4
+
+    return func
+
+
+def N_vv_2_0_0(theta, phi):
+    return 3 * np.cos(theta)
+
+
+def M_vv_2_2_0(kNL):
+    def func(k):
+        return (10000 / 9) * k**2 / kNL**4
+
+    return func
+
+
+def N_vv_2_2_0(theta, phi):
+    return (9 / 2) * np.cos(2 * phi) + (3 / 2) * np.cos(theta)
+
+
+dictionary_terms = {"vv": ["0", "1", "2"]}
+dictionary_lmax = {"vv": [2, 2, 2]}
+dictionary_subterms = {
+    "vv_0_0": 1,
+    "vv_0_1": 0,
+    "vv_0_2": 1,
+    "vv_1_0": 1,
+    "vv_1_1": 0,
+    "vv_1_2": 1,
+    "vv_2_0": 1,
+    "vv_2_1": 0,
+    "vv_2_2": 1,
+}
+multi_index_model = False
+regularize_M_terms = None

flip/covariance/{lai22 → analytical/lai22}/coefficients.py

@@ -1,13 +1,12 @@
 import numpy as np
-
-from flip.covariance.lai22.flip_terms import dictionary_terms
+from flip.covariance.analytical.lai22.flip_terms import dictionary_terms
 
 
 def get_coefficients(
     parameter_values_dict,
     model_kind,
     variant=None,
-    redshift_dict=None,
+    covariance_prefactor_dict=None,
 ):
     coefficients_dict = {}
     if model_kind in ["density", "full", "density_velocity"]:

flip/covariance/{lai22 → analytical/lai22}/fisher_terms.py

@@ -5,7 +5,7 @@ def get_partial_derivative_coefficients(
     model_kind,
     parameter_values_dict,
     variant=None,
-    redshift_dict=None,
+    covariance_prefactor_dict=None,
 ):
     if model_kind == "density":
         return get_partial_derivative_coefficients_density(

flip/covariance/{lai22 → analytical/lai22}/flip_terms.py

@@ -15999,5 +15999,4 @@ dictionary_subterms = {
     "vv_0_0_2": 1,
 }
 multi_index_model = True
-redshift_dependent_model = False
 regularize_M_terms = None