[ENH] add sw_nest, sw_nest_perm_ols, and sw_spice functions to stats …

…module and update documentation

docs/api.rst

@@ -273,6 +273,9 @@ Permutation tests
    permtest_1samp
    permtest_rel
    permtest_pearsonr
+   sw_nest
+   sw_nest_perm_ols
+   sw_spice
 
 Regressions
 

netneurotools/stats/__init__.py

@@ -11,7 +11,10 @@
 from .permutation_test import (
     permtest_1samp,
     permtest_rel,
-    permtest_pearsonr
+    permtest_pearsonr,
+    sw_nest,
+    sw_nest_perm_ols,
+    sw_spice
 )
 
 
@@ -30,6 +33,7 @@
     'efficient_pearsonr', 'weighted_pearsonr', 'make_correlated_xy',
     # permutation_test
     'permtest_1samp', 'permtest_rel', 'permtest_pearsonr',
+    'sw_nest', 'sw_nest_perm_ols', 'sw_spice',
     # regression
     '_add_constant', 'residualize', 'get_dominance_stats',
     # stats_utils

netneurotools/stats/permutation_test.py

@@ -1,13 +1,21 @@
 """Functions for calculating permutation test."""
 
 import numpy as np
+import scipy.stats as sstats
 from sklearn.utils.validation import check_random_state
 
 try:  # scipy >= 1.8.0
     from scipy.stats._stats_py import _chk2_asarray
 except ImportError:  # scipy < 1.8.0
     from scipy.stats.stats import _chk2_asarray
 
+try:
+    import statsmodels.api as sma
+except ImportError:
+    _has_statsmodels = False
+else:
+    _has_statsmodels = True
+
 from .correlation import efficient_pearsonr
 
 
@@ -281,3 +289,275 @@ def permtest_pearsonr(a, b, axis=0, n_perm=1000, resamples=None, seed=0):
     pvals = permutations / (n_perm + 1)  # + 1 in denom accounts for true_corr
 
     return true_corr, pvals
