Lya multiple scattering

[jordanflitter]

This is a draft for the Lya multiple scattering branch.
There are now two new flags, called LYA_MULTIPLE_SCATTERING and TEST_SL_WITH_MS_FILTER, the latter becomes relevant only when the former is True, while the former is relevant only if USE_HALO_FIELD=True.
When LYA_MULTIPLE_SCATTERING=True, if TEST_SL_WITH_MS_FILTER=False, the code uses a modified window function that replaces the window function of Eq. (22) in Davies, Mesinger, Murray 2025 (it replaces only the window function that is applied on the SFRD, not the one that is applied on X-ray luminosity). It is a subtraction of two hypergeometric functions 2F3 (instead of a subtraction of two Bessel functions) where the parameters of 2F3, $\alpha$ and $\beta$, are determined by the current redshift $z$, the filter radius $R$ and the ionization fraction $x_\mathrm{HI}$. Unfortunately GSL doesn't support computing 2F3 so I had to come up with my own implementation, which works pretty well based on the tests that I made.
When LYA_MULTIPLE_SCATTERING=True, and TEST_SL_WITH_MS_FILTER=True, the code uses the same modified window function but with $\alpha\gg\beta$, as in this limit 2F3 approaches the Bessel function. This configuration should be used only as a sanity check that the new window function (multiple scattering = MS) yields the same output as the old one (straight line = SL) in the appropriate limit.
Currently, this branch is designed to work only when USE_MINI_HALOS=False. With mini-halos, the SFRD is also used for the Lyman-Werner feedback, but the applied filter should correspond to the old one (as LW photons travel in straight lines), which is not supported by the current implementation when LYA_MULTIPLE_SCATTERING=True. In the future, two new more fields in XraySourceBox (like filtered_sfr_lw and filtered_sfr_lw_mini) should be introduced in this setting (thereby increasing the amount of used memory), or alternatively we can throw an error when USE_MINI_HALOS and LYA_MULTIPLE_SCATTERING are both True.

[github-actions]

Bencher Report

Branch	mult_scatter
Testbed	ubuntu-latest

Click to view all benchmark results

Benchmark	Latency	Benchmark Result seconds (s) (Result Δ%)	Upper Boundary seconds (s) (Limit %)
tests/test_integration_features.py::test_power_spectra_lightcone[dexm]	📈 view plot 🚷 view threshold	5.08 s (+0.06%) Baseline: 5.08 s	5.59 s (90.97%)
tests/test_integration_features.py::test_power_spectra_lightcone[fftw_wisdom]	📈 view plot 🚷 view threshold	4.14 s (-69.44%) Baseline: 13.54 s	14.89 s (27.78%)
tests/test_integration_features.py::test_power_spectra_lightcone[fixed_halogrids]	📈 view plot 🚷 view threshold	4.25 s (+1.66%) Baseline: 4.18 s	4.60 s (92.42%)
tests/test_integration_features.py::test_power_spectra_lightcone[inhomo]	📈 view plot 🚷 view threshold	13.02 s (-0.11%) Baseline: 13.03 s	14.34 s (90.81%)
tests/test_integration_features.py::test_power_spectra_lightcone[inhomo_ts]	📈 view plot 🚷 view threshold	24.72 s (+1.21%) Baseline: 24.42 s	26.87 s (92.01%)
tests/test_integration_features.py::test_power_spectra_lightcone[mini]	📈 view plot 🚷 view threshold	72.67 s (+1.98%) Baseline: 71.26 s	78.38 s (92.71%)
tests/test_integration_features.py::test_power_spectra_lightcone[mini_gamma_approx]	📈 view plot 🚷 view threshold	24.52 s (+5.20%) Baseline: 23.30 s	25.64 s (95.64%)
tests/test_integration_features.py::test_power_spectra_lightcone[minimize_mem]	📈 view plot 🚷 view threshold	25.64 s (+0.83%) Baseline: 25.42 s	27.97 s (91.66%)
tests/test_integration_features.py::test_power_spectra_lightcone[no-mdz]	📈 view plot 🚷 view threshold	3.68 s (-0.41%) Baseline: 3.69 s	4.06 s (90.54%)
tests/test_integration_features.py::test_power_spectra_lightcone[photoncons-z]	📈 view plot 🚷 view threshold	36.33 s (+0.87%) Baseline: 36.02 s	39.62 s (91.70%)
tests/test_integration_features.py::test_power_spectra_lightcone[sampler]	📈 view plot 🚷 view threshold	9.16 s (+0.49%) Baseline: 9.11 s	10.02 s (91.35%)
tests/test_integration_features.py::test_power_spectra_lightcone[sampler_hires]	📈 view plot 🚷 view threshold	10.31 s (+1.67%) Baseline: 10.14 s	11.15 s (92.43%)
tests/test_integration_features.py::test_power_spectra_lightcone[sampler_ir]	📈 view plot 🚷 view threshold	23.70 s (-7.79%) Baseline: 25.70 s	28.27 s (83.82%)
tests/test_integration_features.py::test_power_spectra_lightcone[sampler_mini]	📈 view plot 🚷 view threshold	69.34 s (+1.17%) Baseline: 68.54 s	75.39 s (91.97%)
tests/test_integration_features.py::test_power_spectra_lightcone[sampler_noncubic]	📈 view plot 🚷 view threshold	9.54 s (+0.39%) Baseline: 9.50 s	10.45 s (91.26%)
tests/test_integration_features.py::test_power_spectra_lightcone[sampler_ts]	📈 view plot 🚷 view threshold	39.44 s (-7.53%) Baseline: 42.65 s	46.91 s (84.07%)
tests/test_integration_features.py::test_power_spectra_lightcone[sampler_ts_ir]	📈 view plot 🚷 view threshold	41.48 s (-6.32%) Baseline: 44.28 s	48.71 s (85.17%)
tests/test_integration_features.py::test_power_spectra_lightcone[sampler_ts_ir_onethread]	📈 view plot 🚷 view threshold	56.87 s (-7.37%) Baseline: 61.40 s	67.54 s (84.21%)
tests/test_integration_features.py::test_power_spectra_lightcone[simple]	📈 view plot 🚷 view threshold	3.97 s (-0.70%) Baseline: 4.00 s	4.40 s (90.27%)
tests/test_integration_features.py::test_power_spectra_lightcone[ts]	📈 view plot 🚷 view threshold	23.40 s (+1.28%) Baseline: 23.10 s	25.41 s (92.07%)
tests/test_integration_features.py::test_power_spectra_lightcone[ts_nomdz]	📈 view plot 🚷 view threshold	21.34 s (+1.68%) Baseline: 20.99 s	23.09 s (92.44%)

