Feature/fkem refactor

[nikosarcevic]

This PR refactors the FKEM non-Limber implementation from the legacy monolithic function into a modular, documented, and unit-tested pipeline, and wires it into angular_cl. It addresses Refactor legacy FKEM non-Limber implementation into modular, documented code #1261.

Closes #1261 and #1260

---

What changed

1. New modular FKEM pipeline

I introduced a structured FKEM module layout:

pyccl/nonlimber_fkem/
    __init__.py
    core.py
    chi_grid.py
    power_spectra.py
    tracers.py
    single_ell.py
    transforms.py
    legacy_fkem.py

• core.py: new nonlimber_fkem driver that:

  • validates the FKEM config (ell grid, auto/manual Limber, etc.),
  • orchesrtates power-spectrum prep, tracer colelctions,chi-grid, and per-ell FKEM eval,
  • implements a robust “auto” Limber transition using n_consec_ell consecutive multipoles below limber_max_error,
  • falls back cleanly to pure Limber (ell_limber = -1, empty FKEM Cls, status = 0) if FKEM fails for a given configuration (and TELLS the user it happened).

• power_spectra.py: prepare_power_spectra:

  • enforces type consistency between p_of_k_a and pk_linear (both string or both Pk2D),
  • parses them into lin / nonlin Pk2D objects and a callable pk,
  • returns None triplets and issues a clear CCLWarning when FKEM cannot be constructed, so that the core can degrade to Limber detreministically

• tracers.py: build_tracer_collections:

  • encapsulates the C-level tracer collection , error/status handling, and keeps this out of the main FKEM loop.

• chi_grid.py: build_chi_grid:

  • derives chi_min, chi_max, and n_chi from tracer kernels (respecting fkem_chi_min / fkem_nchi when provided),
  • handles chi_min == 0 safely,
  • builds the log-spaced chi-grid and returns dlnr + effective params
  • improves chi_min stability by clamping overly small values to the tracer support and adds tests

• single_ell.py + transforms.py:

  • perform the per-ell FKEM computation and FFTLog-related pieces, returning:

    • the FKEM Cl value,
    • the corresponding Limber reference,
    • and the relative FKEM/Limber difference used by the auto-transition logic

2. Legacy FKEM moved and wrapped

• The original _nonlimber_FKEM implementation is now in
  pyccl/nonlimber_fkem/legacy_fkem.py as legacy_nonlimber_fkem.

• The file was renamed and moved out of the top-level pyccl into the dedicated nonlimber_fkem package but kept intact so we can:

  • compare against it in regression tests
  • keep a direct link to the historical implementation referenced in the FKEM paper

3. angular_cl integration and API tweaks

• pyccl/cells.py now calls the new nonlimber_fkem when
  non_limber_integration_method == "FKEM":

  ell_limber_eff, cl_non_limber, status = nonlimber_fkem(
      cosmo=cosmo,
      tracer1=tracer1,
      tracer2=tracer2,
      p_of_k_a=p_of_k_a,
      ell_values=ell_use,
      ell_limber=ell_limber,
      pk_linear=p_of_k_a_lin,
      limber_max_error=limber_max_error,
      n_chi=fkem_nchi,
      chi_min=fkem_chi_min,
      n_consec_ell=3,
  )

• The split between non-Limber and Limber parts is now derived from the length of the FKEM block:

  n_nl = cl_non_limber.size
  if n_nl > 0:
      ell_use_nonlimber = ell_use[:n_nl]
      ell_use_limber = ell_use[n_nl:]
  else:
      ell_use_limber = ell_use

  so the FKEM driver fully controls how many multipoles are treated non-Limber.

• I relaxed and clarified the FKEM sanity checks in angular_cl:

  • ell_limber <= 0 is now treated as “no FKEM, pure Limber” (needed for the reference Limber test).
  • We allow ell_limber == min(ell), meaning FKEM can be active starting from the first multipole.
  • The only invalid configuration is a positive ell_limber smaller than the smallest requested ell:

• Minor API / naming cleanup: changed fkem_Nchi to fkem_nchi to be PEP8-compliant and consistent with other parameters.

4. Tests & docstrings

• New regression tests in pyccl/tests/test_fkem_legacy_vs_new.py:

  1. NC (with and without RSD)

     • nonlimber_fkem matches legacy_nonlimber_fkem to tight tolerance for a simple NC auto-correlation.

  2. WL, known problematic configs (±IA)

     • legacy_nonlimber_fkem raises CCLError for some WL setups.
     • New FKEM core does not crash: it catches the failure, emits a clear [FKEM] CCLWarning, and returns
       ell_limber = -1, empty FKEM Cl array, status = 0 (pure Limber).

  3. WL, successful legacy runs (±IA)

     • when the legacy FKEM runs, the new core matches the legacy Cls with a stringent max relative difference criterion.
     • If the legacy FKEM fails for a configuration, the test is xfailed, as there is no reference Cl.

