API Reference

Cosmological perturbation theory kernels and integration utilities.

This module provides JAX-compatible implementations of standard perturbation theory (SPT) kernels for large-scale structure calculations, including redshift-space distortions and geometric corrections. It also includes numerical integration routines for computing bispectrum multipoles and power spectra.

The module is designed for efficient computation using JAX’s automatic differentiation and just-in-time compilation capabilities.

Main features: - Standard SPT kernels (F2, G2, Z1, Z2) - Redshift-space distortion modeling - Alcock-Paczynski effect corrections - Finger-of-God damping - GEO-FPT geometric corrections - Numerical integration with Gauss-Legendre quadrature - Bispectrum multipole calculations - Sugiyama estimator implementation

Functions are organized into: 1. Numerical integration utilities 2. SPT kernel definitions 3. Bispectrum integrands and multipole calculations 4. Sugiyama estimator implementation 5. Power spectrum 1-loop corrections

All functions are JAX-compatible and support vectorization via jax.vmap.

geofptax.kernels.bk_multip(tr, tr2, tr3, tr4, kp, pk, cosm_par, redshift, num_points=10, fi_vals=Array([[1.0374585, 1.0179875, 1.0033404], [0.0019502, -0.00410578, -0.00397587], [0.0055635, 0.01488085, 0.02396768], [-0.04836375, -0.05466602, -0.05680982], [0.00806689, 0.01114415, 0.01325444]], dtype=float32))

Compute the bispectrum multipoles for a given set of triangles, power spectrum, and cosmological parameters.

High-level wrapper that computes all four multipoles and returns them in a dictionary.

Parameters:

• tr (jnp.ndarray) – Array of triangle side lengths (ka, kb, kc) for the monopole calculation.

• tr2 (jnp.ndarray) – Array of triangle side lengths (ka, kb, kc) for the first multipole calculation.

• tr3 (jnp.ndarray) – Array of triangle side lengths (ka, kb, kc) for the second multipole calculation.

• tr4 (jnp.ndarray) – Array of triangle side lengths (ka, kb, kc) for the third multipole calculation.

• kp (jnp.ndarray) – Array of wavevector magnitudes for the power spectrum.

• pk (jnp.ndarray) – Array of power spectrum values corresponding to kp.

• cosm_par (jnp.ndarray) – Cosmological parameters array.

• redshift (float) – Redshift at which to compute the bispectrum.

• num_points (int, optional) – Number of points for the integration grid. Default is 50.

• fi_vals (jnp.ndarray, optional) – Array of kernel values for interpolation. Default is F_VALS_FULL.

Returns:

A dict containing the bispectrum monopole (000), first multipole (200), second multipole (020), and third multipole (002).

Return type:

dict

Notes

• Keys: ‘000’ (monopole), ‘200’ (quadrupole for ka), ‘020’ (quadrupole for kb), ‘002’ (quadrupole for kc)

• Calls ext_bk_mp internally

• JIT-compiled for performance

Examples

>>> tr = jnp.array([[0.1, 0.1, 0.1]])
>>> tr2 = tr3 = tr4 = tr
>>> kp = jnp.linspace(0.01, 0.3, 100)
>>> pk = linear_power_spectrum(kp, cosm_par)
>>> cosm_par = jnp.array([...])
>>> redshift = 0.5
>>> result = bk_multip(tr, tr2, tr3, tr4, kp, pk, cosm_par, redshift)
>>> bk0 = result['000']
>>> bk200 = result['200']

geofptax.kernels.bk_sugiyama_multip(k1, k2, kp, pk, cosm_par, redshift, num_points=10, fi_vals=Array([[1.0374585, 1.0179875, 1.0033404], [0.0019502, -0.00410578, -0.00397587], [0.0055635, 0.01488085, 0.02396768], [-0.04836375, -0.05466602, -0.05680982], [0.00806689, 0.01114415, 0.01325444]], dtype=float32))

Compute Sugiyama multipole coefficients for given (k1, k2) pairs.

High-level wrapper that returns Sugiyama coefficients as a dictionary.

Parameters:

• k1 (jnp.ndarray) – Array of first wavevector magnitudes.

• k2 (jnp.ndarray) – Array of second wavevector magnitudes.

• kp (jnp.ndarray) – Array of wavevector magnitudes for the power spectrum.

• pk (jnp.ndarray) – Array of power spectrum values corresponding to kp.

• cosm_par (jnp.ndarray) – Cosmological parameters array.

• redshift (float) – Redshift at which to compute the bispectrum.

• num_points (int, optional) – Number of Gauss-Legendre points per angular dimension. Default is 50.

• fi_vals (jnp.ndarray, optional) – Array of kernel values for interpolation. Default is F_VALS_FULL.

Returns:

Dictionary with keys: ‘000’, ‘110’, ‘220’, ‘202’, ‘022’, ‘112’ containing the corresponding Sugiyama coefficients.

Return type:

dict

Examples

>>> k1 = jnp.array([0.1, 0.2, 0.3])
>>> k2 = jnp.array([0.1, 0.2, 0.3])
>>> kp = jnp.linspace(0.01, 0.5, 100)
>>> pk = linear_power_spectrum(kp, cosm_par)
>>> result = bk_sugiyama_multip(k1, k2, kp, pk, cosm_par, redshift=0.5)
>>> b000 = result['000']
>>> b110 = result['110']

geofptax.kernels.bkeff_r_scalar(mua_m, phi, tr, cosm_par, pk_in, sig_fog, log_km, log_pkm, af, mp)

Compute the integrand for the effective bispectrum in redshift space.

This is the core integrand function for computing redshift-space bispectrum multipoles, including all physical effects: RSD, AP, FoG, and GEO-FPT corrections.

Parameters:

• mua_m (float or jnp.ndarray) – Cosine of the angle between ka and the line of sight in real space.

