pystatistics 3.4.0__tar.gz → 3.5.0__tar.gz

pystatistics-3.5.0/.release/UNRELEASED.md

@@ -0,0 +1,39 @@
+# 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
+
+- Categorical imputation for `mice`. Columns can now be declared
+  `binary` / `categorical` (unordered) / `ordered` via `column_kinds`, and are
+  imputed with the matching R-mice method: `logreg` (Bayesian logistic),
+  `polyreg` (multinomial logit, reusing the `multinomial` module), and `polr`
+  (proportional odds, reusing the `ordinal` module). Each draws its parameters
+  from the posterior normal approximation, predicts class probabilities for the
+  missing rows, and samples a category.
+- `method='auto'` (new default) picks the R-default method per column kind
+  (numeric->pmm, binary->logreg, categorical->polyreg, ordered->polr). An
+  explicit method name still applies to all incomplete columns. Categorical
+  columns must be encoded as integer category codes.
+- Categorical predictors are treatment-dummy-encoded in the sweep
+  (`pystatistics/mice/_encode.py`); the numeric-only path keeps a fast slice so
+  its performance is unchanged. The chain maps categorical targets to/from
+  consecutive `0..K-1` class indices.
+- Validated distributionally against R `mice` 3.19.0: imputed category marginal
+  proportions for logreg/polyreg/polr match R within ~0.015 on a shared
+  mixed-type dataset (`tests/mice/references/`).
+- Robustness: if a categorical model fit fails to converge on an awkward
+  intermediate sweep state, the method falls back to a marginal draw for that
+  step with a visible `UserWarning` (matching how R wraps MASS::polr), retrying
+  the full conditional model on the next iteration. Observed ~0.1% of fits.
+- The GPU backend now explicitly refuses data with any categorical column
+  (categorical imputation is CPU-only); previously only the target method was
+  checked.
+- Added a public `vcov` property to the multinomial solution
+  (`pystatistics/multinomial/solution.py`) for parity with the ordinal module,
+  so the imputation method reads only public surfaces.

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

@@ -1,5 +1,44 @@
 # Changelog
 
+## 3.5.0
+
+Categorical imputation for MICE.
+
+- **`mice` now imputes categorical columns**, not just numeric ones. Declare a
+  column's kind with `column_kinds` — `'binary'`, `'categorical'` (unordered),
+  or `'ordered'` — and it is imputed with the matching method: logistic
+  regression for binary columns, multinomial logistic for unordered factors, and
+  proportional-odds (ordinal logistic) for ordered factors. These mirror R
+  `mice`'s `logreg`, `polyreg`, and `polr`. Categorical columns are given as
+  integer category codes.
+- **`method='auto'` is the new default** and picks the appropriate method for
+  each column's kind (predictive mean matching for numeric, logistic for binary,
+  multinomial for unordered, proportional odds for ordered). Passing an explicit
+  method name still applies it to every incomplete column.
+- Categorical columns are handled correctly as predictors too (dummy-encoded),
+  so mixed numeric/categorical datasets impute coherently. Validated against R
+  `mice`: imputed category proportions match closely for all three methods.
+- If a categorical model cannot be fit for a particular step during the
+  iterations, that step falls back to a draw from the observed category
+  distribution (with a warning) and the full model is retried next iteration.
+- GPU acceleration remains numeric-only; running `backend='gpu'` on data with
+  categorical columns raises a clear error (use the CPU backend).
+
+## 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.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pystatistics
-Version: 3.4.0
+Version: 3.5.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 | Multiple imputation by chained equations: PMM and Bayesian linear regression, Rubin's-rules pooling, validated against R mice; CUDA GPU backend |
+| `mice/` | Complete | Multiple imputation by chained equations: numeric (PMM, Bayesian regression) and categorical (logistic, multinomial, proportional-odds), Rubin's-rules pooling, validated against R mice; CUDA GPU backend for numeric |
 
 See [docs/ROADMAP.md](docs/ROADMAP.md) for detailed scope, GPU applicability, and implementation priority for each module.
 