• High-level behaviour test in test_cells_uses_fkem_for_low_ell_number_counts:

  • Ensures that ccl.angular_cl(..., ell_limber=0) produces a pure Limber reference.
  • Confirms that ccl.angular_cl(..., ell_limber=-1 or "auto"/default) uses FKEM at low ℓ for number counts and returns Cls consistent with Limber to a loose tolerance.
  • Includes docstrings explaining the intent of the test (FKEM wiring + behaviour in angular_cl).

• I also added short docstrings to the new FKEM tests to make their intent explicit and easier to navigate.

• Fixed a few relative imports and module-level docstrings in the FKEM package to keep things consistent and importable.

---

Behavioural changes

• FKEM is now modular and testable, with clear separation of:

  • power spectrum prep,
  • tracer collection building,
  • chi-grid construction,
  • per ell FKEM computation,
  • and Limber transition control.

• FKEM never silently disappears:

  • type mismatches or failures now result in explicit [FKEM] warnings and a documented fallback to pure Limber, instead of hidden behaviour.

• The semantics of ell_limber are now:

  • ell_limber <= 0: pure Limber (no FKEM),
  • ell_limber > 0: FKEM is applied up to some ℓ, controlled by nonlimber_fkem (manual or auto),
  • ell_limber == min(ell): allowed; FKEM active starting at the first multipole.

---

Testing status

• FKEM + cells tests (on my arm machine MacOS):

  OMP_NUM_THREADS=2 pytest -vv pyccl/tests/test_fkem* pyccl/tests/test_cells.py

  • All FKEM-related tests pass, including legacy-vs-new comparisons and the high-level angular_cl wiring test.

• Full pyccl test suite:

  OMP_NUM_THREADS=2 pytest -vv pyccl

  The following tests still error, but they are unrelated to this PR and reproduce failures already present on main:

  • pyccl/tests/test_bacco_lbias_power.py
  • pyccl/tests/test_ept_power.py (missing unicode in newer Python)
  • pyccl/tests/test_hmfunc.py
  • pyccl/tests/test_lpt_power.py

These look like scipy / python version-related issues and missing imports in other modules; I plan to address them in a separate PR.

• flake8 complancy fixed in 6fc10a4

• CI is failing now only for MacOS build and this s inherited from main. See this issue I explained here.

[coveralls]

Pull Request Test Coverage Report for Build 20354464860

Details

• 396 of 458 (86.46%) changed or added relevant lines in 12 files are covered.
• 3 unchanged lines in 1 file lost coverage.
• Overall coverage decreased (-1.0%) to 96.514%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
pyccl/nonlimber_fkem/tracers.py	20	24	83.33%
pyccl/cells.py	43	49	87.76%
pyccl/nonlimber_fkem/legacy_fkem.py	13	19	68.42%
pyccl/nonlimber_fkem/chi_grid.py	72	85	84.71%
pyccl/nonlimber_fkem/core.py	80	93	86.02%
pyccl/nonlimber_fkem/transforms.py	77	97	79.38%

Files with Coverage Reduction	New Missed Lines	%
pyccl/cells.py	3	88.24%

Totals
Change from base Build 20350896932:	-1.0%
Covered Lines:	6922
Relevant Lines:	7172

---

💛 - Coveralls

[carlosggarcia]

Just a super quick comment. We cannot do changes to the API without bumping the major version; i.e. without going to CCLv4. Since we don't want that, you need to ensure that the changes respect the API.

[nikosarcevic]

Hmm but this isn't a new feature -- it's a fix

[nikosarcevic]

Because yes this is api but we use ell throughout the package so it's a fix in my opinion

[carlosggarcia]

A fix that changes the current API, so we would have to bump CCL to v4 and maybe break a few working codes. I would open an issue and leave it for CCLv4.

But also in L22 fkem_Nchi=None. I would not change any of these. I haven't looked further so there might be similar issues in other parts of the code, it was a general comment.

[nikosarcevic]

I can keep the old argument names (commit that)
But we really shouldn’t publish a paper on code that has missing docstrings or no unit tests for major components — that part isn’t an API change, it’s just basic scientific software quality. I’d like to at least add the documentation and tests now, even if the API cleanup has to wait for CCL v4.

[carlosggarcia]

That's good. The tests and documentation are more than welcome, of course.

[nikosarcevic]

ok, perfect. I will commit that and then raise an issue to fix it so we keep the ccl standard of arg naming etc

[nikosarcevic]

btw -- just to understand -- you bumped the ccl version when fkem was initially coded in?

[carlosggarcia]

I did not do the merge of that but I think it happened in v2. Since it would be a API compatible change, it would be a minor version change; i.e. from 2.X.Y to 2.X+1.0.

[nikosarcevic]

@carlosggarcia done in 4e83abf