Initialize only once + tracking initializations state

[jordanflitter]

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.

[codecov]

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.

[steven-murray]

Thanks @jordanflitter, this is a really great effort towards something that is very important!

While reading your code (impressive use of decorators, btw!) I think I realized that this can be more simple. I'm going to write psuedo-code below which is untested but that hopefully gives the idea.

First, I would make a new (private) module and define a new singleton object (i.e. an object that should only ever have one instance). Let's say the module name is _c_global_init_manager.py:

from ... import InputParameters

class _GlobalCInitManager:
    def __init__(self):
        self.current_inputs:  InputParameters | None = None

        self._ps_inited = False
        self._sigma_inited = False
        self._heat_inited = False

    def broadcast_inputs(self, inputs: InputParameters):
        if self.current_inputs is None or inputs != self.current_inputs:
            self.free_all()  # free all first, so that we don't mix and match inputs later.
            lib.Broadcast_struct_global_all(<you know the drill>)
            self.current_inputs = inputs

    def init_ps(self, inputs: InputParameters):
        if not self._ps_inited and (self.current_inputs is None or self.current_inputs != inputs):
            # Always broadcast inputs first
            self.broadcast_inputs(inputs)
            lib.init_ps()
            self._ps_inited = True

    def init_heat(self, inputs):
         <same thing>

    ...

    def free_all(self):
        if self._ps_inited:
            lib.free_ps()
            self._ps_inited = False
        if self._sigma_inited:
            lib.freeSigmaMInterpTable()
            self._sigma_inited = False
        ...
        
    def init(self, inputs: InputParameters, *to_init: list[str]):
        if not to_init:
            # be conservative and initialize everything if nothing explicitly given
            to_init = ['ps', 'sigma', 'heat']

        for thing in to_init:
            match thing:
                'ps': self.init_ps()
                'sigma': self.init_sigma()
                ... etc

c_init_manager = GlobalCInitManager()

Now, either in that module or maybe in _param_config like you have now, you just need to define a similar decorator like you have:

from ._c_global_init_manager import c_init_manager

def this_function_does_stuff_in_c(*to_init):
    def wrapper(func):
        def inner(*args, **kwargs):
            inputs = _get_inputs_from_call(*args, **kwargs)  # your function
            c_init_manager.init(inputs, *to_init)
            return func(*args, **kwargs)
        return inner            
    return wrapper

Then, any function that might use the PS or sigma or whatever just needs to have the decorator applied:

from _param_config import this_function_does_stuff_in_c


@this_function_does_stuff_in_c('ps', 'sigma')
def compute_initial_conditions(inputs,...):
    ...

The downside to all of this (which I think is the same downside for you) is that we have to be careful to find and decorate EVERY function that uses the PS/heat/sigma stuff even implicitly, both now and into the future.

It would be great if we could simply mark the C-side functions somehow with a special mark such that we can tell from Python whether they will use these, then we could make automatic judgments. Oh wait, that's called a parameter to the function :-P

[jordanflitter]

Thanks @steven-murray for your feedback and your suggestion!

I went over your suggested pseudo-code. It currently lacks some important logics, like a special case for a generator function (that doesn't return an output), and free memory if exception has occurred, but I'm pretty sure they can be included as well. I am struggling to understand though in what sense this design is simpler compared to mine, at the end I see that in your design we have one decorator that could contain multiple arguments above the user-called functions, whereas in my design there could be multiple decorators of a single argument...
I don't consider btw the appearance of decorators as a downgrade, on the contrary, these decorators give the developer a clear understanding of what must be initialized/freed in order to make the function work.

Ultimately, I don't have a clear preference for your design or mine (again, unless I'm missing something, the main difference boils down to the number of decorators and their arguments). But I will say that this PR was supposed to be a quick patch for solving the bottleneck in run_global_evolution. I think however that, in the end, neither your logic or mine will be used in v5 when we get the power spectrum and sigma from hmf (see for example #657). I imagine we would construct 1D or 2D interpolation tables in CosmoTables and then linearly interpolate those with EvaluateRGTable1D and EvaluateRGTable2D that we already have in the C code, so no initialization/free of power spectrum/sigma would be required! (only broadcasting the parameters to the C code will remain relevant, as well as freeing the memory in CosmoTables, unless we adopt your innovative idea of sending parameters to the functions :)

We should discuss all of this in our next meeting :)

[steven-murray]

Thanks for the your comments @jordanflitter. Yes, it might be easiest to discuss this next week directly. You may be right that there is no real discernible difference between our methods, but I certainly felt mine was much simpler as I was writing it haha.