@@ -413,6 +413,28 @@ pip install pystatistics[dev]
 
 ## What's New
 
+### 3.5.0 — Categorical imputation for MICE
+
+- `mice` now imputes categorical columns, not only numeric ones. Declare each
+  column's kind via `column_kinds` (`'binary'`, `'categorical'`, `'ordered'`)
+  and it is imputed with logistic, multinomial, or proportional-odds regression
+  respectively — mirroring R `mice`'s `logreg`/`polyreg`/`polr`. Categorical
+  columns are integer category codes.
+- `method='auto'` (the new default) selects the right method per column kind;
+  mixed numeric/categorical datasets impute coherently (categorical predictors
+  are dummy-encoded). Imputed category proportions are validated against R
+  `mice`.
+- GPU acceleration stays numeric-only; categorical imputation runs on the CPU.
+
+### 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.5.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 | Multiple imputation by chained equations: PMM and Bayesian linear regression, Rubin's-rules pooling, validated against R mice; CUDA GPU backend |
+| `mice/` | Complete | Multiple imputation by chained equations: numeric (PMM, Bayesian regression) and categorical (logistic, multinomial, proportional-odds), Rubin's-rules pooling, validated against R mice; CUDA GPU backend for numeric |
 
 See [docs/ROADMAP.md](docs/ROADMAP.md) for detailed scope, GPU applicability, and implementation priority for each module.
 
@@ -366,6 +366,28 @@ pip install pystatistics[dev]
 
 ## What's New
 
+### 3.5.0 — Categorical imputation for MICE
+
+- `mice` now imputes categorical columns, not only numeric ones. Declare each
+  column's kind via `column_kinds` (`'binary'`, `'categorical'`, `'ordered'`)
+  and it is imputed with logistic, multinomial, or proportional-odds regression
+  respectively — mirroring R `mice`'s `logreg`/`polyreg`/`polr`. Categorical
+  columns are integer category codes.
+- `method='auto'` (the new default) selects the right method per column kind;
+  mixed numeric/categorical datasets impute coherently (categorical predictors
+  are dummy-encoded). Imputed category proportions are validated against R
+  `mice`.
+- GPU acceleration stays numeric-only; categorical imputation runs on the CPU.
+
+### 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.5.0}/pyproject.toml

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

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

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

{pystatistics-3.4.0 → pystatistics-3.5.0}/pystatistics/mice/_chain.py

@@ -25,6 +25,11 @@ from dataclasses import dataclass
 
 import numpy as np
 
+from pystatistics.mice._encode import (
+    build_predictor_matrix,
+    codes_to_indices,
+    indices_to_codes,
+)
 from pystatistics.mice.design import MICEDesign
 from pystatistics.mice.methods import get_method
 
