pystatistics 3.1.0__tar.gz → 3.3.0__tar.gz

pystatistics-3.3.0/.release/UNRELEASED.md

@@ -0,0 +1,30 @@
+# 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
+
+- New `mice` module: Multiple Imputation by Chained Equations for multivariate
+  missing data. `pystatistics.mice.mice(data, m=5, maxit=5, method='pmm',
+  seed=...)` iteratively imputes each incomplete numeric column from the others
+  and returns `m` completed datasets (`MICESolution`). Numeric methods `pmm`
+  (predictive mean matching, the R default) and `norm` (Bayesian linear
+  regression) are included, with R-faithful defaults (`m=5`, `maxit=5`, 5 PMM
+  donors, left-to-right visit sequence). Imputation is fully deterministic given
+  `seed` (required), with independent per-chain RNG streams.
+- New `pystatistics.mice.pool`: Rubin's-rules pooling of an analysis across the
+  `m` completed datasets, with Barnard-Rubin degrees of freedom and the standard
+  diagnostics (within/between/total variance, relative increase in variance,
+  lambda, fraction of missing information) and confidence intervals.
+- Validated distributionally against R's `mice` 3.19.0: imputed-value
+  distributions (mean, spread, quantiles) and Rubin's-rules pooled regression
+  output match R on a shared incomplete dataset. Reference fixtures and the
+  pinned generation script live in `tests/mice/references/`.
+- CPU implementation only in this release. The module is structured (method
+  registry, single backend entrypoint, per-column type metadata) so GPU
+  acceleration and categorical methods can be added without restructuring.

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

@@ -1,5 +1,64 @@
 # Changelog
 
+## 3.3.0
+
+Multiple imputation by chained equations (MICE) for multivariate missing data.
+
+- **New `mice` module.** `pystatistics.mice.mice(data, m=5, maxit=5,
+  method='pmm', seed=...)` imputes a numeric data matrix that contains missing
+  values (NaN) by iteratively modelling each incomplete column from the others,
+  and returns `m` completed datasets. Two methods are available: `'pmm'`
+  (predictive mean matching — the default, which imputes by copying observed
+  donor values) and `'norm'` (Bayesian linear regression). Defaults follow R's
+  `mice` package (`m=5`, `maxit=5`, 5 PMM donors, left-to-right visit order).
+- **Reproducible by construction.** A `seed` is required; given the same seed,
+  imputations are identical across runs and machines. Each of the `m` chains
+  draws from its own independent random stream.
+- **Rubin's-rules pooling.** `pystatistics.mice.pool(estimates, variances)`
+  combines an analysis fitted on each completed dataset into a single inference,
+  with Barnard–Rubin degrees of freedom, confidence intervals, and the usual
+  diagnostics (within/between/total variance, relative increase in variance,
+  fraction of missing information).
+- **Validated against R.** Imputed-value distributions and pooled regression
+  output are checked against R's `mice` on a shared dataset.
+- This release covers numeric columns on the CPU. Categorical-column methods
+  and GPU acceleration are planned for future releases.
+
+## 3.2.0
+
+Apple Silicon (MPS) GPU support for the FP32 backends, with honest,
+fail-fast boundaries where Metal can't deliver a real GPU path.
+
+- **GPU fitting now runs on Apple Silicon (MPS) for `multinom`, `polr`,
+  `gam`, and `arima` / `arima_batch` (Whittle).** Pass `backend='gpu'` on
+  a Mac with an Apple GPU to run these fits in FP32 on Metal. Every
+  operation on these paths is a native Metal kernel — there is no hidden
+  fallback to the CPU. Results match the CPU backend at the documented
+  GPU/FP32 tolerance tier. (MPS has no double precision, so `use_fp64=True`
+  is rejected on MPS; CUDA FP64 is unchanged.)
+- **`DataSource.to('mps')`** now transfers arrays to the Apple GPU,
+  downcasting float64 to float32 (MPS has no float64). This lets you pay
+  the host→device copy once and reuse a device-resident `DataSource`
+  across multiple GPU fits, as with CUDA.
+- **`backend='auto'` never selects MPS.** On Apple Silicon, `'auto'`
+  uses the CPU (double precision, the R-validated path); the Apple GPU is
+  opt-in only via an explicit `backend='gpu'`. CUDA continues to be
+  auto-selected. This matches how the regression and MVN MLE backends
+  already behaved.
+- **PCA and MVN MLE GPU remain CUDA-only.** `pca(backend='gpu')` and
+  `mlest(algorithm='em', backend='gpu')` now raise a clear error on Apple
+  Silicon instead of running silently on the CPU under a `gpu` label: PCA
+  needs an SVD / eigendecomposition that Metal does not implement, and the
+  EM algorithm's iterative, small-step pattern is far slower on Metal than
+  on the CPU. Use `backend='cpu'` (or `backend='auto'`, which selects the
+  CPU on MPS). CUDA is supported for both. MVN MLE *direct* (BFGS) GPU
+  fitting is unaffected and works on MPS.
+- **Whittle ARIMA GPU is more robust in FP32.** When the optimizer's line
+  search stalls at the FP32 noise floor on an already-converged fit, the
+  result is now accepted (it matches the CPU fit at the FP32 tier) instead
+  of raising a spurious convergence error. Non-stationary fits are still
+  rejected.
+
 ## 3.1.0
 
 Categorical predictors and interaction terms in regression.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pystatistics
