pystatistics 3.3.0__tar.gz → 3.4.0__tar.gz

pystatistics-3.4.0/.release/UNRELEASED.md

@@ -0,0 +1,29 @@
+# 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.

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

@@ -1,5 +1,23 @@
 # Changelog
 
+## 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.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pystatistics
-Version: 3.3.0
+Version: 3.4.0
 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,17 @@ pip install pystatistics[dev]
 
 ## What's New
 
+### 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.0}/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,17 @@ pip install pystatistics[dev]
 
 ## What's New
 
+### 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.0}/pyproject.toml

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

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

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

pystatistics-3.4.0/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.0/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.0}/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.0/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.0}/pystatistics/mice/solvers.py

@@ -20,7 +20,6 @@ import warnings
 from typing import Literal, Mapping, Sequence
 
 from pystatistics.core.compute.device import select_device
-from pystatistics.mice._rng import spawn_streams
 from pystatistics.mice._visit import resolve_visit_sequence
 from pystatistics.mice.design import MICEDesign
 from pystatistics.mice.solution import MICESolution
@@ -38,6 +37,7 @@ def mice(
     methods: Mapping | Sequence[str] | None = None,
     visit_sequence: Sequence[int] | None = None,
     backend: BackendChoice | None = None,
+    use_fp64: bool = False,
     verbose: bool = False,
 ) -> MICESolution:
     """Multiple imputation by chained equations (numeric columns, CPU).
@@ -64,8 +64,14 @@ def mice(
         Column visit order within each iteration. Defaults to incomplete
         columns in ascending index order (R "roman" default).
     backend : {'cpu', 'gpu', 'auto'} or None
-        Compute backend. Default None -> 'cpu'. GPU is not implemented in this
-        release; requesting it raises.
+        Compute backend. Default None -> 'cpu' (the R-validated reference path).
+        'gpu' runs the chains batched on a CUDA GPU; 'auto' uses CUDA when
+        available, else CPU. GPU results match the CPU backend at the GPU/FP32
+        tolerance tier.
+    use_fp64 : bool
+        Only relevant on the GPU. Run the GPU chains in double precision (CUDA
+        only) for closer parity with the CPU reference. Default False (FP32,
+        the fast consumer-GPU path). Ignored on CPU (always double precision).
     verbose : bool
         Print progress information.
 
@@ -105,8 +111,7 @@ def mice(
         )
 
     visit = resolve_visit_sequence(design.incomplete_columns, visit_sequence)
-    backend_impl = _get_backend(backend)
-    streams = spawn_streams(seed, m)
+    backend_impl = _get_backend(backend, use_fp64)
 
     if verbose:
         print(
@@ -116,7 +121,7 @@ def mice(
         )
 
     result = backend_impl.run(
-        design, m=m, maxit=maxit, visit_sequence=visit, streams=streams
+        design, m=m, maxit=maxit, visit_sequence=visit, seed=seed
     )
 
     if verbose:
@@ -125,30 +130,41 @@ def mice(
     return MICESolution(_result=result, _design=design)
 
 
-def _get_backend(choice: BackendChoice | None):
-    """Resolve the compute backend. None -> CPU; GPU not yet implemented."""
+def _get_backend(choice: BackendChoice | None, use_fp64: bool):
+    """Resolve the compute backend.
+
+    None/'cpu' -> CPU. 'gpu' -> CUDA GPU (raises on MPS, which is not yet
+    validated for MICE). 'auto' -> CUDA when available, else CPU; like the
+    other pystatistics modules, 'auto' never selects MPS.
+    """
     if choice is None or choice == "cpu":
         from pystatistics.mice.backends.cpu import CPUMiceBackend
 
         return CPUMiceBackend()
 
     if choice == "auto":
-        # 'auto' currently means CPU: no GPU backend ships in this release.
-        # Kept as an accepted value so callers and the Stage-2 GPU path share
-        # one vocabulary.
+        device = select_device("auto")
+        if device.device_type == "cuda":
+            from pystatistics.mice.backends.gpu import GPUMiceBackend
+
+            return GPUMiceBackend(device="cuda", use_fp64=use_fp64)
+        # CPU, or MPS (never auto-selected): use the CPU reference path.
         from pystatistics.mice.backends.cpu import CPUMiceBackend
 
         return CPUMiceBackend()
 
     if choice == "gpu":
-        # Fail loud rather than silently downgrading (Rule 1). select_device
-        # gives a precise hardware error if no GPU exists; otherwise we still
-        # refuse because the GPU chain backend is Stage 2.
-        select_device("gpu")
-        raise NotImplementedError(
-            "backend='gpu' is not available in this release. MICE ships a "
-            "validated CPU implementation first; GPU acceleration is planned. "
-            "Use backend='cpu' (the default)."
-        )
+        device = select_device("gpu")  # raises RuntimeError if no GPU
+        if device.device_type == "mps":
+            raise NotImplementedError(
+                "backend='gpu' for MICE is validated on CUDA only in this "
+                "release. The chained-equations sweep has not been validated "
+                "on Apple Silicon (MPS), so it is refused rather than run "
+                "unverified. Use backend='cpu' (the default) on a Mac; CUDA "
+                "GPUs are supported."
+            )
+        from pystatistics.mice.backends.gpu import GPUMiceBackend
+
+        return GPUMiceBackend(device=device.device_type, use_fp64=use_fp64)
 
     raise ValueError(f"Unknown backend: {choice!r}")