@@ -48,7 +53,6 @@ def run_chain(
     """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)
@@ -62,14 +66,21 @@ def run_chain(
             mis_rows = mask[:, j]
             obs_rows = ~mis_rows
 
-            predictors = _predictor_columns(data, j, p)
+            predictors = build_predictor_matrix(data, j, design)
             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
+            if design.is_categorical(j):
+                # Methods speak in consecutive 0..K-1 class indices; the chain
+                # translates to/from the column's stored category codes.
+                levels = design.levels_for(j)
+                y_idx = codes_to_indices(y_obs, levels)
+                imputed_idx = method.impute(y_idx, X_obs, X_mis, rng)
+                data[mis_rows, j] = indices_to_codes(imputed_idx, levels)
+            else:
+                data[mis_rows, j] = method.impute(y_obs, X_obs, X_mis, rng)
 
         # Trace: summarise this iteration's imputed cells per incomplete column.
         for j in incomplete:
@@ -99,15 +110,3 @@ def _initialise(
         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]

pystatistics-3.5.0/pystatistics/mice/_encode.py

@@ -0,0 +1,73 @@
+"""
+Predictor encoding and category-code mapping for the chained-equations sweep.
+
+When a column is imputed, the other columns serve as predictors. Numeric columns
+enter the model as-is, but a *categorical* predictor must be dummy-encoded — its
+integer codes are labels, not magnitudes, so feeding the raw codes would impose a
+spurious linear order. This module builds the encoded predictor matrix and maps
+categorical target values between their stored codes and the consecutive
+``0..K-1`` class indices the categorical methods expect.
+
+Contrast choice: categorical predictors use treatment (drop-first) dummies. For
+a model's *predictions* this is equivalent to any other full-rank contrast
+(treatment, polynomial, …) — they span the same column space — so this matches
+R's model matrix for imputation purposes regardless of which contrast R picks.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def build_predictor_matrix(data: NDArray, j: int, design) -> NDArray:
+    """Predictor matrix for imputing column ``j``: every other column, with
+    categorical columns treatment-dummy-encoded. No intercept column (methods
+    add one where they need it).
+
+    Fast path: when no column is categorical this is a plain column slice, so the
+    all-numeric sweep keeps its original cost.
+    """
+    p = design.p
+    if not design.has_categorical:
+        cols = [c for c in range(p) if c != j]
+        return data[:, cols]
+
+    blocks = []
+    for c in range(p):
+        if c == j:
+            continue
+        if design.is_categorical(c):
+            blocks.append(_treatment_dummies(data[:, c], design.levels_for(c)))
+        else:
+            blocks.append(data[:, c].reshape(-1, 1))
+    if not blocks:
+        return np.empty((data.shape[0], 0), dtype=np.float64)
+    return np.hstack(blocks)
+
+
+def _treatment_dummies(col: NDArray, levels: NDArray) -> NDArray:
+    """One-hot encode ``col`` against ``levels``, dropping the first level as the
+    reference. Returns (n, len(levels)-1). Values are exact category codes (drawn
+    from observed values or imputed back as codes), so equality is exact."""
+    # Columns for levels[1:]; the reference level maps to an all-zero row.
+    return (col[:, None] == levels[None, 1:]).astype(np.float64)
+
+
+def codes_to_indices(values: NDArray, levels: NDArray) -> NDArray:
+    """Map category codes to consecutive ``0..K-1`` indices (levels are sorted)."""
+    return np.searchsorted(levels, values).astype(np.intp)
+
+
+def indices_to_codes(indices: NDArray, levels: NDArray) -> NDArray:
+    """Map ``0..K-1`` class indices back to the stored category codes."""
+    return levels[np.asarray(indices, dtype=np.intp)]
+
+
+def add_intercept(X: NDArray) -> NDArray:
+    """Prepend a column of ones (for models that carry an explicit intercept)."""
+    X = np.asarray(X, dtype=np.float64)
+    if X.ndim == 1:
+        X = X.reshape(-1, 1)
+    ones = np.ones((X.shape[0], 1), dtype=np.float64)
+    return np.hstack([ones, X])

{pystatistics-3.4.0 → pystatistics-3.5.0}/pystatistics/mice/backends/gpu.py

@@ -59,6 +59,18 @@ class GPUMiceBackend:
     ) -> Result[MICEParams]:
         import torch
 
+        # The GPU backend handles fully numeric problems only: categorical
+        # columns need dummy-encoding / categorical model fits that this batched
+        # path does not implement. Refuse rather than impute them wrong (a
+        # categorical column can be a *predictor* even when the target is
+        # numeric, so this checks every column, not just the targets).
+        if design.has_categorical:
+            raise ValidationError(
+                "backend='gpu' supports numeric columns only. This data has "
+                "categorical columns; use backend='cpu' (categorical imputation "
+                "is CPU-only)."
+            )
+
         # 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:

{pystatistics-3.4.0 → pystatistics-3.5.0}/pystatistics/mice/design.py

