Add more numba functions, add pre-compilation

Author: Jacobluke-

dabest/__init__.py

@@ -4,4 +4,11 @@
 from ._effsize_objects import TwoGroupsEffectSize, PermutationTest
 from ._dabest_object import Dabest
 
-__version__ = "2024.03.30"
+
+import os
+if os.environ.get('SKIP_NUMBA_COMPILE') != '1':
+    from ._stats_tools.precompile import precompile_all, _NUMBA_COMPILED
+    if not _NUMBA_COMPILED:
+        precompile_all()
+
+__version__ = "2024.03.30"

dabest/_api.py

@@ -1,3 +1,5 @@
+"""Loading data and relevant groups"""
+
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/load.ipynb.
 
 # %% auto 0

dabest/_dabest_object.py

@@ -1,3 +1,5 @@
+"""Main class for estimating statistics and generating plots."""
+
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/dabest_object.ipynb.
 
 # %% auto 0

dabest/_delta_objects.py

@@ -1,3 +1,5 @@
+"""Auxiliary delta classes for estimating statistics and generating plots."""
+
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/delta_objects.ipynb.
 
 # %% auto 0

dabest/_effsize_objects.py

@@ -1,3 +1,5 @@
+"""The auxiliary classes involved in the computations of bootstrapped effect sizes."""
+
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/effsize_objects.ipynb.
 
 # %% auto 0

dabest/_modidx.py

@@ -47,10 +47,14 @@
                                                                                                                               'dabest/_stats_tools/confint_2group_diff.py'),
                                                          'dabest._stats_tools.confint_2group_diff.delta2_bootstrap_loop': ( 'API/confint_2group_diff.html#delta2_bootstrap_loop',
                                                                                                                             'dabest/_stats_tools/confint_2group_diff.py')},