-Version: 3.1.0
+Version: 3.3.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,6 +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 |
 
 See [docs/ROADMAP.md](docs/ROADMAP.md) for detailed scope, GPU applicability, and implementation priority for each module.
 
@@ -373,6 +374,27 @@ y = np.sin(x) + np.random.randn(200) * 0.3
 result = gam(y, smooths=[s('x1')], smooth_data={'x1': x})
 print(result.edf, result.gcv)
 print(result.summary())
+
+# --- Multiple imputation (MICE) ---
+from pystatistics.mice import mice, pool
+
+# data is an (n, p) array with np.nan marking missing values.
+# Predictive mean matching (R default) for numeric columns; seed is required
+# so the imputation is fully reproducible.
+imp = mice(data, m=5, maxit=5, method='pmm', seed=0)
+completed = imp.completed_datasets()        # list of 5 completed (n, p) arrays
+
+# Fit your analysis on each completed dataset, then combine with Rubin's rules:
+estimates, variances = [], []
+for d in completed:
+    X = np.column_stack([np.ones(len(d)), d[:, 1]])
+    beta, *_ = np.linalg.lstsq(X, d[:, 0], rcond=None)
+    resid = d[:, 0] - X @ beta
+    cov = (resid @ resid / (len(d) - 2)) * np.linalg.inv(X.T @ X)
+    estimates.append(beta[1]); variances.append(cov[1, 1])
+
+pooled = pool(estimates, variances, dfcom=len(data) - 2)
+print(pooled.estimate, pooled.se, pooled.ci_low, pooled.ci_high, pooled.fmi)
 ```
 
 ## Installation
@@ -391,6 +413,38 @@ pip install pystatistics[dev]
 
 ## What's New
 
+### 3.3.0 — Multiple imputation (MICE)
+
+- New `mice` module: multiple imputation by chained equations for numeric data
+  with missing values. `mice(data, m=5, method='pmm', seed=...)` returns `m`
+  completed datasets, using predictive mean matching (the R default) or
+  Bayesian linear regression (`method='norm'`). Defaults follow R's `mice`.
+- Imputation is fully reproducible — `seed` is required, and each chain uses an
+  independent random stream.
+- `pool(estimates, variances)` combines per-dataset analyses with Rubin's rules
+  (Barnard–Rubin degrees of freedom, confidence intervals, fraction of missing
+  information).
+- Numeric columns on the CPU in this release; validated against R's `mice`.
+
+### 3.2.0 — Apple Silicon (MPS) GPU support
+
+- `multinom`, `polr`, `gam`, and `arima` / `arima_batch` (Whittle) now run
+  on Apple Silicon GPUs with `backend='gpu'`, in FP32 and entirely on
+  native Metal kernels (no hidden CPU fallback). Results match the CPU
+  backend at the GPU/FP32 tolerance tier.
+- `DataSource.to('mps')` transfers data to the Apple GPU (float64 →
+  float32), so you can pay the host→device copy once and reuse it across
+  fits.
+- `backend='auto'` uses the CPU on Apple Silicon; the Apple GPU is opt-in
+  via an explicit `backend='gpu'`. CUDA is still auto-selected.
+- `pca` and MVN MLE `em` GPU paths remain CUDA-only and now raise a clear
+  error on Apple Silicon rather than silently running on the CPU — PCA's
+  SVD/eigendecomposition and the EM scatter/iteration pattern have no
+  efficient Metal equivalent. Use `backend='cpu'` or `'auto'` on a Mac.
+  (MVN MLE *direct* GPU fitting works on MPS.)
+- Whittle ARIMA GPU fits no longer raise a spurious convergence error when
+  the FP32 line search stalls at an already-converged optimum.
+
 ### 3.1.0 — Categorical predictors & interaction terms
 
 - Regression now supports categorical predictors and interactions via a

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

@@ -39,6 +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 |
 
 See [docs/ROADMAP.md](docs/ROADMAP.md) for detailed scope, GPU applicability, and implementation priority for each module.
 
@@ -326,6 +327,27 @@ y = np.sin(x) + np.random.randn(200) * 0.3
 result = gam(y, smooths=[s('x1')], smooth_data={'x1': x})
 print(result.edf, result.gcv)
 print(result.summary())
+
+# --- Multiple imputation (MICE) ---
+from pystatistics.mice import mice, pool
+
+# data is an (n, p) array with np.nan marking missing values.
+# Predictive mean matching (R default) for numeric columns; seed is required
+# so the imputation is fully reproducible.
+imp = mice(data, m=5, maxit=5, method='pmm', seed=0)
+completed = imp.completed_datasets()        # list of 5 completed (n, p) arrays
+
+# Fit your analysis on each completed dataset, then combine with Rubin's rules:
+estimates, variances = [], []
+for d in completed:
+    X = np.column_stack([np.ones(len(d)), d[:, 1]])
+    beta, *_ = np.linalg.lstsq(X, d[:, 0], rcond=None)
+    resid = d[:, 0] - X @ beta
+    cov = (resid @ resid / (len(d) - 2)) * np.linalg.inv(X.T @ X)
+    estimates.append(beta[1]); variances.append(cov[1, 1])
+
+pooled = pool(estimates, variances, dfcom=len(data) - 2)
+print(pooled.estimate, pooled.se, pooled.ci_low, pooled.ci_high, pooled.fmi)
 ```
 
 ## Installation