@@ -28,10 +28,23 @@ from numpy.typing import NDArray
 from pystatistics.core.exceptions import ValidationError
 from pystatistics.mice.methods import get_method, is_registered
 
-# Stage 1 supports numeric columns only. This single constant + the guard in
-# ``_resolve_kinds`` is the entire "numeric-only" restriction; Stage 3 widens it.
-_SUPPORTED_KINDS = ("numeric",)
-_DEFAULT_METHOD = "pmm"
+# Statistical kinds a column can have. Numeric columns are imputed with
+# regression methods (pmm/norm); the others are categorical.
+_SUPPORTED_KINDS = ("numeric", "binary", "categorical", "ordered")
+
+# Kinds whose values are category codes (integers) rather than measurements.
+_CATEGORICAL_KINDS = ("binary", "categorical", "ordered")
+
+# Default imputation method per kind, mirroring R mice's defaults: PMM for
+# numeric, logistic for binary, multinomial for unordered factors, proportional
+# odds for ordered factors. Used when ``method='auto'`` (the default).
+_DEFAULT_METHOD_BY_KIND = {
+    "numeric": "pmm",
+    "binary": "logreg",
+    "categorical": "polyreg",
+    "ordered": "polr",
+}
+_AUTO_METHOD = "auto"
 
 
 @dataclass(frozen=True)
@@ -43,6 +56,10 @@ class MICEDesign:
     _col_names: tuple[str, ...]
     _col_kinds: tuple[str, ...]
     _methods: tuple[str, ...]  # per column; '' for fully observed columns
+    # Per column: sorted unique category codes for categorical columns, or None
+    # for numeric columns. Drives dummy-encoding of categorical predictors and
+    # code<->index mapping for categorical targets.
+    _levels: tuple[NDArray[np.floating[Any]] | None, ...]
     _n: int
     _p: int
 
@@ -52,7 +69,7 @@ class MICEDesign:
         cls,
         data,
         *,
-        method: str = _DEFAULT_METHOD,
+        method: str = _AUTO_METHOD,
         methods: Mapping[Any, str] | Sequence[str] | None = None,
         column_names: Sequence[str] | None = None,
         column_kinds: Sequence[str] | None = None,
@@ -64,10 +81,13 @@ class MICEDesign:
         data : array-like
             2D matrix (observations x variables); NaN marks missing entries.
             Accepts numpy arrays and anything with a ``.values`` attribute
-            (e.g. pandas DataFrame).
+            (e.g. pandas DataFrame). Categorical columns must be encoded as
+            integer category codes (declare them via ``column_kinds``).
         method : str
-            Default imputation method for every incomplete column (R default
-            ``'pmm'``). Overridden per-column by ``methods``.
+            Default imputation method. ``'auto'`` (default) picks the R-default
+            method for each incomplete column's kind (pmm/logreg/polyreg/polr).
+            An explicit method name is applied to every incomplete column.
+            Overridden per-column by ``methods``.
         methods : mapping or sequence, optional
             Per-column method override. Either a mapping ``{name_or_index:
             method}`` or a length-``p`` sequence. Entries for fully observed
@@ -93,7 +113,7 @@ class MICEDesign:
         source,
         *,
         columns: Sequence[str] | None = None,
-        method: str = _DEFAULT_METHOD,
+        method: str = _AUTO_METHOD,
         methods: Mapping[Any, str] | Sequence[str] | None = None,
     ) -> "MICEDesign":
         """Build a MICEDesign from a DataSource (selected columns, in order)."""
@@ -177,6 +197,7 @@ class MICEDesign:
 
         names = _resolve_names(column_names, p)
         kinds = _resolve_kinds(column_kinds, p)
+        levels = _resolve_levels(data, missing_mask, kinds, names, p)
         per_col_methods = _resolve_methods(
             method, methods, names, kinds, has_missing_per_col, p
         )