-            'dabest._stats_tools.effsize': { 'dabest._stats_tools.effsize._compute_hedges_correction_factor': ( 'API/effsize.html#_compute_hedges_correction_factor',
+            'dabest._stats_tools.effsize': { 'dabest._stats_tools.effsize._cliffs_delta_core': ( 'API/effsize.html#_cliffs_delta_core',
+                                                                                                 'dabest/_stats_tools/effsize.py'),
+                                             'dabest._stats_tools.effsize._compute_hedges_correction_factor': ( 'API/effsize.html#_compute_hedges_correction_factor',
                                                                                                                 'dabest/_stats_tools/effsize.py'),
                                              'dabest._stats_tools.effsize._compute_standardizers': ( 'API/effsize.html#_compute_standardizers',
                                                                                                      'dabest/_stats_tools/effsize.py'),
+                                             'dabest._stats_tools.effsize._mann_whitney_u': ( 'API/effsize.html#_mann_whitney_u',
+                                                                                              'dabest/_stats_tools/effsize.py'),
                                              'dabest._stats_tools.effsize.cliffs_delta': ( 'API/effsize.html#cliffs_delta',
                                                                                            'dabest/_stats_tools/effsize.py'),
                                              'dabest._stats_tools.effsize.cohens_d': ( 'API/effsize.html#cohens_d',
@@ -65,6 +69,8 @@
                                                                                                    'dabest/_stats_tools/effsize.py'),
                                              'dabest._stats_tools.effsize.weighted_delta': ( 'API/effsize.html#weighted_delta',
                                                                                              'dabest/_stats_tools/effsize.py')},
+            'dabest._stats_tools.precompile': { 'dabest._stats_tools.precompile.precompile_all': ( 'API/precompile.html#precompile_all',
+                                                                                                   'dabest/_stats_tools/precompile.py')},
             'dabest.forest_plot': { 'dabest.forest_plot.extract_plot_data': ( 'API/forest_plot.html#extract_plot_data',
                                                                               'dabest/forest_plot.py'),
                                     'dabest.forest_plot.forest_plot': ('API/forest_plot.html#forest_plot', 'dabest/forest_plot.py'),

dabest/_stats_tools/confint_1group.py

@@ -1,3 +1,5 @@
+"""A range of functions to compute bootstraps for a single sample."""
+
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/API/confint_1group.ipynb.
 
 # %% auto 0

dabest/_stats_tools/confint_2group_diff.py

@@ -1,3 +1,5 @@
+"""A range of functions to compute bootstraps for the mean difference"""
+
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/API/confint_2group_diff.ipynb.
 
 # %% auto 0
@@ -268,7 +270,7 @@ def _compute_alpha_from_ci(ci):
     return (100.0 - ci) / 100.0
 
 
-# @njit(cache=True)
+@njit(cache=True)
 def _compute_quantile(z, bias, acceleration):
     numer = bias + z
     denom = 1 - (acceleration * numer)
@@ -319,5 +321,4 @@ def calculate_weighted_delta(group_var, differences):
     num = 0.0
     for i in range(len(weight)):
         num += weight[i] * differences[i]
-
     return num / denom

dabest/_stats_tools/effsize.py

@@ -1,3 +1,5 @@
+"""A range of functions to compute various effect sizes."""
+
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/API/effsize.ipynb.
 
 # %% ../../nbs/API/effsize.ipynb 4
@@ -270,26 +272,39 @@ def hedges_g(control:list|tuple|np.ndarray,
     return correction_factor * d
 
 # %% ../../nbs/API/effsize.ipynb 10
+@njit(cache=True)
+def _mann_whitney_u(x, y):
+    """Numba-optimized Mann-Whitney U calculation"""
+    n1, n2 = len(x), len(y)
+    combined = np.concatenate((x, y))
+    
+    # Use numpy broadcasting for comparison
+    less_than = (combined.reshape(-1, 1) > combined).sum(axis=1)
+    equal_to = (combined.reshape(-1, 1) == combined).sum(axis=1)
+    
+    # Calculate ranks directly
+    ranks = less_than + (equal_to + 1) / 2
+    
+    R1 = np.sum(ranks[:n1])
+    U1 = R1 - (n1 * (n1 + 1)) / 2
+    return U1
+
+@njit(cache=True)
+def _cliffs_delta_core(control, test):
+    """Numba-optimized Cliff's delta calculation"""
+    U = _mann_whitney_u(test, control)
+    return ((2 * U) / (len(control) * len(test))) - 1
+
 def cliffs_delta(control:list|tuple|np.ndarray, 
                  test:list|tuple|np.ndarray
                 )->float:
     """
     Computes Cliff's delta for 2 samples.
     See [here](https://en.wikipedia.org/wiki/Effect_size#Effect_size_for_ordinal_data)
     """
-
-
     c = control[~np.isnan(control)]
     t = test[~np.isnan(test)]
-
-    control_n = len(c)
-    test_n = len(t)
-
-    # Note the order of the control and test arrays.
-    U, _ = mannwhitneyu(t, c, alternative='two-sided') # Not supported by numba.
-    cliffs_delta = ((2 * U) / (control_n * test_n)) - 1
-
-    return cliffs_delta
+    return _cliffs_delta_core(c, t)
 
 
 # %% ../../nbs/API/effsize.ipynb 11

dabest/_stats_tools/precompile.py

@@ -0,0 +1,53 @@
+"""A tool to pre-compile Numba functions for speeding up DABEST bootstrapping"""
+
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/API/precompile.ipynb.
+
+# %% auto 0
+__all__ = ['precompile_all']
+
+# %% ../../nbs/API/precompile.ipynb 4
+import numpy as np
+from tqdm import tqdm
+from . import effsize
+from . import confint_2group_diff
+
+# %% ../../nbs/API/precompile.ipynb 5
+_NUMBA_COMPILED = False
+
+def precompile_all():
+    """Pre-compile all numba functions with dummy data"""
+    global _NUMBA_COMPILED
+    
+    if _NUMBA_COMPILED:
+        return
+        
+    print("Pre-compiling numba functions for DABEST...")
+    
+    # Create dummy data
+    dummy_control = np.array([1.0, 2.0, 3.0])
+    dummy_test = np.array([4.0, 5.0, 6.0])
+    
+    funcs = [
+        # effsize.py functions
+        (effsize.cohens_d, (dummy_control, dummy_test)),
+        (effsize._mann_whitney_u, (dummy_control, dummy_test)),
+        (effsize._cliffs_delta_core, (dummy_control, dummy_test)),
+        (effsize._compute_standardizers, (dummy_control, dummy_test)),
+        (effsize.weighted_delta, (np.array([1.0, 2.0]), np.array([0.1, 0.2]))),
+        
+        # confint_2group_diff.py functions
+        (confint_2group_diff.create_jackknife_indexes, (dummy_control,)),
+        (confint_2group_diff.create_repeated_indexes, (dummy_control,)),
+        (confint_2group_diff.bootstrap_indices, (True, 3, 3, 10, 12345)),
+        (confint_2group_diff.delta2_bootstrap_loop, 
+            (dummy_control, dummy_test, dummy_control, dummy_test, 10, 1.0, 12345, False)),
+        (confint_2group_diff._compute_quantile, (0.5, 0.1, 0.1)),
+        (confint_2group_diff.calculate_group_var, (1.0, 3, 1.0, 3))
+    ]
+    
+    for func, args in tqdm(funcs, desc="Compiling numba functions"):
+        func(*args)
+    
+    _NUMBA_COMPILED = True
+    
+    print("Numba compilation complete!")