Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Factor out irasa funs and provide a consistent interface #37

Merged
merged 4 commits into from
Jul 30, 2024

Conversation

thht
Copy link
Collaborator

@thht thht commented Jul 30, 2024

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

@thht thht linked an issue Jul 30, 2024 that may be closed by this pull request
@thht thht marked this pull request as ready for review July 30, 2024 13:49
@schmidtfa
Copy link
Owner

schmidtfa commented Jul 30, 2024

This sounds and looks great! Merging..

@schmidtfa schmidtfa merged commit 856ed43 into main Jul 30, 2024
4 checks passed
@schmidtfa schmidtfa deleted the Factor-out-_irasa_funs branch July 30, 2024 17:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

move _compute_spectrogram and _psd_wlch outside of irasa_sprint and irasa
2 participants