Source code for ezbootstrap.core

import numba
import warnings
import numpy as np
from numba import TypingError
from typing import Callable, overload, Any


@overload
def _jitted_func_factory(
    func: Callable[[np.ndarray, np.ndarray], Any]
) -> Callable[[np.ndarray, np.ndarray], Any]:
    ...


@overload
def _jitted_func_factory(
    func: Callable[[np.ndarray], Any]
) -> Callable[[np.ndarray], Any]:
    ...


def _jitted_func_factory(func):
    return numba.njit()(func)


def _bs_draw_apply_1sample(
    x: np.ndarray, f: Callable[[np.ndarray], Any], size: int
) -> np.ndarray:
    reps = np.empty(size)
    for i in range(size):
        xs = np.random.choice(x, replace=True, size=size)
        reps[i] = f(xs)
    return reps


def _bs_draw_apply_2sample(
    x: np.ndarray, y: np.ndarray, f: Callable[[np.ndarray, np.ndarray], Any], size: int
) -> np.ndarray:
    reps = np.empty(size)
    for i in range(size):
        xs = np.random.choice(x, replace=True, size=size)
        ys = np.random.choice(y, replace=True, size=size)
        reps[i] = f(xs, ys)
    return reps


def _permutation_draw_apply_2sample(
    x: np.ndarray, y: np.ndarray, f: Callable[[np.ndarray, np.ndarray], Any], size: int
) -> np.ndarray:
    reps = np.empty(size)
    idx = len(x)
    con = np.concatenate((x, y))
    for i in range(size):
        perm = np.random.permutation(con)
        xs = perm[:idx]
        ys = perm[idx:]
        reps[i] = f(xs, ys)
    return reps


def _permute_draw_apply_independent(
    x: np.ndarray, y: np.ndarray, f: Callable[[np.ndarray, np.ndarray], Any], size: int
) -> np.ndarray:
    reps = np.empty(size)
    idx = len(x)
    con = np.concatenate((x, y))
    for i in range(size):
        xs = np.random.permutation(x)
        ys = np.random.permutation(y)
        reps[i] = f(xs, ys)
    return reps


def _bs_pairs_draw_apply(
    x: np.ndarray, y: np.ndarray, f: Callable[[np.ndarray, np.ndarray], Any], size: int
) -> np.ndarray:

    reps = np.empty(size)
    idx = np.arange(len(x))
    for i in range(size):
        idxs = np.random.choice(idx, replace=True, size=len(x))
        xs = x[idxs]
        ys = y[idxs]
        reps[i] = f(xs, ys)
    return reps


[docs]def bs_1sample( x: np.ndarray, func: Callable[[np.ndarray], Any], size: int = 5000 ) -> int: """ One sample bootsrap replicates. Args: x (arraylike): The sample from which replicates will be drawn func (callable[[np.ndarray], float]): The function that returns the statistic to be calculated on x. size (int): The number of bootstrap replicates to generate Returns: A numpy array of bootstrap replicates """ try: jfunc = _jitted_func_factory(func) reps = numba.njit()(_bs_draw_apply_1sample)(x, jfunc, size=size) except TypingError: warnings.warn("Numba compilation failed. Reverting to pure python") reps = _bs_draw_apply_1sample(x, func, size=size) return reps
[docs]def bs_2sample( x: np.ndarray, y: np.ndarray, func: Callable[[np.ndarray, np.ndarray], Any], size: int = 5000, ) -> np.ndarray: """ Two sample bootsrap replicates. Args: x (arraylike): The first sample from which replicates will be drawn y (arraylike): The second sample from which replicates will be drawn func (callable[[np.ndarray, np.ndarray], float]): The function that returns the statistic to be calculated on x and y. size (int): The number of bootstrap replicates to generate Returns: A numpy array of bootstrap replicates """ try: jfunc = _jitted_func_factory(func) reps = numba.njit()(_bs_draw_apply_2sample)(x, y, jfunc, size=size) except TypingError: warnings.warn("Numba compilation failed. Reverting to pure python") reps = _bs_draw_apply_2sample(x, y, func, size=size) return reps
[docs]def bs_pairs( x: np.ndarray, y: np.ndarray, func: Callable[[np.ndarray, np.ndarray], Any], size: int = 5000, ) -> np.ndarray: """ Generate bootstrap replicates from pairs of x and y. Args: x (arraylike): The first variable from which replicates will be drawn y (arraylike): The second variable from which replicates will be drawn func (callable[[np.ndarray, np.ndarray], float]): The function that returns the statistic to be calculated on x and y. size (int): The number of bootstrap replicates to generate Returns: A numpy array of bootstrap replicates """ try: jfunc = _jitted_func_factory(func) reps = numba.njit()(_bs_pairs_draw_apply)(x, y, jfunc, size) except TypingError: warnings.warn("Numba compilation failed. Reverting to pure python") reps = _bs_pairs_draw_apply(x, y, func, size=size) return reps
[docs]def permutation_2sample( x: np.ndarray, y: np.ndarray, func: Callable[[np.ndarray, np.ndarray], Any], size: int = 5000, ) -> np.ndarray: """ Generate permutation replicates from two samples. Args: x (arraylike): The first sample from which replicates will be drawn y (arraylike): The second sample from which replicates will be drawn func (callable[[np.ndarray, np.ndarray], float]): The function that returns the statistic to be calculated on x and y. size (int): The number of permutation replicates to generate Returns: A numpy array of bootstrap replicates """ try: jfunc = _jitted_func_factory(func) reps = numba.njit()(_permutation_draw_apply_2sample)(x, y, jfunc, size) except TypingError: warnings.warn("Numba compilation failed. Reverting to pure python") reps = _permutation_draw_apply_2sample(x, y, func, size) return reps
[docs]def permutation_2sample_independent( x: np.ndarray, y: np.ndarray, func: Callable[[np.ndarray, np.ndarray], Any], size: int = 5000, ) -> np.ndarray: """ Generate permutation replicates from two samples by permuting them independently. Args: x (arraylike): The first sample from which replicates will be drawn y (arraylike): The second sample from which replicates will be drawn func (callable[[np.ndarray, np.ndarray], float]): The function that returns the statistic to be calculated on x and y. size (int): The number of permutation replicates to generate Returns: A numpy array of bootstrap replicates """ try: jfunc = _jitted_func_factory(func) reps = numba.njit()(_permute_draw_apply_independent)(x, y, jfunc, size) except TypingError: warnings.warn("Numba compilation failed. Reverting to pure python") reps = _permute_draw_apply_independent(x, y, func, size) return reps