ml4t-diagnostic 0.1.0a1__py3-none-any.whl

ml4t/diagnostic/evaluation/drift/population_stability_index.py

@@ -0,0 +1,310 @@
+"""Population Stability Index (PSI) for distribution drift detection.
+
+PSI measures the distribution shift between a reference dataset (e.g., training)
+and a test dataset (e.g., production).
+
+PSI Interpretation:
+    - PSI < 0.1: No significant change (green)
+    - 0.1 ≤ PSI < 0.2: Small change, monitor (yellow)
+    - PSI ≥ 0.2: Significant change, investigate (red)
+
+References:
+    - Yurdakul, B. (2018). Statistical Properties of Population Stability Index.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+import polars as pl
+
+
+@dataclass
+class PSIResult:
+    """Result of Population Stability Index calculation.
+
+    Attributes:
+        psi: Overall PSI value (sum of bin-level PSI contributions)
+        bin_psi: PSI contribution per bin
+        bin_edges: Bin boundaries (continuous) or category labels (categorical)
+        reference_counts: Number of samples per bin in reference distribution
+        test_counts: Number of samples per bin in test distribution
+        reference_percents: Percentage of samples per bin in reference
+        test_percents: Percentage of samples per bin in test
+        n_bins: Number of bins used
+        is_categorical: Whether feature is categorical
+        alert_level: Alert level based on PSI thresholds
+            - "green": PSI < 0.1 (no significant change)
+            - "yellow": 0.1 ≤ PSI < 0.2 (small change, monitor)
+            - "red": PSI ≥ 0.2 (significant change, investigate)
+        interpretation: Human-readable interpretation
+    """
+
+    psi: float
+    bin_psi: np.ndarray
+    bin_edges: np.ndarray | list[str]
+    reference_counts: np.ndarray
+    test_counts: np.ndarray
+    reference_percents: np.ndarray
+    test_percents: np.ndarray
+    n_bins: int
+    is_categorical: bool
+    alert_level: Literal["green", "yellow", "red"]
+    interpretation: str
+
+    def summary(self) -> str:
+        """Return formatted summary of PSI results."""
+        lines = [
+            "Population Stability Index (PSI) Report",
+            "=" * 50,
+            f"PSI Value: {self.psi:.4f}",
+            f"Alert Level: {self.alert_level.upper()}",
+            f"Feature Type: {'Categorical' if self.is_categorical else 'Continuous'}",
+            f"Number of Bins: {self.n_bins}",
+            "",
+            f"Interpretation: {self.interpretation}",
+            "",
+            "Bin-Level Analysis:",
+            "-" * 50,
+        ]
+
+        # Add bin-level details
+        for i in range(self.n_bins):
+            if self.is_categorical:
+                bin_label = self.bin_edges[i]
+            else:
+                if i == 0:
+                    bin_label = f"(-inf, {self.bin_edges[i + 1]:.3f}]"
+                elif i == self.n_bins - 1:
+                    bin_label = f"({self.bin_edges[i]:.3f}, +inf)"
+                else:
+                    bin_label = f"({self.bin_edges[i]:.3f}, {self.bin_edges[i + 1]:.3f}]"
+
+            lines.append(
+                f"Bin {i + 1:2d} {bin_label:20s}: "
+                f"Ref={self.reference_percents[i]:6.2%} "
+                f"Test={self.test_percents[i]:6.2%} "
+                f"PSI={self.bin_psi[i]:.4f}"
+            )
+
+        return "\n".join(lines)
+
+
+def compute_psi(
+    reference: np.ndarray | pl.Series,
+    test: np.ndarray | pl.Series,
+    n_bins: int = 10,
+    is_categorical: bool = False,
+    missing_category_handling: Literal["ignore", "separate", "error"] = "separate",
+    psi_threshold_yellow: float = 0.1,
+    psi_threshold_red: float = 0.2,
+) -> PSIResult:
+    """Compute Population Stability Index (PSI) between two distributions.
+
+    PSI measures the distribution shift between a reference dataset (e.g., training)
+    and a test dataset (e.g., production). It quantifies how much the distribution
+    has changed.
+
+    Formula:
+        PSI = Σ (test_% - ref_%) × ln(test_% / ref_%)
+
+    For each bin i:
+        PSI_i = (P_test[i] - P_ref[i]) × ln(P_test[i] / P_ref[i])
+
+    Args:
+        reference: Reference distribution (e.g., training data)
+        test: Test distribution (e.g., production data)
+        n_bins: Number of quantile bins for continuous features (default: 10)
+        is_categorical: Whether feature is categorical (default: False)
+        missing_category_handling: How to handle categories in test not in reference:
+            - "ignore": Skip missing categories (not recommended)
+            - "separate": Create separate bin for missing categories (default)
+            - "error": Raise error if new categories found
+        psi_threshold_yellow: Threshold for yellow alert (default: 0.1)
+        psi_threshold_red: Threshold for red alert (default: 0.2)
+
+    Returns:
+        PSIResult with overall PSI, bin-level contributions, and interpretation
+
+    Raises:
+        ValueError: If inputs are invalid or missing categories found with "error" handling
+
+    Example:
+        >>> # Continuous feature
+        >>> ref = np.random.normal(0, 1, 1000)
+        >>> test = np.random.normal(0.5, 1, 1000)  # Mean shifted
+        >>> result = compute_psi(ref, test, n_bins=10)
+        >>> print(result.summary())
+        >>>
+        >>> # Categorical feature
+        >>> ref_cat = np.array(['A', 'B', 'C'] * 100)
+        >>> test_cat = np.array(['A', 'A', 'B'] * 100)  # Distribution changed
+        >>> result = compute_psi(ref_cat, test_cat, is_categorical=True)
+        >>> print(f"PSI: {result.psi:.4f}, Alert: {result.alert_level}")
+    """
+    # Convert to numpy arrays
+    if isinstance(reference, pl.Series):
+        reference = reference.to_numpy()
+    if isinstance(test, pl.Series):
+        test = test.to_numpy()
+
+    reference = np.asarray(reference)
+    test = np.asarray(test)
+
+    # Validate inputs
+    if len(reference) == 0 or len(test) == 0:
+        raise ValueError("Reference and test arrays must not be empty")
+
+    # Variables with union types for both branches
+    bin_labels: np.ndarray | list[str]
+    bin_edges: np.ndarray | list[str]
+
+    if not is_categorical:
+        # Continuous feature: quantile binning
+        bin_edges, ref_counts, test_counts = _bin_continuous(reference, test, n_bins)
+        bin_labels = bin_edges  # Will be formatted in summary()
+    else:
+        # Categorical feature: category-based binning
+        bin_labels, ref_counts, test_counts = _bin_categorical(
+            reference, test, missing_category_handling
+        )
+        bin_edges = bin_labels
+        n_bins = len(bin_labels)
+
+    # Convert counts to percentages
+    ref_percents = ref_counts / ref_counts.sum()
+    test_percents = test_counts / test_counts.sum()
+
+    # Compute PSI per bin with numerical stability
+    # Add small epsilon to avoid log(0) and division by zero
+    epsilon = 1e-10
+    ref_percents_safe = np.maximum(ref_percents, epsilon)
+    test_percents_safe = np.maximum(test_percents, epsilon)
+
+    # PSI formula: (test% - ref%) * ln(test% / ref%)
+    bin_psi = (test_percents_safe - ref_percents_safe) * np.log(
+        test_percents_safe / ref_percents_safe
+    )
+
+    # Total PSI is sum of bin contributions
+    psi = float(np.sum(bin_psi))
+
+    # Determine alert level
+    alert_level: Literal["green", "yellow", "red"]
+    if psi < psi_threshold_yellow:
+        alert_level = "green"
+        interpretation = (
+            f"No significant distribution change detected (PSI={psi:.4f} < {psi_threshold_yellow}). "
+            "Feature distribution is stable."
+        )
+    elif psi < psi_threshold_red:
+        alert_level = "yellow"
+        interpretation = (
+            f"Small distribution change detected ({psi_threshold_yellow} ≤ PSI={psi:.4f} < {psi_threshold_red}). "
+            "Monitor feature closely but no immediate action required."
+        )
+    else:
+        alert_level = "red"
+        interpretation = (
+            f"Significant distribution change detected (PSI={psi:.4f} ≥ {psi_threshold_red}). "
+            "Investigate cause and consider model retraining."
+        )
+
+    return PSIResult(
+        psi=psi,
+        bin_psi=bin_psi,
+        bin_edges=bin_edges,
+        reference_counts=ref_counts,
+        test_counts=test_counts,
+        reference_percents=ref_percents,
+        test_percents=test_percents,
+        n_bins=n_bins,
+        is_categorical=is_categorical,
+        alert_level=alert_level,
+        interpretation=interpretation,
+    )
+
+
+def _bin_continuous(
+    reference: np.ndarray, test: np.ndarray, n_bins: int
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Bin continuous features using quantiles from reference distribution.
+
+    Uses quantile binning to ensure roughly equal-sized bins in reference distribution.
+    Test distribution is binned using same bin edges.
+
+    Args:
+        reference: Reference data (used to compute quantiles)
+        test: Test data (binned using reference quantiles)
+        n_bins: Number of bins
+
+    Returns:
+        Tuple of (bin_edges, reference_counts, test_counts)
+    """
+    # Compute quantiles from reference distribution
+    # Use (n_bins + 1) to get n_bins bins with n_bins + 1 edges
+    quantiles = np.linspace(0, 100, n_bins + 1)
+    bin_edges = np.percentile(reference, quantiles)
+
+    # Ensure edges are unique (handle constant features)
+    bin_edges = np.unique(bin_edges)
+
+    # If all values are the same, create a single bin
+    if len(bin_edges) == 1:
+        return bin_edges, np.array([len(reference)]), np.array([len(test)])
+
+    # Bin both distributions using same edges
+    # Use digitize for open-interval binning
+    ref_bins = np.digitize(reference, bin_edges[1:-1])
+    test_bins = np.digitize(test, bin_edges[1:-1])
+
+    # Count samples per bin
+    ref_counts = np.bincount(ref_bins, minlength=len(bin_edges) - 1)
+    test_counts = np.bincount(test_bins, minlength=len(bin_edges) - 1)
+
+    return bin_edges, ref_counts, test_counts
+
+
+def _bin_categorical(
+    reference: np.ndarray,
+    test: np.ndarray,
+    missing_handling: Literal["ignore", "separate", "error"],
+) -> tuple[list[str], np.ndarray, np.ndarray]:
+    """Bin categorical features by category labels.
+
+    Args:
+        reference: Reference categories
+        test: Test categories
+        missing_handling: How to handle new categories in test
+
+    Returns:
+        Tuple of (category_labels, reference_counts, test_counts)
+
+    Raises:
+        ValueError: If new categories found and missing_handling="error"
+    """
+    # Get unique categories from reference
+    ref_categories = sorted(set(reference))
+    test_categories = set(test)
+
+    # Check for new categories in test
+    new_categories = test_categories - set(ref_categories)
+
+    if new_categories:
+        if missing_handling == "error":
+            raise ValueError(
+                f"New categories found in test set: {new_categories}. "
+                "These categories were not present in reference distribution."
+            )
+        elif missing_handling == "separate":
+            # Add new categories to the end
+            ref_categories.extend(sorted(new_categories))
+        # else "ignore": new categories will be dropped
+
+    # Count occurrences per category
+    ref_counts = np.array([np.sum(reference == cat) for cat in ref_categories])
+    test_counts = np.array([np.sum(test == cat) for cat in ref_categories])
+
+    return ref_categories, ref_counts, test_counts