@@ -344,6 +366,38 @@ pip install pystatistics[dev]
 
 ## What's New
 
+### 3.3.0 — Multiple imputation (MICE)
+
+- New `mice` module: multiple imputation by chained equations for numeric data
+  with missing values. `mice(data, m=5, method='pmm', seed=...)` returns `m`
+  completed datasets, using predictive mean matching (the R default) or
+  Bayesian linear regression (`method='norm'`). Defaults follow R's `mice`.
+- Imputation is fully reproducible — `seed` is required, and each chain uses an
+  independent random stream.
+- `pool(estimates, variances)` combines per-dataset analyses with Rubin's rules
+  (Barnard–Rubin degrees of freedom, confidence intervals, fraction of missing
+  information).
+- Numeric columns on the CPU in this release; validated against R's `mice`.
+
+### 3.2.0 — Apple Silicon (MPS) GPU support
+
+- `multinom`, `polr`, `gam`, and `arima` / `arima_batch` (Whittle) now run
+  on Apple Silicon GPUs with `backend='gpu'`, in FP32 and entirely on
+  native Metal kernels (no hidden CPU fallback). Results match the CPU
+  backend at the GPU/FP32 tolerance tier.
+- `DataSource.to('mps')` transfers data to the Apple GPU (float64 →
+  float32), so you can pay the host→device copy once and reuse it across
+  fits.
+- `backend='auto'` uses the CPU on Apple Silicon; the Apple GPU is opt-in
+  via an explicit `backend='gpu'`. CUDA is still auto-selected.
+- `pca` and MVN MLE `em` GPU paths remain CUDA-only and now raise a clear
+  error on Apple Silicon rather than silently running on the CPU — PCA's
+  SVD/eigendecomposition and the EM scatter/iteration pattern have no
+  efficient Metal equivalent. Use `backend='cpu'` or `'auto'` on a Mac.
+  (MVN MLE *direct* GPU fitting works on MPS.)
+- Whittle ARIMA GPU fits no longer raise a spurious convergence error when
+  the FP32 line search stalls at an already-converged optimum.
+
 ### 3.1.0 — Categorical predictors & interaction terms
 
 - Regression now supports categorical predictors and interactions via a

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

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

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