• phi (float or jnp.ndarray) – Azimuthal angle (angle between planes defined by ka and kb).

• tr (tuple or jnp.ndarray) – Triangle side lengths (ka_m, kb_m, kc_m) in real space.

• cosm_par (jnp.ndarray) – Cosmological parameters array containing: - cosm_par[2]: α_∥ (parallel AP parameter) - cosm_par[3]: α_⟂ (perpendicular AP parameter) - cosm_par[1]: f (growth rate) - cosm_par[4]: b₁ (linear bias) - cosm_par[5]: b₂ (quadratic bias) - cosm_par[9]: σ_FoG (Finger-of-God damping scale)

• pk_in (tuple or jnp.ndarray) – Power spectrum values at ka_m, kb_m, kc_m.

• sig_fog (float) – Finger-of-God damping factor.

• log_km (jnp.ndarray) – Logarithm of wavevector magnitudes for interpolation.

• log_pkm (jnp.ndarray) – Logarithm of power spectrum values for interpolation.

• af (jnp.ndarray) – Array of coefficients for the GEO-FPT factor.

• mp (int) – Multipole index (0 for monopole, 1 for quadrupole, etc.).

Returns:

Value of the integrand.

Return type:

float or jnp.ndarray

Notes

• Applies Alcock-Paczynski scaling to transform from real to redshift space

• Includes triangle validity checks

• Computes all necessary kernels (Z1, Z2, F2, G2) with GEO-FPT corrections

• Applies FoG damping

• Returns zero for invalid triangles

Examples

>>> mua_m = 0.5
>>> phi = jnp.pi/4
>>> tr = (0.1, 0.1, 0.1)  # Equilateral triangle
>>> cosm_par = jnp.array([...])  # Cosmological parameters
>>> pk_in = (1.0, 1.0, 1.0)  # Power spectrum values
>>> sig_fog = 5.0
>>> log_km = jnp.log10(jnp.linspace(0.01, 0.3, 100))
>>> log_pkm = jnp.log10(pk_interp(10**log_km))
>>> af = jnp.array([1.0, 0.1, 0.1, 0.01, 0.001])
>>> mp = 0  # Monopole
>>> bkeff_r_scalar(mua_m, phi, tr, cosm_par, pk_in, sig_fog, log_km, log_pkm, af, mp)

geofptax.kernels.bkeff_r_vmap(mua_m, phi, tr, cosm_par, pk_in, sig_fog, log_km, log_pkm, af, mp)

Vectorized version of bkeff_r_scalar. Takes similar arguments as bkeff_r_scalar but with additional array axes over which bkeff_r_scalar is mapped.

Original documentation:

Compute the integrand for the effective bispectrum in redshift space.

This is the core integrand function for computing redshift-space bispectrum multipoles, including all physical effects: RSD, AP, FoG, and GEO-FPT corrections.

mua_mfloat or jnp.ndarray

Cosine of the angle between ka and the line of sight in real space.

phifloat or jnp.ndarray

Azimuthal angle (angle between planes defined by ka and kb).

trtuple or jnp.ndarray

Triangle side lengths (ka_m, kb_m, kc_m) in real space.

cosm_parjnp.ndarray

Cosmological parameters array containing: - cosm_par[2]: α_∥ (parallel AP parameter) - cosm_par[3]: α_⟂ (perpendicular AP parameter) - cosm_par[1]: f (growth rate) - cosm_par[4]: b₁ (linear bias) - cosm_par[5]: b₂ (quadratic bias) - cosm_par[9]: σ_FoG (Finger-of-God damping scale)

pk_intuple or jnp.ndarray

Power spectrum values at ka_m, kb_m, kc_m.

sig_fogfloat

Finger-of-God damping factor.

log_kmjnp.ndarray

Logarithm of wavevector magnitudes for interpolation.

log_pkmjnp.ndarray

Logarithm of power spectrum values for interpolation.

afjnp.ndarray

Array of coefficients for the GEO-FPT factor.

mpint

Multipole index (0 for monopole, 1 for quadrupole, etc.).

float or jnp.ndarray

Value of the integrand.

• Applies Alcock-Paczynski scaling to transform from real to redshift space

• Includes triangle validity checks

• Computes all necessary kernels (Z1, Z2, F2, G2) with GEO-FPT corrections

• Applies FoG damping

• Returns zero for invalid triangles

>>> mua_m = 0.5
>>> phi = jnp.pi/4
>>> tr = (0.1, 0.1, 0.1)  # Equilateral triangle
>>> cosm_par = jnp.array([...])  # Cosmological parameters
>>> pk_in = (1.0, 1.0, 1.0)  # Power spectrum values
>>> sig_fog = 5.0
>>> log_km = jnp.log10(jnp.linspace(0.01, 0.3, 100))
>>> log_pkm = jnp.log10(pk_interp(10**log_km))
>>> af = jnp.array([1.0, 0.1, 0.1, 0.01, 0.001])
>>> mp = 0  # Monopole
>>> bkeff_r_scalar(mua_m, phi, tr, cosm_par, pk_in, sig_fog, log_km, log_pkm, af, mp)

geofptax.kernels.bkeff_sugiyama(k1, k2, x12, mu1, phi, cosm_par, pk_interp, log_km, log_pkm, af, mp=None, hh=1.0)

JAX-compatible Sugiyama bispectrum integrand.

Computes the bispectrum integrand for the Sugiyama estimator, including all physical effects (RSD, AP, FoG, GEO-FPT, shot noise).

Parameters:

• k1 (float) – First two wavevector magnitudes (real space)

• k2 (float) – First two wavevector magnitudes (real space)

• x12 (float) – Cosine of angle between k1 and k2

• mu1 (float) – Cosine of angle between k1 and line of sight

• phi (float) – Azimuthal angle

• cosm_par (jnp.ndarray) – Cosmological parameters array

