Skip to content

Initialize only once + tracking initializations state#624

Open
jordanflitter wants to merge 19 commits intomainfrom
init_once
Open

Initialize only once + tracking initializations state#624
jordanflitter wants to merge 19 commits intomainfrom
init_once

Conversation

@jordanflitter
Copy link
Contributor

Fix #591.

This PR deletes all the calls of init_ps, initialiseSigmaMInterpTable and init_heat from the C backend. The initializations (as well as the broadcast to C) are done in the wrapper only ONCE, regardless if we run a high level function, a single field function, or something from cfuncs.py.

I moved the nested wrapper in cfuncs.py to _param_config.py. The update wrapper (whose every line was written with blood, sweat and mostly tears) contains a dynamic _InitManager instance that is passed across different levels of functions via their kwargs, allowing the functions to inspect the state of the initializations and decide whether a given initialization must be performed or not. The decorators above the functions also give an idea for what must be initialized in order for the execution to complete.

Besides the more elegant design (according to me!), this architecture reduces the runtime of run_global_evolution by 50%! (because the sigma interpolation tables are not initialized every time a single field function is called)
Moreover, it cleans some memory leak bugs (init_ps was very often not accompanied with free_ps at the end, sometimes there were even calls to init_ps without actually using the power spectrum).

Other ideas that I had but still haven't implemented them:

  • Update performance.rst.
  • Use the data in _InitManager to raise an error (if something was not initialized, to prevent a segfault) or a warning if the computation was complete but something was left unfreed. While this is a nice feature in theory, the architecture of the updated wrapper will make it very difficult to test these scenarios (since the whole idea is that the wrapper uses _InitManager to decide what to initialize and what not.

… call broadcast_input_struct only one time in the calculation. This forced me to use a generator for run_global_evolution since we use there input_one_cell, which is the real input struct we should pass. A similar logic was applied in the z-photoncons scenario. All tests pass now
…ast_params in cfuncs.py, so that the broadcasting will happen only once when calling functions from this module multiple times (for example, if regenerate=False, in order to check if the appropriate HaloCatalog is in the cache, there is a repeated call to get_halo_catalog_buffer_size). In addition, I removed the call to lib.Free_cosmo_tables_global() at the end of high-level functions, and in return I made sure that free_cosmo_tables is always called (via the decorator @high_level_func), whether the call succeeded or there was an error
…l, renamed broadcast_params->c_wrapper, and simplified the logic in generator_wrapper (since it is always called from a high_level_func)
…hat is decorated with it to state explicitly if it is_generator
…functions, especially single field and high level functions that require initializing the power spectrum. To avoid segfaults, calls to init_ps and free_ps were removed completely by the C backend. This means that power spectrum initialization and free is done ONLY ONCE, at the wrapper level
…of the functions, especially single field and high level functions that require initializing the sigma tables. To avoid segfaults, calls to initialiseSigmaMInterpTable and freeSigmaMInterpTable were removed completely from the C backend. This means that the sigma tables initialization and free is done ONLY ONCE, at the wrapper level
… of the functions, especially single field and high level functions that require initializing the heat tables. To avoid segfaults, calls to init_heat and destruct_heat were removed completely from the C backend. This means that the heat tables initialization and free is done ONLY ONCE, at the wrapper level
…ble and init_heat_tables under a super decorator called _make_lifecycle_decorator
…lization states. All tests have passed, namely no segfaults with this architecture
…init_manager across different levels of functions, thereby increasing its visibility along the calculation. In addition, it simplifies a lot the logic in wrapper, since it doesn't need to deal with a scenario where the called function doesn't accept init_manager as keyword argument. But most importantly, it allows a separation between init_heat_tables and init_sigma_table, since these two wrappers are completely independent (one doesn't need the other, only c_wrapper for broadcasting inputs)
@jordanflitter jordanflitter added context: python wrapper Changes predominantly concerning the python wrapper context: C backend Changes occur predominantly in the C code type: performance: speed Performance improvements that reduce walltime priority: high High priority type: maint: refactoring Refactoring labels Feb 26, 2026
@codecov
Copy link

codecov bot commented Feb 26, 2026

Codecov Report

❌ Patch coverage is 95.40816% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.32%. Comparing base (396b60d) to head (78c1411).
⚠️ Report is 3 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
src/py21cmfast/drivers/_param_config.py 92.85% 6 Missing and 3 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #624      +/-   ##
==========================================
+ Coverage   88.16%   88.32%   +0.15%     
==========================================
  Files          32       32              
  Lines        4741     4829      +88     
  Branches      800      822      +22     
==========================================
+ Hits         4180     4265      +85     
  Misses        401      401              
- Partials      160      163       +3     

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

context: C backend Changes occur predominantly in the C code context: python wrapper Changes predominantly concerning the python wrapper priority: high High priority type: maint: refactoring Refactoring type: performance: speed Performance improvements that reduce walltime

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Move interpolation initializers to python

1 participant