schmidtfa / pyrasa

a package for spectral parametrization in python based on the IRASA algorithm
https://schmidtfa.github.io/pyrasa/
BSD 2-Clause "Simplified" License
7 stars 1 forks source link

Factor out irasa funs and provide a consistent interface #37

Closed thht closed 3 months ago

thht commented 3 months ago

This PR leaves the _compute_psd_welch and _compute_sgramm functions in their irasa_utils.py file but at the same time introduces a consistent and convenient interface for functions that _gen_irasa can use.

The idea is basically that _gen_irasa request a function with a specific set of parameters (which is specified in utils/types.py so mypy catches it). The signature of this function is constraint to be:

def fun(data: np.ndarray, fs: int, h: int | None, up_down: str | None, time_orig: np.ndarray | None = None
    ) -> np.ndarray

The actutal irasa functions (like irasa or irasa_sprint) the generate a lightweight local function following this exact signature and call the respective function that is imported from irasa_utils.py. So, the local function consists only of the definition and the immediatly returned call to the low-level function like this example in irasa:

    def _local_irasa_fun(
        data: np.ndarray,
        fs: int,
        *args: Any,
        **kwargs: Any,
    ) -> np.ndarray:
        return _compute_psd_welch(
            data,
            fs=fs,
            nperseg=psd_kwargs.get('nperseg'),
            win_kwargs=win_kwargs,
            dpss_settings=dpss_settings,
            noverlap=psd_kwargs.get('noverlap'),
            nfft=psd_kwargs.get('nfft'),
        )[1]

    freq, psd = _compute_psd_welch(
        data,
        fs=fs,
        nperseg=psd_kwargs.get('nperseg'),
        win_kwargs=win_kwargs,
        dpss_settings=dpss_settings,
        noverlap=psd_kwargs.get('noverlap'),
        nfft=psd_kwargs.get('nfft'),
    )

    psd, psd_aperiodic, psd_periodic = _gen_irasa(
        data=np.squeeze(data), orig_spectrum=psd, fs=fs, irasa_fun=_local_irasa_fun, hset=hset
    )

As you see, the responsibility to handle the kwargs for the low-level functions is now on the local functions and not anymore on the _gen_irasa because it does not really need to care about these....

This is going to make any efforts of porting it to cython much easier btw.

closes #35

schmidtfa commented 3 months ago

This sounds and looks great! Merging..