pystatistics 3.3.0__tar.gz → 3.4.1__tar.gz

pystatistics-3.4.1/.release/UNRELEASED.md

@@ -0,0 +1,23 @@
+# Unreleased Changes
+
+> This file tracks all changes since the last stable release.
+> Updated by whoever makes a change, on whatever machine.
+> Synced via git so all sessions (Mac, Linux, etc.) see the same state.
+>
+> When ready to release, run: `python .release/release.py --status`
+> and follow the manual release flow in the script docstring.
+
+## Changes
+
+- CPU PMM (`mice` predictive mean matching) now scales to large datasets.
+  Replaced the dense `(n_mis, n_obs)` distance matrix in
+  `pystatistics/mice/methods/pmm.py:_match_donors` with a sorted-array windowed
+  k-NN (sort the observed predictions once, search a width-`2k` window around
+  each missing value's insertion point), matching R `mice`'s `matchindex`
+  approach. Cost drops from `O(n_mis·n_obs)` to `O(n_obs log n_obs + n_mis·k)`
+  in both time and memory. Measured CPU PMM (p=10-15, m=20, maxit=8): n=3000
+  25.4s -> 0.65s (~39x); n=20000 went from minutes/hours to ~6.5s. The window is
+  provably exact (the global k nearest neighbours lie in `[pos-k, pos+k-1]` of
+  the sorted array), so the donor pool is unchanged; imputations remain
+  distributionally identical and still match R. Exact per-seed outputs differ
+  from 3.4.0 because the donor ordering within the k-NN set changed.

{pystatistics-3.3.0 → pystatistics-3.4.1}/CHANGELOG.md

@@ -1,5 +1,38 @@
 # Changelog
 
+## 3.4.1
+
+Performance: CPU predictive mean matching now scales to large datasets.
+
+- The CPU PMM donor search no longer builds a dense distance matrix between
+  every missing and every observed value. It sorts the observed predictions
+  once and searches a small window around each missing value (the same approach
+  R's `mice` uses), reducing both time and memory from quadratic in the number
+  of rows to roughly `n log n`. In practice CPU PMM at n=3000 dropped from ~25s
+  to under 1s, and large problems (n=20000) that were effectively unusable now
+  finish in seconds.
+- Results are unchanged statistically — the donor pool and all imputation
+  distributions are identical, and validation against R's `mice` still passes.
+  Exact per-seed imputed values can differ slightly from 3.4.0.
+
+## 3.4.0
+
+GPU acceleration for MICE.
+
+- **`mice(..., backend='gpu')` runs on CUDA GPUs.** The `m` imputation chains
+  run together on the GPU, with the per-variable linear solves and the
+  predictive-mean-matching donor search batched across chains. `backend='auto'`
+  uses a CUDA GPU when one is available and falls back to the CPU otherwise.
+- **Large speedups on bigger problems.** Because the donor search dominates and
+  parallelises well, the GPU advantage grows with sample size — for example,
+  predictive mean matching was roughly 39× faster at n=1000 and 135× faster at
+  n=3000 than the CPU backend on an RTX 5070 Ti (m=20, maxit=8). GPU results
+  match the CPU backend at the standard GPU/FP32 tolerance.
+- **`use_fp64=True`** runs the GPU path in double precision on CUDA for closer
+  agreement with the CPU reference (default is FP32, the fast path).
+- GPU acceleration requires a CUDA GPU; Apple Silicon (MPS) is not yet supported
+  for MICE and is refused with a clear message rather than run unverified.
+
 ## 3.3.0
 
 Multiple imputation by chained equations (MICE) for multivariate missing data.

{pystatistics-3.3.0 → pystatistics-3.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pystatistics
-Version: 3.3.0
+Version: 3.4.1
 Summary: GPU-accelerated statistical computing for Python
 Project-URL: Homepage, https://sgcx.org/technology/pystatistics/
 Project-URL: Documentation, https://sgcx.org/docs/pystatistics/
@@ -86,7 +86,7 @@ PyStatistics maintains two parallel computational paths with distinct goals:
 | `multivariate/` | Complete | PCA and maximum likelihood factor analysis with varimax/promax rotation |
 | `timeseries/` | Complete | ACF, PACF, ADF, KPSS, ETS, ARIMA, SARIMA, auto_arima, decompose, STL |
 | `gam/` | Complete | Generalized additive models with penalized regression splines matching R mgcv::gam |
-| `mice/` | Numeric (CPU) | Multiple imputation by chained equations: PMM and Bayesian linear regression, Rubin's-rules pooling, validated against R mice |
+| `mice/` | Numeric | Multiple imputation by chained equations: PMM and Bayesian linear regression, Rubin's-rules pooling, validated against R mice; CUDA GPU backend |
 
 See [docs/ROADMAP.md](docs/ROADMAP.md) for detailed scope, GPU applicability, and implementation priority for each module.
 
@@ -413,6 +413,26 @@ pip install pystatistics[dev]
 
 ## What's New
 
+### 3.4.1 — Faster CPU predictive mean matching
+
+- CPU PMM in `mice` now scales to large datasets: the donor search sorts the
+  observed predictions and scans a small window per missing value (as R's
+  `mice` does) instead of forming a full distance matrix, cutting time and
+  memory from quadratic to roughly `n log n`. Large problems that were
+  effectively unusable on the CPU now finish in seconds. Results are
+  statistically unchanged.
+
+### 3.4.0 — GPU acceleration for MICE
+
+- `mice(..., backend='gpu')` runs the imputation chains on a CUDA GPU, batching
+  the per-variable solves and the predictive-mean-matching donor search across
+  chains. `backend='auto'` uses a CUDA GPU when available, else the CPU.
+- The speedup grows with sample size (the donor search dominates and batches
+  well): predictive mean matching ran ~39× faster at n=1000 and ~135× faster at
+  n=3000 than the CPU backend on an RTX 5070 Ti. GPU results match the CPU
+  backend at the GPU/FP32 tolerance; pass `use_fp64=True` for double precision.
+- Requires a CUDA GPU; Apple Silicon (MPS) is not yet supported for MICE.
+
 ### 3.3.0 — Multiple imputation (MICE)
 
 - New `mice` module: multiple imputation by chained equations for numeric data

{pystatistics-3.3.0 → pystatistics-3.4.1}/README.md

@@ -39,7 +39,7 @@ PyStatistics maintains two parallel computational paths with distinct goals:
 | `multivariate/` | Complete | PCA and maximum likelihood factor analysis with varimax/promax rotation |
 | `timeseries/` | Complete | ACF, PACF, ADF, KPSS, ETS, ARIMA, SARIMA, auto_arima, decompose, STL |
 | `gam/` | Complete | Generalized additive models with penalized regression splines matching R mgcv::gam |
-| `mice/` | Numeric (CPU) | Multiple imputation by chained equations: PMM and Bayesian linear regression, Rubin's-rules pooling, validated against R mice |
+| `mice/` | Numeric | Multiple imputation by chained equations: PMM and Bayesian linear regression, Rubin's-rules pooling, validated against R mice; CUDA GPU backend |
 
 See [docs/ROADMAP.md](docs/ROADMAP.md) for detailed scope, GPU applicability, and implementation priority for each module.
 
@@ -366,6 +366,26 @@ pip install pystatistics[dev]
 
 ## What's New
 
+### 3.4.1 — Faster CPU predictive mean matching
+
+- CPU PMM in `mice` now scales to large datasets: the donor search sorts the
+  observed predictions and scans a small window per missing value (as R's
+  `mice` does) instead of forming a full distance matrix, cutting time and
+  memory from quadratic to roughly `n log n`. Large problems that were
+  effectively unusable on the CPU now finish in seconds. Results are
+  statistically unchanged.
+
+### 3.4.0 — GPU acceleration for MICE
+
+- `mice(..., backend='gpu')` runs the imputation chains on a CUDA GPU, batching
+  the per-variable solves and the predictive-mean-matching donor search across
+  chains. `backend='auto'` uses a CUDA GPU when available, else the CPU.
+- The speedup grows with sample size (the donor search dominates and batches
+  well): predictive mean matching ran ~39× faster at n=1000 and ~135× faster at
+  n=3000 than the CPU backend on an RTX 5070 Ti. GPU results match the CPU
+  backend at the GPU/FP32 tolerance; pass `use_fp64=True` for double precision.
+- Requires a CUDA GPU; Apple Silicon (MPS) is not yet supported for MICE.
+
 ### 3.3.0 — Multiple imputation (MICE)
 
 - New `mice` module: multiple imputation by chained equations for numeric data

{pystatistics-3.3.0 → pystatistics-3.4.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pystatistics"
-version = "3.3.0"
+version = "3.4.1"
 description = "GPU-accelerated statistical computing for Python"
 readme = "README.md"
 license = "MIT"

{pystatistics-3.3.0 → pystatistics-3.4.1}/pystatistics/__init__.py

@@ -16,7 +16,7 @@ Usage:
     result = fit(design)
 """
 
-__version__ = "3.3.0"
+__version__ = "3.4.1"
 __author__ = "Hai-Shuo"
 __email__ = "contact@sgcx.org"
 

pystatistics-3.4.1/pystatistics/mice/backends/_gpu_linreg.py

@@ -0,0 +1,138 @@
+"""
+Batched Bayesian linear-regression posterior draw on the GPU.
+
+This is the GPU counterpart of ``methods/_linreg.py``: the same Bayesian draw,
+but vectorized over the ``m`` imputation chains as the leading batch dimension.
+At a given sweep step every chain is fitting the *same* target column on the
+*same* observed rows, but with different predictor values (each chain imputed
+the other columns differently), so we solve ``m`` independent linear systems at
+once with batched cuBLAS/cuSOLVER kernels.
+
+Posterior draw (identical model to the CPU path):
+
+    beta_hat = (X'X + ridge)^{-1} X'y                 batched solve
+    df       = max(n_obs - n_params, 1)               integer, shared by chains
+    sigma*   = sqrt( RSS / chi2(df) )                 chi2 = sum of df squared N(0,1)
+    beta*    = beta_hat + sigma* * L z                L L' = (X'X + ridge)^{-1}
+
+The chi-square is built from standard normals (``df`` is a non-negative integer)
+so the only randomness source is the seeded ``torch.Generator`` — no reliance on
+``torch.distributions``, which ignores explicit generators.
+
+All randomness flows through the passed generator (CLAUDE.md Rule 6). Results
+match the CPU reference distributionally, at the GPU/FP32 tolerance tier — not
+bit-for-bit (different RNG, single precision).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Matches the CPU ridge (methods/_linreg.py) — a tiny relative penalty keeping
+# the batched Gram matrices invertible under FP32.
+_DEFAULT_RIDGE = 1e-5
+
+
+@dataclass
+class BatchedLinRegDraw:
+    """Batched point estimate + posterior draw (leading dim = m chains)."""
+
+    beta_hat: "object"    # (m, q+1) tensor
+    beta_draw: "object"   # (m, q+1) tensor
+    sigma_draw: "object"  # (m,) tensor
+
+
+def add_intercept(X):
+    """Prepend a ones column to a batched predictor tensor (m, n, q) -> (m,n,q+1)."""
+    import torch
+
+    m, n, _ = X.shape
+    ones = torch.ones((m, n, 1), dtype=X.dtype, device=X.device)
+    return torch.cat([ones, X], dim=2)
+
+
+def batched_bayes_linreg_draw(
+    y_obs,
+    X_obs,
+    gen,
+    ridge: float = _DEFAULT_RIDGE,
+) -> BatchedLinRegDraw:
+    """Draw once from the Gaussian linear-model posterior for every chain.
+
+    Parameters
+    ----------
+    y_obs : (m, n_obs) tensor
+        Observed responses (identical across chains, batched for uniform ops).
+    X_obs : (m, n_obs, q) tensor
+        Observed predictors WITHOUT an intercept column (added internally).
+    gen : torch.Generator
+        Sole randomness source, on the same device as the tensors.
+    ridge : float
+        Relative ridge penalty on the diagonal of X'X.
+    """
+    import torch
+
+    Xa = add_intercept(X_obs)                       # (m, n_obs, q+1)
+    m, n_obs, n_params = Xa.shape
+    dtype, device = Xa.dtype, Xa.device
+
+    Xt = Xa.transpose(1, 2)                          # (m, q+1, n_obs)
+    XtX = Xt @ Xa                                    # (m, q+1, q+1)
+
+    diag = torch.diagonal(XtX, dim1=1, dim2=2)       # (m, q+1)
+    diag_mean = diag.mean(dim=1).clamp_min(torch.finfo(dtype).tiny)  # (m,)
+    eye = torch.eye(n_params, dtype=dtype, device=device)
+    XtX_ridge = XtX + ridge * diag_mean[:, None, None] * eye
+
+    Xty = Xt @ y_obs.unsqueeze(-1)                   # (m, q+1, 1)
+    beta_hat = torch.linalg.solve(XtX_ridge, Xty).squeeze(-1)  # (m, q+1)
+
+    resid = y_obs - (Xa @ beta_hat.unsqueeze(-1)).squeeze(-1)  # (m, n_obs)
+    rss = (resid * resid).sum(dim=1)                 # (m,)
+    df = max(n_obs - n_params, 1)
+
+    # chi2(df) = sum of df squared standard normals (df is an integer).
+    z_chi = torch.randn((m, df), generator=gen, dtype=dtype, device=device)
+    chi = (z_chi * z_chi).sum(dim=1).clamp_min(torch.finfo(dtype).tiny)  # (m,)
+    sigma_draw = torch.where(
+        rss > 0, torch.sqrt(rss / chi), torch.zeros_like(rss)
+    )                                                # (m,)
+
+    V = torch.linalg.inv(XtX_ridge)                  # (m, q+1, q+1)
+    L = _batched_safe_cholesky(V)                    # (m, q+1, q+1), L L' = V
+    z_beta = torch.randn((m, n_params), generator=gen, dtype=dtype, device=device)
+    beta_draw = beta_hat + sigma_draw[:, None] * (L @ z_beta.unsqueeze(-1)).squeeze(-1)
+
+    return BatchedLinRegDraw(
+        beta_hat=beta_hat, beta_draw=beta_draw, sigma_draw=sigma_draw
+    )
+
+
+def _batched_safe_cholesky(V):
+    """Lower Cholesky of each (symmetric PSD) matrix in a batch, with jitter.
+
+    ``V`` is a batch of inverse-Gram matrices — positive definite in exact
+    arithmetic, but FP32 rounding can leave one marginally indefinite. We
+    symmetrize and add escalating jitter to the *whole batch* (cheap, keeps the
+    op batched) until ``cholesky_ex`` reports success for every matrix.
+    """
+    import torch
+
+    Vs = 0.5 * (V + V.transpose(1, 2))
+    eye = torch.eye(Vs.shape[1], dtype=Vs.dtype, device=Vs.device)
+    diag = torch.diagonal(Vs, dim1=1, dim2=2)
+    scale = diag.mean().clamp_min(1.0).item()
+
+    jitter = 0.0
+    for _ in range(8):
+        M = Vs if jitter == 0.0 else Vs + jitter * eye
+        L, info = torch.linalg.cholesky_ex(M)
+        if not torch.any(info):
+            return L
+        jitter = scale * (1e-8 if jitter == 0.0 else 10.0 * jitter)
+
+    # Documented last resort (Rule 1): eigenvalue-clipped PSD factor. Reached
+    # only if a matrix is badly degenerate despite ridge + jitter.
+    w, Q = torch.linalg.eigh(Vs)
+    w = w.clamp_min(0.0)
+    return Q @ torch.diag_embed(torch.sqrt(w))

pystatistics-3.4.1/pystatistics/mice/backends/_gpu_methods.py

@@ -0,0 +1,81 @@
+"""
+Batched GPU imputation methods: ``norm`` and ``pmm``.
+
+These mirror ``methods/norm.py`` and ``methods/pmm.py`` but operate on tensors
+batched over the ``m`` chains. Each takes the observed/missing predictor blocks
+for one target column across all chains and returns the imputed values for every
+chain at once, so a whole sweep step is a handful of batched kernels.
+
+The GPU backend dispatches to these by method name; a name with no GPU
+implementation is refused by the backend (fail loud, Rule 1) rather than
+silently downgraded.
+"""
+
+from __future__ import annotations
+
+from pystatistics.mice.backends._gpu_linreg import (
+    add_intercept,
+    batched_bayes_linreg_draw,
+)
+
+
+def gpu_norm_impute(y_obs, X_obs, X_mis, gen, *, donors=None):
+    """Batched Bayesian linear-regression imputation (R ``norm``).
+
+    Parameters mirror the CPU method but are batched: ``y_obs`` (m, n_obs),
+    ``X_obs`` (m, n_obs, q), ``X_mis`` (m, n_mis, q). Returns (m, n_mis).
+    ``donors`` is accepted and ignored so the backend can call every method
+    with one uniform signature.
+    """
+    import torch
+
+    draw = batched_bayes_linreg_draw(y_obs, X_obs, gen)
+    Xa_mis = add_intercept(X_mis)
+    yhat_mis = (Xa_mis @ draw.beta_draw.unsqueeze(-1)).squeeze(-1)  # (m, n_mis)
+    noise = torch.randn(
+        yhat_mis.shape, generator=gen, dtype=yhat_mis.dtype, device=yhat_mis.device
+    )
+    return yhat_mis + draw.sigma_draw[:, None] * noise
+
+
+def gpu_pmm_impute(y_obs, X_obs, X_mis, gen, *, donors=5):
+    """Batched predictive mean matching (R ``pmm``, matchtype=1).
+
+    Observed fitted values use ``beta_hat``; missing fitted values use the
+    posterior draw ``beta*``. For each missing row in each chain, pick one of
+    the ``k`` nearest observed fitted values at random and copy that observed
+    value. Returns (m, n_mis).
+    """
+    import torch
+
+    draw = batched_bayes_linreg_draw(y_obs, X_obs, gen)
+    Xa_obs = add_intercept(X_obs)
+    Xa_mis = add_intercept(X_mis)
+    yhat_obs = (Xa_obs @ draw.beta_hat.unsqueeze(-1)).squeeze(-1)   # (m, n_obs)
+    yhat_mis = (Xa_mis @ draw.beta_draw.unsqueeze(-1)).squeeze(-1)  # (m, n_mis)
+
+    m, n_obs = yhat_obs.shape
+    n_mis = yhat_mis.shape[1]
+    k = min(donors, n_obs)
+
+    # |yhat_mis - yhat_obs| for every (chain, missing row, observed row).
+    dist = (yhat_mis[:, :, None] - yhat_obs[:, None, :]).abs()  # (m, n_mis, n_obs)
+    # k nearest observed fitted values (indices into n_obs).
+    _, cand = torch.topk(dist, k, dim=2, largest=False)         # (m, n_mis, k)
+
+    # One random donor per (chain, missing row) among its k candidates.
+    pick = torch.randint(
+        0, k, (m, n_mis), generator=gen, device=yhat_obs.device
+    )                                                          # (m, n_mis)
+    donor_idx = torch.gather(cand, 2, pick.unsqueeze(-1)).squeeze(-1)  # (m, n_mis)
+
+    # Copy the observed donor values (gather along the observed-row axis).
+    return torch.gather(y_obs, 1, donor_idx)                   # (m, n_mis)
+
+
+# Method-name -> batched implementation. The backend translates the design's
+# per-column method name through this table; unknown names are refused there.
+GPU_METHODS = {
+    "norm": gpu_norm_impute,
+    "pmm": gpu_pmm_impute,
+}

{pystatistics-3.3.0 → pystatistics-3.4.1}/pystatistics/mice/backends/cpu.py

@@ -17,6 +17,7 @@ import numpy as np
 from pystatistics.core.compute.timing import Timer
 from pystatistics.core.result import Result
 from pystatistics.mice._chain import run_chain
+from pystatistics.mice._rng import spawn_streams
 from pystatistics.mice.design import MICEDesign
 from pystatistics.mice.solution import MICEParams
 
@@ -35,7 +36,7 @@ class CPUMiceBackend:
         m: int,
         maxit: int,
         visit_sequence: tuple[int, ...],
-        streams: list[np.random.Generator],
+        seed: int,
     ) -> Result[MICEParams]:
         """Run ``m`` chains and package the imputations.
 
@@ -48,13 +49,15 @@ class CPUMiceBackend:
             Iterations per chain.
         visit_sequence : tuple of int
             Column visit order within each iteration.
-        streams : list of numpy.random.Generator
-            One independent RNG per chain; ``len(streams) == m`` (the solver
-            guarantees this).
+        seed : int
+            Master RNG seed. The backend spawns one independent, reproducible
+            stream per chain from it.
         """
         timer = Timer()
         timer.start()
 
+        streams = spawn_streams(seed, m)
+
         completed = np.empty((m, design.n, design.p), dtype=np.float64)
         chain_means = np.empty((m, maxit, len(design.incomplete_columns)))
         chain_vars = np.empty((m, maxit, len(design.incomplete_columns)))

pystatistics-3.4.1/pystatistics/mice/backends/gpu.py

@@ -0,0 +1,170 @@
+"""
+GPU backend for MICE.
+
+Implements the same ``run`` entrypoint as the CPU backend, but runs all ``m``
+imputation chains concurrently on a CUDA GPU with the chain index as the leading
+tensor batch dimension. Each sweep step — visit column ``j``, fit its method on
+the observed rows, overwrite its missing cells — becomes a handful of batched
+kernels across the ``m`` chains (batched linear solves + Cholesky, and for PMM a
+batched nearest-neighbour donor search).
+
+Why batch over chains: the ``m`` chains are statistically independent and share
+the missingness pattern, so they have identical shapes at every step — an ideal
+batch. The sweep stays sequential *within* a chain (column ``j`` depends on the
+columns imputed before it), which is inherent to chained equations.
+
+Validation tier: GPU results match the CPU reference distributionally at the
+``GPU_FP32`` tolerance, not bit-for-bit (FP32 + a different RNG). FP64 is
+available on CUDA (``use_fp64=True``) for closer parity.
+
+Determinism: a single seeded ``torch.Generator`` on the device is the only
+randomness source. Given the same seed and device, a run reproduces its own
+output to within FP32 kernel non-determinism.
+"""
+
+from __future__ import annotations
+
+from pystatistics.core.compute.timing import Timer
+from pystatistics.core.compute.torch_interop import to_host_f64
+from pystatistics.core.exceptions import ValidationError
+from pystatistics.core.result import Result
+from pystatistics.mice.backends._gpu_methods import GPU_METHODS
+from pystatistics.mice.design import MICEDesign
+from pystatistics.mice.solution import MICEParams
+
+# R mice default donor count for PMM (the design does not carry a per-column
+# donor count, so the GPU path uses the same default as the CPU PMMMethod).
+_PMM_DONORS = 5
+
+
+class GPUMiceBackend:
+    """Runs MICE chains batched on a CUDA GPU."""
+
+    def __init__(self, device: str = "cuda", use_fp64: bool = False):
+        self.device = device
+        self.use_fp64 = use_fp64
+
+    @property
+    def name(self) -> str:
+        return f"gpu_{'fp64' if self.use_fp64 else 'fp32'}"
+
+    def run(
+        self,
+        design: MICEDesign,
+        *,
+        m: int,
+        maxit: int,
+        visit_sequence: tuple[int, ...],
+        seed: int,
+    ) -> Result[MICEParams]:
+        import torch
+
+        # Validate that every incomplete column's method has a GPU kernel before
+        # touching the device (fail loud at the boundary, Rule 2).
+        for j in design.incomplete_columns:
+            meth = design.method_for(j)
+            if meth not in GPU_METHODS:
+                raise ValidationError(
+                    f"Method {meth!r} (column {j}) has no GPU implementation. "
+                    f"GPU methods: {sorted(GPU_METHODS)}. Use backend='cpu'."
+                )
+
+        timer = Timer()
+        timer.start()
+
+        dtype = torch.float64 if self.use_fp64 else torch.float32
+        dev = torch.device(self.device)
+        gen = torch.Generator(device=dev)
+        gen.manual_seed(int(seed))
+
+        n, p = design.n, design.p
+        incomplete = design.incomplete_columns
+        incomplete_pos = {j: k for k, j in enumerate(incomplete)}
+
+        with timer.section("transfer"):
+            base = torch.as_tensor(design.data, device=dev, dtype=dtype)  # (n, p)
+            mask = torch.as_tensor(design.missing_mask, device=dev)        # (n, p)
+            # m independent copies of the data (chains diverge after init).
+            data = base.unsqueeze(0).repeat(m, 1, 1)                       # (m, n, p)
+
+        # Precompute per-incomplete-column index tensors (shared by all chains).
+        obs_idx, mis_idx, pred_idx = {}, {}, {}
+        for j in incomplete:
+            mcol = mask[:, j]
+            obs_idx[j] = torch.nonzero(~mcol, as_tuple=False).squeeze(1)
+            mis_idx[j] = torch.nonzero(mcol, as_tuple=False).squeeze(1)
+            pred_idx[j] = torch.tensor(
+                [c for c in range(p) if c != j], device=dev, dtype=torch.long
+            )
+
+        with timer.section("init"):
+            self._initialise(data, base, obs_idx, mis_idx, incomplete, m, gen, dev)
+
+        chain_mean = torch.empty((m, maxit, len(incomplete)), dtype=dtype, device=dev)
+        chain_var = torch.empty((m, maxit, len(incomplete)), dtype=dtype, device=dev)
+
+        with timer.section("sweep"):
+            for it in range(maxit):
+                for j in visit_sequence:
+                    method_fn = GPU_METHODS[design.method_for(j)]
+                    X = data[:, :, pred_idx[j]]            # (m, n, q)
+                    X_obs = X[:, obs_idx[j], :]            # (m, n_obs, q)
+                    X_mis = X[:, mis_idx[j], :]            # (m, n_mis, q)
+                    y_obs = data[:, obs_idx[j], j]         # (m, n_obs)
+
+                    imputed = method_fn(
+                        y_obs, X_obs, X_mis, gen, donors=_PMM_DONORS
+                    )
+                    data[:, mis_idx[j], j] = imputed
+
+                for j in incomplete:
+                    cells = data[:, mis_idx[j], j]          # (m, n_mis)
+                    k = incomplete_pos[j]
+                    chain_mean[:, it, k] = cells.mean(dim=1)
+                    chain_var[:, it, k] = cells.var(dim=1, unbiased=False)
+
+        with timer.section("download"):
+            completed = to_host_f64(data)
+            chain_mean_host = to_host_f64(chain_mean)
+            chain_var_host = to_host_f64(chain_var)
+
+        timer.stop()
+
+        params = MICEParams(
+            completed=completed,
+            chain_mean=chain_mean_host,
+            chain_var=chain_var_host,
+            incomplete_columns=incomplete,
+            m=m,
+            maxit=maxit,
+            visit_sequence=visit_sequence,
+        )
+        return Result(
+            params=params,
+            info={
+                "method": "chained_equations",
+                "m": m,
+                "maxit": maxit,
+                "device": self.device,
+                "precision": "fp64" if self.use_fp64 else "fp32",
+                "visit_sequence": list(visit_sequence),
+            },
+            timing=timer.result(),
+            backend_name=self.name,
+            warnings=(),
+        )
+
+    @staticmethod
+    def _initialise(data, base, obs_idx, mis_idx, incomplete, m, gen, dev):
+        """Fill missing cells by sampling observed values of the same column,
+        independently per chain. Mutates ``data`` in place."""
+        import torch
+
+        for j in incomplete:
+            observed = base[obs_idx[j], j]          # (n_obs,)
+            n_obs = observed.shape[0]
+            n_mis = mis_idx[j].shape[0]
+            draw_idx = torch.randint(
+                0, n_obs, (m, n_mis), generator=gen, device=dev
+            )
+            data[:, mis_idx[j], j] = observed[draw_idx]

{pystatistics-3.3.0 → pystatistics-3.4.1}/pystatistics/mice/methods/pmm.py

@@ -63,24 +63,45 @@ def _match_donors(
     """For each missing fitted value, pick one of its ``k`` nearest observed
     fitted values at random; return the chosen observed-row indices.
 
-    Vectorized across all missing rows — this is the donor-search kernel a
-    Stage-2 GPU backend would replace with a batched nearest-neighbour search.
+    Scales like R's ``mice`` matcher: sort the observed fitted values once, then
+    for each missing value search a small window around its insertion point
+    instead of forming the full ``(n_mis, n_obs)`` distance matrix. Cost is
+    ``O(n_obs log n_obs + n_mis·k)`` time and ``O(n_mis·k)`` memory — the dense
+    version was ``O(n_mis·n_obs)`` in both, which blows up on large data.
+
+    The window is provably exact. In a sorted array the ``k`` nearest neighbours
+    of a query lie in the contiguous block ``[pos-k, pos+k-1]`` around the
+    insertion point ``pos``; a width-``2k`` window covers that block (shifted to
+    stay in range near the ends), so the windowed k-NN equals the global k-NN.
     """
     n_obs = yhat_obs.shape[0]
+    n_mis = yhat_mis.shape[0]
     k = min(donors, n_obs)
 
-    # |yhat_mis_i - yhat_obs_j| for every (i, j): (n_mis, n_obs).
-    dist = np.abs(yhat_mis[:, None] - yhat_obs[None, :])
+    # Sort donor predictions once; keep the map back to original obs rows.
+    order = np.argsort(yhat_obs, kind="stable")
+    sorted_obs = yhat_obs[order]
 
-    # Indices of the k smallest distances per row (unordered within the k).
-    if k < n_obs:
-        cand = np.argpartition(dist, k - 1, axis=1)[:, :k]
+    # Insertion point of each missing prediction into the sorted donors.
+    pos = np.searchsorted(sorted_obs, yhat_mis)
+
+    # Window of distinct candidate donors guaranteed to contain the k nearest.
+    w = min(2 * k, n_obs)
+    start = np.clip(pos - k, 0, n_obs - w)                  # (n_mis,)
+    win = start[:, None] + np.arange(w)[None, :]            # (n_mis, w) sorted-idx
+    dist = np.abs(sorted_obs[win] - yhat_mis[:, None])      # (n_mis, w)
+
+    # k smallest distances within the window (unordered within the k).
+    if k < w:
+        knn = np.argpartition(dist, k - 1, axis=1)[:, :k]   # (n_mis, k) window cols
     else:
-        cand = np.broadcast_to(np.arange(n_obs), (yhat_mis.shape[0], n_obs))
+        knn = np.broadcast_to(np.arange(w), (n_mis, w))
 
-    # One random donor per row among its k candidates.
-    pick = rng.integers(0, k, size=yhat_mis.shape[0])
-    return cand[np.arange(yhat_mis.shape[0]), pick]
+    rows = np.arange(n_mis)
+    pick = rng.integers(0, k, size=n_mis)                   # one donor per row
+    chosen_win_col = knn[rows, pick]                        # (n_mis,)
+    chosen_sorted = win[rows, chosen_win_col]               # idx into sorted_obs
+    return order[chosen_sorted]                             # back to original rows
 
 
 register(PMMMethod())