🐰 View full continuous benchmarking report in Bencher

[codecov]

Codecov Report

❌ Patch coverage is 98.59155% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 88.42%. Comparing base (396b60d) to head (c52e2f4).
⚠️ Report is 3 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/py21cmfast/wrapper/outputs.py	93.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #534      +/-   ##
==========================================
+ Coverage   88.16%   88.42%   +0.25%     
==========================================
  Files          32       32              
  Lines        4741     4778      +37     
  Branches      800      811      +11     
==========================================
+ Hits         4180     4225      +45     
+ Misses        401      394       -7     
+ Partials      160      159       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[jordanflitter]

Branch is finally ready to be reviewed. Unlike what my previous comment says, it is now possible to have LYA_MULTIPLE_SCATTERING=True together with USE_MINI_HALOS=True. This requires having two extra 4D arrays for the filtered SFRDs for LW photons (as they travel in straight-lines, unlike Lyman alpha photons), this should be fixed when we solve #552. Due to this, and in order to cover the new code lines (since LYA_MULTIPLE_SCATTERING=False by default), I added two new test configurations, multiple_scattering and multiple_scattering_mini (this increases the overall runtime of the tests).
Also, I removed TEST_SL_WITH_MS_FILTER from the branch (this flag was useful during the development stage, it is no longer required).

In the code itself, I refer to equations from here: https://arxiv.org/pdf/2601.14360.

[steven-murray]

Thanks @jordanflitter as far as I can tell, this is well written and makes sense. I had only a few comments. In addition, it would be great to have a quick tutorial on this model (I know there's the paper, but having something very brief with plots would be good).

[jordanflitter]

Thanks @steven-murray for you review and suggestions! I haven't made yet the tutorial you asked since I'm not entirely sure what you have in mind (the results that I can show there are basically already presented in the paper). We could discuss this more next time we meet.

IMPORTANT NOTE: I don't know if you noticed, but I also added in this PR some "unnecessary" outputs (like the Lya flux) since I needed them for my paper. While these outputs were very useful for my project, their presence adds of course more required memory, which is exactly the issue we try to fix in the other PRs. There are several choices we can make here:

1. Merge the PR as it is anyways.
2. Remove these extra "unnecessary" outputs before merging this PR.
3. Wait with merging this PR before we fix Optional quantities #604.
4. Make a new PR, which is similar to what we discussed in Optional quantities #604, but doesn't do anything to the C-code. Instead, it simply purges all the "unnecessary" fields for the calculation of the brightness temperature, unless the user specifically asks for them (and when run_global_evolution is on, the user gets all of these fields/quantities). We can also include in this PR the inclusion of MINIMIZE_MEMORY in ComputeIonizedBox (so we won't necessarily compute and allocate memory for kinetic_temperature and mean_free_path). Once this new PR will be merged, we could safely merge this multiples scattering PR.

Clearly, I favor option (4) :)

[steven-murray]

Thanks @jordanflitter yeah I think it would be a good idea to sort out how we're going to treat "optional" quantities before merging this, so that users don't all of a sudden get a huge increase in required memory. Let's chat tomorrow.

[jordanflitter]

Update of what was done since last review:

• The new arrays (that were useful for my project) were removed from this branch. They'll return after we fix Optional quantities #604.
• Included a preliminary models.rst in the docs. At the moment, I only gave information on the new flag, LYA_MULTIPLE_SCATTERING.
• Added my paper (and other papers that added new features to the code) to acknowledge.rst + show_references. Note that one of the papers that I added was Brad's paper about lightcone + RSD, I think it's an important paper that deserves some credit since before it came out there was no lightcone in 21cmFAST! Since show_references only accepted inputs, it couldn't know if to cite Brad's paper or not, so I added an argument called lightcone with default True.