@@ -187,6 +208,7 @@ class MICEDesign:
             _col_names=names,
             _col_kinds=kinds,
             _methods=per_col_methods,
+            _levels=levels,
             _n=n,
             _p=p,
         )
@@ -244,6 +266,23 @@ class MICEDesign:
         """Method name assigned to column ``col`` ('' if fully observed)."""
         return self._methods[col]
 
+    def kind_for(self, col: int) -> str:
+        """Statistical kind of column ``col``."""
+        return self._col_kinds[col]
+
+    def is_categorical(self, col: int) -> bool:
+        """Whether column ``col`` holds category codes (not measurements)."""
+        return self._col_kinds[col] in _CATEGORICAL_KINDS
+
+    def levels_for(self, col: int):
+        """Sorted unique category codes for column ``col`` (None if numeric)."""
+        return self._levels[col]
+
+    @property
+    def has_categorical(self) -> bool:
+        """Whether any column is categorical."""
+        return any(k in _CATEGORICAL_KINDS for k in self._col_kinds)
+
     def __repr__(self) -> str:
         return (
             f"MICEDesign(n={self._n}, p={self._p}, "
@@ -288,18 +327,48 @@ def _resolve_kinds(column_kinds, p: int) -> tuple[str, ...]:
             raise ValidationError(
                 f"column_kinds has length {len(kinds)}, expected {p}"
             )
-    # Stage-1 numeric-only guard — the single point where the restriction lives.
     bad = [(j, k) for j, k in enumerate(kinds) if k not in _SUPPORTED_KINDS]
     if bad:
         j, k = bad[0]
         raise ValidationError(
-            f"Column {j} has kind {k!r}, but this release supports only "
-            f"numeric columns {_SUPPORTED_KINDS}. Categorical methods "
-            f"(binary/categorical/ordered) are planned for a later release."
+            f"Column {j} has kind {k!r}. Supported kinds: {_SUPPORTED_KINDS}."
         )
     return tuple(kinds)
 
 
+def _resolve_levels(data, missing_mask, kinds, names, p):
+    """Compute and validate category codes for each categorical column.
+
+    For a categorical column the observed values must be integer category codes;
+    the level set is their sorted unique values. Numeric columns get ``None``.
+    Binary columns must have exactly two levels; categorical/ordered need >= 2.
+    """
+    levels: list = []
+    for j in range(p):
+        if kinds[j] not in _CATEGORICAL_KINDS:
+            levels.append(None)
+            continue
+        observed = data[~missing_mask[:, j], j]
+        if not np.all(observed == np.floor(observed)):
+            raise ValidationError(
+                f"Column {names[j]!r} is {kinds[j]!r} but has non-integer "
+                f"values; categorical columns must be integer category codes."
+            )
+        uniq = np.unique(observed)
+        if kinds[j] == "binary" and uniq.size != 2:
+            raise ValidationError(
+                f"Column {names[j]!r} is 'binary' but has {uniq.size} observed "
+                f"levels {uniq.tolist()}; binary columns need exactly 2."
+            )
+        if uniq.size < 2:
+            raise ValidationError(
+                f"Column {names[j]!r} is {kinds[j]!r} but has only "
+                f"{uniq.size} observed level; need at least 2 to impute."
+            )
+        levels.append(uniq.astype(np.float64))
+    return tuple(levels)
+
+
 def _resolve_methods(
     method: str,
     methods,
@@ -308,9 +377,15 @@ def _resolve_methods(
     has_missing_per_col: NDArray[np.bool_],
     p: int,
 ) -> tuple[str, ...]:
-    """Assign a validated method name to each column (per-column override)."""
-    # Start with the default for every column.
-    resolved = [method] * p
+    """Assign a validated method name to each column (per-column override).
+
+    ``method='auto'`` resolves to the R-default method for each column's kind;
+    an explicit method name is applied to every column.
+    """
+    if method == _AUTO_METHOD:
+        resolved = [_DEFAULT_METHOD_BY_KIND[kinds[j]] for j in range(p)]
+    else:
+        resolved = [method] * p
 
     if methods is not None:
         if isinstance(methods, Mapping):

{pystatistics-3.4.0 → pystatistics-3.5.0}/pystatistics/mice/methods/__init__.py

@@ -17,6 +17,9 @@ from pystatistics.mice.methods.registry import (
 # Side-effect imports: each module registers its method on import.
 from pystatistics.mice.methods import norm as _norm  # noqa: F401
 from pystatistics.mice.methods import pmm as _pmm  # noqa: F401
+from pystatistics.mice.methods import logreg as _logreg  # noqa: F401
+from pystatistics.mice.methods import polyreg as _polyreg  # noqa: F401
+from pystatistics.mice.methods import polr as _polr  # noqa: F401
 
 __all__ = [
     "ImputationMethod",

pystatistics-3.5.0/pystatistics/mice/methods/_draw.py

@@ -0,0 +1,72 @@
+"""
+Shared multivariate-normal posterior draw for the categorical methods.
+
+The categorical imputation methods (logreg, polyreg, polr) all follow the same
+Bayesian recipe R's ``mice`` uses: fit the model, then draw the parameter vector
+once from its asymptotic posterior ``N(theta_hat, V)`` before predicting and
+sampling. Drawing the parameters — rather than reusing the point estimate —
+injects the between-imputation variability that makes multiple imputation
+produce valid standard errors.
+
+``V`` is a covariance (or inverse-information) matrix, positive definite in
+theory but occasionally marginally indefinite in finite precision, so the
+Cholesky factor is taken with the same jittered, eigenvalue-clipped fallback as
+the numeric path (reused from ``_linreg``).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from pystatistics.mice.methods._linreg import _safe_cholesky
+
+
+def mvn_draw(mean: np.ndarray, cov: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Draw one sample from ``N(mean, cov)`` via its Cholesky factor.
+
+    Parameters
+    ----------
+    mean : (d,) array
+        Posterior mean (the point estimate).
+    cov : (d, d) array
+        Posterior covariance of the estimate.
+    rng : numpy.random.Generator
+        Sole randomness source.
+    """
+    mean = np.asarray(mean, dtype=np.float64).ravel()
+    L = _safe_cholesky(np.asarray(cov, dtype=np.float64))
+    return mean + L @ rng.standard_normal(mean.shape[0])
+
+
+def sample_categories(probs: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Sample one class index per row from a (n, K) matrix of class probabilities.
+
+    Robust to small negative probabilities (which can arise, e.g., from drawn
+    ordinal thresholds that are not perfectly ordered): negatives are clipped and
+    each row renormalised before inverse-CDF sampling.
+    """
+    probs = np.clip(np.asarray(probs, dtype=np.float64), 0.0, None)
+    row_sums = probs.sum(axis=1, keepdims=True)
+    # Degenerate rows (all-zero) fall back to uniform.
+    row_sums = np.where(row_sums > 0, row_sums, 1.0)
+    probs = probs / row_sums
+
+    cdf = np.cumsum(probs, axis=1)
+    u = rng.random(probs.shape[0])
+    # Smallest index whose cumulative probability reaches u.
+    return (cdf >= u[:, None]).argmax(axis=1).astype(np.intp)
+
+
+def marginal_indices(
+    y_obs: np.ndarray, n_mis: int, rng: np.random.Generator
+) -> np.ndarray:
+    """Sample ``n_mis`` class indices from the observed empirical distribution.
+
+    Used as the documented fallback when a categorical model fit fails to
+    converge mid-sweep (see the categorical methods). It preserves the observed
+    marginal distribution of the column; the next iteration retries the full
+    conditional model, so the fallback is local and self-correcting.
+    """
+    return rng.choice(
+        np.asarray(y_obs, dtype=np.intp), size=int(n_mis), replace=True
+    )