+
+
+def _sw_nest_enrich_score(L, network_ind, one_sided=True):
+    """
+    Calculate the network enrichment score for a given statistic.
+
+    Parameters
+    ----------
+    L : array_like, shape (n_vertices,)
+        Statistics
+    network_ind : array_like, shape (n_vertices,)
+        Network indicator, where 1 indicates membership in the network of
+        interest and 0 otherwise.
+    one_sided : bool, optional
+        Whether to perform a one-sided test. Default: True
+
+    Returns
+    -------
+    enrich_score : float
+        Network enrichment score
+    """
+    L_order = np.argsort(L)[::-1]
+    L_sorted = L[L_order]
+    network_ind_sorted = network_ind[L_order]
+
+    if one_sided:
+        L_sorted_abs = np.abs(L_sorted)
+        P_hit_numerator = np.cumsum(L_sorted_abs * network_ind_sorted)
+    else:
+        P_hit_numerator = np.cumsum(network_ind_sorted)
+    P_hit = P_hit_numerator / P_hit_numerator[-1]
+
+    P_miss_numerator = np.cumsum(1 - network_ind_sorted)
+    P_miss = P_miss_numerator / P_miss_numerator[-1]
+
+    running_sum = P_hit - P_miss
+    enrich_score = np.max(np.abs(running_sum))
+
+    return enrich_score
+
+
+def sw_nest(stat_emp, stat_perm, network_ind, one_sided=True):
+    """
+    Network Enrichment Significance Testing (NEST) from Weinstein et al., 2024.
+
+    Check `original implementation <https://github.com/smweinst/NEST>`_ for more
+    details.
+
+    Parameters
+    ----------
+    stat_emp : array_like, shape (n_vertices,)
+        Empirical statistics
+    stat_perm : array_like, shape (n_permutations, n_vertices)
+        Permuted statistics. Each row corresponds to a permutation calculated by
+        permuting the subjects and re-estimating the statistic.
+    network_ind : array_like, shape (n_vertices,)
+        Network indicator, where 1 indicates membership in the network of
+        interest and 0 otherwise.
+    one_sided : bool, optional
+        Whether to perform a one-sided test. Default: True
+
+    Returns
+    -------
+    p : float
+        Significance level
+
+    References
+    ----------
+    .. [1] Weinstein, S. M., Vandekar, S. N., Li, B., Alexander-Bloch, A. F.,
+       Raznahan, A., Li, M., Gur, R. E., Gur, R. C., Roalf, D. R., Park, M. T.
+       M., Chakravarty, M., Baller, E. B., Linn, K. A., Satterthwaite, T. D., &
+       Shinohara, R. T. (2024). Network enrichment significance testing in
+       brain-phenotype association studies. Human Brain Mapping, 45(8), e26714.
+       https://doi.org/10.1002/hbm.26714
+
+    See Also
+    --------
+    netneurotools.stats.sw_nest_perm_ols
+    """
+    es_emp = _sw_nest_enrich_score(stat_emp, network_ind, one_sided=one_sided)
+    es_perm = np.array([
+        _sw_nest_enrich_score(s, network_ind, one_sided=one_sided) for s in stat_perm
+    ])
+    return (1 + np.sum(es_perm > es_emp)) / (1 + len(es_perm))
+
+
+def sw_nest_perm_ols(
+    observed_vars,  # (N, P)
+    predictor_vars,  # (N,) or (N, 1)
+    covariate_vars=None,  # (N,) or (N, C)
+    freedman_lane=False,
+    n_perm=1000,
+    rng=None
+):
+    """
+    Network Enrichment Significance Testing (NEST) from Weinstein et al., 2024.
+
+    This function implements the permutation test for OLS from the NEST paper.
+    Note that it does not generate the network enrichment score, but rather
+    returns the empirical and permuted statistics for use in the `sw_nest` function.
+
+    Check `original implementation <https://github.com/smweinst/NEST>`_ for more
+    details.
+
+    Parameters
+    ----------
+    observed_vars : array_like, shape (n_subjects, n_vertices)
+        Observed variables
+    predictor_vars : array_like, shape (n_subjects,)
+        Predictor variable
+    covariate_vars : array_like, shape (n_subjects, n_covars), optional
+        Covariate variables. Default: None
+    freedman_lane : bool, optional
+        Whether to use the Freedman-Lane method. Default: False
+    n_perm : int, optional
+        Number of permutations to assess. Default: 1000
+    rng : {int, np.random.Generator, np.random.RandomState}, optional
+        Random number generator. Default: None
+
+    Returns
+    -------
+    empirical : array_like, shape (n_vertices,)
+        Empirical statistics
+    permuted : array_like, shape (n_permutations, n_vertices)
+        Permuted statistics. Each row corresponds to a permutation calculated by
+        permuting the subjects and re-estimating the statistic.
+
+    References
+    ----------
+    .. [1] Weinstein, S. M., Vandekar, S. N., Li, B., Alexander-Bloch, A. F.,
+       Raznahan, A., Li, M., Gur, R. E., Gur, R. C., Roalf, D. R., Park, M. T.
+       M., Chakravarty, M., Baller, E. B., Linn, K. A., Satterthwaite, T. D., &
+       Shinohara, R. T. (2024). Network enrichment significance testing in
+       brain-phenotype association studies. Human Brain Mapping, 45(8), e26714.
+       https://doi.org/10.1002/hbm.26714
+
+    See Also
+    --------
+    netneurotools.stats.sw_nest
+    """
+    if not _has_statsmodels:
+        raise ImportError("statsmodels is required for this function")
+
+    if not isinstance(rng, np.random.Generator):
+        try:
+            rng = np.random.default_rng(rng)
+        except TypeError as e:
+            raise TypeError(
+                f"Cannnot initiate Random Generator from {rng}"
+                ) from e
+
+    if covariate_vars is None:
+        covariate_vars = np.ones_like(predictor_vars)
+    if predictor_vars.ndim == 1:
+        predictor_vars = predictor_vars.reshape((-1, 1))
+    if covariate_vars.ndim == 1:
+        covariate_vars = covariate_vars.reshape((-1, 1))
+
+    n, _ = observed_vars.shape
+    perm_indices = np.array([rng.permutation(n) for _ in range(n_perm)])  # (n_perm, n)
+
+    def _single_ols(_obs, _pred, metric=None):
+        if metric is None:
+            # assumes single predictor and appended intercept
+            return sma.OLS(_obs, _pred).fit()
+        elif metric == "coeff":
+            return sma.OLS(_obs, _pred).fit().params[0, :]
+        else:
+            raise ValueError(f"{metric} is not valid")
+
+    if not freedman_lane:
+        _pred = sma.add_constant(
+            np.c_[predictor_vars, covariate_vars],
+            prepend=False, has_constant="skip"
+        )
+        # observed_vars ~ predictor_vars + covariate_vars + 1
+        empirical = _single_ols(observed_vars, _pred, metric="coeff")
+
+        permuted = []
+        for i_perm in range(n_perm):
+            _pred = sma.add_constant(
+                np.c_[predictor_vars[perm_indices[i_perm], :], covariate_vars],
+                prepend=False, has_constant="skip"
+            )
+            permuted.append(
+                _single_ols(observed_vars, _pred, metric="coeff")
+            )
+        permuted = np.array(permuted)
+    else:
+        #
+        _pred = sma.add_constant(
+            np.c_[predictor_vars, covariate_vars],
+            prepend=False, has_constant="skip"
+        )
+        _pred_cov = sma.add_constant(
+            covariate_vars,
+            prepend=False, has_constant="skip"
+        )
+        # observed_vars ~ predictor_vars + covariate_vars + 1
+        empirical = _single_ols(observed_vars, _pred, metric="coeff")
+
+        # observed_vars ~ covariate_vars + 1 -> residual
+        ols_reduced = _single_ols(observed_vars, _pred_cov)
+        resid = ols_reduced.resid
+
+        permuted = []
+        for i_perm in range(n_perm):
+            # observed_vars = fitted values from ols_reduced + permuted residuals
+            _obs_perm = ols_reduced.predict(_pred_cov) + resid[perm_indices[i_perm], :]
+            # _obs_perm ~ predictor_vars + covariate_vars + 1
+            permuted.append(
+                _single_ols(_obs_perm, _pred, metric="coeff")
+            )
+        permuted = np.array(permuted)
+
+    return empirical, permuted
+
+
+def sw_spice(X, Y, n_perm=10000, rng=None):
+    """
+    Simple Permutation-based Intermodal Correspondence (SPICE) from Weinstein et al., 2021.
+
+    Check `original implementation <https://github.com/smweinst/spice_test>`_ for more details.
+
+    Parameters
+    ----------
+    X : array_like, shape (n_subjects, n_features)
+        Data matrix for the first modality
+    Y : array_like, shape (n_subjects, n_features)
+        Data matrix for the second modality
+    n_perm : int, optional
+        Number of permutations to assess. Default: 10000
+    rng : {int, np.random.Generator, np.random.RandomState}, optional
+        Random number generator. Default: None
+
+    Returns
+    -------
+    p : float
+        Significance level
+
+    References
+    ----------
+    .. [1] Weinstein, S. M., Vandekar, S. N., Adebimpe, A., Tapera, T. M.,
+        Robert‐Fitzgerald, T., Gur, R. C., ... & Shinohara, R. T. (2021). A
+        simple permutation‐based test of intermodal correspondence. Human brain
+        mapping, 42(16), 5175-5187.
+    """  # noqa E501
+    if not isinstance(rng, np.random.Generator):
+        try:
+            rng = np.random.default_rng(rng)
+        except TypeError as e:
+            raise TypeError(
+                f"Cannnot initiate Random Generator from {rng}"
+                ) from e
+
+    if X.shape != Y.shape:
+        raise ValueError("X and Y should be the same shape!")
+    n_subj, n_feat = X.shape  # noqa: F841
+
+    stat_emp = np.mean([
+        sstats.pearsonr(X[i, :], Y[i, :])[0]
+        for i in range(n_subj)
+    ])
+
+    stat_perm = np.empty((n_perm,))
+    for p in range(n_perm):
+        Y_perm = rng.permutation(Y, axis=0)
+        stat_perm[p] = np.mean([
+            sstats.pearsonr(X[i, :], Y_perm[i, :])[0]
+            for i in range(n_subj)
+        ])
+    return (np.sum(np.abs(stat_emp) < np.abs(stat_perm)) + 1) / (n_perm + 1)