@@ -16,7 +16,7 @@ Usage:
     result = fit(design)
 """
 
-__version__ = "3.1.0"
+__version__ = "3.3.0"
 __author__ = "Hai-Shuo"
 __email__ = "contact@sgcx.org"
 
@@ -34,6 +34,7 @@ from pystatistics import ordinal
 from pystatistics import multivariate
 from pystatistics import timeseries
 from pystatistics import gam
+from pystatistics import mice
 
 __all__ = [
     "__version__",
@@ -51,4 +52,5 @@ __all__ = [
     "multivariate",
     "timeseries",
     "gam",
+    "mice",
 ]

{pystatistics-3.1.0 → pystatistics-3.3.0}/pystatistics/core/compute/device.py

@@ -118,12 +118,23 @@ def select_device(prefer: Literal['cpu', 'gpu', 'auto'] = 'auto') -> DeviceInfo:
             - 'cpu': Always use CPU
             - 'gpu': Require GPU (raises if unavailable)
             - 'auto': Use GPU if available, else CPU
-            
+
     Returns:
         DeviceInfo for selected device
-        
+
     Raises:
         RuntimeError: If 'gpu' requested but no GPU available
+
+    Note:
+        This is a hardware *detector*: with ``prefer='auto'`` it returns
+        the best available GPU, which includes Apple Silicon MPS. It does
+        NOT encode the project's dispatch policy. The fitting functions
+        deliberately do NOT auto-select MPS (it is FP32-only and not the
+        R-validated default); they treat ``'auto'`` as "GPU only if
+        ``device_type == 'cuda'``, else CPU", and run on MPS only when the
+        caller passes ``backend='gpu'`` explicitly. Keep that check in the
+        caller, not here, so ``select_device`` stays a pure capability
+        query.
     """
     if prefer == 'cpu':
         return get_cpu_info()

{pystatistics-3.1.0 → pystatistics-3.3.0}/pystatistics/core/compute/linalg/batched.py

@@ -145,9 +145,16 @@ def _batched_ols_gpu(X: NDArray, Y: NDArray, device: str) -> NDArray:
             L.T, Z, upper=True,
         )
     except torch._C._LinAlgError:
-        # Fallback: lstsq
-        result = torch.linalg.lstsq(X_t, Y_t)
-        B = result.solution
+        # Fallback for rank-deficient X: lstsq. MPS has no lstsq, and
+        # the normal-equations Cholesky we just tried is exactly what
+        # failed, so route the rank-revealing solve through CPU LAPACK
+        # (the matrices are small here and this path is rare).
+        if torch_device.type == 'mps':
+            result = torch.linalg.lstsq(X_t.cpu(), Y_t.cpu())
+            B = result.solution.to(torch_device)
+        else:
+            result = torch.linalg.lstsq(X_t, Y_t)
+            B = result.solution
 
     # Transfer back to CPU
     return B.cpu().numpy().astype(np.float64)