• pk_interp (callable) – Power spectrum interpolation function

• log_km (jnp.ndarray) – Logarithm of wavevector magnitudes for interpolation

• log_pkm (jnp.ndarray) – Logarithm of power spectrum values for interpolation

• af (jnp.ndarray) – GEO-FPT coefficients

• mp (int, optional) – Multipole index (not used in Sugiyama, for compatibility)

• hh (float, optional) – Hubble parameter normalization

Returns:

Value of the bispectrum integrand

Return type:

float

Notes

• Based on FOLPSD implementation (https://github.com/alejandroaviles/folpsD)

• Includes shot noise terms with amplitudes A_P and A_B

• All control flow replaced with JAX operations for JIT compatibility

geofptax.kernels.bkeff_sugiyama_vec(k1, k2, x12, mu1, phi, cosm_par, pk_interp, log_km, log_pkm, af, mp=None, hh=1.0)

Vectorized version of bkeff_sugiyama. Takes similar arguments as bkeff_sugiyama but with additional array axes over which bkeff_sugiyama is mapped.

Original documentation:

JAX-compatible Sugiyama bispectrum integrand.

Computes the bispectrum integrand for the Sugiyama estimator, including all physical effects (RSD, AP, FoG, GEO-FPT, shot noise).

k1, k2float

First two wavevector magnitudes (real space)

x12float

Cosine of angle between k1 and k2

mu1float

Cosine of angle between k1 and line of sight

phifloat

Azimuthal angle

cosm_parjnp.ndarray

Cosmological parameters array

pk_interpcallable

Power spectrum interpolation function

log_kmjnp.ndarray

Logarithm of wavevector magnitudes for interpolation

log_pkmjnp.ndarray

Logarithm of power spectrum values for interpolation

afjnp.ndarray

GEO-FPT coefficients

mpint, optional

Multipole index (not used in Sugiyama, for compatibility)

hhfloat, optional

Hubble parameter normalization

float

Value of the bispectrum integrand

• Based on FOLPSD implementation (https://github.com/alejandroaviles/folpsD)

• Includes shot noise terms with amplitudes A_P and A_B

• All control flow replaced with JAX operations for JIT compatibility

geofptax.kernels.compute_basis_grid(x_pts, mu_pts, phi_pts)

Precompute ALL basis functions on the angular grid for Sugiyama estimator.

The Sugiyama estimator expands the bispectrum in a basis of 6 functions. This function precomputes all basis functions on a 3D grid for efficient integration.

Parameters:

• x_pts (jnp.ndarray) – Grid points for x = cosθ₁₂ (angle between k1 and k2)

• mu_pts (jnp.ndarray) – Grid points for μ₁ (cosine of angle between k1 and LOS)

• phi_pts (jnp.ndarray) – Grid points for φ (azimuthal angle)

Returns:

Array of shape (6, N_x, N_mu, N_phi) containing the 6 basis functions evaluated at all grid points.

Return type:

jnp.ndarray

Notes

• Basis functions: 1, P₁(x), P₂(x), P₂(μ₁), P₂(μ₂), P₁(x)P₁(μ₁)

• Normalized according to Sugiyama et al. (2019)

• Efficient when all 6 coefficients are needed simultaneously

geofptax.kernels.compute_sugiyama_multipoles(k1k2_pairs, log_km, log_pkm, cosm_par, redshift, fi_vals=Array([[1.0374585, 1.0179875, 1.0033404], [0.0019502, -0.00410578, -0.00397587], [0.0055635, 0.01488085, 0.02396768], [-0.04836375, -0.05466602, -0.05680982], [0.00806689, 0.01114415, 0.01325444]], dtype=float32), num_points=50)

Compute Sugiyama coefficients using scalar functions and proper vmap.

Main function for computing the 6 Sugiyama multipole coefficients for a set of (k1, k2) pairs.

Parameters:

• k1k2_pairs (jnp.ndarray) – Array of shape (N_pairs, 2) containing (k1, k2) values for each triangle.

• log_km (jnp.ndarray) – Logarithm of wavevector magnitudes for power spectrum interpolation.

• log_pkm (jnp.ndarray) – Logarithm of power spectrum values for interpolation.

• cosm_par (jnp.ndarray) – Cosmological parameters array.

• redshift (float) – Redshift at which to compute the bispectrum.

• fi_vals (jnp.ndarray, optional) – Array of kernel values for interpolation. Default is F_VALS_FULL.

• num_points (int, optional) – Number of Gauss-Legendre points per angular dimension. Default is 50.

Returns:

Array of shape (6, N_pairs) containing the 6 Sugiyama coefficients for each triangle pair.

Return type:

jnp.ndarray

Notes

• Implements the Sugiyama estimator from arXiv:1903.09172

• Uses 3D Gauss-Legendre quadrature over (x, μ, φ)

• Returns coefficients: [B000, B110, B220, B202, B022, B112]

• H-factors from FOLPSD are applied for correct normalization

geofptax.kernels.cosab(ka, kb, kc)

Compute the cosine of the angle between vectors ka and kb using the law of cosines.

Given vectors with magnitudes ka and kb and the magnitude of their difference kc, compute the cosine of the angle between ka and kb.

Parameters:

• ka (float or jnp.ndarray) – Magnitude of the first vector.

• kb (float or jnp.ndarray) – Magnitude of the second vector.

• kc (float or jnp.ndarray) – Magnitude of the third vector (resultant of ka and kb).

Returns:

Cosine of the angle between vectors ka and kb.

Return type:

float or jnp.ndarray

Notes

• Derived from the law of cosines: kc² = ka² + kb² - 2*ka*kb*cos(θ)

• Used extensively in perturbation theory kernels

Examples

>>> ka, kb, kc = 1.0, 1.0, jnp.sqrt(2.0)
>>> cosab(ka, kb, kc)  # Should be cos(90°) = 0
Array(0., dtype=float32)

geofptax.kernels.ext_bk_mp(tr, tr2, tr3, tr4, log_km, log_pkm, cosm_par, redshift, fi_vals=Array([[1.0374585, 1.0179875, 1.0033404], [0.0019502, -0.00410578, -0.00397587], [0.0055635, 0.01488085, 0.02396768], [-0.04836375, -0.05466602, -0.05680982], [0.00806689, 0.01114415, 0.01325444]], dtype=float32), num_points=50)

Compute the effective bispectrum multipoles for a given set of triangles and cosmological parameters.

Computes four multipoles (monopole and three quadrupoles) for four different triangle configurations.

Parameters:

• tr (jnp.ndarray) – Array of triangle side lengths (ka, kb, kc) for the monopole calculation.

• tr2 (jnp.ndarray) – Array of triangle side lengths (ka, kb, kc) for the first multipole calculation.

• tr3 (jnp.ndarray) – Array of triangle side lengths (ka, kb, kc) for the second multipole calculation.

• tr4 (jnp.ndarray) – Array of triangle side lengths (ka, kb, kc) for the third multipole calculation.

• log_km (jnp.ndarray) – Logarithm of wavevector magnitudes for interpolation.

• log_pkm (jnp.ndarray) – Logarithm of power spectrum values for interpolation.

• cosm_par (jnp.ndarray) – Cosmological parameters array.

• redshift (float) – Redshift at which to compute the bispectrum.

• fi_vals (jnp.ndarray, optional) – Array of kernel values for interpolation. Default is F_VALS_FULL.

• num_points (int, optional) – Number of points for the integration grid. Default is 50.

Returns:

A tuple containing the bispectrum monopole (bk0), first multipole (bk200), second multipole (bk020), and third multipole (bk002).

Return type:

tuple

Notes

• Multipole indices: 0=monopole, 1=quadrupole for ka, 2=quadrupole for kb, 3=quadrupole for kc

• All triangles should correspond to the same physical configuration

• Uses vectorized integration vec_integrate_bkeff_r for efficiency

Examples

>>> tr = jnp.array([[0.1, 0.1, 0.1]])  # Equilateral triangle for monopole
>>> tr2 = jnp.array([[0.1, 0.1, 0.1]])  # Same triangle for first quadrupole
>>> tr3 = jnp.array([[0.1, 0.1, 0.1]])  # Same triangle for second quadrupole
>>> tr4 = jnp.array([[0.1, 0.1, 0.1]])  # Same triangle for third quadrupole
>>> log_km = jnp.log10(jnp.linspace(0.01, 0.3, 100))
>>> log_pkm = jnp.log10(pk_interp(10**log_km))
>>> cosm_par = jnp.array([...])
>>> redshift = 0.5
>>> bk0, bk200, bk020, bk002 = ext_bk_mp(tr, tr2, tr3, tr4, log_km, log_pkm, cosm_par, redshift)

geofptax.kernels.f2_ker(ka, kb, kc)

Compute the second-order SPT kernel F2 for matter density perturbations.

The F2 kernel describes the second-order contribution to the density field in standard perturbation theory.

Parameters:

• ka (float or jnp.ndarray) – Magnitude of the first wavevector.

• kb (float or jnp.ndarray) – Magnitude of the second wavevector.

• kc (float or jnp.ndarray) – Magnitude of the resultant wavevector.

Returns:

Value of the F2 kernel.

Return type:

float or jnp.ndarray

Notes

• F2 kernel is symmetric in its arguments: F2(k₁, k₂) = F2(k₂, k₁)

• The kernel is scale-free and depends only on wavevector magnitudes and angles

• Expression: F2 = 5/7 + 1/2*cosθ*(k₁/k₂ + k₂/k₁) + 2/7*cos²θ

Examples

>>> f2_ker(1.0, 1.0, jnp.sqrt(2.0))  # Equal magnitudes, 90° angle
Array(0.5714286, dtype=float32)

geofptax.kernels.g2_ker(ka, kb, kc)

Compute the second-order SPT kernel G2 for velocity divergence perturbations.

The G2 kernel describes the second-order contribution to the velocity divergence field in standard perturbation theory.

Parameters:

• ka (float or jnp.ndarray) – Magnitude of the first wavevector.

• kb (float or jnp.ndarray) – Magnitude of the second wavevector.

• kc (float or jnp.ndarray) – Magnitude of the resultant wavevector.

Returns:

Value of the G2 kernel.

Return type:

float or jnp.ndarray

Notes

• G2 kernel is symmetric in its arguments: G2(k₁, k₂) = G2(k₂, k₁)

• Related to the F2 kernel but with different coefficients

• Expression: G2 = 3/7 + 1/2*cosθ*(k₁/k₂ + k₂/k₁) + 4/7*cos²θ

Examples

>>> g2_ker(1.0, 1.0, jnp.sqrt(2.0))
Array(0.42857143, dtype=float32)

geofptax.kernels.g2_ker_vec(ka, kb, kc)

Vectorized version of g2_ker. Takes similar arguments as g2_ker but with additional array axes over which g2_ker is mapped.

Original documentation:

Compute the second-order SPT kernel G2 for velocity divergence perturbations.

The G2 kernel describes the second-order contribution to the velocity divergence field in standard perturbation theory.

kafloat or jnp.ndarray

Magnitude of the first wavevector.

kbfloat or jnp.ndarray

Magnitude of the second wavevector.

kcfloat or jnp.ndarray

Magnitude of the resultant wavevector.

float or jnp.ndarray

Value of the G2 kernel.

• G2 kernel is symmetric in its arguments: G2(k₁, k₂) = G2(k₂, k₁)

• Related to the F2 kernel but with different coefficients

• Expression: G2 = 3/7 + 1/2*cosθ*(k₁/k₂ + k₂/k₁) + 4/7*cos²θ

>>> g2_ker(1.0, 1.0, jnp.sqrt(2.0))
Array(0.42857143, dtype=float32)

geofptax.kernels.geo_fac(ka, kb, kc, af, hh)

Compute the GEO-FPT factor multiplying Z2_SPT to obtain Z2_GEO.

GEO-FPT (Geometrical Fitting formula for Perturbation Theory) provides a phenomenological correction to improve the accuracy of perturbation theory predictions for the bispectrum.

Parameters:

• ka (float or jnp.ndarray) – Magnitude of the first wavevector.

• kb (float or jnp.ndarray) – Magnitude of the second wavevector.

• kc (float or jnp.ndarray) – Magnitude of the resultant wavevector.

• af (jnp.ndarray) – Array of coefficients for the GEO-FPT factor. Typically has 5 elements: [f₁, f₂, f₃, f₄, f₅]

• hh (float) – Hubble parameter (normalization factor).

Returns:

Value of the GEO-FPT factor.

Return type:

float or jnp.ndarray

Notes

• The factor is computed as: f₁ + f₂*(cos_med/cos_min) + f₃*(cos_max/cos_min) + f₄*area + f₅*area²

• Uses Heron’s formula to compute triangle area

• cos_max, cos_med, cos_min are cosines ordered by magnitude

Examples

>>> ka, kb, kc = 0.1, 0.1, 0.1  # Equilateral triangle
>>> af = jnp.array([1.0, 0.1, 0.1, 0.01, 0.001])
>>> hh = 0.7
>>> geo_fac(ka, kb, kc, af, hh)

geofptax.kernels.geo_fac_damped(ka, kb, kc, af, hh, k_damp=0.12, width=0.03)

GEO-FPT factor with smooth damping beyond a specified scale.

Applies a Gaussian damping to the GEO-FPT corrections beyond k_damp, keeping only the constant term f₁ at high k.

Parameters:

• ka (float or jnp.ndarray) – Magnitude of the first wavevector.

• kb (float or jnp.ndarray) – Magnitude of the second wavevector.

• kc (float or jnp.ndarray) – Magnitude of the resultant wavevector.

• af (jnp.ndarray) – Array of coefficients for the GEO-FPT factor.

• hh (float) – Hubble parameter (normalization factor).

• k_damp (float, optional) – Damping scale (h/Mpc). Default is 0.12.

• width (float, optional) – Width of the damping transition. Default is 0.03.

Returns:

Damped GEO-FPT factor.

Return type:

float or jnp.ndarray

Notes

• Damping factor: exp(-((k_avg - k_damp)/width)²)

• Only the constant term f₁ is kept at high k

• Useful for avoiding unphysical behavior at small scales

Examples

>>> ka, kb, kc = 0.2, 0.2, 0.2  # High k triangle
>>> af = jnp.array([1.0, 0.1, 0.1, 0.01, 0.001])
>>> hh = 0.7
>>> geo_fac_damped(ka, kb, kc, af, hh, k_damp=0.12, width=0.03)

geofptax.kernels.geo_fac_sugiyama_simple(k1, k2, x12, af, hh=1.0)

Compute GEO-FPT factor for given triangle geometry (Sugiyama version).

Simplified version for the Sugiyama estimator that only needs k1, k2, and the cosine between them.

Parameters:

• k1 (float or jnp.ndarray) – First two wavevector magnitudes (real space)

• k2 (float or jnp.ndarray) – First two wavevector magnitudes (real space)

• x12 (float or jnp.ndarray) – Cosine of angle between k1 and k2

• af (jnp.ndarray) – GEO-FPT coefficients

• hh (float, optional) – Hubble parameter normalization

Returns:

eff_fact – GEO-FPT correction factor

Return type:

float or jnp.ndarray

Notes

• Computes k3 internally using law of cosines

• Calls the standard geo_fac function

geofptax.kernels.geo_fac_sugiyama_simple_vec(k1, k2, x12, af, hh=1.0)

Vectorized version of geo_fac_sugiyama_simple. Takes similar arguments as geo_fac_sugiyama_simple but with additional array axes over which geo_fac_sugiyama_simple is mapped.

Original documentation:

Compute GEO-FPT factor for given triangle geometry (Sugiyama version).

Simplified version for the Sugiyama estimator that only needs k1, k2, and the cosine between them.

k1, k2float or jnp.ndarray

First two wavevector magnitudes (real space)

x12float or jnp.ndarray

Cosine of angle between k1 and k2

afjnp.ndarray

GEO-FPT coefficients

hhfloat, optional

Hubble parameter normalization

eff_factfloat or jnp.ndarray

GEO-FPT correction factor

• Computes k3 internally using law of cosines

• Calls the standard geo_fac function

geofptax.kernels.geo_fac_vec(ka, kb, kc, af, hh)

Vectorized version of geo_fac. Takes similar arguments as geo_fac but with additional array axes over which geo_fac is mapped.

Original documentation:

Compute the GEO-FPT factor multiplying Z2_SPT to obtain Z2_GEO.

GEO-FPT (Geometrical Fitting formula for Perturbation Theory) provides a phenomenological correction to improve the accuracy of perturbation theory predictions for the bispectrum.

kafloat or jnp.ndarray

Magnitude of the first wavevector.

kbfloat or jnp.ndarray

Magnitude of the second wavevector.

kcfloat or jnp.ndarray

Magnitude of the resultant wavevector.

afjnp.ndarray

Array of coefficients for the GEO-FPT factor. Typically has 5 elements: [f₁, f₂, f₃, f₄, f₅]

hhfloat

Hubble parameter (normalization factor).

float or jnp.ndarray

Value of the GEO-FPT factor.

• The factor is computed as: f₁ + f₂*(cos_med/cos_min) + f₃*(cos_max/cos_min) + f₄*area + f₅*area²

• Uses Heron’s formula to compute triangle area

• cos_max, cos_med, cos_min are cosines ordered by magnitude

>>> ka, kb, kc = 0.1, 0.1, 0.1  # Equilateral triangle
>>> af = jnp.array([1.0, 0.1, 0.1, 0.01, 0.001])
>>> hh = 0.7
>>> geo_fac(ka, kb, kc, af, hh)

geofptax.kernels.integrate_2d_gauss(f, xmin, xmax, ymin, ymax, nx=20, ny=20)

Compute a 2D integral using vectorized Gauss-Legendre quadrature.

This function performs double integration of a function f(x,y) over the rectangular domain [xmin, xmax] × [ymin, ymax] using Gauss-Legendre quadrature rules. The implementation is fully JAX-compatible and optimized for parallel evaluation.

Parameters:

• f (callable) – A JAX-compatible function of two variables to be integrated. Signature must be f(x, y) where x and y are JAX arrays.

• xmin (float) – Lower bound of integration for the x-axis.

• xmax (float) – Upper bound of integration for the x-axis.

• ymin (float) – Lower bound of integration for the y-axis.

• ymax (float) – Upper bound of integration for the y-axis.

• nx (int, optional) – Number of Gauss-Legendre quadrature points for x-axis. Default is 20.

• ny (int, optional) – Number of Gauss-Legendre quadrature points for y-axis. Default is 20.

Returns:

The computed value of the integral ∫∫ f(x,y) dx dy over the domain.

Return type:

float

Notes

• The implementation uses JAX-native operations and is fully JIT-compilable.

• For smooth functions, nx=ny=20 typically gives machine precision.

• The quadrature weights/nodes are computed using the Golub-Welsch algorithm.

• The function handles vectorized inputs efficiently using JAX’s vmap.

Examples

>>> def f(x, y):
...     return jnp.sin(x) * jnp.cos(y)
>>> integrate_2d_gauss(f, 0, 1, 0, 1)  # Integrate sin(x)cos(y) over [0,1]×[0,1]
Array(0.354036, dtype=float32)

See also

weights_leggauss

The underlying 1D quadrature implementation

jax.vmap

Used for vectorized function evaluation

geofptax.kernels.integrate_bkeff_r(tr, cosm_par, pk_in, sig_fog, log_km, log_pkm, af, mp, xmin, xmax, num_points)

Perform 2D integration of the effective bispectrum integrand using Gauss-Legendre quadrature.

Integrates over μ (cosine of angle with line of sight) and φ (azimuthal angle) to compute bispectrum multipoles.

Parameters:

• tr (tuple or jnp.ndarray) – Triangle side lengths (ka_m, kb_m, kc_m) in real space.

• cosm_par (jnp.ndarray) – Cosmological parameters array.

• pk_in (tuple or jnp.ndarray) – Power spectrum values at ka_m, kb_m, kc_m.

• sig_fog (float) – Finger-of-God damping factor.

• log_km (jnp.ndarray) – Logarithm of wavevector magnitudes for interpolation.

• log_pkm (jnp.ndarray) – Logarithm of power spectrum values for interpolation.

• af (jnp.ndarray) – Array of coefficients for the GEO-FPT factor.

• mp (int) – Multipole index (0 for monopole, 1 for quadrupole, etc.).

• xmin (tuple) – Lower bounds of integration (mua_min, phi_min).

• xmax (tuple) – Upper bounds of integration (mua_max, phi_max).

• num_points (int, optional) – Number of Gauss-Legendre points for mua and phi axes. Default is 20.

Returns:

Result of the 2D integration.

Return type:

float

Notes

• Uses Gauss-Legendre quadrature for high accuracy

• Integration domain: μ ∈ [-1, 1], φ ∈ [0, 2π]

• Calls integrate_2d_gauss for the actual integration

Examples

>>> tr = (0.1, 0.1, 0.1)  # Equilateral triangle
>>> cosm_par = jnp.array([...])
>>> pk_in = (1.0, 1.0, 1.0)
>>> sig_fog = 5.0
>>> log_km = jnp.log10(jnp.linspace(0.01, 0.3, 100))
>>> log_pkm = jnp.log10(pk_interp(10**log_km))
>>> af = jnp.array([1.0, 0.1, 0.1, 0.01, 0.001])
>>> mp = 0
>>> xmin = (-1.0, 0.0)
>>> xmax = (1.0, 2*jnp.pi)
>>> integrate_bkeff_r(tr, cosm_par, pk_in, sig_fog, log_km, log_pkm, af, mp, xmin, xmax, 20)

geofptax.kernels.interpol_ker(a, fi_a_val, a_val=Array([0.33333334, 0.5, 0.6666667], dtype=float32))

Vectorized version of interpol_ker. Takes similar arguments as interpol_ker but with additional array axes over which interpol_ker is mapped.

Original documentation:

Interpolate the kernel values at a given scale factor.

Linearly interpolates kernel values from predefined scale factors to a target scale factor a.

afloat

Scale factor at which to interpolate.

fi_a_valjnp.ndarray

Array of kernel values at predefined scale factors.

a_valjnp.ndarray, optional

Predefined scale factors for interpolation. Default is [1/3, 1/2, 2/3] corresponding to redshifts z=2, 1, 0.5.

jnp.ndarray

Interpolated kernel values at the given scale factor.

>>> a = 0.5
>>> fi_vals = jnp.array([1.0, 1.5, 2.0])
>>> interpol_ker(a, fi_vals)
Array(1.75, dtype=float32)

geofptax.kernels.jax_leggauss(x)

Compute Gauss-Legendre quadrature nodes and weights using JAX-compatible methods.

This function computes the nodes (roots of Legendre polynomials) and weights for Gauss-Legendre quadrature using Newton’s method with Chebyshev nodes as initial guesses.

Parameters:

x (jnp.ndarray) – Array of indices or placeholder values (typically jnp.arange(n)). The length determines the number of quadrature points.

Returns:

• x_nodes (jnp.ndarray) – Quadrature nodes in the interval [-1, 1].

• w_weights (jnp.ndarray) – Quadrature weights corresponding to each node.

Notes

• Uses Newton iteration (5 iterations typically sufficient for convergence)

• Based on the Golub-Welsch algorithm for computing Gaussian quadrature rules

• Fully JAX-compatible and JIT-compilable

Examples

>>> n = 10
>>> x, w = jax_leggauss(jnp.arange(n))
>>> print(x.shape, w.shape)  # Both have shape (10,)

geofptax.kernels.legendre_recurrence(x, n)

Compute Legendre polynomial and its derivative using recurrence relations.

Parameters:

• x (jnp.ndarray) – Points at which to evaluate the Legendre polynomial.

• n (int) – Degree of the Legendre polynomial.

Returns:

• P (jnp.ndarray) – Values of the Legendre polynomial of degree n at points x.

• dP (jnp.ndarray) – Values of the derivative of the Legendre polynomial at points x.

Notes

• Uses the standard three-term recurrence relation for Legendre polynomials

• Efficient for moderate values of n (n < 1000)

• Fully JAX-compatible

Examples

>>> x = jnp.linspace(-1, 1, 100)
>>> P, dP = legendre_recurrence(x, 5)

geofptax.kernels.pt_kernel(k, q, wq)

Compute the 1-loop F3 kernel for power spectrum corrections.

This kernel is used in the 1-loop power spectrum calculation, specifically for the P13 term.

Parameters:

• k (jnp.ndarray) – External wavevector magnitudes (shape: N_k)

• q (jnp.ndarray) – Internal integration wavevector magnitudes (shape: N_q)

• wq (jnp.ndarray) – Integration weights for q (typically from trapezoidal rule)

Returns:

Kernel values of shape (N_k, N_q)

Return type:

jnp.ndarray

Notes

• Implements the angle-averaged F3(q, -q, k) kernel

• Derived from perturbation theory integrals

• Uses piecewise approximations for stability

geofptax.kernels.pt_pk_1loop(k, q, wq, pk_q, kernel13_d)

Compute 1-loop matter power spectrum using standard perturbation theory.

Computes P_δδ(k) at 1-loop order: P11 + P22 + P13.

Parameters:

• k (jnp.ndarray) – External wavevector magnitudes at which to compute P(k).

• q (jnp.ndarray) – Internal integration wavevector magnitudes.

• wq (jnp.ndarray) – Integration weights for q.

• pk_q (jnp.ndarray) – Linear power spectrum at wavevectors q.

• kernel13_d (jnp.ndarray) – Precomputed P13 kernel of shape (len(k), len(q)).

Returns:

1-loop matter power spectrum P_δδ(k).

Return type:

jnp.ndarray

Notes

• P22 term computed by integrating over angle μ

• P13 term uses precomputed kernel

• Could be accelerated with FFTlog techniques (see arXiv:1603.04405)

geofptax.kernels.vec_integrate_bkeff_r(tr, cosm_par, pk_in, sig_fog, log_km, log_pkm, af, mp, xmin, xmax, num_points)

Vectorized version of integrate_bkeff_r. Takes similar arguments as integrate_bkeff_r but with additional array axes over which integrate_bkeff_r is mapped.

Original documentation:

Perform 2D integration of the effective bispectrum integrand using Gauss-Legendre quadrature.

Integrates over μ (cosine of angle with line of sight) and φ (azimuthal angle) to compute bispectrum multipoles.

trtuple or jnp.ndarray

Triangle side lengths (ka_m, kb_m, kc_m) in real space.

cosm_parjnp.ndarray

Cosmological parameters array.

pk_intuple or jnp.ndarray

Power spectrum values at ka_m, kb_m, kc_m.

sig_fogfloat

Finger-of-God damping factor.

log_kmjnp.ndarray

Logarithm of wavevector magnitudes for interpolation.

log_pkmjnp.ndarray

Logarithm of power spectrum values for interpolation.

afjnp.ndarray

Array of coefficients for the GEO-FPT factor.

mpint

Multipole index (0 for monopole, 1 for quadrupole, etc.).

xmintuple

Lower bounds of integration (mua_min, phi_min).

xmaxtuple

Upper bounds of integration (mua_max, phi_max).

num_pointsint, optional

Number of Gauss-Legendre points for mua and phi axes. Default is 20.

float

Result of the 2D integration.

• Uses Gauss-Legendre quadrature for high accuracy

• Integration domain: μ ∈ [-1, 1], φ ∈ [0, 2π]

• Calls integrate_2d_gauss for the actual integration

>>> tr = (0.1, 0.1, 0.1)  # Equilateral triangle
>>> cosm_par = jnp.array([...])
>>> pk_in = (1.0, 1.0, 1.0)
>>> sig_fog = 5.0
>>> log_km = jnp.log10(jnp.linspace(0.01, 0.3, 100))
>>> log_pkm = jnp.log10(pk_interp(10**log_km))
>>> af = jnp.array([1.0, 0.1, 0.1, 0.01, 0.001])
>>> mp = 0
>>> xmin = (-1.0, 0.0)
>>> xmax = (1.0, 2*jnp.pi)
>>> integrate_bkeff_r(tr, cosm_par, pk_in, sig_fog, log_km, log_pkm, af, mp, xmin, xmax, 20)

geofptax.kernels.weights_trapz(x)

Return weights for trapezoidal integration.

Adapted from desilike.utils.

Parameters:

x (jnp.ndarray) – Abscissa points (must be sorted).

Returns:

Integration weights for trapezoidal rule.

Return type:

jnp.ndarray

Notes

• Handles edge cases (empty arrays, single points, two points)

• For N points: weights = (Δx_{i-1} + Δx_i) / 2

geofptax.kernels.z1_ker(mu, cosm_par)

Compute the first-order redshift-space SPT kernel Z1.

The Z1 kernel describes linear redshift-space distortions, incorporating both linear bias and Kaiser effects.

Parameters:

• mu (float or jnp.ndarray) – Cosine of the angle between the wavevector and the line of sight.

• cosm_par (jnp.ndarray) – Cosmological parameters array, where cosm_par[4] is the linear bias (b₁) and cosm_par[1] is the growth rate (f).

Returns:

Value of the Z1 kernel.

Return type:

float or jnp.ndarray

Notes

• Z1 = b₁ + f*μ² where μ = k·ẑ/|k|

• b₁: linear galaxy bias

• f: linear growth rate f ≈ Ωₘ(z)^γ with γ ≈ 0.55

Examples

>>> mu = 0.5
>>> cosm_par = jnp.array([...])  # Contains b₁ at index 4, f at index 1
>>> z1_ker(mu, cosm_par)

geofptax.kernels.z1_ker_vec(mu, cosm_par)

Vectorized version of z1_ker. Takes similar arguments as z1_ker but with additional array axes over which z1_ker is mapped.

Original documentation:

Compute the first-order redshift-space SPT kernel Z1.

The Z1 kernel describes linear redshift-space distortions, incorporating both linear bias and Kaiser effects.

mufloat or jnp.ndarray

Cosine of the angle between the wavevector and the line of sight.

cosm_parjnp.ndarray

Cosmological parameters array, where cosm_par[4] is the linear bias (b₁) and cosm_par[1] is the growth rate (f).

float or jnp.ndarray

Value of the Z1 kernel.

• Z1 = b₁ + f*μ² where μ = k·ẑ/|k|

• b₁: linear galaxy bias

• f: linear growth rate f ≈ Ωₘ(z)^γ with γ ≈ 0.55

>>> mu = 0.5
>>> cosm_par = jnp.array([...])  # Contains b₁ at index 4, f at index 1
>>> z1_ker(mu, cosm_par)

geofptax.kernels.z2_ker(ka, kb, kc, fkern, gkern, mua, mub, cosm_par)

Compute the second-order redshift-space SPT kernel Z2.

The Z2 kernel describes second-order redshift-space distortions, incorporating bias terms, velocity divergence, and nonlinear RSD effects.

Parameters:

• ka (float or jnp.ndarray) – Magnitude of the first wavevector.

• kb (float or jnp.ndarray) – Magnitude of the second wavevector.

• kc (float or jnp.ndarray) – Magnitude of the resultant wavevector.

• fkern (float or jnp.ndarray) – F2 kernel value.

• gkern (float or jnp.ndarray) – G2 kernel value.

• mua (float or jnp.ndarray) – Cosine of the angle between ka and the line of sight.

• mub (float or jnp.ndarray) – Cosine of the angle between kb and the line of sight.

• cosm_par (jnp.ndarray) – Cosmological parameters array, where: - cosm_par[4] is the linear bias (b₁) - cosm_par[1] is the growth rate (f) - cosm_par[5] is the second-order bias (b₂)

Returns:

Value of the Z2 kernel.

Return type:

float or jnp.ndarray

Notes

• Full expression includes terms proportional to b₁, f, f², b₂, and bₛ

• bₛ is the tidal bias term: bₛ = -4/7*(b₁ - 1)

• S₂ kernel: S₂ = cos²θ - 1/3

Examples

>>> ka, kb, kc = 1.0, 1.0, jnp.sqrt(2.0)
>>> fkern = f2_ker(ka, kb, kc)
>>> gkern = g2_ker(ka, kb, kc)
>>> mua, mub = 0.5, -0.5
>>> cosm_par = jnp.array([...])  # Contains b₁, f, b₂ at appropriate indices
>>> z2_ker(ka, kb, kc, fkern, gkern, mua, mub, cosm_par)

geofptax.kernels.z2_ker_vec(ka, kb, kc, fkern, gkern, mua, mub, cosm_par)

Vectorized version of z2_ker. Takes similar arguments as z2_ker but with additional array axes over which z2_ker is mapped.

Original documentation:

Compute the second-order redshift-space SPT kernel Z2.

The Z2 kernel describes second-order redshift-space distortions, incorporating bias terms, velocity divergence, and nonlinear RSD effects.

kafloat or jnp.ndarray

Magnitude of the first wavevector.

kbfloat or jnp.ndarray

Magnitude of the second wavevector.

kcfloat or jnp.ndarray

Magnitude of the resultant wavevector.

fkernfloat or jnp.ndarray

F2 kernel value.

gkernfloat or jnp.ndarray

G2 kernel value.

muafloat or jnp.ndarray

Cosine of the angle between ka and the line of sight.

mubfloat or jnp.ndarray

Cosine of the angle between kb and the line of sight.

cosm_parjnp.ndarray

Cosmological parameters array, where: - cosm_par[4] is the linear bias (b₁) - cosm_par[1] is the growth rate (f) - cosm_par[5] is the second-order bias (b₂)

float or jnp.ndarray

Value of the Z2 kernel.

• Full expression includes terms proportional to b₁, f, f², b₂, and bₛ

• bₛ is the tidal bias term: bₛ = -4/7*(b₁ - 1)

• S₂ kernel: S₂ = cos²θ - 1/3

>>> ka, kb, kc = 1.0, 1.0, jnp.sqrt(2.0)
>>> fkern = f2_ker(ka, kb, kc)
>>> gkern = g2_ker(ka, kb, kc)
>>> mua, mub = 0.5, -0.5
>>> cosm_par = jnp.array([...])  # Contains b₁, f, b₂ at appropriate indices
>>> z2_ker(ka, kb, kc, fkern, gkern, mua, mub, cosm_par)