pystatistics 3.4.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.4.0 → pystatistics-3.4.1}/CHANGELOG.md

@@ -1,5 +1,20 @@
 # 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.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pystatistics
-Version: 3.4.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/
@@ -413,6 +413,15 @@ 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

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

@@ -366,6 +366,15 @@ 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

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

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

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

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

{pystatistics-3.4.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())

{pystatistics-3.4.0 → pystatistics-3.4.1}/tests/mice/test_methods.py

@@ -130,6 +130,35 @@ class TestPMM:
         idx = _match_donors(yhat_mis, yhat_obs, donors=5, rng=rng)
         assert np.all((idx >= 0) & (idx < 2))
 
+    def test_windowed_matcher_returns_true_knn(self):
+        # The chosen donor must always be within the true global k-NN computed
+        # by brute force — the windowed search must not miss nearer donors.
+        rng = make_rng(0)
+        obs_rng = np.random.default_rng(1)
+        yhat_obs = obs_rng.normal(size=300)
+        yhat_mis = obs_rng.normal(size=120)
+        k = 5
+        # Brute-force k nearest observed indices for each missing value.
+        full = np.abs(yhat_mis[:, None] - yhat_obs[None, :])
+        true_knn = set(
+            (i, j)
+            for i in range(len(yhat_mis))
+            for j in np.argpartition(full[i], k - 1)[:k]
+        )
+        chosen = _match_donors(yhat_mis, yhat_obs, donors=k, rng=rng)
+        for i, j in enumerate(chosen):
+            assert (i, int(j)) in true_knn
+
+    def test_matcher_donor_values_are_observed(self):
+        # End-to-end: the matched index points at a real observed value.
+        rng = make_rng(0)
+        obs_rng = np.random.default_rng(2)
+        y_obs = obs_rng.normal(size=200)
+        yhat_obs = obs_rng.normal(size=200)
+        yhat_mis = obs_rng.normal(size=80)
+        idx = _match_donors(yhat_mis, yhat_obs, donors=5, rng=rng)
+        assert np.all((idx >= 0) & (idx < 200))
+
     def test_single_donor_picks_nearest(self):
         # With donors=1, each missing gets the single closest observed value.
         rng = make_rng(0)

pystatistics-3.4.0/.release/UNRELEASED.md

@@ -1,29 +0,0 @@
-# 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
-
-- GPU backend for `mice` (CUDA). `mice(..., backend='gpu')` runs the `m`
-  imputation chains batched on the GPU as the leading tensor dimension: each
-  sweep step is a batched Bayesian linear solve (`torch.linalg.solve` + batched
-  Cholesky) and, for PMM, a batched nearest-neighbour donor search
-  (`(m, n_mis, n_obs)` distances → `topk`). `backend='auto'` selects CUDA when
-  available, else CPU (never MPS, matching the other modules). New `use_fp64`
-  argument runs the GPU path in double precision on CUDA (default FP32).
-- GPU results match the CPU reference distributionally at the GPU/FP32 tolerance
-  tier (PMM donor copies are exact in FP64, exact-to-FP32 otherwise). Validated
-  in `tests/mice/test_gpu.py` against the CPU path rather than R directly.
-- Measured PMM speedups vs the CPU backend on an RTX 5070 Ti: ~39× at n=1000
-  (p=8, m=20, maxit=8) and ~135× at n=3000 (p=10, m=20, maxit=8); the gain grows
-  with n because the donor search dominates and batches well on the GPU.
-- Internal: the MICE backend contract now takes a `seed` (each backend spawns
-  its own RNG streams) instead of pre-spawned streams, so the CPU and GPU
-  backends share one `run()` signature. No public API change.
-- MPS (Apple Silicon) is not yet validated for MICE; `backend='gpu'` raises a
-  clear error on MPS rather than running unverified. CUDA is supported.