Large improvements of flip: mesh, dependencies removal, forward modeling, field comparison likelihood

[corentinravoux]

Mesh improvements.
The idea is to

• separate the pypower module for meshing
• improve meshing integration so that it can be called inside a data_vector instance

[corentinravoux]

Might take this opportunity to remove Hankel transform dependency

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[corentinravoux]

New features added:

• broken-alpha and gamma terms in the peculiar velocity estimator
• GW density mapping
• Module for density comparison

[Copilot]

Pull request overview

Copilot reviewed 7 out of 8 changed files in this pull request and generated 12 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[corentinravoux]

This branch goes beyond the mesh part, it include new developments, in particular the creation of a forward modeling subpackage based on flox.
This first version is linear and will be extended to 2lpt and jaxpm

[corentinravoux]

@copilot Give me a complete review of the flip/forward new implementation, without the flip/forward/flox. Please point at potential bugs or search for unphysical implementation

[Copilot]

@copilot Give me a complete review of the flip/forward new implementation, without the flip/forward/flox. Please point at potential bugs or search for unphysical implementation

Here is a complete review of flip/forward (excluding flip/forward/flox), focusing on potential bugs and unphysical implementations.

---

1. Critical: simulation/__init__.py imports non-existent modules

from . import gaussian, pm, twolpt, simulator, utils

Neither pm nor twolpt directories/files exist inside flip/forward/simulation/. This causes an ImportError the moment anyone does import flip.forward.

---

2. Wrong chi-squared normalization in likelihood.py (unphysical)

File: flip/forward/likelihood.py, line 57

# BUG — divides residual by error^2, then squares → error^4 in denominator
- 0.5 * ((radial_velocity_estimator - radial_velocity_simulator) / error**2) ** 2

Should be:

- 0.5 * ((radial_velocity_estimator - radial_velocity_simulator) / error) ** 2

The current formula computes Δv² / σ⁴ (dimensionally wrong — not a dimensionless χ²). Compare with the correctly written log_likelihood_magnitude_distance_link which uses /observed_distance_modulus_err without squaring.

---

3. get_voxels_in_direction operates on a tuple — runtime TypeError

File: flip/forward/simulation/gaussian/box.py, lines 112–124

spherical_to_cartesian(...) returns a Python tuple (x, y, z), but the very next lines do:

if physical_unit:
    xyz /= (self.box_size / self.number_bins)   # TypeError: can't divide tuple
xyz += self.get_centroid(physical_unit=False)    # TypeError even for physical_unit=False

Augmented assignment (/=, +=) raises TypeError on tuples. The array operations should be applied after stacking, e.g. xyz = jnp.stack(spherical_to_cartesian(...)) — but even then the resulting shape (3, M, ntrial) doesn't match the comment (ntrial, ntargets, 3). This method appears incomplete.

---

4. get_line_of_sight returns a tuple — breaks velocity projection

File: flip/forward/simulation/gaussian/box.py, lines 145–150

def get_line_of_sight(self, ra, dec):
    return spherical_to_cartesian(ra, dec, jnp.ones_like(ra))

Again returns (x, y, z) tuple. In field_utils.py:

return jnp.sum(v * line_of_sight, axis=-1)

v * (x, y, z) fails with a JAX array on the left — line_of_sight must be a numeric array of shape (M, 3) or (3,).

---

5. Duplicate computation of variance in log_prior_density_function

File: flip/forward/likelihood.py, lines 95–101

variance is computed identically twice; the first assignment (lines 95–97) is immediately overwritten by the second (lines 99–101). One copy should be removed.

---

6. sigma8² factor in the density prior may be inconsistent (unphysical risk)

File: flip/forward/likelihood.py, line 96

variance = sigma8**2 * power_spectrum_grid * (number_bins / box_size)**3 * number_bins**3

delta_fourier in the sampler context (paint_modes_on_delta_fourier) are the raw, unnormalised modes — sigma8 enters separately through density_from_delta_fourier and velocity_from_delta_fourier. If the prior is treating the raw modes, the expected variance of |δ_k|² is P(k)·(N/L)³·N³ without an extra sigma8². Including sigma8² means the prior's normalisation depends on a sampled parameter, which changes both the shape and the volume of the prior — potentially biasing the posterior.

---

7. log_prior_position_function can return NaN/-inf

File: flip/forward/likelihood.py, lines 126–130

jnp.log(density_values)

For a Gaussian random field, the real-space density contrast δ can be < −1, giving 1 + b·σ₈·δ ≤ 0 and log(density) = NaN or -inf. The jnp.where guard only applies to positions outside the box, not to negative densities inside it. Consider clamping or using a lognormal model (there is even a commented note for this at line 70–72 of field_model.py).

Also, (comoving_distance_targets > cutoff) + (comoving_distance_targets < 0) uses arithmetic + for logical OR; use | for clarity and correctness when values are not strictly boolean.

The out-of-range penalty -(distance - cutoff)² is not symmetric: for distance < 0 this gives -(distance - cutoff)² = -(d² - 2·cutoff·d + cutoff²), which is positive for large negative d, pointing the sampler in the wrong direction.

---

8. sample_key default argument evaluated at class-definition time

File: flip/forward/sampler.py, line 99

def run(self, sample_key=jax.random.split(jax.random.PRNGKey(0))[1], ...):

Python default argument values are evaluated once when the function is defined. Every call to run() without an explicit sample_key will use the same PRNGKey, producing identical chains. Pass None and generate the key inside the function body.

---

9. Velocity sign in linear theory (potentially unphysical)

File: flip/forward/simulation/gaussian/field_model.py, line 47

vmodes = 1j * wavenumber_ratio * ... * delta_fourier

Standard linear perturbation theory (with the FFT convention δ(k) = Σₓ δ(x)·e^{−2πi·k·x/N}) gives v(k) = −i · (k/k²) · H·f·D·a · δ(k). The code uses +1j instead of -1j. Unless there is a deliberate sign-flip elsewhere in the pipeline (e.g., in how radial_velocity_estimator is defined), all predicted radial velocities have the wrong sign. This should at minimum be documented.

---

10. Hardcoded Planck15 cosmology and redshift range at module import

File: flip/forward/likelihood.py, lines 13–17

_REDSHIFT_LOOKUP_DIST = jcosmo.background.radial_comoving_distance(
    cosmo=jcosmo.Planck15(), a=_REDSHIFT_LOOKUP_A
)

This lookup table is built at import time with a fixed Planck15 cosmology and z ∈ [0, 0.5). Any survey extending beyond z = 0.5 will silently extrapolate. Any cosmology different from Planck15 is ignored in the distance-to-redshift conversion, creating an inconsistency when h or other cosmological parameters are varied.

---

11. dist_range hardcoded in _init_likelihood

File: flip/forward/likelihood.py, line 184

dist_range=[0, 200],

200 Mpc/h is hardcoded with no way to adjust it through the public API of CandleGridGaussianLikelihood. This should be exposed as a likelihood_properties parameter.