pystatistics-3.3.0/pystatistics/core/compute/torch_interop.py

@@ -0,0 +1,48 @@
+"""Host/device transfer helpers for the PyTorch GPU backends.
+
+One job: move tensors off a compute device into host numpy arrays
+correctly across CUDA and Apple Silicon (MPS).
+
+The single invariant this module exists to enforce:
+
+    Cast to float64 only AFTER moving the tensor to the host.
+
+MPS has no float64 dtype, so an on-device ``tensor.to(torch.float64)``
+raises ``TypeError: Cannot convert a MPS Tensor to float64``. The
+download must therefore be ``.cpu().to(torch.float64)``, never
+``.to(torch.float64).cpu()``. Centralising it here means no backend
+can reintroduce the device-side cast by accident (Coding Bible: make
+the wrong thing hard to do accidentally).
+
+This is also correct and lossless on CUDA: for a float64-on-device
+tensor the host-side cast is a no-op; for a float32 tensor the
+resulting float64 values are identical regardless of cast order.
+
+``torch`` is imported lazily inside each function so that importing
+this module never pulls torch into a CPU-only install.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def to_host_f64(tensor: Any) -> NDArray[np.float64]:
+    """Download a torch tensor to a contiguous float64 numpy array.
+
+    Detaches from autograd, moves to the host, then casts to float64
+    (in that order — see the module docstring for why the order is
+    load-bearing on MPS). Safe on CUDA, MPS, and CPU tensors.
+
+    Args:
+        tensor: A ``torch.Tensor`` on any device.
+
+    Returns:
+        A float64 ``numpy.ndarray`` with the tensor's values.
+    """
+    import torch
+
+    return tensor.detach().cpu().to(torch.float64).numpy()

{pystatistics-3.1.0 → pystatistics-3.3.0}/pystatistics/core/datasource.py

@@ -385,11 +385,19 @@ class DataSource:
             else:
                 # Coerce numpy→torch on device; move torch to device.
                 if isinstance(value, torch.Tensor):
-                    new_storage[key] = value.to(target)
+                    tensor = value
                 elif isinstance(value, np.ndarray):
-                    new_storage[key] = torch.as_tensor(value).to(target)
+                    tensor = torch.as_tensor(value)
                 else:
                     new_storage[key] = value
+                    continue
+                # MPS has no float64. Downcast double tensors to float32
+                # before transfer so the device-resident path works on
+                # Apple Silicon; CUDA keeps the source dtype (FP64 is the
+                # R-validated path there).
+                if target.type == "mps" and tensor.dtype == torch.float64:
+                    tensor = tensor.to(torch.float32)
+                new_storage[key] = tensor.to(target)
 
         capabilities = {CAPABILITY_MATERIALIZED, CAPABILITY_REPEATABLE}
         if target_is_gpu:

{pystatistics-3.1.0 → pystatistics-3.3.0}/pystatistics/gam/_gam.py

@@ -169,7 +169,11 @@ def gam(
     if backend != "cpu":
         from pystatistics.core.compute.device import select_device
         dev = select_device("gpu" if backend == "gpu" else "auto")
-        if dev.is_gpu:
+        # backend='auto' must not select MPS: it is FP32-only and not the
+        # R-validated default. MPS runs only on explicit backend='gpu';
+        # 'auto' uses the GPU only for CUDA (matches the regression and
+        # mvnmle dispatch policy).
+        if dev.is_gpu and (backend == "gpu" or dev.device_type == "cuda"):
             from pystatistics.gam.backends.gpu_pirls import GAMGPUFitter
             try:
                 gpu_fitter = GAMGPUFitter(

{pystatistics-3.1.0 → pystatistics-3.3.0}/pystatistics/gam/backends/gpu_pirls.py

@@ -32,6 +32,7 @@ from typing import Any
 import numpy as np
 from numpy.typing import NDArray
 
+from pystatistics.core.compute.torch_interop import to_host_f64
 from pystatistics.gam.backends._gpu_family import (
     GPUFamilyOps,
     resolve_gpu_family,
@@ -173,8 +174,8 @@ class GAMGPUFitter:
         FP64 precision.
         """
         torch = self._torch
