@@ -3,6 +3,7 @@
 
 import numpy as np
 import pandas as pd
+from scipy import stats
 
 __all__ = ["sensitivity_indices"]
 
@@ -14,131 +15,38 @@
     HAS_IPYTHON = False
 
 
-def magic_binning(x: np.ndarray, n_bins_default: int) -> tuple[np.ndarray, int]:
-    """
-    Python equivalent of the MATLAB magic_binning function.
-    Uses equal-frequency binning, groups identical values, and handles NaNs.
-    """
+def _quantile_edges(x: np.ndarray, n_bins: int) -> np.ndarray:
+    """Bin edges holding approximately the same number of points.
 
-    is_nan = np.isnan(x)
-    not_nan_x = x[~is_nan]
+    Bins are defined by value, so identical values always fall in the same
+    bin. Duplicated edges are dropped: a factor with many ties gets fewer
+    than ``n_bins`` bins, and a constant factor gets a single one.
 
-    idx_sorted = np.argsort(not_nan_x)
-    x_sorted = not_nan_x[idx_sorted]
-    n_valid = len(x_sorted)
+    Discrete variables (few unique values relative to n_bins) get one bin
+    per unique value instead of quantile-based edges, to avoid collapsing
+    minority categories into a majority bin.
+    """
+    unique_vals = np.unique(x[~np.isnan(x)])
+    if len(unique_vals) <= n_bins:
+        # midpoints between consecutive unique values, so each value gets
+        # its own bin
+        if unique_vals.size == 1:
+            return np.array([unique_vals[0], np.nextafter(unique_vals[0], np.inf)])
+        midpoints = (unique_vals[:-1] + unique_vals[1:]) / 2
+        return np.concatenate([[unique_vals[0]], midpoints, [unique_vals[-1]]])
 
-    bin_idx_valid = np.zeros(n_valid, dtype=int)
+    edges = np.unique(np.nanquantile(x, np.linspace(0, 1, n_bins + 1)))
+    if edges.size < 2:
+        edges = np.array([edges[0], np.nextafter(edges[0], np.inf)])
+    return edges
 
-    unique_vals = np.unique(x_sorted)
-    if len(unique_vals) <= n_bins_default:
-        _, bin_idx_valid = np.unique(
-            x_sorted, return_inverse=True
-        )  # x_sorted, not not_nan_x
-        bin_idx_valid += 1
-    else:
-        min_bin_size = n_valid // n_bins_default
-        remaining_size = n_valid
-        current_edge_idx = min_bin_size - 1  # 0-based index
-
-        b = 1
-        start_idx = 0
-
-        while b <= n_bins_default:
-            current_bin_size = min_bin_size
-
-            # While the edge is between identical values, move one element further
-            while (current_edge_idx < n_valid - 1) and (
-                x_sorted[current_edge_idx + 1] == x_sorted[current_edge_idx]
-            ):
-                current_edge_idx += 1
-                current_bin_size += 1
-
-            # Assign bin indices
-            bin_idx_valid[start_idx : current_edge_idx + 1] = b
-            remaining_size -= current_bin_size
-
-            # Break if not enough elements left for two distinct bins
-            if remaining_size < min_bin_size * 2:
-                bin_idx_valid[current_edge_idx + 1 :] = b + 1
-                break
-
-            start_idx = current_edge_idx + 1
-            current_edge_idx += min_bin_size
-            b += 1
-
-    bin_idx_valid_orig_order = np.zeros(n_valid, dtype=int)
-    bin_idx_valid_orig_order[idx_sorted] = bin_idx_valid
-
-    # NaNs back in (NaNs get bin 0)
-    bin_idx = np.zeros(len(x), dtype=int)
-    bin_idx[~is_nan] = bin_idx_valid_orig_order
-
-    n_bins_out = np.max(bin_idx) if len(bin_idx) > 0 else 0
-
-    return bin_idx, n_bins_out
-
-
-def bin_data_1d(
-    x: np.ndarray, y: np.ndarray, n_bins_default: int
-) -> tuple[np.ndarray, np.ndarray]:
-    bin_idx, n_bins_x = magic_binning(x, n_bins_default)
-
-    bin_avg = np.full(n_bins_x, np.nan)
-    bin_count = np.full(n_bins_x, np.nan)
-
-    for b in range(1, n_bins_x + 1):
-        mask = bin_idx == b
-        if np.any(mask):
-            bin_avg[b - 1] = np.mean(y[mask])
-            bin_count[b - 1] = np.sum(mask)
-
-    return bin_avg, bin_count
-
-
-def bin_data_2d(xi: np.ndarray, xj: np.ndarray, y: np.ndarray, n_bins_default: int):
-    bin_idx_i, n_bins_i = magic_binning(xi, n_bins_default)
-    bin_idx_j, n_bins_j = magic_binning(xj, n_bins_default)
-
-    bin_avg_ij = np.full((n_bins_i, n_bins_j), np.nan)
-    bin_count_ij = np.full((n_bins_i, n_bins_j), np.nan)
-    bin_avg_i = np.full(n_bins_i, np.nan)
-    bin_count_i = np.full(n_bins_i, np.nan)
-    bin_avg_j = np.full(n_bins_j, np.nan)
-    bin_count_j = np.full(n_bins_j, np.nan)
-
-    for n in range(1, n_bins_i + 1):
-        mask_i = bin_idx_i == n
-        if np.any(mask_i):
-            bin_avg_i[n - 1] = np.mean(y[mask_i])
-            bin_count_i[n - 1] = np.sum(mask_i)
-
-    for m in range(1, n_bins_j + 1):
-        mask_j = bin_idx_j == m
-        if np.any(mask_j):
-            bin_avg_j[m - 1] = np.mean(y[mask_j])
-            bin_count_j[m - 1] = np.sum(mask_j)
-
-    for n in range(1, n_bins_i + 1):
-        for m in range(1, n_bins_j + 1):
-            mask_ij = (bin_idx_i == n) & (bin_idx_j == m)
-            if np.any(mask_ij):
-                bin_avg_ij[n - 1, m - 1] = np.mean(y[mask_ij])
-                bin_count_ij[n - 1, m - 1] = np.sum(mask_ij)
-
-    # Flatten the 2D matrices and exclude NaNs (empty bins)
-    bin_avg_ij_flat = bin_avg_ij.flatten()
-    bin_count_ij_flat = bin_count_ij.flatten()
-
-    valid = ~np.isnan(bin_avg_ij_flat)
-
-    return (
-        bin_avg_ij_flat[valid],
-        bin_count_ij_flat[valid],
-        bin_avg_i,
-        bin_count_i,
-        bin_avg_j,
-        bin_count_j,
-    )
+
+def _conditional_var(sample: np.ndarray, y: np.ndarray, edges: list) -> float:
+    """Var(E[Y | bins]), each bin weighted by its number of points."""
+    mean, *_ = stats.binned_statistic_dd(sample, y, statistic="mean", bins=edges)
+    count, *_ = stats.binned_statistic_dd(sample, y, statistic="count", bins=edges)
+    valid = count > 0
+    return _weighted_var(mean[valid], weights=count[valid])
 
 
 def number_of_bins(n_runs: int, n_factors: int) -> tuple[int, int]:
@@ -264,41 +172,33 @@ def sensitivity_indices(
     foe = np.empty(n_factors)
     soe = np.zeros((n_factors, n_factors))
 
+    edges_foe = [
+        _quantile_edges(inputs[:, k], int(n_bins_foe)) for k in range(n_factors)
+    ]
+    edges_soe = [
+        _quantile_edges(inputs[:, k], int(n_bins_soe)) for k in range(n_factors)
+    ]
+
+    # Marginal Var(E[Y|Xk]) on the SOE binning, identical for every pair
+    var_marginal = np.array(
+        [
+            _conditional_var(inputs[:, [k]], output, [edges_soe[k]])
+            for k in range(n_factors)
+        ]
+    )
+
     for i in range(n_factors):
         # 1. First-order effects (FOE)
-        xi = inputs[:, i]
+        foe[i] = _conditional_var(inputs[:, [i]], output, [edges_foe[i]]) / var_y
 
-        bin_avg, bin_count = bin_data_1d(xi, output, int(n_bins_foe))
-
-        valid_foe = ~np.isnan(bin_avg)
-        foe[i] = _weighted_var(bin_avg[valid_foe], weights=bin_count[valid_foe]) / var_y
-
-        # 2. Second-order effects (SOE)
         for j in range(n_factors):
             if j <= i:
                 continue
 
-            xj = inputs[:, j]
-
-            # Second-order effects (SOE) with magic_binning
-            (
-                bin_avg_ij,
-                bin_count_ij,
-                bin_avg_i,
-                bin_count_i,
-                bin_avg_j,
-                bin_count_j,
-            ) = bin_data_2d(xi, xj, output, int(n_bins_soe))
-
-            var_ij = _weighted_var(bin_avg_ij, weights=bin_count_ij)
-
-            valid_i = ~np.isnan(bin_avg_i)
-            var_i = _weighted_var(bin_avg_i[valid_i], weights=bin_count_i[valid_i])
-
-            valid_j = ~np.isnan(bin_avg_j)
-            var_j = _weighted_var(bin_avg_j[valid_j], weights=bin_count_j[valid_j])
-
-            soe[i, j] = (var_ij - var_i - var_j) / var_y
+            var_ij = _conditional_var(
+                inputs[:, [i, j]], output, [edges_soe[i], edges_soe[j]]
+            )
+            soe[i, j] = (var_ij - var_marginal[i] - var_marginal[j]) / var_y
 
     # Mirror SOE and calculate Combined Effect (SI)
     # SI is FOE + half of all interactions associated with that variable