ml4t/diagnostic/evaluation/drift/wasserstein.py

@@ -0,0 +1,388 @@
+"""Wasserstein distance for continuous distribution drift detection.
+
+The Wasserstein distance (Earth Mover's Distance) measures the minimum cost
+to transform one probability distribution into another.
+
+Properties:
+- True metric: non-negative, symmetric, triangle inequality
+- More sensitive to small shifts than PSI
+- Natural interpretation as "transport cost"
+- No binning artifacts
+
+References:
+    - Villani, C. (2009). Optimal Transport: Old and New. Springer.
+    - Ramdas, A., et al. (2017). On Wasserstein Two-Sample Testing.
+      Entropy, 19(2), 47.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import polars as pl
+from scipy.stats import wasserstein_distance
+
+
+@dataclass
+class WassersteinResult:
+    """Result of Wasserstein distance calculation.
+
+    The Wasserstein distance (also called Earth Mover's Distance) measures the
+    minimum "cost" to transform one distribution into another. It's a true metric
+    and doesn't require binning, making it ideal for continuous features.
+
+    Attributes:
+        distance: Wasserstein distance value (W_p)
+        p: Order of Wasserstein distance (1 or 2)
+        threshold: Calibrated threshold from permutation test (if calibrated)
+        p_value: Statistical significance p-value (if calibrated)
+        drifted: Whether drift was detected (distance > threshold)
+        n_reference: Number of samples in reference distribution
+        n_test: Number of samples in test distribution
+        reference_stats: Summary statistics of reference distribution
+        test_stats: Summary statistics of test distribution
+        threshold_calibration_config: Configuration used for threshold calibration
+        interpretation: Human-readable interpretation
+        computation_time: Time taken to compute (seconds)
+    """
+
+    distance: float
+    p: int
+    threshold: float | None
+    p_value: float | None
+    drifted: bool
+    n_reference: int
+    n_test: int
+    reference_stats: dict[str, float]
+    test_stats: dict[str, float]
+    threshold_calibration_config: dict[str, Any] | None
+    interpretation: str
+    computation_time: float
+
+    def summary(self) -> str:
+        """Return formatted summary of Wasserstein distance results."""
+        lines = [
+            "Wasserstein Distance Drift Detection Report",
+            "=" * 60,
+            f"Wasserstein-{self.p} Distance: {self.distance:.6f}",
+            f"Drift Detected: {'YES' if self.drifted else 'NO'}",
+            "",
+            "Sample Sizes:",
+            f"  Reference: {self.n_reference:,}",
+            f"  Test: {self.n_test:,}",
+            "",
+        ]
+
+        if self.threshold is not None:
+            lines.extend(
+                [
+                    "Threshold Calibration:",
+                    f"  Threshold: {self.threshold:.6f}",
+                    f"  P-value: {self.p_value:.4f}" if self.p_value else "  P-value: N/A",
+                    f"  Config: {self.threshold_calibration_config}",
+                    "",
+                ]
+            )
+
+        lines.extend(
+            [
+                "Distribution Statistics:",
+                "-" * 60,
+                f"Reference: Mean={self.reference_stats['mean']:.4f}, "
+                f"Std={self.reference_stats['std']:.4f}, "
+                f"Min={self.reference_stats['min']:.4f}, "
+                f"Max={self.reference_stats['max']:.4f}",
+                f"Test:      Mean={self.test_stats['mean']:.4f}, "
+                f"Std={self.test_stats['std']:.4f}, "
+                f"Min={self.test_stats['min']:.4f}, "
+                f"Max={self.test_stats['max']:.4f}",
+                "",
+                f"Interpretation: {self.interpretation}",
+                "",
+                f"Computation Time: {self.computation_time:.3f}s",
+            ]
+        )
+
+        return "\n".join(lines)
+
+
+def compute_wasserstein_distance(
+    reference: np.ndarray | pl.Series,
+    test: np.ndarray | pl.Series,
+    p: int = 1,
+    threshold_calibration: bool = True,
+    n_permutations: int = 1000,
+    alpha: float = 0.05,
+    n_samples: int | None = None,
+    random_state: int | None = None,
+) -> WassersteinResult:
+    """Compute Wasserstein distance between reference and test distributions.
+
+    The Wasserstein distance (Earth Mover's Distance) measures the minimum cost
+    to transform one probability distribution into another. Unlike PSI, it doesn't
+    require binning and provides a true metric with desirable properties:
+    - Metric properties: non-negative, symmetric, triangle inequality
+    - More sensitive to small shifts than PSI
+    - Natural interpretation as "transport cost"
+    - No binning artifacts
+
+    The p-Wasserstein distance is defined as:
+        W_p(P, Q) = (∫|F_P^{-1}(u) - F_Q^{-1}(u)|^p du)^{1/p}
+
+    For empirical distributions with sorted samples x_1 ≤ ... ≤ x_n:
+        W_1(P, Q) = (1/n) Σ|x_i^P - x_i^Q|
+
+    Threshold calibration uses a permutation test:
+        H0: reference and test come from the same distribution
+        H1: distributions differ
+
+    Args:
+        reference: Reference distribution (e.g., training data)
+        test: Test distribution (e.g., production data)
+        p: Order of Wasserstein distance (1 or 2). Default: 1
+            - p=1: More robust, easier to interpret
+            - p=2: More sensitive to tail differences
+        threshold_calibration: Whether to calibrate threshold via permutation test
+        n_permutations: Number of permutations for threshold calibration
+        alpha: Significance level for threshold (default: 0.05)
+        n_samples: Subsample to this many samples if provided (for large datasets)
+        random_state: Random seed for reproducibility
+
+    Returns:
+        WassersteinResult with distance, threshold, p-value, and interpretation
+
+    Raises:
+        ValueError: If inputs are invalid or p not in {1, 2}
+
+    Example:
+        >>> # Detect mean shift
+        >>> ref = np.random.normal(0, 1, 1000)
+        >>> test = np.random.normal(0.5, 1, 1000)  # Mean shifted by 0.5
+        >>> result = compute_wasserstein_distance(ref, test)
+        >>> print(result.summary())
+        >>>
+        >>> # Detect variance shift
+        >>> test_var = np.random.normal(0, 2, 1000)  # Variance doubled
+        >>> result = compute_wasserstein_distance(ref, test_var)
+        >>> print(f"Distance: {result.distance:.4f}, Drifted: {result.drifted}")
+        >>>
+        >>> # Without threshold calibration (faster)
+        >>> result = compute_wasserstein_distance(
+        ...     ref, test, threshold_calibration=False
+        ... )
+    """
+    start_time = time.time()
+
+    # Convert to numpy arrays
+    if isinstance(reference, pl.Series):
+        reference = reference.to_numpy()
+    if isinstance(test, pl.Series):
+        test = test.to_numpy()
+
+    reference = np.asarray(reference, dtype=np.float64)
+    test = np.asarray(test, dtype=np.float64)
+
+    # Validate inputs
+    if len(reference) == 0 or len(test) == 0:
+        raise ValueError("Reference and test arrays must not be empty")
+
+    if p not in [1, 2]:
+        raise ValueError(f"p must be 1 or 2, got {p}")
+
+    # Set random state
+    if random_state is not None:
+        np.random.seed(random_state)
+
+    # Subsample if requested
+    if n_samples is not None and len(reference) > n_samples:
+        indices_ref = np.random.choice(len(reference), n_samples, replace=False)
+        reference = reference[indices_ref]
+    if n_samples is not None and len(test) > n_samples:
+        indices_test = np.random.choice(len(test), n_samples, replace=False)
+        test = test[indices_test]
+
+    n_reference = len(reference)
+    n_test = len(test)
+
+    # Compute distribution statistics
+    reference_stats = {
+        "mean": float(np.mean(reference)),
+        "std": float(np.std(reference)),
+        "min": float(np.min(reference)),
+        "max": float(np.max(reference)),
+        "median": float(np.median(reference)),
+        "q25": float(np.percentile(reference, 25)),
+        "q75": float(np.percentile(reference, 75)),
+    }
+
+    test_stats = {
+        "mean": float(np.mean(test)),
+        "std": float(np.std(test)),
+        "min": float(np.min(test)),
+        "max": float(np.max(test)),
+        "median": float(np.median(test)),
+        "q25": float(np.percentile(test, 25)),
+        "q75": float(np.percentile(test, 75)),
+    }
+
+    # Compute Wasserstein distance
+    if p == 1:
+        distance = float(wasserstein_distance(reference, test))
+    else:  # p == 2
+        # scipy's wasserstein_distance computes W_1
+        # For W_2, we need to compute it manually
+        distance = _wasserstein_2(reference, test)
+
+    # Threshold calibration via permutation test
+    threshold = None
+    p_value = None
+    calibration_config = None
+
+    if threshold_calibration:
+        threshold, p_value = _calibrate_wasserstein_threshold(
+            reference, test, distance, n_permutations, alpha, p
+        )
+        calibration_config = {
+            "n_permutations": n_permutations,
+            "alpha": alpha,
+            "method": "permutation",
+        }
+
+    # Determine drift status
+    if threshold is not None:
+        drifted = distance > threshold
+    else:
+        # Without calibration, use heuristic based on distribution statistics
+        # Drift if distance > 0.5 * std of reference
+        drifted = distance > 0.5 * reference_stats["std"]
+        threshold = 0.5 * reference_stats["std"]
+
+    # Generate interpretation
+    if drifted:
+        if p_value is not None:
+            interpretation = (
+                f"Distribution drift detected (W_{p}={distance:.6f} > {threshold:.6f}, "
+                f"p={p_value:.4f}). The test distribution differs significantly from "
+                f"the reference distribution."
+            )
+        else:
+            interpretation = (
+                f"Distribution drift detected (W_{p}={distance:.6f} > {threshold:.6f}). "
+                f"The test distribution differs from the reference distribution."
+            )
+    else:
+        if p_value is not None:
+            interpretation = (
+                f"No significant drift detected (W_{p}={distance:.6f} ≤ {threshold:.6f}, "
+                f"p={p_value:.4f}). Distributions are consistent."
+            )
+        else:
+            interpretation = f"No significant drift detected (W_{p}={distance:.6f} ≤ {threshold:.6f}). Distributions are consistent."
+
+    computation_time = time.time() - start_time
+
+    return WassersteinResult(
+        distance=distance,
+        p=p,
+        threshold=threshold,
+        p_value=p_value,
+        drifted=drifted,
+        n_reference=n_reference,
+        n_test=n_test,
+        reference_stats=reference_stats,
+        test_stats=test_stats,
+        threshold_calibration_config=calibration_config,
+        interpretation=interpretation,
+        computation_time=computation_time,
+    )
+
+
+def _wasserstein_2(u_values: np.ndarray, v_values: np.ndarray) -> float:
+    """Compute Wasserstein-2 distance between two 1D distributions.
+
+    W_2(P, Q) = sqrt(∫|F_P^{-1}(u) - F_Q^{-1}(u)|^2 du)
+
+    For empirical distributions, this is computed as:
+    W_2 = sqrt((1/n) Σ(x_i - y_i)^2) where x, y are sorted samples
+
+    Args:
+        u_values: First distribution samples
+        v_values: Second distribution samples
+
+    Returns:
+        Wasserstein-2 distance
+    """
+    u_sorted = np.sort(u_values)
+    v_sorted = np.sort(v_values)
+
+    # Align to same length via CDF interpolation
+    # Use linear interpolation between sorted samples
+    n = min(len(u_sorted), len(v_sorted))
+    u_quantiles = np.interp(np.linspace(0, 1, n), np.linspace(0, 1, len(u_sorted)), u_sorted)
+    v_quantiles = np.interp(np.linspace(0, 1, n), np.linspace(0, 1, len(v_sorted)), v_sorted)
+
+    # Compute L2 distance
+    return float(np.sqrt(np.mean((u_quantiles - v_quantiles) ** 2)))
+
+
+def _calibrate_wasserstein_threshold(
+    reference: np.ndarray,
+    test: np.ndarray,
+    observed_distance: float,
+    n_permutations: int,
+    alpha: float,
+    p: int,
+) -> tuple[float, float]:
+    """Calibrate Wasserstein distance threshold via permutation test.
+
+    Tests the null hypothesis that reference and test come from the same
+    distribution by computing the null distribution of Wasserstein distances
+    under random permutations.
+
+    H0: P_ref = P_test (no drift)
+    H1: P_ref ≠ P_test (drift detected)
+
+    Args:
+        reference: Reference distribution samples
+        test: Test distribution samples
+        observed_distance: Observed Wasserstein distance
+        n_permutations: Number of permutations
+        alpha: Significance level
+        p: Order of Wasserstein distance
+
+    Returns:
+        Tuple of (threshold, p_value)
+            - threshold: (1-alpha) quantile of null distribution
+            - p_value: Fraction of null distances >= observed
+    """
+    # Pool all samples
+    pooled = np.concatenate([reference, test])
+    n_ref = len(reference)
+
+    # Compute null distribution
+    null_distances = np.zeros(n_permutations)
+
+    for i in range(n_permutations):
+        # Random permutation
+        np.random.shuffle(pooled)
+
+        # Split into two groups
+        ref_perm = pooled[:n_ref]
+        test_perm = pooled[n_ref:]
+
+        # Compute distance
+        if p == 1:
+            null_distances[i] = wasserstein_distance(ref_perm, test_perm)
+        else:  # p == 2
+            null_distances[i] = _wasserstein_2(ref_perm, test_perm)
+
+    # Compute threshold as (1-alpha) quantile
+    threshold = float(np.percentile(null_distances, (1 - alpha) * 100))
+
+    # Compute p-value
+    p_value = float(np.mean(null_distances >= observed_distance))
+
+    return threshold, p_value