-        A_np = A.detach().to(torch.float64).cpu().numpy()
-        b_np = b.detach().to(torch.float64).cpu().numpy()
+        A_np = to_host_f64(A)
+        b_np = to_host_f64(b)
         try:
             L = np.linalg.cholesky(A_np)
             beta_np = np.linalg.solve(L.T, np.linalg.solve(L, b_np))
@@ -205,8 +206,8 @@ class GAMGPUFitter:
         ``X'WX``) stays on device; only the p×p solve goes via host.
         """
         torch = self._torch
-        A_np = A.detach().to(torch.float64).cpu().numpy()
-        XtWX_np = XtWX.detach().to(torch.float64).cpu().numpy()
+        A_np = to_host_f64(A)
+        XtWX_np = to_host_f64(XtWX)
         try:
             F_np = np.linalg.solve(A_np, XtWX_np)
         except np.linalg.LinAlgError:
@@ -330,10 +331,10 @@ class GAMGPUFitter:
         mu = self._fam.linkinv(eta)
 
         return (
-            beta.detach().to(torch.float64).cpu().numpy(),
-            mu.detach().to(torch.float64).cpu().numpy(),
-            eta.detach().to(torch.float64).cpu().numpy(),
-            w.detach().to(torch.float64).cpu().numpy(),
+            to_host_f64(beta),
+            to_host_f64(mu),
+            to_host_f64(eta),
+            to_host_f64(w),
             float(dev_t.detach().cpu().item()),
             n_iter,
             converged,
@@ -426,10 +427,10 @@ class GAMGPUFitter:
         return (
             np.asarray(edf_list, dtype=np.float64),
             total_edf,
-            beta.detach().to(torch.float64).cpu().numpy(),
-            mu.detach().to(torch.float64).cpu().numpy(),
-            eta.detach().to(torch.float64).cpu().numpy(),
-            w.detach().to(torch.float64).cpu().numpy(),
+            to_host_f64(beta),
+            to_host_f64(mu),
+            to_host_f64(eta),
+            to_host_f64(w),
             float(dev_t.detach().cpu().item()),
             n_iter,
             converged,

pystatistics-3.3.0/pystatistics/mice/__init__.py

@@ -0,0 +1,29 @@
+"""
+MICE — Multiple Imputation by Chained Equations.
+
+Iteratively impute each incomplete column from the others until the chain
+stabilises, producing ``m`` completed datasets, then combine analyses across
+them with Rubin's rules. R-pegged: numeric defaults follow R's ``mice`` package
+(predictive mean matching for numeric columns).
+
+Public API (Stage 1, CPU, numeric columns):
+
+    >>> from pystatistics.mice import mice, pool
+    >>> imp = mice(data, m=5, method='pmm', seed=0)
+    >>> completed = imp.completed_datasets()      # list of m arrays
+
+Stage 1 supports numeric columns only; categorical methods are planned.
+"""
+
+from pystatistics.mice.design import MICEDesign
+from pystatistics.mice.pooling import PooledResult, pool
+from pystatistics.mice.solution import MICESolution
+from pystatistics.mice.solvers import mice
+
+__all__ = [
+    "MICEDesign",
+    "MICESolution",
+    "PooledResult",
+    "mice",
+    "pool",
+]

pystatistics-3.3.0/pystatistics/mice/_chain.py

@@ -0,0 +1,113 @@
+"""
+The chained-equations sweep — one imputation chain.
+
+A single chain produces one completed dataset:
+
+  1. *Initialise* every missing cell by a random draw from the observed values of
+     its own column (R mice's default start), so all predictor columns are fully
+     populated from iteration 1.
+  2. *Sweep* ``maxit`` times. In each iteration, visit the incomplete columns in
+     the visit sequence; for each, fit the column's assigned method on the rows
+     where it is observed (using all other columns as predictors) and overwrite
+     its missing cells with fresh draws.
+  3. Record, per iteration and per incomplete column, the mean and variance of
+     the imputed cells — the trace MICE convergence diagnostics are read from.
+
+This module owns orchestration only. The per-variable statistical work lives
+behind the :class:`ImputationMethod` protocol, so the sweep never hard-codes a
+method and a GPU backend can swap in batched method implementations without
+touching this control flow.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from pystatistics.mice.design import MICEDesign
+from pystatistics.mice.methods import get_method
+
+
+@dataclass(frozen=True)
+class ChainResult:
+    """Output of one chain: completed data plus convergence traces."""
+
+    completed: np.ndarray            # (n, p) — no NaN
+    chain_mean: np.ndarray           # (maxit, n_incomplete)
+    chain_var: np.ndarray            # (maxit, n_incomplete)
+    incomplete_columns: tuple[int, ...]
+
+
+def run_chain(
+    design: MICEDesign,
+    rng: np.random.Generator,
+    maxit: int,
+    visit_sequence: tuple[int, ...],
+) -> ChainResult:
+    """Run one chained-equations chain. See module docstring."""
+    data = design.data.copy()
+    mask = design.missing_mask
+    p = design.p
+    incomplete = design.incomplete_columns
+
+    _initialise(data, mask, incomplete, rng)
+
+    chain_mean = np.empty((maxit, len(incomplete)), dtype=np.float64)
+    chain_var = np.empty((maxit, len(incomplete)), dtype=np.float64)
+    incomplete_pos = {j: k for k, j in enumerate(incomplete)}
+
+    for it in range(maxit):
+        for j in visit_sequence:
+            mis_rows = mask[:, j]
+            obs_rows = ~mis_rows
+
+            predictors = _predictor_columns(data, j, p)
+            X_obs = predictors[obs_rows]
+            X_mis = predictors[mis_rows]
+            y_obs = data[obs_rows, j]
+
+            method = get_method(design.method_for(j))
+            imputed = method.impute(y_obs, X_obs, X_mis, rng)
+            data[mis_rows, j] = imputed
+
+        # Trace: summarise this iteration's imputed cells per incomplete column.
+        for j in incomplete:
+            cells = data[mask[:, j], j]
+            k = incomplete_pos[j]
+            chain_mean[it, k] = float(np.mean(cells))
+            chain_var[it, k] = float(np.var(cells))
+
+    return ChainResult(
+        completed=data,
+        chain_mean=chain_mean,
+        chain_var=chain_var,
+        incomplete_columns=incomplete,
+    )
+
+
+def _initialise(
+    data: np.ndarray,
+    mask: np.ndarray,
+    incomplete: tuple[int, ...],
+    rng: np.random.Generator,
+) -> None:
+    """Fill missing cells by sampling (with replacement) from observed values of
+    the same column. Mutates ``data`` in place."""
+    for j in incomplete:
+        mis_rows = mask[:, j]
+        observed = data[~mis_rows, j]
+        n_missing = int(np.count_nonzero(mis_rows))
+        data[mis_rows, j] = rng.choice(observed, size=n_missing, replace=True)
+
+
+def _predictor_columns(data: np.ndarray, j: int, p: int) -> np.ndarray:
+    """All columns except ``j`` (the current target), as a 2D predictor matrix.
+
+    Columns are fully populated at this point (initialised on entry, kept
+    complete by every visit), so the returned matrix is finite. Fancy-indexing
+    with a list always yields a 2D matrix, including the single-predictor
+    (p == 2) case.
+    """
+    cols = [c for c in range(p) if c != j]
+    return data[:, cols]